refactor: move prevent set state in render logic to useIsFetching hook

Author: boschni

src/core/queryCache.ts

@@ -5,6 +5,7 @@ import {
   Console,
   isObject,
   Updater,
+  functionalUpdate,
 } from './utils'
 import { defaultConfigRef, ReactQueryConfigRef } from './config'
 import { Query } from './query'
@@ -268,17 +269,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)
       }
     }
 
@@ -388,13 +379,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: !config.staleTime,
+      initialData: functionalUpdate(updater, undefined),
+      ...config,
+    })
   }
 }
 

src/react/tests/useInfiniteQuery.test.tsx

@@ -182,7 +182,7 @@ describe('useInfiniteQuery', () => {
       rendered.getByText('Page 0: 0')
     })
 
-    await fireEvent.click(rendered.getByText('Load More'))
+    fireEvent.click(rendered.getByText('Load More'))
 
     await waitFor(() => rendered.getByText('Loading more...'))
 

src/react/tests/useIsFetching.test.tsx

@@ -3,6 +3,7 @@ import * as React from 'react'
 
 import { useQuery, useIsFetching } from '../index'
 import { sleep, queryKey } from './utils'
+import { useQueryCache } from '../ReactQueryCacheProvider'
 
 describe('useIsFetching', () => {
   // See https://github.com/tannerlinsley/react-query/issues/105
@@ -40,4 +41,47 @@ 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 key3 = queryKey()
+
+    const isFetchings: number[] = []
+
+    function Child() {
+      const isFetching = useIsFetching()
+      isFetchings.push(isFetching)
+      return null
+    }
+
+    function Page() {
+      const queryCache = useQueryCache()
+
+      useQuery(key1, async () => {
+        await sleep(100)
+        return 'data'
+      })
+
+      useQuery(key2, async () => {
+        await sleep(10)
+        return 'data'
+      })
+
+      if (isFetchings.length === 1) {
+        queryCache.setQueryData(key3, () => 'data')
+      }
+
+      return <Child />
+    }
+
+    render(<Page />)
+    await waitFor(() => expect(isFetchings).toEqual([2, 2, 1, 1, 1, 0]))
+    expect(spy).not.toHaveBeenCalled()
+    expect(spy.mock.calls[0]?.[0] ?? '').not.toMatch('setState')
+
+    spy.mockRestore()
+  })
 })

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,43 +1,72 @@
 import React from 'react'
 
-import { uid, isServer } from '../core/utils'
+import { isServer } from '../core/utils'
 import { QueryResultBase, BaseQueryConfig, QueryStatus } from '../core/types'
 
-export function useUid(): number {
-  const ref = React.useRef(0)
-
-  if (ref.current === null) {
-    ref.current = uid()
-  }
-
-  return ref.current
-}
-
 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[]) => (isMounted() ? callback(...args) : void 0),
+    [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 parent 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 prevent updating state while rendering.
+ */
+function scheduleMicrotask(fn: () => void): void {
+  Promise.resolve().then(fn)
 }
 
 export function handleSuspense(