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

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,6 +77,8 @@ function InnerApp() {
           <React.Fragment
             // Resets the entire tree below when it changes:
             key={currentAccount?.did}>
+            <QueryProvider currentDid={currentAccount?.did}>
+              <PushNotificationsListener>
                 <StatsigProvider>
                   <LabelDefsProvider>
                     <LoggedOutViewProvider>
@@ -98,6 +98,8 @@ function InnerApp() {
                     </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,7 +129,6 @@ function App() {
    * that is set up in the InnerApp component above.
    */
   return (
-    <QueryProvider>
     <SessionProvider>
       <ShellStateProvider>
         <PrefsStateProvider>
@@ -143,7 +150,6 @@ function App() {
         </PrefsStateProvider>
       </ShellStateProvider>
     </SessionProvider>
-    </QueryProvider>
   )
 }
 

src/App.web.tsx

@@ -54,6 +54,7 @@ function InnerApp() {
       <React.Fragment
         // Resets the entire tree below when it changes:
         key={currentAccount?.did}>
+        <QueryProvider currentDid={currentAccount?.did}>
           <StatsigProvider>
             <LabelDefsProvider>
               <LoggedOutViewProvider>
@@ -73,6 +74,7 @@ function InnerApp() {
               </LoggedOutViewProvider>
             </LabelDefsProvider>
           </StatsigProvider>
+        </QueryProvider>
       </React.Fragment>
     </Alf>
   )
@@ -94,7 +96,6 @@ function App() {
    * that is set up in the InnerApp component above.
    */
   return (
-    <QueryProvider>
     <SessionProvider>
       <ShellStateProvider>
         <PrefsStateProvider>
@@ -116,7 +117,6 @@ function App() {
         </PrefsStateProvider>
       </ShellStateProvider>
     </SessionProvider>
-    </QueryProvider>
   )
 }
 

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,7 +39,8 @@ focusManager.setEventListener(onFocus => {
   }
 })
 
-const queryClient = new QueryClient({
+const createQueryClient = () =>
+  new QueryClient({
     defaultOptions: {
       queries: {
         // NOTE
@@ -60,11 +61,6 @@ const queryClient = new QueryClient({
     },
   })
 
-const asyncStoragePersister = createAsyncStoragePersister({
-  storage: AsyncStorage,
-  key: 'queryCache',
-})
-
 const dehydrateOptions: PersistQueryClientProviderProps['persistOptions']['dehydrateOptions'] =
   {
     shouldDehydrateMutation: (_: any) => false,
@@ -73,12 +69,50 @@ const dehydrateOptions: PersistQueryClientProviderProps['persistOptions']['dehyd
     },
   }
 
-const persistOptions = {
+export function QueryProvider({
-  persister: asyncStoragePersister,
+  children,
-  dehydrateOptions,
+  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}

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']>(