Use dispose uncommitted implementation

__tests__/feature-detection-no-finalization-registry.test.ts

@@ -0,0 +1,7 @@
+import "./kill-finalization-registry";
+import useDisposeUncommittedDefault, { useDisposeUncommitted } from "../src/index";
+
+test("The timers impl is exported", () => {
+  expect(useDisposeUncommittedDefault.name).toBe("useDisposeUncommittedTimerBased");
+  expect(useDisposeUncommitted.name).toBe("useDisposeUncommittedTimerBased");
+});

__tests__/feature-detection-supports-finalization-registry.test.ts

@@ -0,0 +1,6 @@
+import useDisposeUncommittedDefault, { useDisposeUncommitted } from "../src/index";
+
+test("The finalization impl is exported", () => {
+  expect(useDisposeUncommittedDefault.name).toBe("useDisposeUncommittedFinalizationRegistry");
+  expect(useDisposeUncommitted.name).toBe("useDisposeUncommittedFinalizationRegistry");
+});

__tests__/finalization-registry-based-impl.test.tsx

@@ -0,0 +1,119 @@
+import { cleanup, render } from "@testing-library/react";
+// @ts-expect-error lib have no @types package, and the api is super simple
+import gc from "expose-gc/function";
+import React, { useLayoutEffect, useRef } from "react";
+import {
+  createUseDisposeUncommitted,
+  userCallbacksMap,
+} from "../src/finalization-registry-based-impl";
+import { FinalizationRegistry } from "../src/FinalizationRegistryWrapper";
+import { assertNotNullable } from "../src/utils";
+
+interface CounterGroups {
+  instances: number;
+  mounts: number;
+  unmounts: number;
+  disposes: number;
+  renders: number;
+}
+
+describe("finalization registry based implementation", () => {
+  assertNotNullable(FinalizationRegistry);
+
+  const useDisposeUncommitted = createUseDisposeUncommitted(FinalizationRegistry);
+
+  test("Dispose is called for uncommitted component and not called for committed component", async () => {
+    const countersComp1: CounterGroups = {
+      instances: 0,
+      mounts: 0,
+      unmounts: 0,
+      disposes: 0,
+      renders: 0,
+    };
+
+    const countersComp2: CounterGroups = {
+      instances: 0,
+      mounts: 0,
+      unmounts: 0,
+      disposes: 0,
+      renders: 0,
+    };
+
+    function useUpdateCounters(countersGroup: CounterGroups) {
+      countersGroup.renders += 1;
+      const isFirstRender = useRef<{ firstRender: boolean }>({ firstRender: true });
+
+      if (isFirstRender.current.firstRender) {
+        isFirstRender.current.firstRender = false;
+        countersGroup.instances += 1;
+      }
+
+      useDisposeUncommitted(() => {
+        countersGroup.disposes += 1;
+      });
+
+      useLayoutEffect(() => {
+        countersGroup.mounts += 1;
+        return () => {
+          countersGroup.unmounts += 1;
+        };
+      }, [countersGroup]);
+    }
+
+    function TestComponent1() {
+      useUpdateCounters(countersComp1);
+      return <div />;
+    }
+
+    function TestComponent2() {
+      useUpdateCounters(countersComp2);
+      return <div />;
+    }
+
+    // Render, then remove only #2
+    const rendering = render(
+      <React.StrictMode>
+        <TestComponent1 />
+        <TestComponent2 />
+      </React.StrictMode>
+    );
+    rendering.rerender(
+      <React.StrictMode>
+        <TestComponent1 />
+      </React.StrictMode>
+    );
+
+    cleanup();
+    // eslint-disable-next-line @typescript-eslint/no-unsafe-call
+    gc();
+
+    await sleep(50);
+
+    expect(countersComp1).toMatchInlineSnapshot(`
+      Object {
+        "disposes": 1,
+        "instances": 2,
+        "mounts": 1,
+        "renders": 4,
+        "unmounts": 1,
+      }
+    `);
+    expect(countersComp2).toMatchInlineSnapshot(`
+      Object {
+        "disposes": 1,
+        "instances": 2,
+        "mounts": 1,
+        "renders": 2,
+        "unmounts": 1,
+      }
+    `);
+
+    expect(userCallbacksMap.size).toBe(0);
+  });
+});
+
+function sleep(time: number): Promise<void> {
+  return new Promise<void>((res) => {
+    setTimeout(res, time);
+  });
+}

__tests__/kill-finalization-registry.ts

@@ -0,0 +1,2 @@
+// @ts-expect-error FinalizationRegistry is missing from types
+global.FinalizationRegistry = undefined;

