fix useBackgroundQuery: dispose ref after unmount and not used

.changeset/fast-roses-kick.md

@@ -0,0 +1,5 @@
+---
+"@apollo/client": patch
+---
+
+useBackgroundQuery, handle ref disposal if unmount before used

src/react/hooks/__tests__/useBackgroundQuery.test.tsx

@@ -446,6 +446,122 @@ it("auto resubscribes when mounting useReadQuery after naturally disposed by use
   await expect(Profiler).not.toRerender({ timeout: 50 });
 });
 
+it("disposes of the queryRef when unmounting before it is used by useReadQuery", async () => {
+  const { query, mocks } = setupSimpleCase();
+  const client = new ApolloClient({
+    link: new MockLink(mocks),
+    cache: new InMemoryCache(),
+  });
+
+  const Profiler = createDefaultProfiler<SimpleCaseData>();
+
+  function App() {
+    useTrackRenders();
+    useBackgroundQuery(query);
+
+    return null;
+  }
+
+  const { unmount } = renderWithClient(<App />, { client, wrapper: Profiler });
+
+  expect(client.getObservableQueries().size).toBe(1);
+  expect(client).toHaveSuspenseCacheEntryUsing(query);
+
+  {
+    const { renderedComponents } = await Profiler.takeRender();
+
+    expect(renderedComponents).toStrictEqual([App]);
+  }
+
+  unmount();
+  await wait(0);
+
+  expect(client.getObservableQueries().size).toBe(0);
+  expect(client).not.toHaveSuspenseCacheEntryUsing(query);
+});
+
+it("disposes of old queryRefs when changing variables before the queryRef is used by useReadQuery", async () => {
+  const { query, mocks } = setupVariablesCase();
+  const client = new ApolloClient({
+    link: new MockLink(mocks),
+    cache: new InMemoryCache(),
+  });
+
+  const Profiler = createDefaultProfiler<SimpleCaseData>();
+
+  function App({ id }: { id: string }) {
+    useTrackRenders();
+    useBackgroundQuery(query, { variables: { id } });
+
+    return null;
+  }
+
+  const { rerender } = renderWithClient(<App id="1" />, {
+    client,
+    wrapper: Profiler,
+  });
+
+  expect(client.getObservableQueries().size).toBe(1);
+  expect(client).toHaveSuspenseCacheEntryUsing(query, {
+    variables: { id: "1" },
+  });
+
+  {
+    const { renderedComponents } = await Profiler.takeRender();
+
+    expect(renderedComponents).toStrictEqual([App]);
+  }
+
+  rerender(<App id="2" />);
+
+  await wait(0);
+
+  expect(client.getObservableQueries().size).toBe(1);
+  expect(client).toHaveSuspenseCacheEntryUsing(query, {
+    variables: { id: "2" },
+  });
+  expect(client).not.toHaveSuspenseCacheEntryUsing(query, {
+    variables: { id: "1" },
+  });
+});
+
+it("does not prematurely dispose of the queryRef when using strict mode", async () => {
+  const { query, mocks } = setupSimpleCase();
+  const client = new ApolloClient({
+    link: new MockLink(mocks),
+    cache: new InMemoryCache(),
+  });
+
+  const Profiler = createDefaultProfiler<SimpleCaseData>();
+
+  function App() {
+    useTrackRenders();
+    useBackgroundQuery(query);
+
+    return null;
+  }
+
+  renderWithClient(<App />, {
+    client,
+    wrapper: ({ children }) => (
+      <React.StrictMode>
+        <Profiler>{children}</Profiler>
+      </React.StrictMode>
+    ),
+  });
+
+  {
+    const { renderedComponents } = await Profiler.takeRender();
+
+    expect(renderedComponents).toStrictEqual([App]);
+  }
+
+  await wait(10);
+
+  expect(client.getObservableQueries().size).toBe(1);
+  expect(client).toHaveSuspenseCacheEntryUsing(query);
+});
+
 it("allows the client to be overridden", async () => {
   const { query } = setupSimpleCase();
 

src/react/hooks/useBackgroundQuery.ts

@@ -15,7 +15,7 @@ import {
 } from "../internal/index.js";
 import type { CacheKey, QueryReference } from "../internal/index.js";
 import type { BackgroundQueryHookOptions, NoInfer } from "../types/types.js";
-import { __use, wrapHook } from "./internal/index.js";
+import { wrapHook } from "./internal/index.js";
 import { useWatchQueryOptions } from "./useSuspenseQuery.js";
 import type { FetchMoreFunction, RefetchFunction } from "./useSuspenseQuery.js";
 import { canonicalStringify } from "../../cache/index.js";
@@ -224,8 +224,10 @@ function _useBackgroundQuery<
     ...([] as any[]).concat(queryKey),
   ];
 
-  const queryRef = suspenseCache.getQueryRef(cacheKey, () =>
-    client.watchQuery(watchQueryOptions as WatchQueryOptions<any, any>)
+  const queryRef = suspenseCache.getQueryRef(
+    cacheKey,
+    () => client.watchQuery(watchQueryOptions as WatchQueryOptions<any, any>),
+    true
   );
 
   const [wrappedQueryRef, setWrappedQueryRef] = React.useState(
@@ -261,6 +263,12 @@ function _useBackgroundQuery<
     [queryRef]
   );
 
+  React.useEffect(() => {
+    return () => {
+      queryRef.disposeOnUnmount();
+    };
+  }, [queryRef]);
+
   return [
     didFetchResult.current ? wrappedQueryRef : void 0,
     { fetchMore, refetch },

src/react/internal/cache/QueryReference.ts

@@ -154,20 +154,25 @@ export class InternalQueryReference<TData = unknown> {
   private subscription!: ObservableSubscription;
   private listeners = new Set<Listener<TData>>();
   private autoDisposeTimeoutId?: NodeJS.Timeout;
+  private readonly autoDisposeTimeoutMs?: number;
 
   private resolve: ((result: ApolloQueryResult<TData>) => void) | undefined;
   private reject: ((error: unknown) => void) | undefined;
 
   private references = 0;
+  private nbOfUse = 0;
 
   constructor(
     observable: ObservableQuery<TData, any>,
-    options: InternalQueryReferenceOptions
+    options: InternalQueryReferenceOptions,
+    trackUses: boolean = false
   ) {
     this.handleNext = this.handleNext.bind(this);
     this.handleError = this.handleError.bind(this);
     this.dispose = this.dispose.bind(this);
     this.observable = observable;
+    this.autoDisposeTimeoutMs = options.autoDisposeTimeoutMs ?? 30_000;
+    this.nbOfUse = trackUses ? 1 : 0;
 
     if (options.onDispose) {
       this.onDispose = options.onDispose;
@@ -176,23 +181,13 @@ export class InternalQueryReference<TData = unknown> {
     this.setResult();
     this.subscribeToQuery();
 
-    // Start a timer that will automatically dispose of the query if the
-    // suspended resource does not use this queryRef in the given time. This
-    // helps prevent memory leaks when a component has unmounted before the
-    // query has finished loading.
-    const startDisposeTimer = () => {
-      if (!this.references) {
-        this.autoDisposeTimeoutId = setTimeout(
-          this.dispose,
-          options.autoDisposeTimeoutMs ?? 30_000
-        );
-      }
-    };
-
     // We wait until the request has settled to ensure we don't dispose of the
     // query ref before the request finishes, otherwise we would leave the
     // promise in a pending state rendering the suspense boundary indefinitely.
-    this.promise.then(startDisposeTimer, startDisposeTimer);
+    this.promise.then(
+      this.startDisposeTimer.bind(this),
+      this.startDisposeTimer.bind(this)
+    );
   }
 
   get disposed() {
@@ -203,6 +198,19 @@ export class InternalQueryReference<TData = unknown> {
     return this.observable.options;
   }
 
+  // Start a timer that will automatically dispose of the query if the
+  // suspended resource does not use this queryRef in the given time. This
+  // helps prevent memory leaks when a component has unmounted before the
+  // query has finished loading.
+  startDisposeTimer() {
+    if (!this.references) {
+      this.autoDisposeTimeoutId = setTimeout(
+        this.dispose,
+        this.autoDisposeTimeoutMs ?? 30_000
+      );
+    }
+  }
+
   reinitialize() {
     const { observable } = this;
 
@@ -251,6 +259,19 @@ export class InternalQueryReference<TData = unknown> {
     };
   }
 
+  newUsage() {
+    this.nbOfUse++;
+    clearTimeout(this.autoDisposeTimeoutId);
+    this.startDisposeTimer();
+  }
+
+  disposeOnUnmount() {
+    this.nbOfUse--;
+    if (!this.nbOfUse && !this.references) {
+      this.dispose();
+    }
+  }
+
   didChangeOptions(watchQueryOptions: ObservedOptions) {
     return OBSERVED_CHANGED_OPTIONS.some(
       (option) =>

src/react/internal/cache/SuspenseCache.ts

@@ -30,19 +30,26 @@ export class SuspenseCache {
 
   getQueryRef<TData = any>(
     cacheKey: CacheKey,
-    createObservable: () => ObservableQuery<TData>
+    createObservable: () => ObservableQuery<TData>,
+    trackUses: boolean = false
   ) {
     const ref = this.queryRefs.lookupArray(cacheKey) as {
       current?: InternalQueryReference<TData>;
     };
 
     if (!ref.current) {
-      ref.current = new InternalQueryReference(createObservable(), {
-        autoDisposeTimeoutMs: this.options.autoDisposeTimeoutMs,
-        onDispose: () => {
-          delete ref.current;
+      ref.current = new InternalQueryReference(
+        createObservable(),
+        {
+          autoDisposeTimeoutMs: this.options.autoDisposeTimeoutMs,
+          onDispose: () => {
+            delete ref.current;
+          },
         },
-      });
+        trackUses
+      );
+    } else if (trackUses) {
+      ref.current.newUsage();
     }
 
     return ref.current;