useLocalStorageKey

useLocalStorageKey.tsx

import { useSyncExternalStore, useCallback } from "react";

/**
 * Use for getting the value of given `key` in `localStorage`, or setting it.
 *
 * @param key - `localStorage` key we want to get its value, or change it
 * @example
 *
 * ```tsx
 * const [hello, setHello] = useLocalStorageKey("hello");
 *
 * return (
 *   <button
 *     onClick={() => {
 *       if (hello === null) {
 *         // Sets value to `"world"` in `localStorage` and synchronized where
 *         // `useLocalStorageKey` is used with the same `key`
 *         setHello("world");
 *       } else {
 *         // `null` will remove the `key` from `localStorage`
 *         setHello(null);
 *       }
 *     }}
 *   >
 *     hello: {hello}
 *   </button>
 * );
 * ```
 */
export function useLocalStorageKey(
  key: string,
): readonly [
  value: string | null,
  setValue: (value: string | null) => undefined,
] {
  const value = useSyncExternalStore(
    // There is no real gain from making `subscribe` "stable reference", and
    // React Compiler will do the trick anyway
    (onStoreChange) => subscribeKey(key, onStoreChange),
    // - `getSnapshot()` is always called, and "optimizing" is meaningless
    // - `string` or `null` don't have references. So, they are stable between
    //   re-renders
    // - Modern browsers don't use threads lock for `localStorage.getItem()`, so
    //   it's safe reading it every render without worry about performance
    () => localStorage.getItem(key),
    // `localStorage` (of the user) isn't relevant inside the server, so we
    // return `null` same as when `key` doesn't exist in `localStorage`
    () => null,
  );

  // `useCallback` isn't needed with React Compiler. But if not used, the setter
  // better be a stable reference
  const setValue = useCallback(
    (value: string | null) => setKey(key, value),
    [key],
  );

  return [value, setValue];
}

// Listeners are divided by keys for better efficiency when a value change
// happens for a specific key
const listeners = new Map<string, Set<() => void>>();

// `Set` creation is cheap, but not cheap enough if not needed. With React
// Compiler and a stable `subscribe` in `useSyncExternalStore`, it's meaningless
function createSet<T>() {
  return new Set<T>();
}

// This function is used by `useSyncExternalStore` to know when `localStorage`
// value was changed
function subscribeKey(key: string, onStoreChange: () => void) {
  // Adding the listener
  listeners.getOrInsertComputed(key, createSet).add(onStoreChange);

  // Adding listener to global `storage` event. It's important for receiving
  // updates from other tabs
  window.addEventListener(
    "storage",
    // `addEventListener` stores listeners in `Set` like structure, so you can
    // add the same reference multiple times, and it would store it only once.
    // No need to worry about multiple calls when `storage` event
    handleStorage,
    // `passive = true` because it's not important enough to block main thread.
    // It's probably meaningless in real life usage
    { passive: true },
  );

  return () => {
    const keyListeners = listeners.get(key);
    // If not listeners left to a given key, we can remove them all
    if (keyListeners?.delete(onStoreChange) && keyListeners.size === 0) {
      listeners.delete(key);

      // If no any listeners left to any key, no reason to keep a listener for
      // the global `storage` event
      if (listeners.size === 0) {
        // Only if `capture = true` we need to pass third parameter
        window.removeEventListener("storage", handleStorage);
      }
    }
  };
}

// The is a wrapper for `localStorage` mutation methods, so we can inform other
// usages for the same `key`
function setKey(key: string, value: string | null): undefined {
  if (value === null) {
    localStorage.removeItem(key);
  } else {
    localStorage.setItem(key, value);
  }

  // We inform changes by ourselves, since `storage` event is only fired for
  // changes in other tabs. So we must inform in tab changes
  listeners
    .get(key)
    // It might look like a bad idea not to check if value changed before, but:
    // - React doesn't re-render if the value from `getSnapshot()` isn't change
    // - usually the listeners number per key is low
    // - we can't make sure no updated was done to `localStorage` directly, and
    //   currently we show a stale value
    ?.forEach((cb) => cb());
}

// The handler for `storage` event on `window`
function handleStorage(event: StorageEvent) {
  // This check probably isn't needed since this event only relevant to
  // `localStorage`
  if (event.storageArea !== localStorage) {
    return;
  }

  // `event.key === null` means some other tab has called `localStorage.clear()`
  if (event.key === null) {
    // We iterate listeners from all keys for `localStorage.clear()`, because
    // it might effect them all
    listeners.forEach((keyListeners) => {
      keyListeners.forEach((cb) => cb());
    });
  } else {
    // We inform listener of the key that was changed. More notes in `setKey`
    listeners.get(event.key)?.forEach((cb) => cb());
  }
}