__tests__/timers-based-impl.test.tsx

@@ -0,0 +1,129 @@
+import { cleanup, render } from "@testing-library/react";
+import React, { useLayoutEffect, useRef } from "react";
+import { DEFAULT_CLEANUP_TIMER_LOOP_MILLIS } from "../src/consts";
+import {
+  useDisposeUncommittedTimerBased,
+  waitingToBeCommittedOrDisposed,
+} from "../src/timers-based-impl";
+
+interface CounterGroups {
+  instances: number;
+  mounts: number;
+  unmounts: number;
+  disposes: number;
+  renders: number;
+  revivesOnCommit: number;
+  revivesOnRender: number;
+}
+
+describe("timers based impl", () => {
+  test("Dispose is called for uncommitted component and not called for committed component", async () => {
+    const countersComp1: CounterGroups = {
+      instances: 0,
+      mounts: 0,
+      unmounts: 0,
+      disposes: 0,
+      renders: 0,
+      revivesOnCommit: 0,
+      revivesOnRender: 0,
+    };
+
+    const countersComp2: CounterGroups = {
+      instances: 0,
+      mounts: 0,
+      unmounts: 0,
+      disposes: 0,
+      renders: 0,
+      revivesOnCommit: 0,
+      revivesOnRender: 0,
+    };
+
+    function useUpdateCounters(countersGroup: CounterGroups) {
+      countersGroup.renders += 1;
+      const isFirstRender = useRef<{ firstRender: boolean }>({ firstRender: true });
+
+      if (isFirstRender.current.firstRender) {
+        isFirstRender.current.firstRender = false;
+        countersGroup.instances += 1;
+      }
+
+      useDisposeUncommittedTimerBased(
+        () => {
+          countersGroup.disposes += 1;
+        },
+        (revivedOnRender) => {
+          if (revivedOnRender) {
+            countersGroup.revivesOnRender += 1;
+          } else {
+            countersGroup.revivesOnCommit += 1;
+          }
+        }
+      );
+
+      useLayoutEffect(() => {
+        countersGroup.mounts += 1;
+        return () => {
+          countersGroup.unmounts += 1;
+        };
+      }, [countersGroup]);
+    }
+
+    function TestComponent1() {
+      useUpdateCounters(countersComp1);
+      return <div />;
+    }
+
+    function TestComponent2() {
+      useUpdateCounters(countersComp2);
+      return <div />;
+    }
+
+    // Render, then remove only #2
+    const rendering = render(
+      <React.StrictMode>
+        <TestComponent1 />
+        <TestComponent2 />
+      </React.StrictMode>
+    );
+    rendering.rerender(
+      <React.StrictMode>
+        <TestComponent1 />
+      </React.StrictMode>
+    );
+
+    cleanup();
+
+    await sleep(DEFAULT_CLEANUP_TIMER_LOOP_MILLIS + 1000);
+
+    expect(waitingToBeCommittedOrDisposed.size).toBe(0);
+
+    expect(countersComp1).toMatchInlineSnapshot(`
+      Object {
+        "disposes": 1,
+        "instances": 2,
+        "mounts": 1,
+        "renders": 4,
+        "revivesOnCommit": 0,
+        "revivesOnRender": 0,
+        "unmounts": 1,
+      }
+    `);
+    expect(countersComp2).toMatchInlineSnapshot(`
+      Object {
+        "disposes": 1,
+        "instances": 2,
+        "mounts": 1,
+        "renders": 2,
+        "revivesOnCommit": 0,
+        "revivesOnRender": 0,
+        "unmounts": 1,
+      }
+    `);
+  }, 20_000);
+});
+
+function sleep(time: number): Promise<void> {
+  return new Promise<void>((res) => {
+    setTimeout(res, time);
+  });
+}

__tests__/tsconfig.json

@@ -0,0 +1,9 @@
+{
+  "extends": "../tsconfig",
+  "compilerOptions": {
+    "types": ["jest", "node"],
+    "jsx": "react",
+    "rootDir": "../"
+  },
+  "include": [".", "../src"]
+}

src/FinalizationRegistryWrapper.ts

@@ -0,0 +1,12 @@
+declare class FinalizationRegistryType<T, R = unknown, U = unknown> {
+  constructor(cleanup: (cleanupToken: T) => void);
+  register(object: R, cleanupToken: T, unregisterToken?: U): void;
+  unregister(unregisterToken: U): void;
+}
+
+declare const FinalizationRegistry: typeof FinalizationRegistryType | undefined;
+
+const FinalizationRegistryLocal =
+  typeof FinalizationRegistry === "undefined" ? undefined : FinalizationRegistry;
+
+export { FinalizationRegistryLocal as FinalizationRegistry };

