Ensure that PreloadedQueryRef instances are unsubscribed when garbage collected

.api-reports/api-report-utilities_internal_ponyfills.api.md

@@ -0,0 +1,12 @@
+## API Report File for "@apollo/client"
+
+> Do not edit this file. It is a report generated by [API Extractor](https://api-extractor.com/).
+
+```ts
+
+// @public (undocumented)
+export const FinalizationRegistry: FinalizationRegistryConstructor;
+
+// (No @packageDocumentation comment for this package)
+
+```

.changeset/clever-students-guess.md

@@ -0,0 +1,5 @@
+---
+"@apollo/client": patch
+---
+
+Ensure that `PreloadedQueryRef` instances are unsubscribed when garbage collected

package.json

@@ -66,6 +66,10 @@
       "production": "./src/utilities/internal/index.production.ts",
       "default": "./src/utilities/internal/index.ts"
     },
+    "./utilities/internal/ponyfills": {
+      "react-native": "./src/utilities/internal/ponyfills/index.react-native.ts",
+      "default": "./src/utilities/internal/ponyfills/index.ts"
+    },
     "./utilities/internal/globals": "./src/utilities/internal/globals/index.ts",
     "./utilities/subscriptions/relay": "./src/utilities/subscriptions/relay/index.ts",
     "./utilities/invariant": {

src/__tests__/__snapshots__/exports.ts.snap

@@ -480,6 +480,12 @@ Array [
 ]
 `;
 
+exports[`exports of public entry points @apollo/client/utilities/internal/ponyfills 1`] = `
+Array [
+  "FinalizationRegistry",
+]
+`;
+
 exports[`exports of public entry points @apollo/client/utilities/invariant 1`] = `
 Array [
   "ApolloErrorMessageHandler",

src/__tests__/exports.ts

@@ -42,6 +42,7 @@ import * as utilities from "@apollo/client/utilities";
 import * as utilitiesEnvironment from "@apollo/client/utilities/environment";
 import * as utilitiesInternal from "@apollo/client/utilities/internal";
 import * as utilitiesInternalGlobals from "@apollo/client/utilities/internal/globals";
+import * as utilitiesInternalPonyfills from "@apollo/client/utilities/internal/ponyfills";
 import * as utilitiesInvariant from "@apollo/client/utilities/invariant";
 import * as v4_migration from "@apollo/client/v4-migration";
 
@@ -104,6 +105,10 @@ describe("exports of public entry points", () => {
   check("@apollo/client/utilities", utilities);
   check("@apollo/client/utilities/internal", utilitiesInternal);
   check("@apollo/client/utilities/internal/globals", utilitiesInternalGlobals);
+  check(
+    "@apollo/client/utilities/internal/ponyfills",
+    utilitiesInternalPonyfills
+  );
   check("@apollo/client/utilities/invariant", utilitiesInvariant);
   check("@apollo/client/utilities/environment", utilitiesEnvironment);
   check("@apollo/client/v4-migration", v4_migration);

src/react/query-preloader/__tests__/createQueryPreloader.test.tsx

@@ -1,4 +1,4 @@
-import { act, screen } from "@testing-library/react";
+import { act, screen, waitFor } from "@testing-library/react";
 import {
   createRenderStream,
   disableActEnvironment,
@@ -2007,6 +2007,102 @@ test("does not mask results by default", async () => {
   }
 });
 
+describe("PreloadedQueryRef` disposal", () => {
+  test("when the `PreloadedQueryRef` is disposed of, the ObservableQuery is unsubscribed", async () => {
+    const { query, mocks } = setupVariablesCase();
+    const client = new ApolloClient({
+      cache: new InMemoryCache(),
+      link: new MockLink(mocks),
+    });
+    const preloadQuery = createQueryPreloader(client);
+
+    let queryRef: PreloadedQueryRef | null = preloadQuery(query, {
+      variables: { id: "1" },
+    });
+    const internalQueryRef = unwrapQueryRef(queryRef)!;
+
+    expect(internalQueryRef.observable.hasObservers()).toBe(true);
+    expect(internalQueryRef["softReferences"]).toBe(1);
+    queryRef = null;
+
+    await waitFor(() => {
+      global.gc!();
+      expect(internalQueryRef.observable.hasObservers()).toBe(false);
+    });
+    expect(internalQueryRef["softReferences"]).toBe(0);
+  });
+
+  test("when the `PreloadedQueryRef` is disposed of, while the initial request is still ongoing the ObservableQuery stays subscribed to", async () => {
+    const { query } = setupVariablesCase();
+    const link = new MockSubscriptionLink();
+    const client = new ApolloClient({
+      cache: new InMemoryCache(),
+      link,
+    });
+    const preloadQuery = createQueryPreloader(client);
+
+    let queryRef: PreloadedQueryRef | null = preloadQuery(query, {
+      variables: { id: "1" },
+    });
+    const internalQueryRef = unwrapQueryRef(queryRef)!;
+
+    expect(internalQueryRef.observable.hasObservers()).toBe(true);
+    expect(internalQueryRef["softReferences"]).toBe(1);
+    queryRef = null;
+
+    await expect(
+      waitFor(() => {
+        global.gc!();
+        expect(internalQueryRef.observable.hasObservers()).toBe(false);
+      })
+    ).rejects.toThrow();
+    expect(internalQueryRef["softReferences"]).toBe(1);
+
+    link.simulateResult(
+      {
+        result: {
+          data: {
+            character: { __typename: "Character", id: "1", name },
+          },
+        },
+      },
+      true
+    );
+
+    await waitFor(() => {
+      global.gc!();
+      expect(internalQueryRef.observable.hasObservers()).toBe(false);
+    });
+    expect(internalQueryRef["softReferences"]).toBe(0);
+  });
+
+  test("when retained by a component, the soft retain lets go", async () => {
+    const { query, mocks } = setupVariablesCase();
+    const client = new ApolloClient({
+      cache: new InMemoryCache(),
+      link: new MockLink(mocks),
+    });
+    const preloadQuery = createQueryPreloader(client);
+
+    const queryRef = preloadQuery(query, {
+      variables: { id: "1" },
+    });
+    const internalQueryRef = unwrapQueryRef(queryRef)!;
+
+    expect(internalQueryRef["softReferences"]).toBe(1);
+
+    using _disabledAct = disableActEnvironment();
+    const { renderStream } = await renderDefaultTestApp<VariablesCaseData>({
+      client,
+      queryRef,
+    });
+    await renderStream.takeRender();
+    await renderStream.takeRender();
+
+    expect(internalQueryRef["softReferences"]).toBe(0);
+  });
+});
+
 describe.skip("type tests", () => {
   const client = new ApolloClient({
     cache: new InMemoryCache(),

src/react/query-preloader/createQueryPreloader.ts

@@ -19,6 +19,7 @@ import type {
   NoInfer,
   VariablesOption,
 } from "@apollo/client/utilities/internal";
+import { FinalizationRegistry } from "@apollo/client/utilities/internal/ponyfills";
 
 import { wrapHook } from "../hooks/internal/index.js";
 
@@ -197,10 +198,12 @@ const _createQueryPreloader: typeof createQueryPreloader = (client) => {
       }
     );
 
-    return wrapQueryRef(queryRef) as unknown as PreloadedQueryRef<
+    const wrapped = wrapQueryRef(queryRef) as unknown as PreloadedQueryRef<
       TData,
       TVariables
     >;
+    softRetainWhileReferenced(wrapped, queryRef);
+    return wrapped;
   }
 
   return Object.assign(preloadQuery, {
@@ -212,3 +215,54 @@ const _createQueryPreloader: typeof createQueryPreloader = (client) => {
     },
   });
 };
+
+/**
+ * Soft-retains the underlying `InternalQueryReference` while the `PreloadedQueryRef`
+ * is still reachable.
+ * When the `PreloadedQueryRef` is garbage collected, the soft retain is
+ * disposed of, but only after the initial query has finished loading.
+ * Once the `InternalQueryReference` is properly retained, the check for garbage
+ * collection is unregistered and the soft retain is disposed of immediately.
+ */
+// this is an individual function to avoid closing over any values more than necessary
+function softRetainWhileReferenced(
+  wrapped: PreloadedQueryRef<any, any, any>,
+  queryRef: InternalQueryReference
+) {
+  const { softDispose, delayedSoftDispose } = getCleanup(queryRef);
+  registry.register(wrapped, delayedSoftDispose, queryRef);
+  // This will unregister the cleanup from the finalization registry when
+  // the queryRef is properly retained.
+  // This is mostly done to keep the FinalizationRegistry from holding too many
+  // cleanup functions, as our React Native polyfill has to iterate all of them regularly.
+  queryRef.retain = unregisterOnRetain(queryRef.retain, softDispose);
+}
+
+// this is an individual function to avoid closing over any values more than necessary
+function unregisterOnRetain(
+  originalRetain: InternalQueryReference["retain"],
+  softDispose: () => void
+) {
+  return function (
+    this: InternalQueryReference,
+    ...args: Parameters<InternalQueryReference["retain"]>
+  ) {
+    registry.unregister(this);
+    const dispose = originalRetain.apply(this, args);
+    softDispose();
+    return dispose;
+  };
+}
+
+// this is an individual function to avoid closing over any values more than necessary
+function getCleanup(queryRef: InternalQueryReference) {
+  const softDispose = queryRef.softRetain();
+  const initialPromise = queryRef.promise;
+
+  return {
+    softDispose,
+    delayedSoftDispose: () => initialPromise.finally(softDispose),
+  };
+}
+
+const registry = new FinalizationRegistry<() => void>((cleanup) => cleanup());

src/utilities/internal/ponyfills/FinalizationRegistry.ts

@@ -0,0 +1,55 @@
+import { invariant } from "@apollo/client/utilities/invariant";
+
+interface Entry<T> {
+  targetRef: WeakRef<WeakKey>;
+  value: T;
+}
+
+/**
+ * An approximation of `FinalizationRegistry` based on `WeakRef`.
+ * Checks every 500ms if registered values have been garbage collected.
+ */
+export const FinalizationRegistry: typeof globalThis.FinalizationRegistry = class FinalizationRegistry<
+  T,
+> {
+  private intervalLength = 500;
+  private callback: (value: T) => void;
+  private references = new Set<Entry<T>>();
+  private unregisterTokens = new WeakMap<WeakKey, Entry<T>>();
+  private interval: ReturnType<typeof setInterval> | null = null;
+  constructor(callback: (value: T) => void) {
+    this.callback = callback;
+    this.handler = this.handler.bind(this);
+  }
+  handler() {
+    if (this.references.size === 0) {
+      clearInterval(this.interval!);
+      this.interval = null;
+      return;
+    }
+    this.references.forEach((entry) => {
+      if (entry.targetRef.deref() === undefined) {
+        this.references.delete(entry);
+        this.callback(entry.value);
+      }
+    });
+  }
+  register(target: WeakKey, value: T, unregisterToken?: WeakKey): void {
+    const entry = { targetRef: new WeakRef(target), value };
+    this.references.add(entry);
+    if (unregisterToken) {
+      // some simplifications here as it's an internal polyfill
+      // we don't allow the same unregisterToken to be reused
+      invariant(!this.unregisterTokens.has(unregisterToken));
+      this.unregisterTokens.set(unregisterToken, entry);
+    }
+    if (!this.interval) {
+      this.interval = setInterval(this.handler, this.intervalLength);
+    }
+  }
+  unregister(unregisterToken: WeakKey): boolean {
+    this.references.delete(this.unregisterTokens.get(unregisterToken)!);
+    return this.unregisterTokens.delete(unregisterToken);
+  }
+  [Symbol.toStringTag] = "FinalizationRegistry" as const;
+};

src/utilities/internal/ponyfills/__tests__/FinalizationRegistry.test.ts

@@ -0,0 +1,84 @@
+import { waitFor } from "@testing-library/react";
+
+// eslint-disable-next-line
+import { FinalizationRegistry } from "../FinalizationRegistry.js";
+
+test("register", async () => {
+  const cleanedUp: number[] = [];
+  const registry = new FinalizationRegistry<number>((value) => {
+    cleanedUp.push(value);
+  });
+  // @ts-ignore we want to speed this up a bit
+  registry["intervalLength"] = 1;
+
+  let obj1: {} | null = {};
+  let obj2: {} | null = {};
+  let obj3: {} | null = {};
+
+  registry.register(obj1, 1);
+  registry.register(obj2, 2);
+  registry.register(obj3, 3);
+
+  expect(cleanedUp).toStrictEqual([]);
+
+  obj1 = null;
+  await waitFor(() => {
+    global.gc!();
+    expect(cleanedUp).toStrictEqual([1]);
+  });
+
+  obj3 = null;
+  await waitFor(() => {
+    global.gc!();
+    expect(cleanedUp).toStrictEqual([1, 3]);
+  });
+
+  obj2 = null;
+  await waitFor(() => {
+    global.gc!();
+    expect(cleanedUp).toStrictEqual([1, 3, 2]);
+  });
+});
+
+test("unregister", async () => {
+  const cleanedUp: number[] = [];
+  const registry = new FinalizationRegistry<number>((value) => {
+    cleanedUp.push(value);
+  });
+  // @ts-ignore we want to speed this up a bit
+  registry["intervalLength"] = 1;
+
+  let obj1: {} | null = {};
+  const token1 = {};
+  let obj2: {} | null = {};
+  const token2 = {};
+  let obj3: {} | null = {};
+  const token3 = {};
+
+  registry.register(obj1, 1, token1);
+  registry.register(obj2, 2, token2);
+  registry.register(obj3, 3, token3);
+
+  expect(cleanedUp).toStrictEqual([]);
+
+  obj1 = null;
+  await waitFor(() => {
+    global.gc!();
+    expect(cleanedUp).toStrictEqual([1]);
+  });
+
+  registry.unregister(token3);
+  obj3 = null;
+  await expect(
+    waitFor(() => {
+      global.gc!();
+      expect(cleanedUp).toStrictEqual([1, 3]);
+    })
+  ).rejects.toThrow();
+
+  obj2 = null;
+  await waitFor(() => {
+    global.gc!();
+    expect(cleanedUp).toStrictEqual([1, 2]);
+  });
+});

src/utilities/internal/ponyfills/index.react-native.ts

@@ -0,0 +1 @@
+export { FinalizationRegistry } from "./FinalizationRegistry.js";

src/utilities/internal/ponyfills/index.ts

@@ -0,0 +1,2 @@
+const F = FinalizationRegistry;
+export { F as FinalizationRegistry };