perf: ⚡ improve persist client write operations

packages/web/package.json

@@ -102,6 +102,7 @@
     "fuse.js": "^6.5.3",
     "highcharts": "^9.1.0",
     "highcharts-react-official": "^3.0.0",
+    "idb-keyval": "^6.2.1",
     "launchdarkly-react-client-sdk": "^3.0.6",
     "lightweight-charts": "^4.1.4",
     "lottie-react": "^2.4.0",

packages/web/utils/trpc.ts

@@ -1,13 +1,17 @@
 import { superjson } from "@osmosis-labs/server";
-import { makeIndexedKVStore } from "@osmosis-labs/stores";
 import {
   createTRPCRouter,
   localLink,
   makeSkipBatchLink,
 } from "@osmosis-labs/trpc";
-import { createAsyncStoragePersister } from "@tanstack/query-async-storage-persister";
+import { noop } from "@osmosis-labs/utils";
 import { QueryClient } from "@tanstack/react-query";
-import { persistQueryClient } from "@tanstack/react-query-persist-client";
+import {
+  PersistedClient,
+  Persister,
+  persistQueryClient,
+  Promisable,
+} from "@tanstack/react-query-persist-client";
 import { loggerLink } from "@trpc/client";
 import { createTRPCNext } from "@trpc/next";
 import type {
@@ -16,6 +20,7 @@ import type {
   inferRouterInputs,
   inferRouterOutputs,
 } from "@trpc/server";
+import { get, set, values } from "idb-keyval";
 
 import { AssetLists } from "~/config/generated/asset-lists";
 import { ChainList } from "~/config/generated/chain-list";
@@ -42,25 +47,186 @@ const trpcLocalRouter = createTRPCRouter({
   local: localRouter,
 });
 
+interface AsyncThrottleOptions {
+  interval?: number;
+  onError?: (error: unknown) => void;
+}
+
+function asyncThrottle<Args extends readonly unknown[]>(
+  func: (...args: Args) => Promise<void>,
+  { interval = 1000, onError = noop }: AsyncThrottleOptions = {}
+) {
+  if (typeof func !== "function") throw new Error("argument is not function.");
+
+  let running = false;
+  let lastTime = 0;
+  let timeout: ReturnType<typeof setTimeout>;
+  let currentArgs: Args | null = null;
+
+  const execFunc = async () => {
+    if (currentArgs) {
+      const args = currentArgs;
+      currentArgs = null;
+      try {
+        running = true;
+        await func(...args);
+      } catch (error) {
+        onError(error);
+      } finally {
+        lastTime = Date.now(); // this line must after 'func' executed to avoid two 'func' running in concurrent.
+        running = false;
+      }
+    }
+  };
+
+  const delayFunc = async () => {
+    clearTimeout(timeout);
+    timeout = setTimeout(() => {
+      if (running) {
+        delayFunc(); // Will come here when 'func' execution time is greater than the interval.
+      } else {
+        execFunc();
+      }
+    }, interval);
+  };
+
+  return (...args: Args) => {
+    currentArgs = args;
+
+    const tooSoon = Date.now() - lastTime < interval;
+    if (running || tooSoon) {
+      delayFunc();
+    } else {
+      execFunc();
+    }
+  };
+}
+
+interface AsyncStorage {
+  getItem: (key: string) => Promise<string | null>;
+  setItem: (key: string, value: string) => Promise<unknown>;
+  removeItem: (key: string) => Promise<void>;
+  values: () => Promise<string[]>;
+}
+
+type AsyncPersistRetryer = (props: {
+  persistedClient: PersistedClient;
+  error: Error;
+  errorCount: number;
+}) => Promisable<PersistedClient | undefined>;
+
+interface CreateAsyncStoragePersisterOptions {
+  /** The storage client used for setting and retrieving items from cache.
+   * For SSR pass in `undefined`. Note that window.localStorage can be
+   * `null` in Android WebViews depending on how they are configured.
+   */
+  storage: AsyncStorage | undefined | null;
+  /** The key to use when storing the cache */
+  key?: string;
+  /** To avoid spamming,
+   * pass a time in ms to throttle saving the cache to disk */
+  throttleTime?: number;
+  retry?: AsyncPersistRetryer;
+}
+
+const createAsyncStoragePersister = ({
+  storage,
+  key = `REACT_QUERY_OFFLINE_CACHE`,
+  throttleTime = 1000,
+  retry,
+}: CreateAsyncStoragePersisterOptions): Persister => {
+  if (storage) {
+    const trySave = async (
+      persistedClient: PersistedClient
+    ): Promise<Error | undefined> => {
+      try {
+        for (const query of persistedClient.clientState.queries) {
+          storage.setItem(query.queryKey.join("-"), superjson.stringify(query));
+        }
+
+        const clientState = {
+          buster: persistedClient.buster,
+          timestamp: persistedClient.timestamp,
+        };
+
+        await storage.setItem(key, JSON.stringify(clientState));
+        return;
+      } catch (error) {
+        return error as Error;
+      }
+    };
+
+    return {
+      persistClient: asyncThrottle(
+        async (persistedClient) => {
+          let client: PersistedClient | undefined = persistedClient;
+          let error = await trySave(client);
+          let errorCount = 0;
+          while (error && client) {
+            errorCount++;
+            client = await retry?.({
+              persistedClient: client,
+              error,
+              errorCount,
+            });
+
+            if (client) {
+              error = await trySave(client);
+            }
+          }
+        },
+        { interval: throttleTime }
+      ),
+      restoreClient: async () => {
+        const cacheClientStateString = await storage.getItem(key);
+
+        if (!cacheClientStateString) {
+          return;
+        }
+
+        const queriesString = await storage.values();
+        const queries = queriesString.map(
+          superjson.parse<PersistedClient["clientState"]["queries"][0]>
+        );
+        const client: PersistedClient = JSON.parse(cacheClientStateString);
+
+        return {
+          ...client,
+          clientState: {
+            queries,
+            mutations: [],
+          },
+        } as PersistedClient;
+      },
+      removeClient: () => storage.removeItem(key),
+    };
+  }
+
+  return {
+    persistClient: noop,
+    restoreClient: () => Promise.resolve(undefined),
+    removeClient: noop,
+  };
+};
+
 /** A set of type-safe react-query hooks for your tRPC API. */
 export const api = createTRPCNext<AppRouter>({
   config() {
-    const storage = makeIndexedKVStore("tanstack-query-cache");
+    /* const storage = makeIndexedKVStore("tanstack-query-cache"); */
 
     const localStoragePersister = createAsyncStoragePersister({
       storage:
         typeof window !== "undefined"
           ? {
               getItem: async (key) => {
-                const item: string | null | undefined = await storage.get(key);
+                const item: string | null | undefined = await get(key);
                 return item ?? null;
               },
-              setItem: (key, value) => storage.set(key, value),
-              removeItem: (key) => storage.set(key, undefined),
+              setItem: (key, value) => set(key, value),
+              removeItem: (key) => set(key, undefined),
+              values: () => values(),
             }
           : undefined,
-      serialize: (client) => superjson.stringify(client),
-      deserialize: (cachedString) => superjson.parse(cachedString),
     });
 
     const queryClient = new QueryClient({