fix(react-query-persist-client): decouple subscribe and unsubscribe (#6028)

* fix(react-query-persist-client): decouple subscribe and unsubscribe

to make sure subscriptions are always done at the end, even if we unsubscribe while we're restoring

* chore: fix types

* chore: I think this is better

Author: TkDodo

packages/react-query-persist-client/src/PersistQueryClientProvider.tsx

@@ -1,7 +1,10 @@
 'use client'
 import * as React from 'react'
 
-import { persistQueryClient } from '@tanstack/query-persist-client-core'
+import {
+  persistQueryClientRestore,
+  persistQueryClientSubscribe,
+} from '@tanstack/query-persist-client-core'
 import { IsRestoringProvider, QueryClientProvider } from '@tanstack/react-query'
 import type { PersistQueryClientOptions } from '@tanstack/query-persist-client-core'
 import type { QueryClientProviderProps } from '@tanstack/react-query'
@@ -27,25 +30,23 @@ export const PersistQueryClientProvider = ({
   })
 
   React.useEffect(() => {
+    const options = {
+      ...refs.current.persistOptions,
+      queryClient: client,
+    }
     if (!didRestore.current) {
       didRestore.current = true
       setIsRestoring(true)
-      const [unsubscribe, promise] = persistQueryClient({
-        ...refs.current.persistOptions,
-        queryClient: client,
-      })
-
-      promise.then(() => {
-        refs.current.onSuccess?.()
-        setIsRestoring(false)
+      persistQueryClientRestore(options).then(async () => {
+        try {
+          await refs.current.onSuccess?.()
+        } finally {
+          setIsRestoring(false)
+        }
       })
-
-      return () => {
-        unsubscribe()
-      }
     }
-    return undefined
-  }, [client])
+    return isRestoring ? undefined : persistQueryClientSubscribe(options)
+  }, [client, isRestoring])
 
   return (
     <QueryClientProvider client={client} {...props}>

packages/react-query-persist-client/src/__tests__/PersistQueryClientProvider.test.tsx

@@ -1,5 +1,5 @@
 import * as React from 'react'
-import { render, waitFor } from '@testing-library/react'
+import { fireEvent, render, waitFor } from '@testing-library/react'
 
 import { QueryClient, useQueries, useQuery } from '@tanstack/react-query'
 import { persistQueryClientSave } from '@tanstack/query-persist-client-core'
@@ -121,6 +121,70 @@ describe('PersistQueryClientProvider', () => {
     })
   })
 
+  test('should subscribe correctly in StrictMode', async () => {
+    const key = queryKey()
+
+    const queryClient = createQueryClient()
+    await queryClient.prefetchQuery({
+      queryKey: key,
+      queryFn: () => Promise.resolve('hydrated'),
+    })
+
+    const persister = createMockPersister()
+
+    await persistQueryClientSave({ queryClient, persister })
+
+    queryClient.clear()
+
+    function Page() {
+      const state = useQuery({
+        queryKey: key,
+        queryFn: async () => {
+          await sleep(10)
+          return 'fetched'
+        },
+      })
+
+      return (
+        <div>
+          <h1>{state.data}</h1>
+          <h2>fetchStatus: {state.fetchStatus}</h2>
+          <button
+            onClick={() => {
+              queryClient.setQueryData(key, 'updated')
+            }}
+          >
+            update
+          </button>
+        </div>
+      )
+    }
+
+    const rendered = render(
+      <React.StrictMode>
+        <PersistQueryClientProvider
+          client={queryClient}
+          persistOptions={{ persister }}
+        >
+          <Page />
+        </PersistQueryClientProvider>
+        ,
+      </React.StrictMode>,
+    )
+
+    await waitFor(() => rendered.getByText('fetchStatus: idle'))
+    await waitFor(() => rendered.getByText('hydrated'))
+    await waitFor(() => rendered.getByText('fetched'))
+
+    fireEvent.click(rendered.getByRole('button', { name: /update/i }))
+
+    await waitFor(() => rendered.getByText('updated'))
+
+    const state = await persister.restoreClient()
+
+    expect(state?.clientState.queries[0]?.state.data).toBe('updated')
+  })
+
   test('should also put useQueries into idle state', async () => {
     const key = queryKey()
     const states: UseQueryResult[] = []