Add support for autoDisposeTimeoutMs in useSuspenseFragment

Author: jerelmiller

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

@@ -20,11 +20,12 @@ import {
   renderHookToSnapshotStream,
   useTrackRenders,
 } from "@testing-library/react-render-stream";
-import { spyOnConsole } from "../../../testing/internal";
-import { renderHook } from "@testing-library/react";
+import { renderAsync, spyOnConsole } from "../../../testing/internal";
+import { act, renderHook, screen, waitFor } from "@testing-library/react";
 import { InvariantError } from "ts-invariant";
-import { MockedProvider, wait } from "../../../testing";
+import { MockedProvider, MockSubscriptionLink, wait } from "../../../testing";
 import { expectTypeOf } from "expect-type";
+import userEvent from "@testing-library/user-event";
 
 function createDefaultRenderStream<TData = unknown>() {
   return createRenderStream({
@@ -1566,6 +1567,220 @@ test("tears down all watches when rendering multiple records", async () => {
   expect(cache["watches"].size).toBe(0);
 });
 
+test("tears down watches after default autoDisposeTimeoutMs if component never renders again after suspending", async () => {
+  jest.useFakeTimers();
+  interface ItemFragment {
+    __typename: "Item";
+    id: number;
+    text: string;
+  }
+
+  const fragment: TypedDocumentNode<ItemFragment> = gql`
+    fragment ItemFragment on Item {
+      id
+      text
+    }
+  `;
+
+  const cache = new InMemoryCache();
+  const user = userEvent.setup({ advanceTimers: jest.advanceTimersByTime });
+  const link = new MockSubscriptionLink();
+  const client = new ApolloClient({ link, cache });
+
+  function App() {
+    const [showItem, setShowItem] = React.useState(true);
+
+    return (
+      <ApolloProvider client={client}>
+        <button onClick={() => setShowItem(false)}>Hide item</button>
+        {showItem && (
+          <Suspense fallback="Loading item...">
+            <Item />
+          </Suspense>
+        )}
+      </ApolloProvider>
+    );
+  }
+
+  function Item() {
+    const { data } = useSuspenseFragment({
+      fragment,
+      from: { __typename: "Item", id: 1 },
+    });
+
+    return <span>{data.text}</span>;
+  }
+
+  await renderAsync(<App />);
+
+  // Ensure <Greeting /> suspends immediately
+  expect(screen.getByText("Loading item...")).toBeInTheDocument();
+
+  // Hide the greeting before it finishes loading data
+  await act(() => user.click(screen.getByText("Hide item")));
+
+  expect(screen.queryByText("Loading item...")).not.toBeInTheDocument();
+
+  client.writeFragment({
+    fragment,
+    data: { __typename: "Item", id: 1, text: "Item #1" },
+  });
+
+  // clear the microtask queue
+  await act(() => Promise.resolve());
+
+  expect(cache["watches"].size).toBe(1);
+
+  jest.advanceTimersByTime(30_000);
+
+  expect(cache["watches"].size).toBe(0);
+
+  jest.useRealTimers();
+});
+
+test("tears down watches after configured autoDisposeTimeoutMs if component never renders again after suspending", async () => {
+  jest.useFakeTimers();
+  interface ItemFragment {
+    __typename: "Item";
+    id: number;
+    text: string;
+  }
+
+  const fragment: TypedDocumentNode<ItemFragment> = gql`
+    fragment ItemFragment on Item {
+      id
+      text
+    }
+  `;
+
+  const user = userEvent.setup({ advanceTimers: jest.advanceTimersByTime });
+  const link = new MockSubscriptionLink();
+  const cache = new InMemoryCache();
+  const client = new ApolloClient({
+    link,
+    cache,
+    defaultOptions: {
+      react: {
+        suspense: {
+          autoDisposeTimeoutMs: 5000,
+        },
+      },
+    },
+  });
+
+  function App() {
+    const [showItem, setShowItem] = React.useState(true);
+
+    return (
+      <ApolloProvider client={client}>
+        <button onClick={() => setShowItem(false)}>Hide item</button>
+        {showItem && (
+          <Suspense fallback="Loading item...">
+            <Item />
+          </Suspense>
+        )}
+      </ApolloProvider>
+    );
+  }
+
+  function Item() {
+    const { data } = useSuspenseFragment({
+      fragment,
+      from: { __typename: "Item", id: 1 },
+    });
+
+    return <span>{data.text}</span>;
+  }
+
+  await renderAsync(<App />);
+
+  // Ensure <Greeting /> suspends immediately
+  expect(screen.getByText("Loading item...")).toBeInTheDocument();
+
+  // Hide the greeting before it finishes loading data
+  await act(() => user.click(screen.getByText("Hide item")));
+
+  expect(screen.queryByText("Loading item...")).not.toBeInTheDocument();
+
+  client.writeFragment({
+    fragment,
+    data: { __typename: "Item", id: 1, text: "Item #1" },
+  });
+
+  // clear the microtask queue
+  await act(() => Promise.resolve());
+
+  expect(cache["watches"].size).toBe(1);
+
+  jest.advanceTimersByTime(5000);
+
+  expect(cache["watches"].size).toBe(0);
+
+  jest.useRealTimers();
+});
+
+test("cancels autoDisposeTimeoutMs if the component renders before timer finishes", async () => {
+  jest.useFakeTimers();
+  interface ItemFragment {
+    __typename: "Item";
+    id: number;
+    text: string;
+  }
+
+  const fragment: TypedDocumentNode<ItemFragment> = gql`
+    fragment ItemFragment on Item {
+      id
+      text
+    }
+  `;
+
+  const link = new MockSubscriptionLink();
+  const cache = new InMemoryCache();
+  const client = new ApolloClient({ link, cache });
+
+  function App() {
+    return (
+      <ApolloProvider client={client}>
+        <Suspense fallback="Loading item...">
+          <Item />
+        </Suspense>
+      </ApolloProvider>
+    );
+  }
+
+  function Item() {
+    const { data } = useSuspenseFragment({
+      fragment,
+      from: { __typename: "Item", id: 1 },
+    });
+
+    return <span>{data.text}</span>;
+  }
+
+  await renderAsync(<App />);
+
+  // Ensure <Greeting /> suspends immediately
+  expect(screen.getByText("Loading item...")).toBeInTheDocument();
+
+  client.writeFragment({
+    fragment,
+    data: { __typename: "Item", id: 1, text: "Item #1" },
+  });
+
+  // clear the microtask queue
+  await act(() => Promise.resolve());
+
+  await waitFor(() => {
+    expect(screen.getByText("Item #1")).toBeInTheDocument();
+  });
+
+  jest.advanceTimersByTime(30_000);
+
+  expect(cache["watches"].size).toBe(1);
+
+  jest.useRealTimers();
+});
+
 describe.skip("type tests", () => {
   test("returns TData when from is a non-null value", () => {
     const fragment: TypedDocumentNode<{ foo: string }> = gql``;

src/react/internal/cache/FragmentReference.ts

@@ -20,6 +20,7 @@ type FragmentRefPromise<TData> = PromiseWithState<TData>;
 type Listener<TData> = (promise: FragmentRefPromise<TData>) => void;
 
 interface FragmentReferenceOptions {
+  autoDisposeTimeoutMs?: number;
   onDispose?: () => void;
 }
 
@@ -36,6 +37,7 @@ export class FragmentReference<
 
   private subscription!: ObservableSubscription;
   private listeners = new Set<Listener<MaybeMasked<TData>>>();
+  private autoDisposeTimeoutId?: NodeJS.Timeout;
 
   private references = 0;
 
@@ -58,11 +60,26 @@ export class FragmentReference<
 
     const diff = this.getDiff(client, watchFragmentOptions);
 
+    // Start a timer that will automatically dispose of the query if the
+    // suspended resource does not use this fragmentRef 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
+        );
+      }
+    };
+
     this.promise =
       diff.complete ?
         createFulfilledPromise(diff.result)
       : this.createPendingPromise();
     this.subscribeToFragment();
+
+    this.promise.then(startDisposeTimer, startDisposeTimer);
   }
 
   listen(listener: Listener<MaybeMasked<TData>>) {
@@ -75,6 +92,7 @@ export class FragmentReference<
 
   retain() {
     this.references++;
+    clearTimeout(this.autoDisposeTimeoutId);
     let disposed = false;
 
     return () => {

src/react/internal/cache/SuspenseCache.ts

@@ -68,6 +68,7 @@ export class SuspenseCache {
 
     if (!ref.current) {
       ref.current = new FragmentReference(client, options, {
+        autoDisposeTimeoutMs: this.options.autoDisposeTimeoutMs,
         onDispose: () => {
           delete ref.current;
         },