base.ts
1 import type { IsNever } from 'type-fest'; 2 import type { StorageItemKey, Unwatch, WxtStorageItem } from 'wxt/storage'; 3 4 import { EventDispatcher, EventListener } from '@/events'; 5 import { filterMap } from '@/helpers/filterMap'; 6 import type { MaybePromise, ToReadonly } from '@/types/util'; 7 import type { Logger } from '@/utils/logger'; 8 9 import { NoShardsError } from './errors'; 10 import { 11 type StorageEvent, 12 StorageEventTrigger, 13 type StorageMeta, 14 type StorageSyncMeta, 15 storageListenerArgs, 16 storageStateType, 17 } from './types'; 18 19 export type AnyStorageBase = StorageBase<any, any, any, any, any, any>; 20 export type AnySyncableStorage = SyncableStorage<any, any, any, any, any, any, any>; 21 22 export type StorageWatchCallback<State, ListenerArgs extends unknown[] = []> = ( 23 state: ToReadonly<State>, 24 ...args: ListenerArgs 25 ) => void; 26 27 export interface StorageItem<Value> { 28 key: StorageItemKey; 29 value: Value; 30 } 31 32 export interface StorageShard<Prefix extends string = never, Value = unknown> { 33 key: IsNever<Prefix> extends true ? string | undefined : `${Prefix}:${string}`; 34 value: Value; 35 } 36 37 export abstract class StorageBase< 38 Key extends string, 39 MetaKey extends `${Key}$`, 40 State, 41 TMetadata extends StorageMeta = StorageMeta, 42 RawState = State, 43 ListenerArgs extends unknown[] = [], 44 > implements Disposable 45 { 46 protected abstract readonly logger: Logger; 47 protected abstract readonly version: number; 48 49 readonly key: Key; 50 readonly metaKey: MetaKey; 51 readonly [storageStateType]?: NoInfer<State>; 52 readonly [storageListenerArgs]?: NoInfer<ListenerArgs>; 53 readonly #storage; 54 readonly #source; 55 readonly #tabId; 56 readonly #watchers = new Set<StorageWatchCallback<State, ListenerArgs>>(); 57 readonly #unwatch: Unwatch; 58 readonly #eventListener: EventListener; 59 protected statePromise: Promise<ToReadonly<State>> | undefined; 60 #state!: ToReadonly<State>; 61 #isVersionMetaSaved = false; 62 63 constructor( 64 tabId: number | undefined, 65 source: string, 66 logger: Logger, 67 key: Key, 68 metaKey: MetaKey, 69 storage: WxtStorageItem<RawState, TMetadata>, 70 ) { 71 this.#tabId = tabId; 72 this.#source = source; 73 this.key = key; 74 this.metaKey = metaKey; 75 this.#storage = storage; 76 77 this.initialize(); 78 79 this.#unwatch = this.#storage.watch((_state, oldState) => { 80 const state = this.parseRawValue(_state as ToReadonly<RawState>); 81 this.#state = state; 82 this.#notifyWatchers(state, oldState, false); 83 }); 84 85 this.#eventListener = new EventListener(logger, msg => { 86 if (msg.type !== 'StorageUpdatedEvent' || msg.data.key !== this.key) { 87 return; 88 } 89 90 const isOutsideUpdate = msg.data.source !== this.#source || msg.data.tabId !== this.#tabId; 91 92 if (isOutsideUpdate || msg.data.trigger === StorageEventTrigger.SetValue) { 93 if (isOutsideUpdate) { 94 // we don't have to clone `state` as it's already `structuredClone`d (and also it's immutable) 95 // https://developer.mozilla.org/en-US/docs/Mozilla/Add-ons/WebExtensions/Chrome_incompatibilities#data_cloning_algorithm 96 this.#state = msg.data.state as ToReadonly<State>; 97 } 98 99 this.#notifyWatchers(this.#state, msg.data.oldState as RawState | null, false); 100 } 101 }); 102 } 103 104 initialize() { 105 return (this.statePromise = Promise.all([this.#storage.getValue(), this.#storage.getMeta()]) 106 .then(([rawState, meta]) => { 107 this.#isVersionMetaSaved = 'v' in meta; 108 const state = (this.#state = this.parseRawValue(rawState as ToReadonly<RawState>)); 109 this.logger.debug('initialized with state', state, 'from raw', rawState); 110 return state; 111 }) 112 .finally(() => (this.statePromise = undefined))); 113 } 114 115 async reinitialize() { 116 const oldState = this.toRawValue(await this.getValue()); 117 const state = await this.initialize(); 118 this.#notifyWatchers(state, oldState, true, StorageEventTrigger.SetValue); 119 } 120 121 getValue() { 122 if (this.statePromise) { 123 return this.statePromise; 124 } 125 126 return this.#state; 127 } 128 129 async getItems(): Promise<[state: StorageItem<RawState>, meta: StorageItem<TMetadata>]> { 130 const [state, meta] = await storage.getItems([this.#storage.key, `${this.#storage.key}$`]); 131 return [state, meta]; 132 } 133 134 watch(cb: StorageWatchCallback<State, ListenerArgs>): Unwatch { 135 this.#watchers.add(cb); 136 return () => void this.#watchers.delete(cb); 137 } 138 139 clear() { 140 return this.setValue(structuredClone(this.#storage.fallback) as ToReadonly<RawState>); 141 } 142 143 #notifyWatchers( 144 state: ToReadonly<State>, 145 oldState: RawState | null, 146 isDispatchEvent: boolean, 147 eventTrigger?: StorageEventTrigger, 148 ) { 149 using logger = this.logger.scopedGroupAuto('notifying watchers'); 150 logger.debug({ state, oldState, watchers: this.#watchers }); 151 152 if (this.#watchers.size) { 153 for (const cb of this.#watchers) { 154 this.notifyWatcher(cb, state, oldState); 155 } 156 } else { 157 this.notifyWatcher(undefined, state, oldState); 158 } 159 160 isDispatchEvent && 161 EventDispatcher.dispatchStorageUpdate({ 162 tabId: this.#tabId, 163 source: this.#source, 164 key: this.key, 165 state, 166 oldState, 167 trigger: eventTrigger, 168 } as StorageEvent); 169 } 170 171 #saveValue(value: RawState) { 172 const setValuePromise = this.#storage.setValue(value); 173 174 if (this.#isVersionMetaSaved) { 175 return setValuePromise; 176 } 177 178 // WXT does not save the version when saving the value, 179 // which causes the current non-1 version (e.g., 2) 180 // to be interpreted as 1 when migrations are run, 181 // resulting in the error 182 // 183 // https://github.com/wxt-dev/wxt/issues/1775 184 185 return Promise.all([ 186 setValuePromise, 187 this.#storage.setMeta({ v: this.version } as Partial<TMetadata>), 188 ]); 189 } 190 191 protected async setValue(value: ToReadonly<RawState>) { 192 this.logger.debug('new value:', value); 193 194 const oldState = await this.#storage.getValue(); 195 196 await this.#saveValue(value as RawState); 197 const state = (this.#state = this.parseRawValue(value)); 198 199 this.#notifyWatchers(state, oldState, true, StorageEventTrigger.SetValue); 200 } 201 202 protected async setParsedValue(state: ToReadonly<State>) { 203 this.logger.debug('new parsed value:', state); 204 205 const oldState = this.toRawValue(this.#state); 206 207 await this.#saveValue(this.toRawValue(state)); 208 this.#state = state; 209 210 this.#notifyWatchers(state, oldState, true, StorageEventTrigger.SetValue); 211 } 212 213 protected parseRawValue(raw: ToReadonly<RawState>): ToReadonly<State> { 214 return raw as unknown as ToReadonly<State>; 215 } 216 217 protected toRawValue(parsed: ToReadonly<State>): RawState { 218 return parsed as unknown as RawState; 219 } 220 221 protected notifyWatcher( 222 cb: StorageWatchCallback<State, ListenerArgs> | undefined, 223 state: ToReadonly<State>, 224 _oldState: RawState | null, 225 ): void { 226 (cb as unknown as StorageWatchCallback<State, []> | undefined)?.(state as ToReadonly<State>); 227 } 228 229 [Symbol.dispose]() { 230 this.#eventListener[Symbol.dispose](); 231 this.#unwatch(); 232 } 233 } 234 235 export abstract class SyncableStorage< 236 Key extends string, 237 MetaKey extends `${Key}$`, 238 State, 239 TMetadata extends StorageMeta & StorageSyncMeta = StorageMeta & StorageSyncMeta, 240 RawState = State, 241 SyncRawState = RawState, 242 ListenerArgs extends unknown[] = [], 243 > extends StorageBase<Key, MetaKey, State, TMetadata, RawState, ListenerArgs> { 244 async getSyncItems( 245 syncMeta: StorageSyncMeta, 246 ): Promise<[shards: StorageShard<'sync'>[], removeKeys: `sync:${string}`[]]> { 247 const [state, meta] = await this.getItems(); 248 const stateValue = this.getSyncValueFromRaw(state.value); 249 250 const syncShards: StorageShard<'sync'>[] = [ 251 { 252 key: `sync:${this.metaKey}`, 253 value: { 254 ...meta.value, 255 ...syncMeta, 256 }, 257 }, 258 ]; 259 const removeKeys: `sync:${string}`[] = []; 260 261 if (this instanceof ShardedStorage) { 262 const rawShards = await this.shardRawValue(this.key, stateValue); 263 this.logger.debug('generated shards', rawShards); 264 265 for (const shard of rawShards) { 266 const itemKey = `sync:${shard.key}` as const; 267 268 if (typeof shard.value === 'undefined') { 269 removeKeys.push(itemKey); 270 } else { 271 syncShards.push({ 272 key: itemKey, 273 value: shard.value, 274 }); 275 } 276 } 277 } else { 278 syncShards.push({ 279 key: `sync:${this.key}`, 280 value: stateValue, 281 }); 282 } 283 284 return [syncShards, removeKeys]; 285 } 286 287 getSyncValueFromRaw(state: RawState): SyncRawState { 288 return state as unknown as SyncRawState; 289 } 290 291 async getShardsFromSync( 292 keepStoragePrefix?: boolean, 293 ): Promise<[shards: StorageShard<never>[], state: Record<string, unknown>]> { 294 const state = await browser.storage.sync.get(); 295 296 const shards = filterMap( 297 Object.entries(state), 298 ([key]) => key === this.key, 299 ([key, value]): StorageShard<never> => ({ 300 key: keepStoragePrefix ? key : undefined, 301 value, 302 }), 303 ); 304 305 return [shards, state]; 306 } 307 308 async restoreFromSync() { 309 const [state, meta] = 310 this instanceof ShardedStorage 311 ? await this.recoverRawValueFromShards() 312 : await this.defaultGetRawValueFromSync(); 313 314 return { 315 [this.key]: state, 316 [this.metaKey]: meta, 317 }; 318 } 319 320 getRawValueFromSync(state: SyncRawState): MaybePromise<RawState> { 321 return state as unknown as RawState; 322 } 323 324 private async defaultGetRawValueFromSync(): Promise<[RawState, TMetadata]> { 325 const [state, meta] = await storage.getItems([`sync:${this.key}`, `sync:${this.metaKey}`]); 326 327 const missingKeys: string[] = []; 328 329 for (const { key, value } of [state, meta]) { 330 if (typeof value === 'undefined' || value === null) { 331 missingKeys.push(key); 332 } 333 } 334 335 if (missingKeys.length > 0) { 336 throw new NoShardsError(missingKeys); 337 } 338 339 const stateValue = await this.getRawValueFromSync(state.value); 340 341 return [stateValue, meta.value as TMetadata]; 342 } 343 } 344 345 export abstract class ShardedStorage< 346 Key extends string, 347 MetaKey extends `${Key}$`, 348 State, 349 TMetadata extends StorageMeta & StorageSyncMeta = StorageMeta & StorageSyncMeta, 350 RawState = State, 351 SyncRawState = RawState, 352 ListenerArgs extends unknown[] = [], 353 > extends SyncableStorage<Key, MetaKey, State, TMetadata, RawState, SyncRawState, ListenerArgs> { 354 /** 355 * shards value into the smaller pieces 356 * 357 * can be used to mitigate the standard 8 KB/item sync storage limit 358 * 359 * f.e., on the realistic list of blocked channels I exceeded the limit 360 * on 192 channels with the one-key-per-object approach, and 361 * on 365 channels with the one-key-per-field approach 362 */ 363 abstract shardRawValue(keyPrefix: string, raw: SyncRawState): MaybePromise<StorageShard<never>[]>; 364 abstract getShardsFromSync( 365 keepStoragePrefix?: boolean, 366 ): Promise<[shards: StorageShard<never>[], state: Record<string, unknown>]>; 367 abstract recoverRawValueFromShards(): Promise<[RawState, TMetadata]>; 368 }