fix: update synchronously to make sure observers always have the latest state (#1757)

Author: boschni

src/core/index.ts

@@ -1,3 +1,4 @@
+export { CancelledError } from './retryer'
 export { QueryCache } from './queryCache'
 export { QueryClient } from './queryClient'
 export { QueryObserver } from './queryObserver'
@@ -14,6 +15,5 @@ export { isCancelledError } from './retryer'
 
 // Types
 export * from './types'
-export type { CancelledError } from './retryer'
 export type { Query } from './query'
 export type { Logger } from './logger'

src/core/query.ts

@@ -135,7 +135,7 @@ export class Query<
   private cache: QueryCache
   private promise?: Promise<TData>
   private gcTimeout?: number
-  private retryer?: Retryer<unknown, TError>
+  private retryer?: Retryer<TData, TError>
   private observers: QueryObserver<any, any, any, any>[]
   private defaultOptions?: QueryOptions<TQueryFnData, TError, TData>
 
@@ -387,28 +387,21 @@ export class Query<
 
     // Try to fetch the data
     this.retryer = new Retryer({
-      fn: context.fetchFn,
-      onFail: () => {
-        this.dispatch({ type: 'failed' })
-      },
-      onPause: () => {
-        this.dispatch({ type: 'pause' })
-      },
-      onContinue: () => {
-        this.dispatch({ type: 'continue' })
-      },
-      retry: context.options.retry,
-      retryDelay: context.options.retryDelay,
-    })
+      fn: context.fetchFn as () => TData,
+      onSuccess: data => {
+        this.setData(data as TData)
 
-    this.promise = this.retryer.promise
-      .then(data => this.setData(data as TData))
-      .catch(error => {
-        // Set error state if needed
+        // Remove query after fetching if cache time is 0
+        if (this.cacheTime === 0) {
+          this.optionalRemove()
+        }
+      },
+      onError: error => {
+        // Optimistically update state if needed
         if (!(isCancelledError(error) && error.silent)) {
           this.dispatch({
             type: 'error',
-            error,
+            error: error as TError,
           })
         }
 
@@ -426,18 +419,21 @@ export class Query<
         if (this.cacheTime === 0) {
           this.optionalRemove()
         }
+      },
+      onFail: () => {
+        this.dispatch({ type: 'failed' })
+      },
+      onPause: () => {
+        this.dispatch({ type: 'pause' })
+      },
+      onContinue: () => {
+        this.dispatch({ type: 'continue' })
+      },
+      retry: context.options.retry,
+      retryDelay: context.options.retryDelay,
+    })
 
-        // Propagate error
-        throw error
-      })
-      .then(data => {
-        // Remove query after fetching if cache time is 0
-        if (this.cacheTime === 0) {
-          this.optionalRemove()
-        }
-
-        return data
-      })
+    this.promise = this.retryer.promise
 
     return this.promise
   }

src/core/retryer.ts

