fix: make sure all callbacks are called in suspense mode (TanStack#828)

Author: boschni

src/core/query.ts

@@ -37,6 +37,7 @@ export interface QueryState<TResult, TError> {
   error: TError | null
   failureCount: number
   isError: boolean
+  isFetched: boolean
   isFetching: boolean
   isFetchingMore?: IsFetchingMoreValue
   isIdle: boolean
@@ -57,7 +58,7 @@ export interface FetchMoreOptions {
   previous: boolean
 }
 
-enum ActionType {
+export enum ActionType {
   Failed = 'Failed',
   MarkStale = 'MarkStale',
   Fetch = 'Fetch',
@@ -95,7 +96,7 @@ interface SetStateAction<TResult, TError> {
   updater: Updater<QueryState<TResult, TError>, QueryState<TResult, TError>>
 }
 
-type Action<TResult, TError> =
+export type Action<TResult, TError> =
   | ErrorAction<TError>
   | FailedAction
   | FetchAction
@@ -112,8 +113,6 @@ export class Query<TResult, TError> {
   config: QueryConfig<TResult, TError>
   instances: QueryInstance<TResult, TError>[]
   state: QueryState<TResult, TError>
-  fallbackInstance?: QueryInstance<TResult, TError>
-  wasSuspended?: boolean
   shouldContinueRetryOnFocus?: boolean
   promise?: Promise<TResult | undefined>
 
@@ -163,7 +162,7 @@ export class Query<TResult, TError> {
     // Only update state if something has changed
     if (!shallowEqual(this.state, newState)) {
       this.state = newState
-      this.instances.forEach(d => d.onStateUpdate?.(this.state))
+      this.instances.forEach(d => d.onStateUpdate(newState, action))
       this.notifyGlobalListeners(this)
     }
   }
@@ -502,15 +501,6 @@ export class Query<TResult, TError> {
       // If there are any retries pending for this query, kill them
       this.cancelled = null
 
-      const getCallbackInstances = () => {
-        const callbackInstances = [...this.instances]
-
-        if (this.wasSuspended && this.fallbackInstance) {
-          callbackInstances.unshift(this.fallbackInstance)
-        }
-        return callbackInstances
-      }
-
       try {
         // Set up the query refreshing state
         this.dispatch({ type: ActionType.Fetch })
@@ -520,14 +510,6 @@ export class Query<TResult, TError> {
 
         this.setData(old => (this.config.isDataEqual!(old, data) ? old! : data))
 
-        getCallbackInstances().forEach(instance => {
-          instance.config.onSuccess?.(this.state.data!)
-        })
-
-        getCallbackInstances().forEach(instance =>
-          instance.config.onSettled?.(this.state.data, null)
-        )
-
         delete this.promise
 
         return data
@@ -541,14 +523,6 @@ export class Query<TResult, TError> {
         delete this.promise
 
         if (error !== this.cancelled) {
-          getCallbackInstances().forEach(instance =>
-            instance.config.onError?.(error)
-          )
-
-          getCallbackInstances().forEach(instance =>
-            instance.config.onSettled?.(undefined, error)
-          )
-
           throw error
         }
 
@@ -597,6 +571,7 @@ function getDefaultState<TResult, TError>(
   return {
     ...getStatusProps(initialStatus),
     error: null,
+    isFetched: false,
     isFetching: initialStatus === QueryStatus.Loading,
     failureCount: 0,
     isStale,
@@ -638,6 +613,7 @@ export function queryReducer<TResult, TError>(
         data: functionalUpdate(action.updater, state.data),
         error: null,
         isStale: action.isStale,
+        isFetched: true,
         isFetching: false,
         updatedAt: Date.now(),
         failureCount: 0,
@@ -646,6 +622,7 @@ export function queryReducer<TResult, TError>(
       return {
         ...state,
         failureCount: state.failureCount + 1,
+        isFetched: true,
         isFetching: false,
         isStale: true,
         ...(!action.cancelled && {

src/core/queryCache.ts

@@ -18,7 +18,6 @@ import {
   TupleQueryFunction,
   TupleQueryKey,
 } from './types'
-import { QueryInstance } from './queryInstance'
 
 // TYPES
 
@@ -283,14 +282,6 @@ export class QueryCache {
       }
     }
 
-    query.fallbackInstance = {
-      config: {
-        onSuccess: query.config.onSuccess,
-        onError: query.config.onError,
-        onSettled: query.config.onSettled,
-      },
-    } as QueryInstance<TResult, TError>
-
     return query
   }
 

src/core/queryInstance.ts

@@ -1,5 +1,5 @@
 import { uid, isServer, isDocumentVisible, Console } from './utils'
-import { Query, QueryState } from './query'
+import { Query, QueryState, Action, ActionType } from './query'
 import { BaseQueryConfig } from './types'
 
 // TYPES
@@ -13,17 +13,17 @@ export type OnStateUpdateFunction<TResult, TError> = (
 export class QueryInstance<TResult, TError> {
   id: number
   config: BaseQueryConfig<TResult, TError>
-  onStateUpdate?: OnStateUpdateFunction<TResult, TError>
 
   private query: Query<TResult, TError>
   private refetchIntervalId?: number
+  private stateUpdateListener?: OnStateUpdateFunction<TResult, TError>
 
   constructor(
     query: Query<TResult, TError>,
     onStateUpdate?: OnStateUpdateFunction<TResult, TError>
   ) {
     this.id = uid()
-    this.onStateUpdate = onStateUpdate
+    this.stateUpdateListener = onStateUpdate
     this.query = query
     this.config = {}
   }
@@ -77,30 +77,45 @@ export class QueryInstance<TResult, TError> {
       // Perform the refetch for this query if necessary
       if (
         this.query.instances.some(d => d.config.enabled) && // Don't auto refetch if disabled
-        !this.query.wasSuspended && // Don't double refetch for suspense
+        !(this.config.suspense && this.query.state.isFetched) && // Don't refetch if in suspense mode and the data is already fetched
         this.query.state.isStale && // Only refetch if stale
-        (this.query.config.refetchOnMount || this.query.instances.length === 1)
+        (this.config.refetchOnMount || this.query.instances.length === 1)
       ) {
         await this.query.fetch()
       }
-
-      this.query.wasSuspended = false
     } catch (error) {
       Console.error(error)
     }
   }
 
-  unsubscribe(): void {
+  unsubscribe(preventGC?: boolean): void {
     this.query.instances = this.query.instances.filter(d => d.id !== this.id)
 
     if (!this.query.instances.length) {
       this.clearInterval()
       this.query.cancel()
 
-      if (!isServer) {
+      if (!preventGC && !isServer) {
         // Schedule garbage collection
         this.query.scheduleGarbageCollection()
       }
     }
   }
+
+  onStateUpdate(
+    state: QueryState<TResult, TError>,
+    action: Action<TResult, TError>
+  ): void {
+    if (action.type === ActionType.Success && state.isSuccess) {
+      this.config.onSuccess?.(state.data!)
+      this.config.onSettled?.(state.data!, null)
+    }
+
+    if (action.type === ActionType.Error && state.isError) {
+      this.config.onError?.(state.error!)
+      this.config.onSettled?.(undefined, state.error!)
+    }
+
+    this.stateUpdateListener?.(state)
+  }
 }

src/react/tests/suspense.test.tsx

@@ -90,6 +90,43 @@ describe("useQuery's in Suspense mode", () => {
     expect(successFn).toHaveBeenCalledTimes(1)
   })
 
+  it('should call every onSuccess handler within a suspense boundary', async () => {
+    const key = queryKey()
+
+    const successFn1 = jest.fn()
+    const successFn2 = jest.fn()
+
+    function FirstComponent() {
+      useQuery(key, () => sleep(10), {
+        suspense: true,
+        onSuccess: successFn1,
+      })
+
+      return <span>first</span>
+    }
+
+    function SecondComponent() {
+      useQuery(key, () => sleep(20), {
+        suspense: true,
+        onSuccess: successFn2,
+      })
+
+      return <span>second</span>
+    }
+
+    const rendered = render(
+      <React.Suspense fallback="loading">
+        <FirstComponent />
+        <SecondComponent />
+      </React.Suspense>
+    )
+
+    await waitFor(() => rendered.getByText('second'))
+
+    expect(successFn1).toHaveBeenCalledTimes(1)
+    expect(successFn1).toHaveBeenCalledTimes(1)
+  })
+
   // https://github.com/tannerlinsley/react-query/issues/468
   it('should reset error state if new component instances are mounted', async () => {
     const key = queryKey()

src/react/tests/useQuery.test.tsx

@@ -546,21 +546,16 @@ describe('useQuery', () => {
       staleTime: 10,
     })
 
-    await sleep(100)
+    await sleep(11)
 
     function Page() {
-      const query = useQuery(key, queryFn)
-
-      return (
-        <div>
-          <div>{query.data}</div>
-        </div>
-      )
+      useQuery(key, queryFn)
+      return null
     }
 
     render(<Page />)
 
-    await sleep(100)
+    await act(() => sleep(0))
 
     expect(prefetchQueryFn).toHaveBeenCalledTimes(1)
     expect(queryFn).toHaveBeenCalledTimes(1)
@@ -579,21 +574,16 @@ describe('useQuery', () => {
       staleTime: 1000,
     })
 
-    sleep(100)
+    await sleep(0)
 
     function Page() {
-      const query = useQuery(key, queryFn)
-
-      return (
-        <div>
-          <div>{query.data}</div>
-        </div>
-      )
+      useQuery(key, queryFn)
+      return null
     }
 
     render(<Page />)
 
-    sleep(100)
+    await sleep(0)
 
     expect(prefetchQueryFn).toHaveBeenCalledTimes(1)
     expect(queryFn).toHaveBeenCalledTimes(0)

src/react/useInfiniteQuery.ts

@@ -80,7 +80,7 @@ export function useInfiniteQuery<TResult, TError>(
   const query = result.query
   const state = result.query.state
 
-  handleSuspense(result)
+  handleSuspense(config, result)
 
   const fetchMore = React.useMemo(() => query.fetchMore.bind(query), [query])
 

src/react/usePaginatedQuery.ts

@@ -124,7 +124,7 @@ export function usePaginatedQuery<TResult, TError>(
     Object.assign(result, overrides)
   }
 
-  handleSuspense(result)
+  handleSuspense(config, result)
 
   return {
     ...result,

src/react/useQuery.ts

@@ -66,7 +66,7 @@ export function useQuery<TResult, TError>(
   const [queryKey, config] = useQueryArgs<TResult, TError>(args)
   const result = useBaseQuery<TResult, TError>(queryKey, config)
 
-  handleSuspense(result)
+  handleSuspense(config, result)
 
   return {
     ...result,

src/react/utils.ts

@@ -1,7 +1,7 @@
 import React from 'react'
 
 import { uid, isServer } from '../core/utils'
-import { QueryResultBase, QueryStatus } from '../core/types'
+import { QueryResultBase, BaseQueryConfig, QueryStatus } from '../core/types'
 
 export function useUid(): number {
   const ref = React.useRef(0)
@@ -40,9 +40,12 @@ export function useRerenderer() {
   return React.useCallback(() => rerender({}), [rerender])
 }
 
-export function handleSuspense(result: QueryResultBase<any, any>) {
+export function handleSuspense(
+  config: BaseQueryConfig<any, any>,
+  result: QueryResultBase<any, any>
+) {
   const { error, query } = result
-  const { config, state } = query
+  const { state } = query
 
   if (config.suspense || config.useErrorBoundary) {
     if (state.status === QueryStatus.Error && state.throwInErrorBoundary) {
@@ -54,7 +57,16 @@ export function handleSuspense(result: QueryResultBase<any, any>) {
       state.status !== QueryStatus.Success &&
       config.enabled
     ) {
-      query.wasSuspended = true
+      const instance = query.subscribe()
+
+      instance.updateConfig({
+        ...config,
+        onSettled: (data, error) => {
+          instance.unsubscribe(true)
+          config.onSettled?.(data, error)
+        },
+      })
+
       throw query.fetch()
     }
   }