Repository: tannerlinsley/react-query
Pull request #933 (2020-08-27 16:01:45)

feat: determine staleness locally instead of globally
This MR allows users to define different stale times for different situations. Instead of marking a whole query as stale, the staleness of a query is determined when consuming the query. The rationale for this can be found here: https://github.com/tannerlinsley/react-query/discussions/770#discussioncomment-42352 . This also resolves some of the confusion about what happens if different hooks define different stale times.

vercel[bot]: This pull request is being automatically deployed with Vercel ([learn more](https://vercel.link/github-learn-more)).  
To see the status of your deployment, click below or on the icon next to each commit.


tannerlinsley: This is genius.
tannerlinsley: Do you know if this is something we should highlight in the comparison docs? Or does everyone else have this too...?
tannerlinsley: This good to merge?
boschni: Added a comparison to the docs and integrated the hydration changes. Should be good to go now!

Patch:
diff --git a/docs/src/pages/docs/comparison.md b/docs/src/pages/docs/comparison.md
index 05670a2d8..68c7985a0 100644
--- a/docs/src/pages/docs/comparison.md
+++ b/docs/src/pages/docs/comparison.md
@@ -19,6 +19,7 @@ Feature/Capability Key:
 | Supported Query Keys                         | JSON                                   | JSON                       | GraphQL Query                       |
 | Query Key Change Detection                   | Deep Compare (Serialization)           | Referential Equality (===) | Deep Compare (Serialization)        |
 | Query Data Memoization Level                 | Query + Structural Sharing             | Query                      | Query + Entity + Structural Sharing |
+| Stale While Revalidate                       | Server-Side + Client-Side              | Server-Side                | None                                |
 | Bundle Size                                  | [![][bp-react-query]][bpl-react-query] | [![][bp-swr]][bpl-swr]     | [![][bp-apollo]][bpl-apollo]        |
 | Queries                                      | ✅                                     | ✅                         | ✅                                  |
 | Caching                                      | ✅                                     | ✅                         | ✅                                  |
diff --git a/src/core/query.ts b/src/core/query.ts
index 12a01f00d..eaad23a61 100644
--- a/src/core/query.ts
+++ b/src/core/query.ts
@@ -44,7 +44,6 @@ export interface QueryState<TResult, TError> {
   isFetchingMore: IsFetchingMoreValue
   isIdle: boolean
   isLoading: boolean
-  isStale: boolean
   isSuccess: boolean
   status: QueryStatus
   throwInErrorBoundary?: boolean
@@ -66,7 +65,6 @@ export interface RefetchOptions {
 
 export enum ActionType {
   Failed = 'Failed',
-  MarkStale = 'MarkStale',
   Fetch = 'Fetch',
   Success = 'Success',
   Error = 'Error',
@@ -76,10 +74,6 @@ interface FailedAction {
   type: ActionType.Failed
 }
 
-interface MarkStaleAction {
-  type: ActionType.MarkStale
-}
-
 interface FetchAction {
   type: ActionType.Fetch
   isFetchingMore?: IsFetchingMoreValue
@@ -89,7 +83,6 @@ interface SuccessAction<TResult> {
   type: ActionType.Success
   data: TResult | undefined
   canFetchMore?: boolean
-  isStale: boolean
 }
 
 interface ErrorAction<TError> {
@@ -101,7 +94,6 @@ export type Action<TResult, TError> =
   | ErrorAction<TError>
   | FailedAction
   | FetchAction
-  | MarkStaleAction
   | SuccessAction<TResult>
 
 // CLASS
@@ -115,14 +107,11 @@ export class Query<TResult, TError> {
 
   private queryCache: QueryCache
   private promise?: Promise<TResult | undefined>
-  private cacheTimeout?: number
-  private staleTimeout?: number
+  private gcTimeout?: number
   private cancelFetch?: () => void
   private continueFetch?: () => void
   private isTransportCancelable?: boolean
   private notifyGlobalListeners: (query: Query<TResult, TError>) => void
-  private enableStaleTimeout: boolean
-  private enableGarbageCollectionTimeout: boolean
 
   constructor(init: QueryInitConfig<TResult, TError>) {
     this.config = init.config
@@ -132,23 +121,7 @@ export class Query<TResult, TError> {
     this.notifyGlobalListeners = init.notifyGlobalListeners
     this.observers = []
     this.state = getDefaultState(init.config)
-    this.enableStaleTimeout = false
-    this.enableGarbageCollectionTimeout = false
-  }
-
-  activateStaleTimeout(): void {
-    this.enableStaleTimeout = true
-    this.rescheduleStaleTimeout()
-  }
-
-  activateGarbageCollectionTimeout(): void {
-    this.enableGarbageCollectionTimeout = true
-    this.rescheduleGarbageCollection()
-  }
-
-  activateTimeouts(): void {
-    this.activateStaleTimeout()
-    this.activateGarbageCollectionTimeout()
+    this.scheduleGc()
   }
 
   updateConfig(config: QueryConfig<TResult, TError>): void {
@@ -161,61 +134,18 @@ export class Query<TResult, TError> {
     this.notifyGlobalListeners(this)
   }
 
-  private rescheduleStaleTimeout(): void {
+  private scheduleGc(): void {
     if (isServer) {
       return
     }
 
-    this.clearStaleTimeout()
-
-    if (
-      !this.enableStaleTimeout ||
-      this.state.isStale ||
-      this.state.status !== QueryStatus.Success ||
-      this.config.staleTime === Infinity
-    ) {
-      return
-    }
-
-    const staleTime = this.config.staleTime || 0
-    let timeout = staleTime
-    if (this.state.updatedAt) {
-      const timeElapsed = Date.now() - this.state.updatedAt
-      const timeUntilStale = staleTime - timeElapsed
-      timeout = Math.max(timeUntilStale, 0)
-    }
-
-    this.staleTimeout = setTimeout(() => {
-      this.invalidate()
-    }, timeout)
-  }
+    this.clearGcTimeout()
 
-  invalidate(): void {
-    this.clearStaleTimeout()
-
-    if (this.state.isStale) {
+    if (this.config.cacheTime === Infinity || this.observers.length > 0) {
       return
     }
 
-    this.dispatch({ type: ActionType.MarkStale })
-  }
-
-  private rescheduleGarbageCollection(): void {
-    if (isServer) {
-      return
-    }
-
-    this.clearCacheTimeout()
-
-    if (
-      !this.enableGarbageCollectionTimeout ||
-      this.config.cacheTime === Infinity ||
-      this.observers.length > 0
-    ) {
-      return
-    }
-
-    this.cacheTimeout = setTimeout(() => {
+    this.gcTimeout = setTimeout(() => {
       this.clear()
     }, this.config.cacheTime)
   }
@@ -241,21 +171,14 @@ export class Query<TResult, TError> {
 
   private clearTimersObservers(): void {
     this.observers.forEach(observer => {
-      observer.clearRefetchInterval()
+      observer.clearTimers()
     })
   }
 
-  private clearStaleTimeout() {
-    if (this.staleTimeout) {
-      clearTimeout(this.staleTimeout)
-      this.staleTimeout = undefined
-    }
-  }
-
-  private clearCacheTimeout() {
-    if (this.cacheTimeout) {
-      clearTimeout(this.cacheTimeout)
-      this.cacheTimeout = undefined
+  private clearGcTimeout() {
+    if (this.gcTimeout) {
+      clearTimeout(this.gcTimeout)
+      this.gcTimeout = undefined
     }
   }
 
@@ -275,8 +198,6 @@ export class Query<TResult, TError> {
       data = prevData
     }
 
-    const isStale = this.config.staleTime === 0
-
     // Try to determine if more data can be fetched
     const canFetchMore = hasMorePages(this.config, data)
 
@@ -284,16 +205,12 @@ export class Query<TResult, TError> {
     this.dispatch({
       type: ActionType.Success,
       data,
-      isStale,
       canFetchMore,
     })
-
-    this.rescheduleStaleTimeout()
   }
 
   clear(): void {
-    this.clearStaleTimeout()
-    this.clearCacheTimeout()
+    this.clearGcTimeout()
     this.clearTimersObservers()
     this.cancel()
     delete this.queryCache.queries[this.queryHash]
@@ -304,12 +221,23 @@ export class Query<TResult, TError> {
     return this.observers.some(observer => observer.config.enabled)
   }
 
+  isStale(): boolean {
+    return this.observers.some(observer => observer.isStale())
+  }
+
+  isStaleByTime(staleTime = 0): boolean {
+    return (
+      !this.state.isSuccess || this.state.updatedAt + staleTime <= Date.now()
+    )
+  }
+
   onWindowFocus(): void {
     if (
-      this.state.isStale &&
       this.observers.some(
         observer =>
-          observer.config.enabled && observer.config.refetchOnWindowFocus
+          observer.isStale() &&
+          observer.config.enabled &&
+          observer.config.refetchOnWindowFocus
       )
     ) {
       this.fetch()
@@ -319,10 +247,11 @@ export class Query<TResult, TError> {
 
   onOnline(): void {
     if (
-      this.state.isStale &&
       this.observers.some(
         observer =>
-          observer.config.enabled && observer.config.refetchOnReconnect
+          observer.isStale() &&
+          observer.config.enabled &&
+          observer.config.refetchOnReconnect
       )
     ) {
       this.fetch()
@@ -348,7 +277,7 @@ export class Query<TResult, TError> {
     this.observers.push(observer)
 
     // Stop the query from being garbage collected
-    this.clearCacheTimeout()
+    this.clearGcTimeout()
   }
 
   unsubscribeObserver(observer: QueryObserver<TResult, TError>): void {
@@ -362,7 +291,7 @@ export class Query<TResult, TError> {
       }
     }
 
-    this.rescheduleGarbageCollection()
+    this.scheduleGc()
   }
 
   private async tryFetchData<T>(
@@ -639,12 +568,6 @@ function getDefaultState<TResult, TError>(
 
   const hasInitialData = typeof initialData !== 'undefined'
 
-  const isStale =
-    !config.enabled ||
-    (typeof config.initialStale === 'function'
-      ? config.initialStale()
-      : config.initialStale ?? !hasInitialData)
-
   const initialStatus = hasInitialData
     ? QueryStatus.Success
     : config.enabled
@@ -658,9 +581,8 @@ function getDefaultState<TResult, TError>(
     isFetching: initialStatus === QueryStatus.Loading,
     isFetchingMore: false,
     failureCount: 0,
-    isStale,
     data: initialData,
-    updatedAt: hasInitialData ? Date.now() : 0,
+    updatedAt: Date.now(),
     canFetchMore: hasMorePages(config, initialData),
   }
 }
@@ -675,11 +597,6 @@ export function queryReducer<TResult, TError>(
         ...state,
         failureCount: state.failureCount + 1,
       }
-    case ActionType.MarkStale:
-      return {
-        ...state,
-        isStale: true,
-      }
     case ActionType.Fetch:
       const status =
         typeof state.data !== 'undefined'
@@ -698,7 +615,6 @@ export function queryReducer<TResult, TError>(
         ...getStatusProps(QueryStatus.Success),
         data: action.data,
         error: null,
-        isStale: action.isStale,
         isFetched: true,
         isFetching: false,
         isFetchingMore: false,
@@ -714,7 +630,6 @@ export function queryReducer<TResult, TError>(
         isFetched: true,
         isFetching: false,
         isFetchingMore: false,
-        isStale: true,
         failureCount: state.failureCount + 1,
         throwInErrorBoundary: true,
       }
diff --git a/src/core/queryCache.ts b/src/core/queryCache.ts
index d6ee1eac6..3c8c63dde 100644
--- a/src/core/queryCache.ts
+++ b/src/core/queryCache.ts
@@ -189,7 +189,7 @@ export class QueryCache {
             }
           }
 
-          return query.invalidate()
+          return undefined
         })
       )
     } catch (err) {
@@ -305,7 +305,7 @@ export class QueryCache {
     let query
     try {
       query = this.buildQuery<TResult, TError>(queryKey, configWithoutRetry)
-      if (options?.force || query.state.isStale) {
+      if (options?.force || query.isStaleByTime(config.staleTime)) {
         await query.fetch()
       }
       return query.state.data
@@ -314,14 +314,6 @@ export class QueryCache {
         throw error
       }
       return
-    } finally {
-      if (query) {
-        // When prefetching, no observer is tied to the query,
-        // so to avoid immediate garbage collection of the still
-        // empty query, we wait with activating timeouts until
-        // the prefetch is done
-        query.activateTimeouts()
-      }
     }
   }
 
@@ -337,13 +329,11 @@ export class QueryCache {
       return
     }
 
-    const newQuery = this.buildQuery<TResult, TError>(queryKey, {
+    this.buildQuery<TResult, TError>(queryKey, {
       initialStale: typeof config?.staleTime === 'undefined',
       initialData: functionalUpdate(updater, undefined),
       ...config,
     })
-
-    newQuery.activateTimeouts()
   }
 }
 
diff --git a/src/core/queryObserver.ts b/src/core/queryObserver.ts
index bed80161a..505308ed0 100644
--- a/src/core/queryObserver.ts
+++ b/src/core/queryObserver.ts
@@ -19,6 +19,7 @@ export class QueryObserver<TResult, TError> {
   private currentResult!: QueryResult<TResult, TError>
   private previousResult?: QueryResult<TResult, TError>
   private updateListener?: UpdateListener<TResult, TError>
+  private staleTimeoutId?: number
   private refetchIntervalId?: number
   private started?: boolean
 
@@ -39,14 +40,14 @@ export class QueryObserver<TResult, TError> {
     this.updateListener = listener
     this.currentQuery.subscribeObserver(this)
     this.optionalFetch()
-    this.updateRefetchInterval()
+    this.updateTimers()
     return this.unsubscribe.bind(this)
   }
 
   unsubscribe(): void {
     this.started = false
     this.updateListener = undefined
-    this.clearRefetchInterval()
+    this.clearTimers()
     this.currentQuery.unsubscribeObserver(this)
   }
 
@@ -64,7 +65,7 @@ export class QueryObserver<TResult, TError> {
     // If we subscribed to a new query, optionally fetch and update refetch
     if (updated) {
       this.optionalFetch()
-      this.updateRefetchInterval()
+      this.updateTimers()
       return
     }
 
@@ -73,6 +74,14 @@ export class QueryObserver<TResult, TError> {
       this.optionalFetch()
     }
 
+    // Update stale interval if needed
+    if (
+      config.enabled !== prevConfig.enabled ||
+      config.staleTime !== prevConfig.staleTime
+    ) {
+      this.updateStaleTimeout()
+    }
+
     // Update refetch interval if needed
     if (
       config.enabled !== prevConfig.enabled ||
@@ -84,6 +93,10 @@ export class QueryObserver<TResult, TError> {
     }
   }
 
+  isStale(): boolean {
+    return this.currentResult.isStale
+  }
+
   getCurrentResult(): QueryResult<TResult, TError> {
     return this.currentResult
   }
@@ -125,6 +138,37 @@ export class QueryObserver<TResult, TError> {
     }
   }
 
+  private updateIsStale(): void {
+    const isStale = this.currentQuery.isStaleByTime(this.config.staleTime)
+    if (isStale !== this.currentResult.isStale) {
+      this.currentResult = this.createResult()
+      this.updateListener?.(this.currentResult)
+    }
+  }
+
+  private updateStaleTimeout(): void {
+    if (isServer) {
+      return
+    }
+
+    this.clearStaleTimeout()
+
+    const staleTime = this.config.staleTime || 0
+    const { isStale, updatedAt } = this.currentResult
+
+    if (isStale || staleTime === Infinity) {
+      return
+    }
+
+    const timeElapsed = Date.now() - updatedAt
+    const timeUntilStale = staleTime - timeElapsed
+    const timeout = Math.max(timeUntilStale, 0)
+
+    this.staleTimeoutId = setTimeout(() => {
+      this.updateIsStale()
+    }, timeout)
+  }
+
   private updateRefetchInterval(): void {
     if (isServer) {
       return
@@ -148,7 +192,24 @@ export class QueryObserver<TResult, TError> {
     }, this.config.refetchInterval)
   }
 
-  clearRefetchInterval(): void {
+  updateTimers(): void {
+    this.updateStaleTimeout()
+    this.updateRefetchInterval()
+  }
+
+  clearTimers(): void {
+    this.clearStaleTimeout()
+    this.clearRefetchInterval()
+  }
+
+  private clearStaleTimeout(): void {
+    if (this.staleTimeoutId) {
+      clearInterval(this.staleTimeoutId)
+      this.staleTimeoutId = undefined
+    }
+  }
+
+  private clearRefetchInterval(): void {
     if (this.refetchIntervalId) {
       clearInterval(this.refetchIntervalId)
       this.refetchIntervalId = undefined
@@ -156,7 +217,7 @@ export class QueryObserver<TResult, TError> {
   }
 
   private createResult(): QueryResult<TResult, TError> {
-    const { currentQuery, previousResult, config } = this
+    const { currentResult, currentQuery, previousResult, config } = this
 
     const {
       canFetchMore,
@@ -166,7 +227,6 @@ export class QueryObserver<TResult, TError> {
       isFetching,
       isFetchingMore,
       isLoading,
-      isStale,
     } = currentQuery.state
 
     let { data, status, updatedAt } = currentQuery.state
@@ -178,6 +238,22 @@ export class QueryObserver<TResult, TError> {
       status = previousResult.status
     }
 
+    let isStale = false
+
+    // When the query has not been fetched yet and this is the initial render,
+    // determine the staleness based on the initialStale or existence of initial data.
+    if (!currentResult && !currentQuery.state.isFetched) {
+      if (typeof config.initialStale === 'function') {
+        isStale = config.initialStale()
+      } else if (typeof config.initialStale === 'boolean') {
+        isStale = config.initialStale
+      } else {
+        isStale = typeof currentQuery.state.data === 'undefined'
+      }
+    } else {
+      isStale = currentQuery.isStaleByTime(config.staleTime)
+    }
+
     return {
       ...getStatusProps(status),
       canFetchMore,
@@ -211,8 +287,6 @@ export class QueryObserver<TResult, TError> {
       return false
     }
 
-    newQuery.activateTimeouts()
-
     this.previousResult = this.currentResult
     this.currentQuery = newQuery
     this.currentResult = this.createResult()
@@ -236,11 +310,11 @@ export class QueryObserver<TResult, TError> {
     if (action.type === 'Success' && isSuccess) {
       this.config.onSuccess?.(data!)
       this.config.onSettled?.(data!, null)
-      this.updateRefetchInterval()
+      this.updateTimers()
     } else if (action.type === 'Error' && isError) {
       this.config.onError?.(error!)
       this.config.onSettled?.(undefined, error!)
-      this.updateRefetchInterval()
+      this.updateTimers()
     }
 
     this.updateListener?.(this.currentResult)
diff --git a/src/hydration/hydration.ts b/src/hydration/hydration.ts
index d73cf6c08..a3c41ba33 100644
--- a/src/hydration/hydration.ts
+++ b/src/hydration/hydration.ts
@@ -1,10 +1,9 @@
-import { DEFAULT_STALE_TIME, DEFAULT_CACHE_TIME } from '../core/config'
+import { DEFAULT_CACHE_TIME } from '../core/config'
 
 import type { Query, QueryCache, QueryKey, QueryConfig } from 'react-query'
 
 export interface DehydratedQueryConfig {
   queryKey: QueryKey
-  staleTime?: number
   cacheTime?: number
   initialData?: unknown
 }
@@ -41,9 +40,6 @@ function dehydrateQuery<TResult, TError = unknown>(
   // in the html-payload, but not consume it on the initial render.
   // We still schedule stale and garbage collection right away, which means
   // we need to specifically include staleTime and cacheTime in dehydration.
-  if (query.config.staleTime !== DEFAULT_STALE_TIME) {
-    dehydratedQuery.config.staleTime = query.config.staleTime
-  }
   if (query.config.cacheTime !== DEFAULT_CACHE_TIME) {
     dehydratedQuery.config.cacheTime = query.config.cacheTime
   }
@@ -93,6 +89,5 @@ export function hydrate<TResult>(
 
     const query = queryCache.buildQuery(queryKey, queryConfig)
     query.state.updatedAt = dehydratedQuery.updatedAt
-    query.activateGarbageCollectionTimeout()
   }
 }


Test patch:
diff --git a/src/core/tests/queryCache.test.tsx b/src/core/tests/queryCache.test.tsx
index c84f3b2f2..1a7df97ad 100644
--- a/src/core/tests/queryCache.test.tsx
+++ b/src/core/tests/queryCache.test.tsx
@@ -71,6 +71,7 @@ describe('queryCache', () => {
       fetchFn,
       {
         initialData: 'initial',
+        staleTime: 100,
       },
       {
         throwOnError: true,
@@ -79,6 +80,33 @@ describe('queryCache', () => {
     expect(first).toBe('og')
   })
 
+  test('prefetchQuery should only fetch if the data is older then the given stale time', async () => {
+    const key = queryKey()
+
+    let count = 0
+    const fetchFn = () => ++count
+
+    defaultQueryCache.setQueryData(key, count)
+    const first = await defaultQueryCache.prefetchQuery(key, fetchFn, {
+      staleTime: 100,
+    })
+    await sleep(11)
+    const second = await defaultQueryCache.prefetchQuery(key, fetchFn, {
+      staleTime: 10,
+    })
+    const third = await defaultQueryCache.prefetchQuery(key, fetchFn, {
+      staleTime: 10,
+    })
+    await sleep(11)
+    const fourth = await defaultQueryCache.prefetchQuery(key, fetchFn, {
+      staleTime: 10,
+    })
+    expect(first).toBe(0)
+    expect(second).toBe(1)
+    expect(third).toBe(1)
+    expect(fourth).toBe(2)
+  })
+
   test('prefetchQuery should throw error when throwOnError is true', async () => {
     const consoleMock = mockConsoleError()
 
@@ -200,43 +228,6 @@ describe('queryCache', () => {
     expect(defaultQueryCache.getQuery(key)).toBeFalsy()
   })
 
-  test('setQueryData should schedule stale timeout, if staleTime is set', async () => {
-    const key = queryKey()
-
-    defaultQueryCache.setQueryData(key, 'test data', { staleTime: 10 })
-    // @ts-expect-error
-    expect(defaultQueryCache.getQuery(key)!.staleTimeout).not.toBeUndefined()
-  })
-
-  test('setQueryData should not schedule stale timeout by default', async () => {
-    const key = queryKey()
-
-    defaultQueryCache.setQueryData(key, 'test data')
-    // @ts-expect-error
-    expect(defaultQueryCache.getQuery(key)!.staleTimeout).toBeUndefined()
-  })
-
-  test('setQueryData should not schedule stale timeout, if staleTime is set to `Infinity`', async () => {
-    const key = queryKey()
-
-    defaultQueryCache.setQueryData(key, 'test data', { staleTime: Infinity })
-    // @ts-expect-error
-    expect(defaultQueryCache.getQuery(key)!.staleTimeout).toBeUndefined()
-  })
-
-  test('setQueryData schedules stale timeouts appropriately', async () => {
-    const key = queryKey()
-
-    defaultQueryCache.setQueryData(key, 'test data', { staleTime: 100 })
-
-    expect(defaultQueryCache.getQuery(key)!.state.data).toEqual('test data')
-    expect(defaultQueryCache.getQuery(key)!.state.isStale).toEqual(false)
-
-    await new Promise(resolve => setTimeout(resolve, 100))
-
-    expect(defaultQueryCache.getQuery(key)!.state.isStale).toEqual(true)
-  })
-
   test('setQueryData updater function works as expected', () => {
     const key = queryKey()
 
@@ -266,20 +257,6 @@ describe('queryCache', () => {
     expect(data).toEqual(['data1', 'data2'])
   })
 
-  test('stale timeout dispatch is not called if query is no longer in the query cache', async () => {
-    const key = queryKey()
-
-    const fetchData = () => Promise.resolve('data')
-    await defaultQueryCache.prefetchQuery(key, fetchData, {
-      staleTime: 100,
-    })
-    const query = defaultQueryCache.getQuery(key)
-    expect(query!.state.isStale).toBe(false)
-    defaultQueryCache.removeQueries(key)
-    await sleep(50)
-    expect(query!.state.isStale).toBe(false)
-  })
-
   test('query interval is cleared when unsubscribed to a refetchInterval query', async () => {
     const key = queryKey()
 
diff --git a/src/hydration/tests/hydration.test.tsx b/src/hydration/tests/hydration.test.tsx
index 2b5d44b2d..64e475bce 100644
--- a/src/hydration/tests/hydration.test.tsx
+++ b/src/hydration/tests/hydration.test.tsx
@@ -2,8 +2,6 @@ import { sleep } from '../../react/tests/utils'
 import { makeQueryCache } from '../..'
 import { dehydrate, hydrate } from '../hydration'
 
-import type { QueryCache } from '../..'
-
 const fetchData: <TResult>(
   value: TResult,
   ms?: number
@@ -12,11 +10,6 @@ const fetchData: <TResult>(
   return value
 }
 
-// It is up to the queryObserver to schedule staleness, this simulates that
-function simulateQueryObserver(queryCache: QueryCache, queryKey: string): void {
-  queryCache.getQuery(queryKey)?.activateStaleTimeout()
-}
-
 describe('dehydration and rehydration', () => {
   test('should work with serializeable values', async () => {
     const queryCache = makeQueryCache()
@@ -49,106 +42,27 @@ describe('dehydration and rehydration', () => {
     })
 
     const fetchDataAfterHydration = jest.fn()
-    await hydrationQueryCache.prefetchQuery('string', fetchDataAfterHydration)
-    await hydrationQueryCache.prefetchQuery('number', fetchDataAfterHydration)
-    await hydrationQueryCache.prefetchQuery('boolean', fetchDataAfterHydration)
-    await hydrationQueryCache.prefetchQuery('null', fetchDataAfterHydration)
-    await hydrationQueryCache.prefetchQuery('array', fetchDataAfterHydration)
-    await hydrationQueryCache.prefetchQuery('nested', fetchDataAfterHydration)
-    expect(fetchDataAfterHydration).toHaveBeenCalledTimes(0)
-
-    queryCache.clear({ notify: false })
-    hydrationQueryCache.clear({ notify: false })
-  })
-
-  test('should not schedule staleness unless observed', async () => {
-    const queryCache = makeQueryCache()
-    await queryCache.prefetchQuery('string', () => fetchData('string'))
-    const dehydrated = dehydrate(queryCache)
-    const stringified = JSON.stringify(dehydrated)
-
-    // ---
-
-    const parsed = JSON.parse(stringified)
-    const hydrationQueryCache = makeQueryCache()
-    hydrate(hydrationQueryCache, parsed)
-    expect(hydrationQueryCache.getQuery('string')?.state.data).toBe('string')
-    expect(hydrationQueryCache.getQuery('string')?.state.isStale).toBe(false)
-    await sleep(10)
-    expect(hydrationQueryCache.getQuery('string')?.state.isStale).toBe(false)
-
-    queryCache.clear({ notify: false })
-    hydrationQueryCache.clear({ notify: false })
-  })
-
-  test('should default to scheduling staleness immediately', async () => {
-    const queryCache = makeQueryCache()
-    await queryCache.prefetchQuery('string', () => fetchData('string'))
-    const dehydrated = dehydrate(queryCache)
-    const stringified = JSON.stringify(dehydrated)
-
-    // ---
-
-    const parsed = JSON.parse(stringified)
-    const hydrationQueryCache = makeQueryCache()
-    hydrate(hydrationQueryCache, parsed)
-    simulateQueryObserver(hydrationQueryCache, 'string')
-    expect(hydrationQueryCache.getQuery('string')?.state.data).toBe('string')
-    expect(hydrationQueryCache.getQuery('string')?.state.isStale).toBe(false)
-    await sleep(10)
-    expect(hydrationQueryCache.getQuery('string')?.state.isStale).toBe(true)
-
-    queryCache.clear({ notify: false })
-    hydrationQueryCache.clear({ notify: false })
-  })
-
-  test('should respect staleTime, measured from when data was fetched', async () => {
-    const queryCache = makeQueryCache()
-    await queryCache.prefetchQuery('string', () => fetchData('string'), {
-      staleTime: 50,
+    await hydrationQueryCache.prefetchQuery('string', fetchDataAfterHydration, {
+      staleTime: 100,
     })
-    const dehydrated = dehydrate(queryCache)
-    const stringified = JSON.stringify(dehydrated)
-
-    await sleep(20)
-
-    // ---
-
-    const parsed = JSON.parse(stringified)
-    const hydrationQueryCache = makeQueryCache()
-    hydrate(hydrationQueryCache, parsed)
-    simulateQueryObserver(hydrationQueryCache, 'string')
-    expect(hydrationQueryCache.getQuery('string')?.state.data).toBe('string')
-    expect(hydrationQueryCache.getQuery('string')?.state.isStale).toBe(false)
-    await sleep(10)
-    expect(hydrationQueryCache.getQuery('string')?.state.isStale).toBe(false)
-    await sleep(30)
-    expect(hydrationQueryCache.getQuery('string')?.state.isStale).toBe(true)
-
-    queryCache.clear({ notify: false })
-    hydrationQueryCache.clear({ notify: false })
-  })
-
-  test('should schedule stale immediately if enough time has elapsed between dehydrate and hydrate', async () => {
-    const queryCache = makeQueryCache()
-    await queryCache.prefetchQuery('string', () => fetchData('string'), {
-      staleTime: 20,
+    await hydrationQueryCache.prefetchQuery('number', fetchDataAfterHydration, {
+      staleTime: 100,
     })
-    const dehydrated = dehydrate(queryCache)
-    const stringified = JSON.stringify(dehydrated)
-
-    await sleep(30)
-
-    // ---
-
-    const parsed = JSON.parse(stringified)
-    const hydrationQueryCache = makeQueryCache()
-    hydrate(hydrationQueryCache, parsed)
-    simulateQueryObserver(hydrationQueryCache, 'string')
-    expect(hydrationQueryCache.getQuery('string')?.state.data).toBe('string')
-    expect(hydrationQueryCache.getQuery('string')?.state.isStale).toBe(false)
-    await sleep(10)
-    expect(hydrationQueryCache.getQuery('string')?.state.isStale).toBe(true)
+    await hydrationQueryCache.prefetchQuery(
+      'boolean',
+      fetchDataAfterHydration,
+      { staleTime: 100 }
+    )
+    await hydrationQueryCache.prefetchQuery('null', fetchDataAfterHydration, {
+      staleTime: 100,
+    })
+    await hydrationQueryCache.prefetchQuery('array', fetchDataAfterHydration, {
+      staleTime: 100,
+    })
+    await hydrationQueryCache.prefetchQuery('nested', fetchDataAfterHydration, {
+      staleTime: 100,
+    })
+    expect(fetchDataAfterHydration).toHaveBeenCalledTimes(0)
 
     queryCache.clear({ notify: false })
     hydrationQueryCache.clear({ notify: false })
@@ -201,7 +115,8 @@ describe('dehydration and rehydration', () => {
     const fetchDataAfterHydration = jest.fn()
     await hydrationQueryCache.prefetchQuery(
       ['string', { key: ['string'], key2: 0 }],
-      fetchDataAfterHydration
+      fetchDataAfterHydration,
+      { staleTime: 10 }
     )
     expect(fetchDataAfterHydration).toHaveBeenCalledTimes(0)
 
@@ -223,7 +138,6 @@ describe('dehydration and rehydration', () => {
         (dehydratedQuery?.config?.queryKey as Array<string>)[0] === 'string'
     )
     expect(dehydratedQuery).toBeTruthy()
-    expect(dehydratedQuery?.config.staleTime).toBe(undefined)
     expect(dehydratedQuery?.config.cacheTime).toBe(undefined)
   })
 
diff --git a/src/hydration/tests/react.test.tsx b/src/hydration/tests/react.test.tsx
index 39b725e27..21945693c 100644
--- a/src/hydration/tests/react.test.tsx
+++ b/src/hydration/tests/react.test.tsx
@@ -1,11 +1,13 @@
 import React from 'react'
-import { render, waitFor } from '@testing-library/react'
+import { render } from '@testing-library/react'
+
 import {
   ReactQueryCacheProvider as OriginalCacheProvider,
   makeQueryCache,
   useQuery,
 } from '../..'
 import { dehydrate, useHydrate, ReactQueryCacheProvider } from '../'
+import { waitForMs } from '../../react/tests/utils'
 
 describe('React hydration', () => {
   const fetchData: (value: string) => Promise<string> = value =>
@@ -37,6 +39,7 @@ describe('React hydration', () => {
 
       const rendered = render(<Page />)
 
+      await waitForMs(10)
       rendered.getByText('string')
     })
 
@@ -60,12 +63,8 @@ describe('React hydration', () => {
         </OriginalCacheProvider>
       )
 
+      await waitForMs(10)
       rendered.getByText('string')
-      expect(clientQueryCache.getQuery('string')?.state.isStale).toBe(false)
-      await waitFor(() =>
-        expect(clientQueryCache.getQuery('string')?.state.isStale).toBe(true)
-      )
-
       clientQueryCache.clear({ notify: false })
     })
   })
@@ -93,11 +92,8 @@ describe('React hydration', () => {
         </ReactQueryCacheProvider>
       )
 
+      await waitForMs(10)
       rendered.getByText('string')
-      expect(clientQueryCache.getQuery('string')?.state.isStale).toBe(false)
-      await waitFor(() =>
-        expect(clientQueryCache.getQuery('string')?.state.isStale).toBe(true)
-      )
 
       const intermediateCache = makeQueryCache()
       await intermediateCache.prefetchQuery('string', () =>
@@ -119,17 +115,10 @@ describe('React hydration', () => {
 
       // Existing query data should not be overwritten,
       // so this should still be the original data
+      await waitForMs(10)
       rendered.getByText('string')
       // But new query data should be available immediately
       rendered.getByText('added string')
-      expect(clientQueryCache.getQuery('added string')?.state.isStale).toBe(
-        false
-      )
-      await waitFor(() =>
-        expect(clientQueryCache.getQuery('added string')?.state.isStale).toBe(
-          true
-        )
-      )
 
       clientQueryCache.clear({ notify: false })
     })
@@ -156,11 +145,8 @@ describe('React hydration', () => {
         </ReactQueryCacheProvider>
       )
 
+      await waitForMs(10)
       rendered.getByText('string')
-      expect(clientQueryCache.getQuery('string')?.state.isStale).toBe(false)
-      await waitFor(() =>
-        expect(clientQueryCache.getQuery('string')?.state.isStale).toBe(true)
-      )
 
       const newClientQueryCache = makeQueryCache()
 
@@ -173,11 +159,8 @@ describe('React hydration', () => {
         </ReactQueryCacheProvider>
       )
 
+      await waitForMs(10)
       rendered.getByText('string')
-      expect(newClientQueryCache.getQuery('string')?.state.isStale).toBe(false)
-      await waitFor(() =>
-        expect(newClientQueryCache.getQuery('string')?.state.isStale).toBe(true)
-      )
 
       clientQueryCache.clear({ notify: false })
       newClientQueryCache.clear({ notify: false })
diff --git a/src/react/tests/useQuery.test.tsx b/src/react/tests/useQuery.test.tsx
index 6e4966eb5..2307b4efe 100644
--- a/src/react/tests/useQuery.test.tsx
+++ b/src/react/tests/useQuery.test.tsx
@@ -461,6 +461,88 @@ describe('useQuery', () => {
     })
   })
 
+  it('should be able to set different stale times for a query', async () => {
+    const key = queryKey()
+    const states1: QueryResult<string>[] = []
+    const states2: QueryResult<string>[] = []
+
+    await queryCache.prefetchQuery(key, () => 'prefetch')
+
+    await sleep(10)
+
+    function FirstComponent() {
+      const state = useQuery(key, () => 'one', {
+        staleTime: 100,
+      })
+      states1.push(state)
+      return null
+    }
+
+    function SecondComponent() {
+      const state = useQuery(key, () => 'two', {
+        staleTime: 5,
+      })
+      states2.push(state)
+      return null
+    }
+
+    function Page() {
+      return (
+        <>
+          <FirstComponent />
+          <SecondComponent />
+        </>
+      )
+    }
+
+    render(<Page />)
+
+    await waitFor(() => expect(states1.length).toBe(4))
+    await waitFor(() => expect(states2.length).toBe(4))
+
+    // First render
+    expect(states1[0]).toMatchObject({
+      data: 'prefetch',
+      isStale: false,
+    })
+    // Second useQuery started fetching
+    expect(states1[1]).toMatchObject({
+      data: 'prefetch',
+      isStale: false,
+    })
+    // Second useQuery data came in
+    expect(states1[2]).toMatchObject({
+      data: 'two',
+      isStale: false,
+    })
+    // Data became stale after 100ms
+    expect(states1[3]).toMatchObject({
+      data: 'two',
+      isStale: true,
+    })
+
+    // First render, data is stale
+    expect(states2[0]).toMatchObject({
+      data: 'prefetch',
+      isStale: true,
+    })
+    // Second useQuery started fetching
+    expect(states2[1]).toMatchObject({
+      data: 'prefetch',
+      isStale: true,
+    })
+    // Second useQuery data came in
+    expect(states2[2]).toMatchObject({
+      data: 'two',
+      isStale: false,
+    })
+    // Data became stale after 5ms
+    expect(states2[3]).toMatchObject({
+      data: 'two',
+      isStale: true,
+    })
+  })
+
   // See https://github.com/tannerlinsley/react-query/issues/137
   it('should not override initial data in dependent queries', async () => {
     const key1 = queryKey()
@@ -767,6 +849,7 @@ describe('useQuery', () => {
   // See https://github.com/tannerlinsley/react-query/issues/195
   it('should refetch if stale after a prefetch', async () => {
     const key = queryKey()
+    const states: QueryResult<string>[] = []
 
     const queryFn = jest.fn()
     queryFn.mockImplementation(() => 'data')
@@ -781,13 +864,14 @@ describe('useQuery', () => {
     await sleep(11)
 
     function Page() {
-      useQuery(key, queryFn)
+      const state = useQuery(key, queryFn)
+      states.push(state)
       return null
     }
 
     render(<Page />)
 
-    await act(() => sleep(0))
+    await waitFor(() => expect(states.length).toBe(3))
 
     expect(prefetchQueryFn).toHaveBeenCalledTimes(1)
     expect(queryFn).toHaveBeenCalledTimes(1)
diff --git a/src/react/tests/utils.tsx b/src/react/tests/utils.tsx
index 4e00ff1a3..d1b273cf9 100644
--- a/src/react/tests/utils.tsx
+++ b/src/react/tests/utils.tsx
@@ -1,3 +1,5 @@
+import { waitFor } from '@testing-library/react'
+
 let queryKeyCount = 0
 
 export function mockVisibilityState(value: string) {
@@ -31,6 +33,15 @@ export function sleep(timeout: number): Promise<void> {
   })
 }
 
+export function waitForMs(ms: number) {
+  const end = Date.now() + ms
+  return waitFor(() => {
+    if (Date.now() < end) {
+      throw new Error('Time not elapsed yet')
+    }
+  })
+}
+
 /**
  * Checks that `T` is of type `U`.
  */