refactor: only batch react callbacks

Author: boschni

src/core/mutationCache.ts

@@ -66,9 +66,7 @@ export class MutationCache extends Subscribable<MutationCacheListener> {
   notify(mutation?: Mutation<any, any, any, any>) {
     notifyManager.batch(() => {
       this.listeners.forEach(listener => {
-        notifyManager.schedule(() => {
-          listener(mutation)
-        })
+        listener(mutation)
       })
     })
   }

src/core/mutationObserver.ts

@@ -118,12 +118,9 @@ export class MutationObserver<
   }
 
   private notify() {
-    const { currentResult } = this
     notifyManager.batch(() => {
       this.listeners.forEach(listener => {
-        notifyManager.schedule(() => {
-          listener(currentResult)
-        })
+        listener(this.currentResult)
       })
     })
   }

src/core/notifyManager.ts

@@ -49,6 +49,17 @@ class NotifyManager {
     }
   }
 
+  /**
+   * All calls to the wrapped function will be batched.
+   */
+  batchCalls<T extends Function>(callback: T): T {
+    return ((...args: any[]) => {
+      this.schedule(() => {
+        callback(...args)
+      })
+    }) as any
+  }
+
   flush(): void {
     const queue = this.queue
     this.queue = []

src/core/queriesObserver.ts

@@ -119,12 +119,9 @@ export class QueriesObserver extends Subscribable<QueriesObserverListener> {
   }
 
   private notify(): void {
-    const { result } = this
     notifyManager.batch(() => {
       this.listeners.forEach(listener => {
-        notifyManager.schedule(() => {
-          listener(result)
-        })
+        listener(this.result)
       })
     })
   }

src/core/queryCache.ts

@@ -112,9 +112,7 @@ export class QueryCache extends Subscribable<QueryCacheListener> {
   notify(query?: Query<any, any>) {
     notifyManager.batch(() => {
       this.listeners.forEach(listener => {
-        notifyManager.schedule(() => {
-          listener(query)
-        })
+        listener(query)
       })
     })
   }

src/core/queryObserver.ts

@@ -462,47 +462,26 @@ export class QueryObserver<
   }
 
   private notify(notifyOptions: NotifyOptions): void {
-    const { currentResult, currentQuery, listeners } = this
-    const { onSuccess, onSettled, onError } = this.options
-
     notifyManager.batch(() => {
       // First trigger the configuration callbacks
       if (notifyOptions.onSuccess) {
-        if (onSuccess) {
-          notifyManager.schedule(() => {
-            onSuccess(currentResult.data!)
-          })
-        }
-        if (onSettled) {
-          notifyManager.schedule(() => {
-            onSettled(currentResult.data!, null)
-          })
-        }
+        this.options.onSuccess?.(this.currentResult.data!)
+        this.options.onSettled?.(this.currentResult.data!, null)
       } else if (notifyOptions.onError) {
-        if (onError) {
-          notifyManager.schedule(() => {
-            onError(currentResult.error!)
-          })
-        }
-        if (onSettled) {
-          notifyManager.schedule(() => {
-            onSettled(undefined, currentResult.error!)
-          })
-        }
+        this.options.onError?.(this.currentResult.error!)
+        this.options.onSettled?.(undefined, this.currentResult.error!)
       }
 
       // Then trigger the listeners
       if (notifyOptions.listeners) {
-        listeners.forEach(listener => {
-          notifyManager.schedule(() => {
-            listener(currentResult)
-          })
+        this.listeners.forEach(listener => {
+          listener(this.currentResult)
         })
       }
 
       // Then the cache listeners
       if (notifyOptions.cache) {
-        this.client.getQueryCache().notify(currentQuery)
+        this.client.getQueryCache().notify(this.currentQuery)
       }
     })
   }

src/core/subscribable.ts

@@ -1,6 +1,6 @@
 type Listener = () => void
 
-export class Subscribable<TListener = Listener> {
+export class Subscribable<TListener extends Function = Listener> {
   protected listeners: TListener[]
 
   constructor() {

src/react/tests/useQuery.test.tsx

@@ -340,6 +340,41 @@ describe('useQuery', () => {
     expect(onSuccess).toHaveBeenCalledWith('data')
   })
 
+  it('should not call onSuccess if a component has unmounted', async () => {
+    const key = queryKey()
+    const states: UseQueryResult<string>[] = []
+    const onSuccess = jest.fn()
+
+    function Page() {
+      const [show, setShow] = React.useState(true)
+
+      React.useEffect(() => {
+        setShow(false)
+      }, [setShow])
+
+      return show ? <Component /> : null
+    }
+
+    function Component() {
+      const state = useQuery(
+        key,
+        async () => {
+          await sleep(10)
+          return 'data'
+        },
+        { onSuccess }
+      )
+      states.push(state)
+      return null
+    }
+
+    renderWithClient(queryClient, <Page />)
+
+    await sleep(50)
+    expect(states.length).toBe(1)
+    expect(onSuccess).toHaveBeenCalledTimes(0)
+  })
+
   it('should call onError after a query has been fetched with an error', async () => {
     const key = queryKey()
     const states: UseQueryResult<unknown>[] = []

src/react/useBaseQuery.ts

@@ -1,6 +1,6 @@
 import React from 'react'
 
-import { useIsMounted } from './utils'
+import { notifyManager } from '../core/notifyManager'
 import { QueryObserver } from '../core/queryObserver'
 import { useQueryErrorResetBoundary } from './QueryErrorResetBoundary'
 import { useQueryClient } from './QueryClientProvider'
@@ -11,10 +11,30 @@ export function useBaseQuery<TData, TError, TQueryFnData, TQueryData>(
   Observer: typeof QueryObserver
 ) {
   const queryClient = useQueryClient()
-  const isMounted = useIsMounted()
   const errorResetBoundary = useQueryErrorResetBoundary()
   const defaultedOptions = queryClient.defaultQueryObserverOptions(options)
 
+  // Batch calls to callbacks if not in suspense mode
+  if (!defaultedOptions.suspense) {
+    if (defaultedOptions.onError) {
+      defaultedOptions.onError = notifyManager.batchCalls(
+        defaultedOptions.onError
+      )
+    }
+
+    if (defaultedOptions.onSuccess) {
+      defaultedOptions.onSuccess = notifyManager.batchCalls(
+        defaultedOptions.onSuccess
+      )
+    }
+
+    if (defaultedOptions.onSettled) {
+      defaultedOptions.onSettled = notifyManager.batchCalls(
+        defaultedOptions.onSettled
+      )
+    }
+  }
+
   // Always set stale time when using suspense
   if (defaultedOptions.suspense && !defaultedOptions.staleTime) {
     defaultedOptions.staleTime = 2000
@@ -38,12 +58,8 @@ export function useBaseQuery<TData, TError, TQueryFnData, TQueryData>(
   // Subscribe to the observer
   React.useEffect(() => {
     errorResetBoundary.clearReset()
-    return observer.subscribe(result => {
-      if (isMounted()) {
-        setCurrentResult(result)
-      }
-    })
-  }, [isMounted, observer, errorResetBoundary])
+    return observer.subscribe(notifyManager.batchCalls(setCurrentResult))
+  }, [observer, errorResetBoundary])
 
   // Handle suspense
   if (observer.options.suspense || observer.options.useErrorBoundary) {

src/react/useIsFetching.ts

@@ -1,9 +1,9 @@
 import React from 'react'
 
-import { QueryKey } from '../core'
+import { notifyManager } from '../core/notifyManager'
+import { QueryKey } from '../core/types'
 import { parseFilterArgs, QueryFilters } from '../core/utils'
 import { useQueryClient } from './QueryClientProvider'
-import { useIsMounted } from './utils'
 
 export function useIsFetching(filters?: QueryFilters): number
 export function useIsFetching(
@@ -15,7 +15,6 @@ export function useIsFetching(
   arg2?: QueryFilters
 ): number {
   const queryClient = useQueryClient()
-  const isMounted = useIsMounted()
   const [filters] = parseFilterArgs(arg1, arg2)
   const [isFetching, setIsFetching] = React.useState(
     queryClient.isFetching(filters)
@@ -28,15 +27,15 @@ export function useIsFetching(
 
   React.useEffect(
     () =>
-      queryClient.getQueryCache().subscribe(() => {
-        if (isMounted()) {
+      queryClient.getQueryCache().subscribe(
+        notifyManager.batchCalls(() => {
           const newIsFetching = queryClient.isFetching(filtersRef.current)
           if (isFetchingRef.current !== newIsFetching) {
             setIsFetching(newIsFetching)
           }
-        }
-      }),
-    [queryClient, isMounted]
+        })
+      ),
+    [queryClient]
   )
 
   return isFetching