dispose of queryRefs faster on rejection

Author: phryneas

src/react/internal/cache/QueryReference.ts

@@ -79,7 +79,7 @@ export interface PreloadedQueryRef<
 }
 
 interface InternalQueryReferenceOptions {
-  onDispose?: () => void;
+  onDispose: (ref: InternalQueryReference<any, any>) => void;
   autoDisposeTimeoutMs?: number;
 }
 
@@ -209,9 +209,10 @@ export class InternalQueryReference<
     this.dispose = this.dispose.bind(this);
     this.observable = observable;
 
-    if (options.onDispose) {
-      this.onDispose = options.onDispose;
-    }
+    this.onDispose = () => {
+      clearTimeout(this.autoDisposeTimeoutId);
+      options.onDispose?.(this);
+    };
 
     this.setResult();
     this.subscribeToQuery();
@@ -351,7 +352,7 @@ export class InternalQueryReference<
     return this.initiateFetch(this.observable.fetchMore<TData>(options));
   }
 
-  private dispose() {
+  public dispose() {
     this.subscription.unsubscribe();
   }
 

src/react/internal/cache/SuspenseCache.ts

@@ -43,15 +43,30 @@ export class SuspenseCache {
   >(cacheKey: CacheKey, createObservable: () => ObservableQuery<TData>) {
     const ref = this.queryRefs.lookupArray(cacheKey) as {
       current?: InternalQueryReference<TData, TStates>;
+      disposeTimeout?: ReturnType<typeof setTimeout>;
     };
 
     if (!ref.current) {
       ref.current = new InternalQueryReference(createObservable(), {
         autoDisposeTimeoutMs: this.options.autoDisposeTimeoutMs,
-        onDispose: () => {
+        onDispose: (internalRef) => {
+          if (internalRef !== ref.current) {
+            return;
+          }
           delete ref.current;
+          if (ref.disposeTimeout) {
+            clearTimeout(ref.disposeTimeout);
+            delete ref.disposeTimeout;
+          }
         },
       });
+    } else if (ref.current.promise.status === "rejected") {
+      if (ref.disposeTimeout) {
+        clearTimeout(ref.disposeTimeout);
+      }
+      ref.disposeTimeout = setTimeout(() => {
+        ref.current?.dispose();
+      }, 1000);
     }
 
     return ref.current;