refactor: move prevent set state in render logic to useIsFetching hook (TanStack#874)

Author: boschni

src/core/queryCache.ts

@@ -5,6 +5,7 @@ import {
   Console,
   isObject,
   Updater,
+  functionalUpdate,
 } from './utils'
 import { getDefaultedQueryConfig } from './config'
 import { Query } from './query'
@@ -231,17 +232,7 @@ export class QueryCache {
 
       if (!this.config.frozen) {
         this.queries[queryHash] = query
-
-        if (isServer) {
-          this.notifyGlobalListeners()
-        } else {
-          // Here, we setTimeout so as to not trigger
-          // any setState's in parent components in the
-          // middle of the render phase.
-          setTimeout(() => {
-            this.notifyGlobalListeners()
-          })
-        }
+        this.notifyGlobalListeners(query)
       }
     }
 
@@ -333,13 +324,18 @@ export class QueryCache {
     updater: Updater<TResult | undefined, TResult>,
     config?: QueryConfig<TResult, TError>
   ) {
-    let query = this.getQuery<TResult, TError>(queryKey)
+    const query = this.getQuery<TResult, TError>(queryKey)
 
-    if (!query) {
-      query = this.buildQuery<TResult, TError>(queryKey, config)
+    if (query) {
+      query.setData(updater)
+      return
     }
 
-    query.setData(updater)
+    this.buildQuery<TResult, TError>(queryKey, {
+      initialStale: typeof config?.staleTime === 'undefined',
+      initialData: functionalUpdate(updater, undefined),
+      ...config,
+    })
   }
 }
 

src/react/tests/useIsFetching.test.tsx

@@ -40,4 +40,60 @@ describe('useIsFetching', () => {
     await waitFor(() => rendered.getByText('isFetching: 1'))
     await waitFor(() => rendered.getByText('isFetching: 0'))
   })
+
+  it('should not update state while rendering', async () => {
+    const spy = jest.spyOn(console, 'error')
+
+    const key1 = queryKey()
+    const key2 = queryKey()
+
+    const isFetchings: number[] = []
+
+    function IsFetching() {
+      const isFetching = useIsFetching()
+      isFetchings.push(isFetching)
+      return null
+    }
+
+    function FirstQuery() {
+      useQuery(key1, async () => {
+        await sleep(100)
+        return 'data'
+      })
+      return null
+    }
+
+    function SecondQuery() {
+      useQuery(key2, async () => {
+        await sleep(100)
+        return 'data'
+      })
+      return null
+    }
+
+    function Page() {
+      const [renderSecond, setRenderSecond] = React.useState(false)
+
+      React.useEffect(() => {
+        setTimeout(() => {
+          setRenderSecond(true)
+        }, 10)
+      }, [])
+
+      return (
+        <>
+          <FirstQuery />
+          {renderSecond && <SecondQuery />}
+          <IsFetching />
+        </>
+      )
+    }
+
+    render(<Page />)
+    await waitFor(() => expect(isFetchings).toEqual([1, 1, 2, 1, 0]))
+    expect(spy).not.toHaveBeenCalled()
+    expect(spy.mock.calls[0]?.[0] ?? '').not.toMatch('setState')
+
+    spy.mockRestore()
+  })
 })

src/react/useBaseQuery.ts

@@ -23,7 +23,7 @@ export function useBaseQuery<TResult, TError>(
   React.useEffect(
     () =>
       observer.subscribe(() => {
-        Promise.resolve().then(rerender)
+        rerender()
       }),
     [observer, rerender]
   )

src/react/useIsFetching.ts

@@ -1,23 +1,19 @@
 import React from 'react'
 
-import { useRerenderer, useGetLatest } from './utils'
 import { useQueryCache } from './ReactQueryCacheProvider'
+import { useSafeState } from './utils'
 
 export function useIsFetching(): number {
   const queryCache = useQueryCache()
-  const rerender = useRerenderer()
-  const isFetching = queryCache.isFetching
 
-  const getIsFetching = useGetLatest(isFetching)
+  const [isFetching, setIsFetching] = useSafeState(queryCache.isFetching)
 
   React.useEffect(
     () =>
-      queryCache.subscribe(newCache => {
-        if (getIsFetching() !== newCache.isFetching) {
-          rerender()
-        }
+      queryCache.subscribe(() => {
+        setIsFetching(queryCache.isFetching)
       }),
-    [getIsFetching, queryCache, rerender]
+    [queryCache, setIsFetching]
   )
 
   return isFetching

src/react/utils.ts

@@ -1,40 +1,79 @@
 import React from 'react'
 
-import { uid, isServer } from '../core/utils'
-
-export function useUid(): number {
-  const ref = React.useRef(0)
-
-  if (ref.current === null) {
-    ref.current = uid()
-  }
-
-  return ref.current
-}
+import { isServer } from '../core/utils'
 
 export function useGetLatest<T>(obj: T): () => T {
   const ref = React.useRef<T>(obj)
   ref.current = obj
   return React.useCallback(() => ref.current, [])
 }
 
-export function useMountedCallback<T extends Function>(callback: T): T {
-  const mounted = React.useRef(false)
+function useIsMounted(): () => boolean {
+  const mountedRef = React.useRef(false)
+  const isMounted = React.useCallback(() => mountedRef.current, [])
 
   React[isServer ? 'useEffect' : 'useLayoutEffect'](() => {
-    mounted.current = true
+    mountedRef.current = true
     return () => {
-      mounted.current = false
+      mountedRef.current = false
     }
   }, [])
 
+  return isMounted
+}
+
+export function useMountedCallback<T extends Function>(callback: T): T {
+  const isMounted = useIsMounted()
   return (React.useCallback(
-    (...args: any[]) => (mounted.current ? callback(...args) : void 0),
-    [callback]
+    (...args: any[]) => {
+      if (isMounted()) {
+        return callback(...args)
+      }
+    },
+    [callback, isMounted]
   ) as any) as T
 }
 
+/**
+ * This hook is a safe useState version which schedules state updates in microtasks
+ * to prevent updating a component state while React is rendering different components
+ * or when the component is not mounted anymore.
+ */
+export function useSafeState<S>(
+  initialState: S | (() => S)
+): [S, React.Dispatch<React.SetStateAction<S>>] {
+  const isMounted = useIsMounted()
+  const [state, setState] = React.useState(initialState)
+
+  const safeSetState = React.useCallback(
+    (value: React.SetStateAction<S>) => {
+      scheduleMicrotask(() => {
+        if (isMounted()) {
+          setState(value)
+        }
+      })
+    },
+    [isMounted]
+  )
+
+  return [state, safeSetState]
+}
+
 export function useRerenderer() {
-  const rerender = useMountedCallback(React.useState<unknown>()[1])
-  return React.useCallback(() => rerender({}), [rerender])
+  const [, setState] = useSafeState({})
+  return React.useCallback(() => setState({}), [setState])
+}
+
+/**
+ * Schedules a microtask.
+ * This can be useful to schedule state updates after rendering.
+ */
+function scheduleMicrotask(callback: () => void): void {
+  Promise.resolve()
+    .then(callback)
+    .catch(error =>
+      setTimeout(() => {
+        throw error
+      })
+    )
 }