src/consts.ts

@@ -0,0 +1,9 @@
+/**
+ * The amount of time since render an uncommitted component instance will be considered disposed
+ */
+export const DEFAULT_ASSUMED_UNCOMMITTED_AFTER_MILLIS = 1_000;
+
+/**
+ * The frequency with which we'll check for assumed uncommitted component
+ */
+export const DEFAULT_CLEANUP_TIMER_LOOP_MILLIS = 1_000;

src/finalization-registry-based-impl.ts

@@ -0,0 +1,87 @@
+import { useLayoutEffect, useRef, useState } from "react";
+import type { FinalizationRegistry as FinalizationRegistryType } from "./FinalizationRegistryWrapper";
+import type { UserCallback, UseDisposeUncommittedFinalizationRegistry } from "./types";
+
+/**
+ * We use class so it will be easy to find these objects in heap snapshots explorers by the class name
+ */
+class ObjectToBeRetainedByReact {}
+
+interface ContainerRef {
+  cleanupToken: number;
+  committed: boolean;
+}
+
+export const userCallbacksMap = new Map<number, UserCallback>();
+
+/**
+ * This is behind create function, because FinalizationRegistry might not be defined
+ * We want to create the actual impl only if FinalizationRegistry available
+ * revisit
+ * @param FinalizationRegistry
+ */
+export function createUseDisposeUncommitted(
+  FinalizationRegistry: NonNullable<typeof FinalizationRegistryType>
+): UseDisposeUncommittedFinalizationRegistry {
+  // Global state, all of the instances of finalization-registry UseDisposeUncommitted shares these vars
+  let cleanupTokensCounter = 0;
+
+  const registry = new FinalizationRegistry<
+    number,
+    ObjectToBeRetainedByReact,
+    React.MutableRefObject<ContainerRef | null>
+  >(function cleanup(cleanupToken: number) {
+    const callback = userCallbacksMap.get(cleanupToken);
+    userCallbacksMap.delete(cleanupToken);
+    // This is not expected to be false if we got here...
+    if (callback) {
+      callback();
+      userCallbacksMap.delete(cleanupToken);
+    }
+  });
+  // end of global state
+
+  return function useDisposeUncommittedFinalizationRegistry(onUncommittedCallback: UserCallback) {
+    /**
+     * This is the "magic"
+     * As long as the calling component instance is not disposed by react, and still queued for committing somewhere in the future,
+     * React must retain this state object
+     *
+     * When this object is no longer retained, our finalization callback will be called,
+     * and we know that react is no longer to commit our component instance
+     */
+    const [objectRetainedByReact] = useState(new ObjectToBeRetainedByReact());
+    const cleanupTokenRef = useRef<ContainerRef | null>(null);
+
+    if (cleanupTokenRef.current === null) {
+      cleanupTokenRef.current = {
+        cleanupToken: cleanupTokensCounter++,
+        committed: false,
+      };
+
+      // console.info(`component registered ${cleanupTokenRef.current.cleanupToken}`);
+
+      registry.register(
+        objectRetainedByReact,
+        cleanupTokenRef.current.cleanupToken,
+        cleanupTokenRef
+      );
+      userCallbacksMap.set(cleanupTokenRef.current.cleanupToken, onUncommittedCallback);
+    } else if (!cleanupTokenRef.current.committed) {
+      // The user might pass different onUncommittedCallback on additional renders but before commit,
+      // Make sure we point to the most current passed function
+      userCallbacksMap.set(cleanupTokenRef.current.cleanupToken, onUncommittedCallback);
+    }
+
+    /**
+     * We can tell that component is committed only when react runs the side effects
+     */
+    useLayoutEffect(() => {
+      if (cleanupTokenRef.current !== null) {
+        cleanupTokenRef.current.committed = true;
+        registry.unregister(cleanupTokenRef);
+        userCallbacksMap.delete(cleanupTokenRef.current.cleanupToken);
+      }
+    }, []);
+  };
+}

src/index.ts

@@ -0,0 +1,10 @@
+import { createUseDisposeUncommitted } from "./finalization-registry-based-impl";
+import { FinalizationRegistry } from "./FinalizationRegistryWrapper";
+import { useDisposeUncommittedTimerBased } from "./timers-based-impl";
+import { UseDisposeUncommitted } from "./types";
+
+export const useDisposeUncommitted: UseDisposeUncommitted = FinalizationRegistry
+  ? createUseDisposeUncommitted(FinalizationRegistry)
+  : useDisposeUncommittedTimerBased;
+
+export default useDisposeUncommitted;