/ src / storage / base.ts
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  }