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

@@ -111,6 +111,7 @@ function createErrorProfiler<TData = unknown>() {
     },
   });
 }
+
 function createDefaultProfiler<TData = unknown>() {
   return createProfiler({
     initialSnapshot: {
@@ -446,6 +447,169 @@ 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("disposes of the queryRef when unmounting before it is used by useReadQuery even if it has been rerendered", async () => {
+  const { query, mocks } = setupSimpleCase();
+  const client = new ApolloClient({
+    link: new MockLink(mocks),
+    cache: new InMemoryCache(),
+  });
+  const user = userEvent.setup();
+
+  const Profiler = createDefaultProfiler<SimpleCaseData>();
+
+  function App() {
+    useTrackRenders();
+    useBackgroundQuery(query);
+
+    const [a, setA] = React.useState(0);
+
+    return (
+      <>
+        <button onClick={() => setA(a + 1)}>Increment</button>
+      </>
+    );
+  }
+
+  const { unmount } = renderWithClient(<App />, {
+    client,
+    wrapper: Profiler,
+  });
+  const button = screen.getByText("Increment");
+
+  await act(() => user.click(button));
+
+  {
+    const { renderedComponents } = await Profiler.takeRender();
+
+    expect(renderedComponents).toStrictEqual([App]);
+  }
+
+  expect(client.getObservableQueries().size).toBe(1);
+  expect(client).toHaveSuspenseCacheEntryUsing(query);
+
+  await wait(0);
+
+  unmount();
+  await wait(0);
+  expect(client.getObservableQueries().size).toBe(0);
+});
+
 it("allows the client to be overridden", async () => {
   const { query } = setupSimpleCase();
 
@@ -985,6 +1149,7 @@ it("works with startTransition to change variables", async () => {
       completed: boolean;
     };
   }
+
   const user = userEvent.setup();
 
   const query: TypedDocumentNode<Data, Variables> = gql`
@@ -4189,6 +4354,7 @@ describe("refetch", () => {
         completed: boolean;
       };
     }
+
     const user = userEvent.setup();
 
     const query: TypedDocumentNode<Data, Variables> = gql`
@@ -4437,6 +4603,7 @@ describe("refetch", () => {
         completed: boolean;
       };
     }
+
     const user = userEvent.setup();
 
     const query: TypedDocumentNode<Data, Variables> = gql`
@@ -5046,9 +5213,11 @@ describe("fetchMore", () => {
       name: string;
       completed: boolean;
     }
+
     interface Data {
       todos: Todo[];
     }
+
     const user = userEvent.setup();
 
     const query: TypedDocumentNode<Data, Variables> = gql`

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";
@@ -261,6 +261,13 @@ function _useBackgroundQuery<
     [queryRef]
   );
 
+  React.useEffect(() => {
+    queryRef.newUsage();
+    return () => {
+      queryRef.disposeOnUnmount();
+    };
+  }, [queryRef]);
+
   return [
     didFetchResult.current ? wrappedQueryRef : void 0,
     { fetchMore, refetch },

src/react/internal/cache/QueryReference.ts

@@ -154,11 +154,13 @@ 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 numberOfUse = 0;
 
   constructor(
     observable: ObservableQuery<TData, any>,
@@ -168,6 +170,7 @@ export class InternalQueryReference<TData = unknown> {
     this.handleError = this.handleError.bind(this);
     this.dispose = this.dispose.bind(this);
     this.observable = observable;
+    this.autoDisposeTimeoutMs = options.autoDisposeTimeoutMs ?? 30_000;
 
     if (options.onDispose) {
       this.onDispose = options.onDispose;
@@ -176,23 +179,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 +196,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 +257,24 @@ export class InternalQueryReference<TData = unknown> {
     };
   }
 
+  newUsage() {
+    this.numberOfUse++;
+    if (!this.references) {
+      clearTimeout(this.autoDisposeTimeoutId);
+      this.startDisposeTimer();
+    }
+  }
+
+  disposeOnUnmount() {
+    this.numberOfUse--;
+    // Wait before fully disposing in case the app is running in strict mode.
+    setTimeout(() => {
+      if (!this.numberOfUse && !this.references) {
+        this.dispose();
+      }
+    });
+  }
+
   didChangeOptions(watchQueryOptions: ObservedOptions) {
     return OBSERVED_CHANGED_OPTIONS.some(
       (option) =>