fix useBackgroundQuery: dispose ref after unmount and not used (#11696)

Co-authored-by: Jerel Miller <jerelmiller@gmail.com>
apollographql · Mar 21, 2024 · 466ef82 · 466ef82
1 parent 8e5c66b
commit 466ef82
Show file tree

Hide file tree

Showing 9 changed files with 209 additions and 2 deletions.
diff --git a/.api-reports/api-report-react.md b/.api-reports/api-report-react.md
@@ -953,6 +953,8 @@ class InternalQueryReference<TData = unknown> {
     // (undocumented)
     retain(): () => void;
     // (undocumented)
+    softRetain(): () => void;
+    // (undocumented)
     get watchQueryOptions(): WatchQueryOptions<OperationVariables, TData>;
 }
 

diff --git a/.api-reports/api-report-react_hooks.md b/.api-reports/api-report-react_hooks.md
@@ -895,6 +895,8 @@ class InternalQueryReference<TData = unknown> {
     // (undocumented)
     retain(): () => void;
     // (undocumented)
+    softRetain(): () => void;
+    // (undocumented)
     get watchQueryOptions(): WatchQueryOptions<OperationVariables, TData>;
 }
 

diff --git a/.api-reports/api-report-react_internal.md b/.api-reports/api-report-react_internal.md
@@ -885,6 +885,8 @@ export class InternalQueryReference<TData = unknown> {
     // (undocumented)
     retain(): () => void;
     // (undocumented)
+    softRetain(): () => void;
+    // (undocumented)
     get watchQueryOptions(): WatchQueryOptions<OperationVariables, TData>;
 }
 

diff --git a/.api-reports/api-report.md b/.api-reports/api-report.md
@@ -1328,6 +1328,8 @@ class InternalQueryReference<TData = unknown> {
     // (undocumented)
     retain(): () => void;
     // (undocumented)
+    softRetain(): () => void;
+    // (undocumented)
     get watchQueryOptions(): WatchQueryOptions<OperationVariables, TData>;
 }
 

diff --git a/.changeset/fast-roses-kick.md b/.changeset/fast-roses-kick.md
@@ -0,0 +1,5 @@
+---
+"@apollo/client": patch
+---
+
+Immediately dispose of the `queryRef` if `useBackgroundQuery` unmounts before the auto dispose timeout kicks in.
diff --git a/.size-limits.json b/.size-limits.json
@@ -1,4 +1,4 @@
 {
-  "dist/apollo-client.min.cjs": 39277,
+  "dist/apollo-client.min.cjs": 39319,
   "import { ApolloClient, InMemoryCache, HttpLink } from \"dist/index.js\" (production)": 32630
 }
diff --git a/src/react/hooks/__tests__/useBackgroundQuery.test.tsx b/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`
@@ -5055,9 +5222,11 @@ describe("fetchMore", () => {
       name: string;
       completed: boolean;
     }
+
     interface Data {
       todos: Todo[];
     }
+
     const user = userEvent.setup();
 
     const query: TypedDocumentNode<Data, Variables> = gql`

diff --git a/src/react/hooks/useBackgroundQuery.ts b/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,8 @@ function _useBackgroundQuery<
     [queryRef]
   );
 
+  React.useEffect(() => queryRef.softRetain(), [queryRef]);
+
   return [
     didFetchResult.current ? wrappedQueryRef : void 0,
     { fetchMore, refetch },

diff --git a/src/react/internal/cache/QueryReference.ts b/src/react/internal/cache/QueryReference.ts
@@ -158,6 +158,7 @@ export class InternalQueryReference<TData = unknown> {
   private reject: ((error: unknown) => void) | undefined;
 
   private references = 0;
+  private softReferences = 0;
 
   constructor(
     observable: ObservableQuery<TData, any>,
@@ -250,6 +251,28 @@ export class InternalQueryReference<TData = unknown> {
     };
   }
 
+  softRetain() {
+    this.softReferences++;
+    let disposed = false;
+
+    return () => {
+      // Tracking if this has already been called helps ensure that
+      // multiple calls to this function won't decrement the reference
+      // counter more than it should. Subsequent calls just result in a noop.
+      if (disposed) {
+        return;
+      }
+
+      disposed = true;
+      this.softReferences--;
+      setTimeout(() => {
+        if (!this.softReferences && !this.references) {
+          this.dispose();
+        }
+      });
+    };
+  }
+
   didChangeOptions(watchQueryOptions: ObservedOptions) {
     return OBSERVED_CHANGED_OPTIONS.some(
       (option) =>