@@ -6,6 +6,8 @@ import { functionalUpdate, sleep } from './utils'
 
 interface RetryerConfig<TData = unknown, TError = unknown> {
   fn: () => TData | Promise<TData>
+  onError?: (error: TError) => void
+  onSuccess?: (data: TData) => void
   onFail?: (failureCount: number, error: TError) => void
   onPause?: () => void
   onContinue?: () => void
@@ -88,15 +90,21 @@ export class Retryer<TData = unknown, TError = unknown> {
     })
 
     const resolve = (value: any) => {
-      this.isResolved = true
-      continueFn?.()
-      promiseResolve(value)
+      if (!this.isResolved) {
+        this.isResolved = true
+        config.onSuccess?.(value)
+        continueFn?.()
+        promiseResolve(value)
+      }
     }
 
     const reject = (value: any) => {
-      this.isResolved = true
-      continueFn?.()
-      promiseReject(value)
+      if (!this.isResolved) {
+        this.isResolved = true
+        config.onError?.(value)
+        continueFn?.()
+        promiseReject(value)
+      }
     }
 
     const pause = () => {
@@ -129,13 +137,15 @@ export class Retryer<TData = unknown, TError = unknown> {
 
       // Create callback to cancel this fetch
       cancelFn = cancelOptions => {
-        reject(new CancelledError(cancelOptions))
-
-        // Cancel transport if supported
-        if (isCancelable(promiseOrValue)) {
-          try {
-            promiseOrValue.cancel()
-          } catch {}
+        if (!this.isResolved) {
+          reject(new CancelledError(cancelOptions))
+
+          // Cancel transport if supported
+          if (isCancelable(promiseOrValue)) {
+            try {
+              promiseOrValue.cancel()
+            } catch {}
+          }
         }
       }
 

src/react/tests/useQuery.test.tsx

@@ -13,6 +13,7 @@ import {
 } from './utils'
 import {
   useQuery,
+  CancelledError,
   QueryClient,
   UseQueryResult,
   QueryCache,
@@ -900,7 +901,7 @@ describe('useQuery', () => {
 
     renderWithClient(queryClient, <Page />)
 
-    await waitFor(() => renderCount > 1)
+    await sleep(10)
     expect(renderCount).toBe(2)
     expect(states.length).toBe(2)
     expect(states[0]).toMatchObject({ data: undefined })
@@ -1051,7 +1052,14 @@ describe('useQuery', () => {
     queryClient.setQueryData(key, 'set')
 
     function Page() {
-      const result = useQuery(key, () => 'fetched', { enabled: false })
+      const result = useQuery(
+        key,
+        async () => {
+          await sleep(1)
+          return 'fetched'
+        },
+        { enabled: false }
+      )
 
       results.push(result)
 
@@ -1082,7 +1090,8 @@ describe('useQuery', () => {
     function Page() {
       const state = useQuery(
         key,
-        () => {
+        async () => {
+          await sleep(1)
           count++
           return count
         },
@@ -2214,10 +2223,17 @@ describe('useQuery', () => {
     let count = 0
 
     function Page() {
-      const state = useQuery(key, () => count++, {
-        staleTime: Infinity,
-        refetchOnWindowFocus: 'always',
-      })
+      const state = useQuery(
+        key,
+        async () => {
+          await sleep(1)
+          return count++
+        },
+        {
+          staleTime: Infinity,
+          refetchOnWindowFocus: 'always',
+        }
+      )
       states.push(state)
       return null
     }
@@ -2725,7 +2741,10 @@ describe('useQuery', () => {
     queryClient.setQueryData(key, 'prefetched')
 
     function Page() {
-      const state = useQuery(key, () => 'data')
+      const state = useQuery(key, async () => {
+        await sleep(1)
+        return 'data'
+      })
       states.push(state)
       return null
     }
@@ -3269,6 +3288,70 @@ describe('useQuery', () => {
     expect(cancelFn).toHaveBeenCalled()
   })
 
+  it('should refetch when quickly switching to a failed query', async () => {
+    const key = queryKey()
+    const states: UseQueryResult<string>[] = []
+
+    const queryFn = () => {
+      let cancelFn = jest.fn()
+
+      const promise = new Promise<string>((resolve, reject) => {
+        cancelFn = jest.fn(() => reject('Cancelled'))
+        sleep(50).then(() => resolve('OK'))
+      })
+
+      ;(promise as any).cancel = cancelFn
+
+      return promise
+    }
+
+    function Page() {
+      const [id, setId] = React.useState(1)
+      const [hasChanged, setHasChanged] = React.useState(false)
+
+      const state = useQuery([key, id], queryFn)
+
+      states.push(state)
+
+      React.useEffect(() => {
+        setId(prevId => (prevId === 1 ? 2 : 1))
+        setHasChanged(true)
+      }, [hasChanged])
+
+      return null
+    }
+
+    renderWithClient(queryClient, <Page />)
+
+    await sleep(100)
+    expect(states.length).toBe(5)
+    // Load query 1
+    expect(states[0]).toMatchObject({
+      status: 'loading',
+      error: null,
+    })
+    // Load query 2
+    expect(states[1]).toMatchObject({
+      status: 'loading',
+      error: null,
+    })
+    // Load query 1
+    expect(states[2]).toMatchObject({
+      status: 'loading',
+      error: expect.any(CancelledError),
+    })
+    // State update
+    expect(states[3]).toMatchObject({
+      status: 'loading',
+      error: expect.any(CancelledError),
+    })
+    // Loaded query 1
+    expect(states[4]).toMatchObject({
+      status: 'success',
+      error: null,
+    })
+  })
+
   it('should update query state and refetch when reset with resetQueries', async () => {
     const key = queryKey()
     const states: UseQueryResult<number>[] = []
@@ -3277,7 +3360,8 @@ describe('useQuery', () => {
     function Page() {
       const state = useQuery(
         key,
-        () => {
+        async () => {
+          await sleep(1)
           count++
           return count
         },