15.8 KB524 lines
Blame
1/**
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 *
4 * This source code is licensed under the MIT license found in the
5 * LICENSE file in the root directory of this source tree.
6 */
7
8import type {Atom, Getter, WritableAtom} from 'jotai';
9import type {Json} from 'shared/typeUtils';
10import type {Platform} from './platform';
11import type {ConfigName, LocalStorageName, SettableConfigName} from './types';
12
13import {atom, getDefaultStore, useAtomValue} from 'jotai';
14import {loadable} from 'jotai/utils';
15import {useMemo} from 'react';
16import {RateLimiter} from 'shared/RateLimiter';
17import {isPromise} from 'shared/utils';
18import serverAPI from './ClientToServerAPI';
19import platform from './platform';
20import {assert} from './utils';
21
22/** A mutable atom that stores type `T`. */
23export type MutAtom<T> = WritableAtom<T, [T | ((prev: T) => T)], void>;
24
25/**
26 * The store being used. Do not use this directly. Alternatives are:
27 * - use `readAtom` instead of `store.get`.
28 * - use `writeAtom` instead of `store.set`.
29 * - use `atomWithOnChange` instead of `store.sub`.
30 */
31let store = getDefaultStore();
32
33/**
34 * Replace the current Jotai store used by this module.
35 * Practically, this is only useful for tests to reset states.
36 */
37export function setJotaiStore(newStore: typeof store) {
38 store = newStore;
39}
40
41/** Define a read-write atom backed by a config. */
42export function configBackedAtom<T extends Json>(
43 name: SettableConfigName,
44 defaultValue: T,
45 readonly?: false,
46): MutAtom<T>;
47
48/**
49 * Define a read-only atom backed by a config.
50 *
51 * This can be useful for staged rollout features
52 * where the config is not supposed to be set by the user.
53 * (user config will override the staged rollout config)
54 */
55export function configBackedAtom<T extends Json>(
56 name: ConfigName,
57 defaultValue: T,
58 readonly: true,
59 useRawValue?: boolean,
60): Atom<T>;
61
62export function configBackedAtom<T extends Json>(
63 name: ConfigName | SettableConfigName,
64 defaultValue: T,
65 readonly = false,
66 useRawValue?: boolean,
67): MutAtom<T> | Atom<T> {
68 // https://jotai.org/docs/guides/persistence
69 const primitiveAtom = atom<T>(defaultValue);
70
71 let lastStrValue: undefined | string = undefined;
72 serverAPI.onMessageOfType('gotConfig', event => {
73 if (event.name !== name) {
74 return;
75 }
76 lastStrValue = event.value;
77 writeAtom(
78 primitiveAtom,
79 event.value === undefined
80 ? defaultValue
81 : useRawValue === true
82 ? event.value
83 : JSON.parse(event.value),
84 );
85 });
86 serverAPI.onConnectOrReconnect(() => {
87 serverAPI.postMessage({
88 type: 'getConfig',
89 name: name as ConfigName,
90 });
91 });
92
93 return readonly
94 ? atom<T>(get => get(primitiveAtom))
95 : atom<T, [T | ((prev: T) => T)], void>(
96 get => get(primitiveAtom),
97 (get, set, update) => {
98 const newValue = typeof update === 'function' ? update(get(primitiveAtom)) : update;
99 set(primitiveAtom, newValue);
100 const strValue = useRawValue ? String(newValue) : JSON.stringify(newValue);
101 if (strValue !== lastStrValue) {
102 lastStrValue = strValue;
103 serverAPI.postMessage({
104 type: 'setConfig',
105 name: name as SettableConfigName,
106 value: strValue,
107 });
108 }
109 },
110 );
111}
112
113/**
114 * Loads this atom from a local persistent cache (usually browser local storage),
115 * and persists any changes back to it.
116 * Useful for some customizations that don't warrant a user-visible sl config,
117 * for example UI expansion state.
118 */
119export function localStorageBackedAtom<T extends Json>(
120 name: LocalStorageName,
121 defaultValue: T,
122): MutAtom<T> {
123 const primitiveAtom = atom<T>(platform.getPersistedState<T>(name) ?? defaultValue);
124
125 return atom(
126 get => get(primitiveAtom),
127 (get, set, update) => {
128 const newValue = typeof update === 'function' ? update(get(primitiveAtom)) : update;
129 set(primitiveAtom, newValue);
130 platform.setPersistedState(name, newValue);
131 },
132 );
133}
134
135/**
136 * Wraps an atom with an "onChange" callback.
137 * Changing the returned atom will trigger the callback.
138 * Calling this function will trigger `onChange` with the current value except when `skipInitialCall` is `true`.
139 */
140export function atomWithOnChange<T>(
141 originalAtom: MutAtom<T>,
142 onChange: (value: T) => void,
143 skipInitialCall?: boolean,
144): MutAtom<T> {
145 if (skipInitialCall !== true) {
146 onChange(readAtom(originalAtom));
147 }
148 return atom(
149 get => get(originalAtom),
150 (get, set, args) => {
151 const oldValue = get(originalAtom);
152 set(originalAtom, args);
153 const newValue = get(originalAtom);
154 if (oldValue !== newValue) {
155 onChange(newValue);
156 }
157 },
158 );
159}
160
161/**
162 * Creates a lazily initialized atom.
163 * On first read, trigger `load` to get the actual value.
164 * `fallback` provides the value when the async `load` is running.
165 * `original` is an optional nullable atom to provide the value.
166 */
167export function lazyAtom<T>(
168 load: (get: Getter) => Promise<T> | T,
169 fallback: T,
170 original?: MutAtom<T | undefined>,
171): MutAtom<T> {
172 const originalAtom = original ?? atom<T | undefined>(undefined);
173 const limiter = new RateLimiter(1);
174 return atom(
175 get => {
176 const value = get(originalAtom);
177 if (value !== undefined) {
178 return value;
179 }
180 const loaded = load(get);
181 if (!isPromise(loaded)) {
182 writeAtom(originalAtom, loaded);
183 return loaded;
184 }
185 // Kick off the "load" but rate limit it.
186 limiter.enqueueRun(async () => {
187 if (get(originalAtom) !== undefined) {
188 // A previous "load" was completed.
189 return;
190 }
191 const newValue = await loaded;
192 writeAtom(originalAtom, newValue);
193 });
194 // Use the fallback value while waiting for the promise.
195 return fallback;
196 },
197 (get, set, args) => {
198 const newValue =
199 typeof args === 'function' ? (args as (prev: T) => T)(get(originalAtom) ?? fallback) : args;
200 set(originalAtom, newValue);
201 },
202 );
203}
204
205export function readAtom<T>(atom: Atom<T>): T {
206 return store.get(atom);
207}
208
209export function writeAtom<T>(atom: MutAtom<T>, value: T | ((prev: T) => T)) {
210 store.set(atom, value);
211}
212
213export function refreshAtom<T>(atom: WritableAtom<T, [], void>) {
214 store.set(atom);
215}
216
217/** Create an atom that is automatically reset when the depAtom is changed. */
218export function atomResetOnDepChange<T>(defaultValue: T, depAtom: Atom<unknown>): MutAtom<T> {
219 assert(
220 typeof depAtom !== 'undefined',
221 'depAtom should not be undefined (is there a circular dependency?)',
222 );
223 const primitiveAtom = atom<T>(defaultValue);
224 let lastDep = readAtom(depAtom);
225 return atom(
226 get => {
227 const dep = get(depAtom);
228 if (dep !== lastDep) {
229 lastDep = dep;
230 writeAtom(primitiveAtom, defaultValue);
231 }
232 return get(primitiveAtom);
233 },
234 (_get, set, update) => {
235 set(primitiveAtom, update);
236 },
237 );
238}
239
240/**
241 * Creates a derived atom that can be force-refreshed, by using the update function.
242 * Uses Suspense for async update functions.
243 * ```
244 * const postsAtom = atomWithRefresh(get => fetchPostsAsync());
245 * ...
246 * const [posts, refreshPosts] = useAtom(postsAtom);
247 * ```
248 */
249export function atomWithRefresh<T>(fn: (get: Getter) => T) {
250 const refreshCounter = atom(0);
251
252 return atom(
253 get => {
254 get(refreshCounter);
255 return fn(get);
256 },
257 (_, set) => set(refreshCounter, i => i + 1),
258 );
259}
260
261/**
262 * Creates a derived atom that can be force-refreshed, by using the update function.
263 * The underlying async state is given as a Loadable atom instead of one that suspends.
264 * ```
265 * const postsAtom = atomWithRefresh(get => fetchPostsAsync());
266 * ...
267 * const [postsLoadable, refreshPosts] = useAtom(postsAtom);
268 * if (postsLoadable.state === 'hasData') {
269 * const posts = postsLoadable.data;
270 * }
271 * ```
272 */
273export function atomLoadableWithRefresh<T>(fn: (get: Getter) => Promise<T>) {
274 const refreshCounter = atom(0);
275 const loadableAtom = loadable(
276 atom(get => {
277 get(refreshCounter);
278 return fn(get);
279 }),
280 );
281
282 return atom(
283 get => get(loadableAtom),
284 (_, set) => set(refreshCounter, i => i + 1),
285 );
286}
287
288/**
289 * Drop-in replacement of `atomFamily` that tries to book-keep internal cache
290 * periodically to avoid memory leak.
291 *
292 * There are 2 caches:
293 * - "strong" cache: keep atoms alive even if all references are gone.
294 * - "weak" cache: keep atoms alive as long as there is a reference to it.
295 *
296 * Periodically, when the weak cache size reaches a threshold, a "cleanup"
297 * process runs to:
298 * - Clear the "strong" cache to mitigate memory leak.
299 * - Drop dead entries in the "weak" cache.
300 * - Update the threshold to 2x the "weak" cache size.
301 *
302 * Therefore the memory waste is hopefully within 2x of the needed memory.
303 *
304 * Setting `options.useStrongCache` to `false` disables the "strong" cache
305 * to further limit memory usage.
306 */
307export function atomFamilyWeak<K, A extends Atom<unknown>>(
308 fn: (key: K) => A,
309 options?: AtomFamilyWeakOptions,
310): AtomFamilyWeak<K, A> {
311 const {useStrongCache = true, initialCleanupThreshold = 4} = options ?? {};
312
313 // This cache persists through component unmount / remount, therefore
314 // it can be memory leaky.
315 const strongCache = new Map<K, A>();
316
317 // This cache ensures atoms in use are returned as-is during re-render,
318 // to avoid infinite re-render with React StrictMode.
319 const weakCache = new Map<K, WeakRef<A>>();
320
321 const cleanup = () => {
322 // Clear the strong cache. This allows GC to drop weakRefs.
323 strongCache.clear();
324 // Clean up weak cache - remove dead entries.
325 weakCache.forEach((weakRef, key) => {
326 if (weakRef.deref() == null) {
327 weakCache.delete(key);
328 }
329 });
330 // Adjust threshold to trigger the next cleanup.
331 resultFunc.threshold = weakCache.size * 2;
332 };
333
334 const resultFunc: AtomFamilyWeak<K, A> = (key: K) => {
335 const cached = strongCache.get(key);
336 if (cached != null) {
337 // This state was accessed recently.
338 return cached;
339 }
340 const weakCached = weakCache.get(key)?.deref();
341 if (weakCached != null) {
342 // State is not dead yet.
343 return weakCached;
344 }
345 // Not in cache. Need re-calculate.
346 const atom = fn(key);
347 if (useStrongCache) {
348 strongCache.set(key, atom);
349 }
350 weakCache.set(key, new WeakRef(atom));
351 if (weakCache.size > resultFunc.threshold) {
352 cleanup();
353 }
354 if (resultFunc.debugLabel != null && atom.debugLabel == null) {
355 atom.debugLabel = `${resultFunc.debugLabel}:${key}`;
356 }
357 return atom;
358 };
359
360 resultFunc.cleanup = cleanup;
361 resultFunc.threshold = initialCleanupThreshold;
362 resultFunc.strongCache = strongCache;
363 resultFunc.weakCache = weakCache;
364 resultFunc.clear = () => {
365 weakCache.clear();
366 strongCache.clear();
367 resultFunc.threshold = initialCleanupThreshold;
368 };
369
370 return resultFunc;
371}
372
373type AtomFamilyWeakOptions = {
374 /**
375 * Enable the "strong" cache so unmount / remount can try to reuse the cache.
376 *
377 * If this is disabled, then only the weakRef cache is used, states that
378 * are no longer referred by components might be lost to GC.
379 *
380 * Default: true.
381 */
382 useStrongCache?: boolean;
383
384 /**
385 * Number of items before triggering an initial cleanup.
386 * Default: 4.
387 */
388 initialCleanupThreshold?: number;
389};
390
391export interface AtomFamilyWeak<K, A extends Atom<unknown>> {
392 (key: K): A;
393 /** The "strong" cache (can be empty). */
394 strongCache: Map<K, A>;
395 /** The weakRef cache (must contain entries that are still referred elsewhere). */
396 weakCache: Map<K, WeakRef<A>>;
397 /** Trigger a cleanup ("GC"). */
398 cleanup(): void;
399 /** Clear the cache */
400 clear(): void;
401 /** Auto cleanup threshold on weakCache size. */
402 threshold: number;
403 /** Prefix of debugLabel. */
404 debugLabel?: string;
405}
406
407function getAllPersistedStateWithPrefix<T>(
408 prefix: LocalStorageName,
409 islPlatform: Platform,
410): Record<string, T> {
411 const all = islPlatform.getAllPersistedState();
412 if (all == null) {
413 return {};
414 }
415 return Object.fromEntries(
416 Object.entries(all)
417 .filter(([key]) => key.startsWith(prefix))
418 .map(([key, value]) => [key.slice(prefix.length), value]),
419 );
420}
421
422/**
423 * An atom family that loads and persists data from local storage.
424 * Each key is stored in a separate local storage entry, using the `storageKeyPrefix`.
425 * Each stored value includes a timestamp so that stale data can be evicted,
426 * on next startup.
427 * Data is loaded once on startup, but written to local storage on every change.
428 * Write `undefined` to any atom to explicitly remove it from local storage.
429 */
430export function localStorageBackedAtomFamily<K extends string, T extends Json | Partial<Json>>(
431 storageKeyPrefix: LocalStorageName,
432 getDefault: (key: K) => T,
433 maxAgeDays = 14,
434 islPlatform = platform,
435): AtomFamilyWeak<K, WritableAtom<T, [T | undefined | ((prev: T) => T | undefined)], void>> {
436 type StoredData = {
437 data: T;
438 date: number;
439 };
440 const initialData = getAllPersistedStateWithPrefix<StoredData>(storageKeyPrefix, islPlatform);
441
442 const ONE_DAY_MS = 1000 * 60 * 60 * 24;
443 // evict previously stored old data
444 for (const key in initialData) {
445 const data = initialData[key];
446 if (data?.date != null && Date.now() - data?.date > ONE_DAY_MS * maxAgeDays) {
447 islPlatform.setPersistedState(storageKeyPrefix + key, undefined);
448 delete initialData[key];
449 }
450 }
451
452 return atomFamilyWeak((key: K) => {
453 // We use the full getPersistedState instead of initialData, as this atom may have been evicted from the weak cache,
454 // and is now being recreated and requires checking the actual cache to get any changes after initialization.
455 const data = islPlatform.getPersistedState(storageKeyPrefix + key) as StoredData | null;
456 const initial = data?.data ?? getDefault(key);
457 const storageKey = storageKeyPrefix + key;
458
459 const inner = atom<T>(initial);
460 return atom(
461 get => get(inner),
462 (get, set, value) => {
463 const oldValue = get(inner);
464 const result = typeof value === 'function' ? value(oldValue) : value;
465 set(inner, result === undefined ? getDefault(key) : result);
466 const newValue = get(inner);
467 if (oldValue !== newValue) {
468 // TODO: debounce?
469 islPlatform.setPersistedState(
470 storageKey,
471 result == null
472 ? undefined
473 : ({
474 data: newValue as Json,
475 date: Date.now(),
476 } as StoredData as Json),
477 );
478 }
479 },
480 );
481 });
482}
483
484function setDebugLabelForDerivedAtom<A extends Atom<unknown>>(
485 original: Atom<unknown>,
486 derived: A,
487 key: unknown,
488): A {
489 derived.debugLabel = `${original.debugLabel ?? original.toString()}:${key}`;
490 return derived;
491}
492
493/**
494 * Similar to `useAtomValue(mapAtom).get(key)` but avoids re-render if the map
495 * is changed but the `get(key)` does not change.
496 *
497 * This might be an appealing alternative to `atomFamilyWeak` in some cases.
498 * The `atomFamilyWeak` keeps caching state within itself and it has
499 * undesirable memory overhead regardless of settings. This function makes
500 * the hook own the caching state so states can be released cleanly on unmount.
501 */
502export function useAtomGet<K, V>(
503 mapAtom: Atom<{get(k: K): V | undefined}>,
504 key: K,
505): Awaited<V | undefined> {
506 const derivedAtom = useMemo(() => {
507 const derived = atom(get => get(mapAtom).get(key));
508 return setDebugLabelForDerivedAtom(mapAtom, derived, key);
509 }, [key, mapAtom]);
510 return useAtomValue(derivedAtom);
511}
512
513/**
514 * Similar to `useAtomValue(setAtom).has(key)` but avoids re-render if the set
515 * is changed but the `has(key)` does not change.
516 *
517 * This might be an appealing alternative to `atomFamilyWeak`. See `useAtomGet`
518 * for explanation.
519 */
520export function useAtomHas<K>(setAtom: Atom<{has(k: K): boolean}>, key: K): Awaited<boolean> {
521 const derivedAtom = useMemo(() => atom(get => get(setAtom).has(key)), [key, setAtom]);
522 return useAtomValue(derivedAtom);
523}
524