addons/isl/src/jotaiUtils.tsblame
View source
b69ab311/**
b69ab312 * Copyright (c) Meta Platforms, Inc. and affiliates.
b69ab313 *
b69ab314 * This source code is licensed under the MIT license found in the
b69ab315 * LICENSE file in the root directory of this source tree.
b69ab316 */
b69ab317
b69ab318import type {Atom, Getter, WritableAtom} from 'jotai';
b69ab319import type {Json} from 'shared/typeUtils';
b69ab3110import type {Platform} from './platform';
b69ab3111import type {ConfigName, LocalStorageName, SettableConfigName} from './types';
b69ab3112
b69ab3113import {atom, getDefaultStore, useAtomValue} from 'jotai';
b69ab3114import {loadable} from 'jotai/utils';
b69ab3115import {useMemo} from 'react';
b69ab3116import {RateLimiter} from 'shared/RateLimiter';
b69ab3117import {isPromise} from 'shared/utils';
b69ab3118import serverAPI from './ClientToServerAPI';
b69ab3119import platform from './platform';
b69ab3120import {assert} from './utils';
b69ab3121
b69ab3122/** A mutable atom that stores type `T`. */
b69ab3123export type MutAtom<T> = WritableAtom<T, [T | ((prev: T) => T)], void>;
b69ab3124
b69ab3125/**
b69ab3126 * The store being used. Do not use this directly. Alternatives are:
b69ab3127 * - use `readAtom` instead of `store.get`.
b69ab3128 * - use `writeAtom` instead of `store.set`.
b69ab3129 * - use `atomWithOnChange` instead of `store.sub`.
b69ab3130 */
b69ab3131let store = getDefaultStore();
b69ab3132
b69ab3133/**
b69ab3134 * Replace the current Jotai store used by this module.
b69ab3135 * Practically, this is only useful for tests to reset states.
b69ab3136 */
b69ab3137export function setJotaiStore(newStore: typeof store) {
b69ab3138 store = newStore;
b69ab3139}
b69ab3140
b69ab3141/** Define a read-write atom backed by a config. */
b69ab3142export function configBackedAtom<T extends Json>(
b69ab3143 name: SettableConfigName,
b69ab3144 defaultValue: T,
b69ab3145 readonly?: false,
b69ab3146): MutAtom<T>;
b69ab3147
b69ab3148/**
b69ab3149 * Define a read-only atom backed by a config.
b69ab3150 *
b69ab3151 * This can be useful for staged rollout features
b69ab3152 * where the config is not supposed to be set by the user.
b69ab3153 * (user config will override the staged rollout config)
b69ab3154 */
b69ab3155export function configBackedAtom<T extends Json>(
b69ab3156 name: ConfigName,
b69ab3157 defaultValue: T,
b69ab3158 readonly: true,
b69ab3159 useRawValue?: boolean,
b69ab3160): Atom<T>;
b69ab3161
b69ab3162export function configBackedAtom<T extends Json>(
b69ab3163 name: ConfigName | SettableConfigName,
b69ab3164 defaultValue: T,
b69ab3165 readonly = false,
b69ab3166 useRawValue?: boolean,
b69ab3167): MutAtom<T> | Atom<T> {
b69ab3168 // https://jotai.org/docs/guides/persistence
b69ab3169 const primitiveAtom = atom<T>(defaultValue);
b69ab3170
b69ab3171 let lastStrValue: undefined | string = undefined;
b69ab3172 serverAPI.onMessageOfType('gotConfig', event => {
b69ab3173 if (event.name !== name) {
b69ab3174 return;
b69ab3175 }
b69ab3176 lastStrValue = event.value;
b69ab3177 writeAtom(
b69ab3178 primitiveAtom,
b69ab3179 event.value === undefined
b69ab3180 ? defaultValue
b69ab3181 : useRawValue === true
b69ab3182 ? event.value
b69ab3183 : JSON.parse(event.value),
b69ab3184 );
b69ab3185 });
b69ab3186 serverAPI.onConnectOrReconnect(() => {
b69ab3187 serverAPI.postMessage({
b69ab3188 type: 'getConfig',
b69ab3189 name: name as ConfigName,
b69ab3190 });
b69ab3191 });
b69ab3192
b69ab3193 return readonly
b69ab3194 ? atom<T>(get => get(primitiveAtom))
b69ab3195 : atom<T, [T | ((prev: T) => T)], void>(
b69ab3196 get => get(primitiveAtom),
b69ab3197 (get, set, update) => {
b69ab3198 const newValue = typeof update === 'function' ? update(get(primitiveAtom)) : update;
b69ab3199 set(primitiveAtom, newValue);
b69ab31100 const strValue = useRawValue ? String(newValue) : JSON.stringify(newValue);
b69ab31101 if (strValue !== lastStrValue) {
b69ab31102 lastStrValue = strValue;
b69ab31103 serverAPI.postMessage({
b69ab31104 type: 'setConfig',
b69ab31105 name: name as SettableConfigName,
b69ab31106 value: strValue,
b69ab31107 });
b69ab31108 }
b69ab31109 },
b69ab31110 );
b69ab31111}
b69ab31112
b69ab31113/**
b69ab31114 * Loads this atom from a local persistent cache (usually browser local storage),
b69ab31115 * and persists any changes back to it.
b69ab31116 * Useful for some customizations that don't warrant a user-visible sl config,
b69ab31117 * for example UI expansion state.
b69ab31118 */
b69ab31119export function localStorageBackedAtom<T extends Json>(
b69ab31120 name: LocalStorageName,
b69ab31121 defaultValue: T,
b69ab31122): MutAtom<T> {
b69ab31123 const primitiveAtom = atom<T>(platform.getPersistedState<T>(name) ?? defaultValue);
b69ab31124
b69ab31125 return atom(
b69ab31126 get => get(primitiveAtom),
b69ab31127 (get, set, update) => {
b69ab31128 const newValue = typeof update === 'function' ? update(get(primitiveAtom)) : update;
b69ab31129 set(primitiveAtom, newValue);
b69ab31130 platform.setPersistedState(name, newValue);
b69ab31131 },
b69ab31132 );
b69ab31133}
b69ab31134
b69ab31135/**
b69ab31136 * Wraps an atom with an "onChange" callback.
b69ab31137 * Changing the returned atom will trigger the callback.
b69ab31138 * Calling this function will trigger `onChange` with the current value except when `skipInitialCall` is `true`.
b69ab31139 */
b69ab31140export function atomWithOnChange<T>(
b69ab31141 originalAtom: MutAtom<T>,
b69ab31142 onChange: (value: T) => void,
b69ab31143 skipInitialCall?: boolean,
b69ab31144): MutAtom<T> {
b69ab31145 if (skipInitialCall !== true) {
b69ab31146 onChange(readAtom(originalAtom));
b69ab31147 }
b69ab31148 return atom(
b69ab31149 get => get(originalAtom),
b69ab31150 (get, set, args) => {
b69ab31151 const oldValue = get(originalAtom);
b69ab31152 set(originalAtom, args);
b69ab31153 const newValue = get(originalAtom);
b69ab31154 if (oldValue !== newValue) {
b69ab31155 onChange(newValue);
b69ab31156 }
b69ab31157 },
b69ab31158 );
b69ab31159}
b69ab31160
b69ab31161/**
b69ab31162 * Creates a lazily initialized atom.
b69ab31163 * On first read, trigger `load` to get the actual value.
b69ab31164 * `fallback` provides the value when the async `load` is running.
b69ab31165 * `original` is an optional nullable atom to provide the value.
b69ab31166 */
b69ab31167export function lazyAtom<T>(
b69ab31168 load: (get: Getter) => Promise<T> | T,
b69ab31169 fallback: T,
b69ab31170 original?: MutAtom<T | undefined>,
b69ab31171): MutAtom<T> {
b69ab31172 const originalAtom = original ?? atom<T | undefined>(undefined);
b69ab31173 const limiter = new RateLimiter(1);
b69ab31174 return atom(
b69ab31175 get => {
b69ab31176 const value = get(originalAtom);
b69ab31177 if (value !== undefined) {
b69ab31178 return value;
b69ab31179 }
b69ab31180 const loaded = load(get);
b69ab31181 if (!isPromise(loaded)) {
b69ab31182 writeAtom(originalAtom, loaded);
b69ab31183 return loaded;
b69ab31184 }
b69ab31185 // Kick off the "load" but rate limit it.
b69ab31186 limiter.enqueueRun(async () => {
b69ab31187 if (get(originalAtom) !== undefined) {
b69ab31188 // A previous "load" was completed.
b69ab31189 return;
b69ab31190 }
b69ab31191 const newValue = await loaded;
b69ab31192 writeAtom(originalAtom, newValue);
b69ab31193 });
b69ab31194 // Use the fallback value while waiting for the promise.
b69ab31195 return fallback;
b69ab31196 },
b69ab31197 (get, set, args) => {
b69ab31198 const newValue =
b69ab31199 typeof args === 'function' ? (args as (prev: T) => T)(get(originalAtom) ?? fallback) : args;
b69ab31200 set(originalAtom, newValue);
b69ab31201 },
b69ab31202 );
b69ab31203}
b69ab31204
b69ab31205export function readAtom<T>(atom: Atom<T>): T {
b69ab31206 return store.get(atom);
b69ab31207}
b69ab31208
b69ab31209export function writeAtom<T>(atom: MutAtom<T>, value: T | ((prev: T) => T)) {
b69ab31210 store.set(atom, value);
b69ab31211}
b69ab31212
b69ab31213export function refreshAtom<T>(atom: WritableAtom<T, [], void>) {
b69ab31214 store.set(atom);
b69ab31215}
b69ab31216
b69ab31217/** Create an atom that is automatically reset when the depAtom is changed. */
b69ab31218export function atomResetOnDepChange<T>(defaultValue: T, depAtom: Atom<unknown>): MutAtom<T> {
b69ab31219 assert(
b69ab31220 typeof depAtom !== 'undefined',
b69ab31221 'depAtom should not be undefined (is there a circular dependency?)',
b69ab31222 );
b69ab31223 const primitiveAtom = atom<T>(defaultValue);
b69ab31224 let lastDep = readAtom(depAtom);
b69ab31225 return atom(
b69ab31226 get => {
b69ab31227 const dep = get(depAtom);
b69ab31228 if (dep !== lastDep) {
b69ab31229 lastDep = dep;
b69ab31230 writeAtom(primitiveAtom, defaultValue);
b69ab31231 }
b69ab31232 return get(primitiveAtom);
b69ab31233 },
b69ab31234 (_get, set, update) => {
b69ab31235 set(primitiveAtom, update);
b69ab31236 },
b69ab31237 );
b69ab31238}
b69ab31239
b69ab31240/**
b69ab31241 * Creates a derived atom that can be force-refreshed, by using the update function.
b69ab31242 * Uses Suspense for async update functions.
b69ab31243 * ```
b69ab31244 * const postsAtom = atomWithRefresh(get => fetchPostsAsync());
b69ab31245 * ...
b69ab31246 * const [posts, refreshPosts] = useAtom(postsAtom);
b69ab31247 * ```
b69ab31248 */
b69ab31249export function atomWithRefresh<T>(fn: (get: Getter) => T) {
b69ab31250 const refreshCounter = atom(0);
b69ab31251
b69ab31252 return atom(
b69ab31253 get => {
b69ab31254 get(refreshCounter);
b69ab31255 return fn(get);
b69ab31256 },
b69ab31257 (_, set) => set(refreshCounter, i => i + 1),
b69ab31258 );
b69ab31259}
b69ab31260
b69ab31261/**
b69ab31262 * Creates a derived atom that can be force-refreshed, by using the update function.
b69ab31263 * The underlying async state is given as a Loadable atom instead of one that suspends.
b69ab31264 * ```
b69ab31265 * const postsAtom = atomWithRefresh(get => fetchPostsAsync());
b69ab31266 * ...
b69ab31267 * const [postsLoadable, refreshPosts] = useAtom(postsAtom);
b69ab31268 * if (postsLoadable.state === 'hasData') {
b69ab31269 * const posts = postsLoadable.data;
b69ab31270 * }
b69ab31271 * ```
b69ab31272 */
b69ab31273export function atomLoadableWithRefresh<T>(fn: (get: Getter) => Promise<T>) {
b69ab31274 const refreshCounter = atom(0);
b69ab31275 const loadableAtom = loadable(
b69ab31276 atom(get => {
b69ab31277 get(refreshCounter);
b69ab31278 return fn(get);
b69ab31279 }),
b69ab31280 );
b69ab31281
b69ab31282 return atom(
b69ab31283 get => get(loadableAtom),
b69ab31284 (_, set) => set(refreshCounter, i => i + 1),
b69ab31285 );
b69ab31286}
b69ab31287
b69ab31288/**
b69ab31289 * Drop-in replacement of `atomFamily` that tries to book-keep internal cache
b69ab31290 * periodically to avoid memory leak.
b69ab31291 *
b69ab31292 * There are 2 caches:
b69ab31293 * - "strong" cache: keep atoms alive even if all references are gone.
b69ab31294 * - "weak" cache: keep atoms alive as long as there is a reference to it.
b69ab31295 *
b69ab31296 * Periodically, when the weak cache size reaches a threshold, a "cleanup"
b69ab31297 * process runs to:
b69ab31298 * - Clear the "strong" cache to mitigate memory leak.
b69ab31299 * - Drop dead entries in the "weak" cache.
b69ab31300 * - Update the threshold to 2x the "weak" cache size.
b69ab31301 *
b69ab31302 * Therefore the memory waste is hopefully within 2x of the needed memory.
b69ab31303 *
b69ab31304 * Setting `options.useStrongCache` to `false` disables the "strong" cache
b69ab31305 * to further limit memory usage.
b69ab31306 */
b69ab31307export function atomFamilyWeak<K, A extends Atom<unknown>>(
b69ab31308 fn: (key: K) => A,
b69ab31309 options?: AtomFamilyWeakOptions,
b69ab31310): AtomFamilyWeak<K, A> {
b69ab31311 const {useStrongCache = true, initialCleanupThreshold = 4} = options ?? {};
b69ab31312
b69ab31313 // This cache persists through component unmount / remount, therefore
b69ab31314 // it can be memory leaky.
b69ab31315 const strongCache = new Map<K, A>();
b69ab31316
b69ab31317 // This cache ensures atoms in use are returned as-is during re-render,
b69ab31318 // to avoid infinite re-render with React StrictMode.
b69ab31319 const weakCache = new Map<K, WeakRef<A>>();
b69ab31320
b69ab31321 const cleanup = () => {
b69ab31322 // Clear the strong cache. This allows GC to drop weakRefs.
b69ab31323 strongCache.clear();
b69ab31324 // Clean up weak cache - remove dead entries.
b69ab31325 weakCache.forEach((weakRef, key) => {
b69ab31326 if (weakRef.deref() == null) {
b69ab31327 weakCache.delete(key);
b69ab31328 }
b69ab31329 });
b69ab31330 // Adjust threshold to trigger the next cleanup.
b69ab31331 resultFunc.threshold = weakCache.size * 2;
b69ab31332 };
b69ab31333
b69ab31334 const resultFunc: AtomFamilyWeak<K, A> = (key: K) => {
b69ab31335 const cached = strongCache.get(key);
b69ab31336 if (cached != null) {
b69ab31337 // This state was accessed recently.
b69ab31338 return cached;
b69ab31339 }
b69ab31340 const weakCached = weakCache.get(key)?.deref();
b69ab31341 if (weakCached != null) {
b69ab31342 // State is not dead yet.
b69ab31343 return weakCached;
b69ab31344 }
b69ab31345 // Not in cache. Need re-calculate.
b69ab31346 const atom = fn(key);
b69ab31347 if (useStrongCache) {
b69ab31348 strongCache.set(key, atom);
b69ab31349 }
b69ab31350 weakCache.set(key, new WeakRef(atom));
b69ab31351 if (weakCache.size > resultFunc.threshold) {
b69ab31352 cleanup();
b69ab31353 }
b69ab31354 if (resultFunc.debugLabel != null && atom.debugLabel == null) {
b69ab31355 atom.debugLabel = `${resultFunc.debugLabel}:${key}`;
b69ab31356 }
b69ab31357 return atom;
b69ab31358 };
b69ab31359
b69ab31360 resultFunc.cleanup = cleanup;
b69ab31361 resultFunc.threshold = initialCleanupThreshold;
b69ab31362 resultFunc.strongCache = strongCache;
b69ab31363 resultFunc.weakCache = weakCache;
b69ab31364 resultFunc.clear = () => {
b69ab31365 weakCache.clear();
b69ab31366 strongCache.clear();
b69ab31367 resultFunc.threshold = initialCleanupThreshold;
b69ab31368 };
b69ab31369
b69ab31370 return resultFunc;
b69ab31371}
b69ab31372
b69ab31373type AtomFamilyWeakOptions = {
b69ab31374 /**
b69ab31375 * Enable the "strong" cache so unmount / remount can try to reuse the cache.
b69ab31376 *
b69ab31377 * If this is disabled, then only the weakRef cache is used, states that
b69ab31378 * are no longer referred by components might be lost to GC.
b69ab31379 *
b69ab31380 * Default: true.
b69ab31381 */
b69ab31382 useStrongCache?: boolean;
b69ab31383
b69ab31384 /**
b69ab31385 * Number of items before triggering an initial cleanup.
b69ab31386 * Default: 4.
b69ab31387 */
b69ab31388 initialCleanupThreshold?: number;
b69ab31389};
b69ab31390
b69ab31391export interface AtomFamilyWeak<K, A extends Atom<unknown>> {
b69ab31392 (key: K): A;
b69ab31393 /** The "strong" cache (can be empty). */
b69ab31394 strongCache: Map<K, A>;
b69ab31395 /** The weakRef cache (must contain entries that are still referred elsewhere). */
b69ab31396 weakCache: Map<K, WeakRef<A>>;
b69ab31397 /** Trigger a cleanup ("GC"). */
b69ab31398 cleanup(): void;
b69ab31399 /** Clear the cache */
b69ab31400 clear(): void;
b69ab31401 /** Auto cleanup threshold on weakCache size. */
b69ab31402 threshold: number;
b69ab31403 /** Prefix of debugLabel. */
b69ab31404 debugLabel?: string;
b69ab31405}
b69ab31406
b69ab31407function getAllPersistedStateWithPrefix<T>(
b69ab31408 prefix: LocalStorageName,
b69ab31409 islPlatform: Platform,
b69ab31410): Record<string, T> {
b69ab31411 const all = islPlatform.getAllPersistedState();
b69ab31412 if (all == null) {
b69ab31413 return {};
b69ab31414 }
b69ab31415 return Object.fromEntries(
b69ab31416 Object.entries(all)
b69ab31417 .filter(([key]) => key.startsWith(prefix))
b69ab31418 .map(([key, value]) => [key.slice(prefix.length), value]),
b69ab31419 );
b69ab31420}
b69ab31421
b69ab31422/**
b69ab31423 * An atom family that loads and persists data from local storage.
b69ab31424 * Each key is stored in a separate local storage entry, using the `storageKeyPrefix`.
b69ab31425 * Each stored value includes a timestamp so that stale data can be evicted,
b69ab31426 * on next startup.
b69ab31427 * Data is loaded once on startup, but written to local storage on every change.
b69ab31428 * Write `undefined` to any atom to explicitly remove it from local storage.
b69ab31429 */
b69ab31430export function localStorageBackedAtomFamily<K extends string, T extends Json | Partial<Json>>(
b69ab31431 storageKeyPrefix: LocalStorageName,
b69ab31432 getDefault: (key: K) => T,
b69ab31433 maxAgeDays = 14,
b69ab31434 islPlatform = platform,
b69ab31435): AtomFamilyWeak<K, WritableAtom<T, [T | undefined | ((prev: T) => T | undefined)], void>> {
b69ab31436 type StoredData = {
b69ab31437 data: T;
b69ab31438 date: number;
b69ab31439 };
b69ab31440 const initialData = getAllPersistedStateWithPrefix<StoredData>(storageKeyPrefix, islPlatform);
b69ab31441
b69ab31442 const ONE_DAY_MS = 1000 * 60 * 60 * 24;
b69ab31443 // evict previously stored old data
b69ab31444 for (const key in initialData) {
b69ab31445 const data = initialData[key];
b69ab31446 if (data?.date != null && Date.now() - data?.date > ONE_DAY_MS * maxAgeDays) {
b69ab31447 islPlatform.setPersistedState(storageKeyPrefix + key, undefined);
b69ab31448 delete initialData[key];
b69ab31449 }
b69ab31450 }
b69ab31451
b69ab31452 return atomFamilyWeak((key: K) => {
b69ab31453 // We use the full getPersistedState instead of initialData, as this atom may have been evicted from the weak cache,
b69ab31454 // and is now being recreated and requires checking the actual cache to get any changes after initialization.
b69ab31455 const data = islPlatform.getPersistedState(storageKeyPrefix + key) as StoredData | null;
b69ab31456 const initial = data?.data ?? getDefault(key);
b69ab31457 const storageKey = storageKeyPrefix + key;
b69ab31458
b69ab31459 const inner = atom<T>(initial);
b69ab31460 return atom(
b69ab31461 get => get(inner),
b69ab31462 (get, set, value) => {
b69ab31463 const oldValue = get(inner);
b69ab31464 const result = typeof value === 'function' ? value(oldValue) : value;
b69ab31465 set(inner, result === undefined ? getDefault(key) : result);
b69ab31466 const newValue = get(inner);
b69ab31467 if (oldValue !== newValue) {
b69ab31468 // TODO: debounce?
b69ab31469 islPlatform.setPersistedState(
b69ab31470 storageKey,
b69ab31471 result == null
b69ab31472 ? undefined
b69ab31473 : ({
b69ab31474 data: newValue as Json,
b69ab31475 date: Date.now(),
b69ab31476 } as StoredData as Json),
b69ab31477 );
b69ab31478 }
b69ab31479 },
b69ab31480 );
b69ab31481 });
b69ab31482}
b69ab31483
b69ab31484function setDebugLabelForDerivedAtom<A extends Atom<unknown>>(
b69ab31485 original: Atom<unknown>,
b69ab31486 derived: A,
b69ab31487 key: unknown,
b69ab31488): A {
b69ab31489 derived.debugLabel = `${original.debugLabel ?? original.toString()}:${key}`;
b69ab31490 return derived;
b69ab31491}
b69ab31492
b69ab31493/**
b69ab31494 * Similar to `useAtomValue(mapAtom).get(key)` but avoids re-render if the map
b69ab31495 * is changed but the `get(key)` does not change.
b69ab31496 *
b69ab31497 * This might be an appealing alternative to `atomFamilyWeak` in some cases.
b69ab31498 * The `atomFamilyWeak` keeps caching state within itself and it has
b69ab31499 * undesirable memory overhead regardless of settings. This function makes
b69ab31500 * the hook own the caching state so states can be released cleanly on unmount.
b69ab31501 */
b69ab31502export function useAtomGet<K, V>(
b69ab31503 mapAtom: Atom<{get(k: K): V | undefined}>,
b69ab31504 key: K,
b69ab31505): Awaited<V | undefined> {
b69ab31506 const derivedAtom = useMemo(() => {
b69ab31507 const derived = atom(get => get(mapAtom).get(key));
b69ab31508 return setDebugLabelForDerivedAtom(mapAtom, derived, key);
b69ab31509 }, [key, mapAtom]);
b69ab31510 return useAtomValue(derivedAtom);
b69ab31511}
b69ab31512
b69ab31513/**
b69ab31514 * Similar to `useAtomValue(setAtom).has(key)` but avoids re-render if the set
b69ab31515 * is changed but the `has(key)` does not change.
b69ab31516 *
b69ab31517 * This might be an appealing alternative to `atomFamilyWeak`. See `useAtomGet`
b69ab31518 * for explanation.
b69ab31519 */
b69ab31520export function useAtomHas<K>(setAtom: Atom<{has(k: K): boolean}>, key: K): Awaited<boolean> {
b69ab31521 const derivedAtom = useMemo(() => atom(get => get(setAtom).has(key)), [key, setAtom]);
b69ab31522 return useAtomValue(derivedAtom);
b69ab31523}