refactor(hooks): refactor storage with useSyncExternalStore

Author: xiejay97

packages/hooks/src/storage/store.ts

@@ -0,0 +1,48 @@
+import type { AbstractStorage } from './storage';
+
+import { useMemo } from 'react';
+
+class Store {
+  private _listeners: (() => void)[] = [];
+  private _key: any;
+  private _storage: AbstractStorage<any, any>;
+
+  constructor(key: any, storage: AbstractStorage<any, any>) {
+    this._key = key;
+    this._storage = storage;
+  }
+
+  subscribe(onStoreChange: () => void) {
+    this._listeners = this._listeners.concat([onStoreChange]);
+    return () => {
+      this._listeners = this._listeners.filter((f) => f !== onStoreChange);
+    };
+  }
+
+  getSnapshot() {
+    return this._storage.getItem(this._key);
+  }
+
+  emitChange() {
+    for (const listener of this._listeners) {
+      listener();
+    }
+  }
+}
+
+export const stores = new Map<any, Store>();
+
+export function useStore(key: any, storage: AbstractStorage<any, any>) {
+  return useMemo(() => {
+    let store = stores.get(key);
+    if (!store) {
+      store = new Store(key, storage);
+      stores.set(key, store);
+    }
+    return {
+      subscribe: store.subscribe.bind(store),
+      getSnapshot: store.getSnapshot.bind(store),
+      emitChange: store.emitChange.bind(store),
+    };
+  }, [key, storage]);
+}

packages/hooks/src/storage/useStorage.ts

@@ -2,13 +2,11 @@ import type { AbstractParserOptions } from './parser';
 import type { AbstractStorage } from './storage';
 
 import { isNull, isUndefined } from 'lodash';
-import { useEffect, useMemo, useState } from 'react';
+import { useMemo, useSyncExternalStore } from 'react';
 
-import { useIsomorphicLayoutEffect } from '../useIsomorphicLayoutEffect';
 import { LocalStorageService } from './localStorage';
 import { STRING_PARSER } from './parser';
-
-const updates = new Map<any, Set<(...args: any[]) => any>>();
+import { useStore, stores } from './store';
 
 interface UseStorageMethod<V> {
   set: (value: V) => void;
@@ -34,61 +32,37 @@ export function useStorage<V, K = string>(
 
   const { serializer, deserializer } = isUndefined(parser) ? PARSER.plain : PARSER[parser];
 
-  const res = useMemo<{ readonly value: any } & UseStorageMethod<V>>(() => {
-    const updateKey = (ov: any) => {
-      const cbs = updates.get(key);
-      if (cbs) {
-        for (const cb of cbs) {
-          cb(ov);
-        }
-      }
-    };
-    return {
+  const store = useStore(key, SERVICE);
+  const value = useSyncExternalStore(store.subscribe, store.getSnapshot);
+
+  return useMemo<{ readonly value: any } & UseStorageMethod<V>>(
+    () => ({
       get value() {
-        const originValue = SERVICE.getItem(key);
-        if (isNull(originValue)) {
+        if (isNull(value)) {
           return defaultValue ?? null;
         }
 
-        return deserializer(originValue);
+        return deserializer(value);
       },
       set: (value) => {
         const originValue = serializer(value);
         SERVICE.setItem(key, originValue);
-        updateKey(originValue);
+        store.emitChange();
       },
       remove: () => {
         SERVICE.removeItem(key);
-        updateKey(null);
+        store.emitChange();
       },
-    };
-  }, [SERVICE, defaultValue, deserializer, key, serializer]);
-  const [, setOriginValue] = useState(SERVICE.getItem(key));
-
-  // eslint-disable-next-line react-hooks/exhaustive-deps
-  useIsomorphicLayoutEffect(() => {
-    const updatesOfKey = updates.get(key);
-    if (isUndefined(updatesOfKey)) {
-      updates.set(key, new Set([setOriginValue]));
-    } else if (!updatesOfKey.has(setOriginValue)) {
-      updatesOfKey.add(setOriginValue);
-    }
-  });
-
-  useEffect(() => {
-    updates.get(key)?.delete(setOriginValue);
-  }, [key]);
-
-  return res;
+    }),
+    [SERVICE, defaultValue, deserializer, key, serializer, store, value]
+  );
 }
 
 useStorage.SERVICE = new LocalStorageService() as AbstractStorage<any, any>;
 useStorage.PARSER = STRING_PARSER as AbstractParserOptions<any>;
 useStorage.clear = () => {
   useStorage.SERVICE.clear();
-  for (const [, cbs] of updates) {
-    for (const cb of cbs) {
-      cb(null);
-    }
+  for (const [, store] of stores) {
+    store.emitChange();
   }
 };