From 2e2e554f079bab5d88c89d393cb9177dfbd71e90 Mon Sep 17 00:00:00 2001
From: Dan Abramov <dan.abramov@gmail.com>
Date: Fri, 22 Mar 2024 00:06:25 +0000
Subject: [PATCH] Scope query client per DID

---
 src/App.native.tsx      |  2 +-
 src/App.web.tsx         |  2 +-
 src/lib/react-query.tsx | 92 ++++++++++++++++++++++++++++-------------
 3 files changed, 65 insertions(+), 31 deletions(-)

diff --git a/src/App.native.tsx b/src/App.native.tsx
index b9e9a5eca0..ede587c899 100644
--- a/src/App.native.tsx
+++ b/src/App.native.tsx
@@ -75,7 +75,7 @@ function InnerApp() {
           <React.Fragment
             // Resets the entire tree below when it changes:
             key={currentAccount?.did}>
-            <QueryProvider>
+            <QueryProvider currentDid={currentAccount?.did}>
               <PushNotificationsListener>
                 <StatsigProvider>
                   <LabelDefsProvider>
diff --git a/src/App.web.tsx b/src/App.web.tsx
index a876fc894a..ccf7ecb491 100644
--- a/src/App.web.tsx
+++ b/src/App.web.tsx
@@ -54,7 +54,7 @@ function InnerApp() {
       <React.Fragment
         // Resets the entire tree below when it changes:
         key={currentAccount?.did}>
-        <QueryProvider>
+        <QueryProvider currentDid={currentAccount?.did}>
           <StatsigProvider>
             <LabelDefsProvider>
               <LoggedOutViewProvider>
diff --git a/src/lib/react-query.tsx b/src/lib/react-query.tsx
index 08b61ee205..2fcd469427 100644
--- a/src/lib/react-query.tsx
+++ b/src/lib/react-query.tsx
@@ -1,4 +1,4 @@
-import React from 'react'
+import React, {useRef, useState} from 'react'
 import {AppState, AppStateStatus} from 'react-native'
 import AsyncStorage from '@react-native-async-storage/async-storage'
 import {createAsyncStoragePersister} from '@tanstack/query-async-storage-persister'
@@ -39,31 +39,27 @@ focusManager.setEventListener(onFocus => {
   }
 })
 
-const queryClient = new QueryClient({
-  defaultOptions: {
-    queries: {
-      // NOTE
-      // refetchOnWindowFocus breaks some UIs (like feeds)
-      // so we only selectively want to enable this
-      // -prf
-      refetchOnWindowFocus: false,
-      // Structural sharing between responses makes it impossible to rely on
-      // "first seen" timestamps on objects to determine if they're fresh.
-      // Disable this optimization so that we can rely on "first seen" timestamps.
-      structuralSharing: false,
-      // We don't want to retry queries by default, because in most cases we
-      // want to fail early and show a response to the user. There are
-      // exceptions, and those can be made on a per-query basis. For others, we
-      // should give users controls to retry.
-      retry: false,
+const createQueryClient = () =>
+  new QueryClient({
+    defaultOptions: {
+      queries: {
+        // NOTE
+        // refetchOnWindowFocus breaks some UIs (like feeds)
+        // so we only selectively want to enable this
+        // -prf
+        refetchOnWindowFocus: false,
+        // Structural sharing between responses makes it impossible to rely on
+        // "first seen" timestamps on objects to determine if they're fresh.
+        // Disable this optimization so that we can rely on "first seen" timestamps.
+        structuralSharing: false,
+        // We don't want to retry queries by default, because in most cases we
+        // want to fail early and show a response to the user. There are
+        // exceptions, and those can be made on a per-query basis. For others, we
+        // should give users controls to retry.
+        retry: false,
+      },
     },
-  },
-})
-
-const asyncStoragePersister = createAsyncStoragePersister({
-  storage: AsyncStorage,
-  key: 'queryCache',
-})
+  })
 
 const dehydrateOptions: PersistQueryClientProviderProps['persistOptions']['dehydrateOptions'] =
   {
@@ -73,12 +69,50 @@ const dehydrateOptions: PersistQueryClientProviderProps['persistOptions']['dehyd
     },
   }
 
-const persistOptions = {
-  persister: asyncStoragePersister,
-  dehydrateOptions,
+export function QueryProvider({
+  children,
+  currentDid,
+}: {
+  children: React.ReactNode
+  currentDid: string | undefined
+}) {
+  return (
+    <QueryProviderInner
+      // Enforce we never reuse cache between users.
+      // These two props MUST stay in sync.
+      key={currentDid}
+      currentDid={currentDid}>
+      {children}
+    </QueryProviderInner>
+  )
 }
 
-export function QueryProvider({children}: {children: React.ReactNode}) {
+function QueryProviderInner({
+  children,
+  currentDid,
+}: {
+  children: React.ReactNode
+  currentDid: string | undefined
+}) {
+  const initialDid = useRef(currentDid)
+  if (currentDid !== initialDid.current) {
+    throw Error(
+      'Something is very wrong. Expected did to be stable due to key above.',
+    )
+  }
+  // We create the query client here so that it's scoped to a specific DID.
+  // Do not move the query client creation outside of this component.
+  const [queryClient, _setQueryClient] = useState(() => createQueryClient())
+  const [persistOptions, _setPersistOptions] = useState(() => {
+    const asyncPersister = createAsyncStoragePersister({
+      storage: AsyncStorage,
+      key: 'queryClient-' + (currentDid ?? 'logged-out'),
+    })
+    return {
+      persister: asyncPersister,
+      dehydrateOptions,
+    }
+  })
   return (
     <PersistQueryClientProvider
       client={queryClient}