Refactor query timeouts

src/core/query.ts

@@ -120,6 +120,7 @@ export class Query<TResult, TError> {
   private continueFetch?: () => void
   private isTransportCancelable?: boolean
   private notifyGlobalListeners: (query: Query<TResult, TError>) => void
+  private enableTimeouts: boolean
 
   constructor(init: QueryInitConfig<TResult, TError>) {
     this.config = init.config
@@ -129,16 +130,13 @@ export class Query<TResult, TError> {
     this.notifyGlobalListeners = init.notifyGlobalListeners
     this.observers = []
     this.state = getDefaultState(init.config)
+    this.enableTimeouts = false
+  }
 
-    // If the query started with data, schedule
-    // a stale timeout
-    if (!isServer && this.state.data) {
-      this.scheduleStaleTimeout()
-
-      // Schedule for garbage collection in case
-      // nothing subscribes to this query
-      this.scheduleCacheTimeout()
-    }
+  activateTimeouts(): void {
+    this.enableTimeouts = true
+    this.rescheduleStaleTimeout()
+    this.rescheduleGarbageCollection()
   }
 
   updateConfig(config: QueryConfig<TResult, TError>): void {
@@ -151,20 +149,33 @@ export class Query<TResult, TError> {
     this.notifyGlobalListeners(this)
   }
 
-  private scheduleStaleTimeout(): void {
+  private rescheduleStaleTimeout(): void {
     if (isServer) {
       return
     }
 
     this.clearStaleTimeout()
 
-    if (this.state.isStale || this.config.staleTime === Infinity) {
+    if (
+      !this.enableTimeouts ||
+      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()
-    }, this.config.staleTime)
+    }, timeout)
   }
 
   invalidate(): void {
@@ -177,14 +188,18 @@ export class Query<TResult, TError> {
     this.dispatch({ type: ActionType.MarkStale })
   }
 
-  private scheduleCacheTimeout(): void {
+  private rescheduleGarbageCollection(): void {
     if (isServer) {
       return
     }
 
     this.clearCacheTimeout()
 
-    if (this.config.cacheTime === Infinity) {
+    if (
+      !this.enableTimeouts ||
+      this.config.cacheTime === Infinity ||
+      this.observers.length > 0
+    ) {
       return
     }
 
@@ -261,10 +276,7 @@ export class Query<TResult, TError> {
       canFetchMore,
     })
 
-    if (!isStale) {
-      // Schedule a fresh invalidation!
-      this.scheduleStaleTimeout()
-    }
+    this.rescheduleStaleTimeout()
   }
 
   clear(): void {
@@ -318,10 +330,9 @@ export class Query<TResult, TError> {
       if (this.isTransportCancelable) {
         this.cancel()
       }
-
-      // Schedule garbage collection
-      this.scheduleCacheTimeout()
     }
+
+    this.rescheduleGarbageCollection()
   }
 
   private async tryFetchData<T>(

src/core/queryCache.ts

@@ -300,11 +300,9 @@ export class QueryCache {
     // https://github.com/tannerlinsley/react-query/issues/652
     const configWithoutRetry = { retry: false, ...config }
 
+    let query
     try {
-      const query = this.buildQuery<TResult, TError>(
-        queryKey,
-        configWithoutRetry
-      )
+      query = this.buildQuery<TResult, TError>(queryKey, configWithoutRetry)
       if (options?.force || query.state.isStale) {
         await query.fetch()
       }
@@ -314,6 +312,14 @@ 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()
+      }
     }
   }
 
@@ -329,11 +335,13 @@ export class QueryCache {
       return
     }
 
-    this.buildQuery<TResult, TError>(queryKey, {
+    const newQuery = this.buildQuery<TResult, TError>(queryKey, {
       initialStale: typeof config?.staleTime === 'undefined',
       initialData: functionalUpdate(updater, undefined),
       ...config,
     })
+
+    newQuery.activateTimeouts()
   }
 }
 

src/core/queryObserver.ts

@@ -211,6 +211,8 @@ export class QueryObserver<TResult, TError> {
       return false
     }
 
+    newQuery.activateTimeouts()
+
     this.previousResult = this.currentResult
     this.currentQuery = newQuery
     this.currentResult = this.createResult()

src/react/tests/useQuery.test.tsx

@@ -800,7 +800,10 @@ describe('useQuery', () => {
     queryFn.mockImplementation(() => 'data')
 
     const prefetchQueryFn = jest.fn()
-    prefetchQueryFn.mockImplementation(() => 'not yet...')
+    prefetchQueryFn.mockImplementation(async () => {
+      await sleep(10)
+      return 'not yet...'
+    })
 
     await queryCache.prefetchQuery(key, prefetchQueryFn, {
       staleTime: 1000,
@@ -809,7 +812,9 @@ describe('useQuery', () => {
     await sleep(0)
 
     function Page() {
-      useQuery(key, queryFn)
+      useQuery(key, queryFn, {
+        staleTime: 1000,
+      })
       return null
     }