Scope query client per DID (#3333)

* Move QueryProvider inside the key

* Pull useQueryClient-dependent code down in App.native

* Remove useQueryClient dependency from session provider

* Scope query client per DID

4 files changed, 159 insertions, 126 deletions

diff --git a/src/App.native.tsx b/src/App.native.tsx
index 57ebe4951..9abe4a559 100644
--- a/src/App.native.tsx
+++ b/src/App.native.tsx
@@ -54,12 +54,10 @@ SplashScreen.preventAutoHideAsync()
 function InnerApp() {
   const {isInitialLoad, currentAccount} = useSession()
   const {resumeSession} = useSessionApi()
-  const queryClient = useQueryClient()
   const theme = useColorModeTheme()
   const {_} = useLingui()
 
   useIntentHandler()
-  useNotificationsListener(queryClient)
   useOTAUpdates()
 
   // init
@@ -79,25 +77,29 @@ function InnerApp() {
           <React.Fragment
             // Resets the entire tree below when it changes:
             key={currentAccount?.did}>
-            <StatsigProvider>
-              <LabelDefsProvider>
-                <LoggedOutViewProvider>
-                  <SelectedFeedProvider>
-                    <UnreadNotifsProvider>
-                      <ThemeProvider theme={theme}>
-                        {/* All components should be within this provider */}
-                        <RootSiblingParent>
-                          <GestureHandlerRootView style={s.h100pct}>
-                            <TestCtrls />
-                            <Shell />
-                          </GestureHandlerRootView>
-                        </RootSiblingParent>
-                      </ThemeProvider>
-                    </UnreadNotifsProvider>
-                  </SelectedFeedProvider>
-                </LoggedOutViewProvider>
-              </LabelDefsProvider>
-            </StatsigProvider>
+            <QueryProvider currentDid={currentAccount?.did}>
+              <PushNotificationsListener>
+                <StatsigProvider>
+                  <LabelDefsProvider>
+                    <LoggedOutViewProvider>
+                      <SelectedFeedProvider>
+                        <UnreadNotifsProvider>
+                          <ThemeProvider theme={theme}>
+                            {/* All components should be within this provider */}
+                            <RootSiblingParent>
+                              <GestureHandlerRootView style={s.h100pct}>
+                                <TestCtrls />
+                                <Shell />
+                              </GestureHandlerRootView>
+                            </RootSiblingParent>
+                          </ThemeProvider>
+                        </UnreadNotifsProvider>
+                      </SelectedFeedProvider>
+                    </LoggedOutViewProvider>
+                  </LabelDefsProvider>
+                </StatsigProvider>
+              </PushNotificationsListener>
+            </QueryProvider>
           </React.Fragment>
         </Splash>
       </Alf>
@@ -105,6 +107,12 @@ function InnerApp() {
   )
 }
 
+function PushNotificationsListener({children}: {children: React.ReactNode}) {
+  const queryClient = useQueryClient()
+  useNotificationsListener(queryClient)
+  return children
+}
+
 function App() {
   const [isReady, setReady] = useState(false)
 
@@ -121,29 +129,27 @@ function App() {
    * that is set up in the InnerApp component above.
    */
   return (
-    <QueryProvider>
-      <SessionProvider>
-        <ShellStateProvider>
-          <PrefsStateProvider>
-            <MutedThreadsProvider>
-              <InvitesStateProvider>
-                <ModalStateProvider>
-                  <DialogStateProvider>
-                    <LightboxStateProvider>
-                      <I18nProvider>
-                        <PortalProvider>
-                          <InnerApp />
-                        </PortalProvider>
-                      </I18nProvider>
-                    </LightboxStateProvider>
-                  </DialogStateProvider>
-                </ModalStateProvider>
-              </InvitesStateProvider>
-            </MutedThreadsProvider>
-          </PrefsStateProvider>
-        </ShellStateProvider>
-      </SessionProvider>
-    </QueryProvider>
+    <SessionProvider>
+      <ShellStateProvider>
+        <PrefsStateProvider>
+          <MutedThreadsProvider>
+            <InvitesStateProvider>
+              <ModalStateProvider>
+                <DialogStateProvider>
+                  <LightboxStateProvider>
+                    <I18nProvider>
+                      <PortalProvider>
+                        <InnerApp />
+                      </PortalProvider>
+                    </I18nProvider>
+                  </LightboxStateProvider>
+                </DialogStateProvider>
+              </ModalStateProvider>
+            </InvitesStateProvider>
+          </MutedThreadsProvider>
+        </PrefsStateProvider>
+      </ShellStateProvider>
+    </SessionProvider>
   )
 }
 
diff --git a/src/App.web.tsx b/src/App.web.tsx
index 2910bbbae..ccf7ecb49 100644
--- a/src/App.web.tsx
+++ b/src/App.web.tsx
@@ -54,25 +54,27 @@ function InnerApp() {
       <React.Fragment
         // Resets the entire tree below when it changes:
         key={currentAccount?.did}>
-        <StatsigProvider>
-          <LabelDefsProvider>
-            <LoggedOutViewProvider>
-              <SelectedFeedProvider>
-                <UnreadNotifsProvider>
-                  <ThemeProvider theme={theme}>
-                    {/* All components should be within this provider */}
-                    <RootSiblingParent>
-                      <SafeAreaProvider>
-                        <Shell />
-                      </SafeAreaProvider>
-                    </RootSiblingParent>
-                    <ToastContainer />
-                  </ThemeProvider>
-                </UnreadNotifsProvider>
-              </SelectedFeedProvider>
-            </LoggedOutViewProvider>
-          </LabelDefsProvider>
-        </StatsigProvider>
+        <QueryProvider currentDid={currentAccount?.did}>
+          <StatsigProvider>
+            <LabelDefsProvider>
+              <LoggedOutViewProvider>
+                <SelectedFeedProvider>
+                  <UnreadNotifsProvider>
+                    <ThemeProvider theme={theme}>
+                      {/* All components should be within this provider */}
+                      <RootSiblingParent>
+                        <SafeAreaProvider>
+                          <Shell />
+                        </SafeAreaProvider>
+                      </RootSiblingParent>
+                      <ToastContainer />
+                    </ThemeProvider>
+                  </UnreadNotifsProvider>
+                </SelectedFeedProvider>
+              </LoggedOutViewProvider>
+            </LabelDefsProvider>
+          </StatsigProvider>
+        </QueryProvider>
       </React.Fragment>
     </Alf>
   )
@@ -94,29 +96,27 @@ function App() {
    * that is set up in the InnerApp component above.
    */
   return (
-    <QueryProvider>
-      <SessionProvider>
-        <ShellStateProvider>
-          <PrefsStateProvider>
-            <MutedThreadsProvider>
-              <InvitesStateProvider>
-                <ModalStateProvider>
-                  <DialogStateProvider>
-                    <LightboxStateProvider>
-                      <I18nProvider>
-                        <PortalProvider>
-                          <InnerApp />
-                        </PortalProvider>
-                      </I18nProvider>
-                    </LightboxStateProvider>
-                  </DialogStateProvider>
-                </ModalStateProvider>
-              </InvitesStateProvider>
-            </MutedThreadsProvider>
-          </PrefsStateProvider>
-        </ShellStateProvider>
-      </SessionProvider>
-    </QueryProvider>
+    <SessionProvider>
+      <ShellStateProvider>
+        <PrefsStateProvider>
+          <MutedThreadsProvider>
+            <InvitesStateProvider>
+              <ModalStateProvider>
+                <DialogStateProvider>
+                  <LightboxStateProvider>
+                    <I18nProvider>
+                      <PortalProvider>
+                        <InnerApp />
+                      </PortalProvider>
+                    </I18nProvider>
+                  </LightboxStateProvider>
+                </DialogStateProvider>
+              </ModalStateProvider>
+            </InvitesStateProvider>
+          </MutedThreadsProvider>
+        </PrefsStateProvider>
+      </ShellStateProvider>
+    </SessionProvider>
   )
 }
 
diff --git a/src/lib/react-query.tsx b/src/lib/react-query.tsx
index 08b61ee20..2fcd46942 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}
diff --git a/src/state/session/index.tsx b/src/state/session/index.tsx
index c7dba3089..5c7cc1591 100644
--- a/src/state/session/index.tsx
+++ b/src/state/session/index.tsx
@@ -4,7 +4,6 @@ import {
   BSKY_LABELER_DID,
   BskyAgent,
 } from '@atproto/api'
-import {useQueryClient} from '@tanstack/react-query'
 import {jwtDecode} from 'jwt-decode'
 
 import {track} from '#/lib/analytics/analytics'
@@ -178,7 +177,6 @@ function createPersistSessionHandler(
 }
 
 export function Provider({children}: React.PropsWithChildren<{}>) {
-  const queryClient = useQueryClient()
   const isDirty = React.useRef(false)
   const [state, setState] = React.useState<SessionState>({
     isInitialLoad: true,
@@ -211,12 +209,11 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
   const clearCurrentAccount = React.useCallback(() => {
     logger.warn(`session: clear current account`)
     __globalAgent = PUBLIC_BSKY_AGENT
-    queryClient.clear()
     setStateAndPersist(s => ({
       ...s,
       currentAccount: undefined,
     }))
-  }, [setStateAndPersist, queryClient])
+  }, [setStateAndPersist])
 
   const createAccount = React.useCallback<ApiContext['createAccount']>(
     async ({
@@ -286,14 +283,13 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       )
 
       __globalAgent = agent
-      queryClient.clear()
       upsertAccount(account)
 
       logger.debug(`session: created account`, {}, logger.DebugContext.session)
       track('Create Account')
       logEvent('account:create:success', {})
     },
-    [upsertAccount, queryClient, clearCurrentAccount],
+    [upsertAccount, clearCurrentAccount],
   )
 
   const login = React.useCallback<ApiContext['login']>(
@@ -334,7 +330,6 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       __globalAgent = agent
       // @ts-ignore
       if (IS_DEV && isWeb) window.agent = agent
-      queryClient.clear()
       upsertAccount(account)
 
       logger.debug(`session: logged in`, {}, logger.DebugContext.session)
@@ -342,7 +337,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       track('Sign In', {resumedSession: false})
       logEvent('account:loggedIn', {logContext, withPassword: true})
     },
-    [upsertAccount, queryClient, clearCurrentAccount],
+    [upsertAccount, clearCurrentAccount],
   )
 
   const logout = React.useCallback<ApiContext['logout']>(
@@ -411,7 +406,6 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
 
         agent.session = prevSession
         __globalAgent = agent
-        queryClient.clear()
         upsertAccount(account)
 
         if (prevSession.deactivated) {
@@ -448,7 +442,6 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
         try {
           const freshAccount = await resumeSessionWithFreshAccount()
           __globalAgent = agent
-          queryClient.clear()
           upsertAccount(freshAccount)
         } catch (e) {
           /*
@@ -489,7 +482,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
         }
       }
     },
-    [upsertAccount, queryClient, clearCurrentAccount],
+    [upsertAccount, clearCurrentAccount],
   )
 
   const resumeSession = React.useCallback<ApiContext['resumeSession']>(