fix(@typegpu/react): Cleanup on-demand root on unmount (#2554)

Author: iwoplaza

.github/workflows/pkg-pr.yml

@@ -37,7 +37,7 @@ jobs:
         run: pnpm nightly-build
 
       - name: Publish (pkg.pr.new)
-        run: pnpm exec pkg-pr-new publish './packages/typegpu' './packages/typegpu-noise' './packages/typegpu-cli' './packages/unplugin-typegpu' --json output.json --comment=off --pnpm --no-compact
+        run: pnpm exec pkg-pr-new publish './packages/typegpu' './packages/typegpu-noise' './packages/typegpu-react' './packages/typegpu-cli' './packages/unplugin-typegpu' --json output.json --comment=off --pnpm --no-compact
       - name: Post or update comment
         uses: actions/github-script@v6
         with:

package.json

@@ -32,7 +32,7 @@
     "test:browser": "vitest run --browser.enabled --project browser",
     "test:browser:watch": "vitest --browser.enabled --project browser",
     "test:coverage": "vitest --coverage run",
-    "nightly-build": "SKIP_TESTS=true pnpm --filter typegpu --filter @typegpu/noise --filter @typegpu/cli --filter unplugin-typegpu prepublishOnly --skip-publish-tag-check",
+    "nightly-build": "SKIP_TESTS=true pnpm --filter typegpu --filter @typegpu/noise --filter @typegpu/react --filter @typegpu/cli --filter unplugin-typegpu prepublishOnly --skip-publish-tag-check",
     "changes": "tgpu-dev-cli changes"
   },
   "devDependencies": {

packages/typegpu-react/src/core/root-context.tsx

@@ -7,6 +7,7 @@ import React, {
   useState,
 } from 'react';
 import tgpu, { type TgpuRoot } from 'typegpu';
+import { useDeferredCleanup } from './helper-hooks.ts';
 
 function attachPromiseStatus<T>(
   promise: PromiseLike<T> & {
@@ -59,18 +60,31 @@ type RootContextResult =
 
 interface RootContext {
   initOrGetRoot(): RootContextResult;
+  unmount(): void;
 }
 
 class OwnRootContext implements RootContext {
   #result: RootContextResult | undefined;
+  #destroyed: boolean = false;
 
   initOrGetRoot(): RootContextResult {
+    if (this.#destroyed) {
+      console.warn(`[@typegpu/react]: Tried to init an already destroyed root.`);
+      return { status: 'rejected', error: new Error('Already destroyed') };
+    }
+
     if (!this.#result) {
       this.#result = {
         status: 'pending',
         promise: tgpu.init().then(
           (root) => {
-            this.#result = { status: 'resolved', value: root };
+            if (this.#destroyed) {
+              root.destroy();
+              this.#result = undefined;
+            } else {
+              this.#result = { status: 'resolved', value: root };
+            }
+
             return root;
           },
           (error) => {
@@ -83,6 +97,15 @@ class OwnRootContext implements RootContext {
 
     return this.#result;
   }
+
+  unmount(): void {
+    this.#destroyed = true;
+
+    if (this.#result?.status === 'resolved') {
+      this.#result.value.destroy();
+    }
+    this.#result = undefined;
+  }
 }
 
 class ExistingRootContext implements RootContext {
@@ -95,6 +118,10 @@ class ExistingRootContext implements RootContext {
   initOrGetRoot(): RootContextResult {
     return this.result;
   }
+
+  unmount(): void {
+    // Nothing to do on unmount
+  }
 }
 
 /**
@@ -124,6 +151,10 @@ export const Root = ({ children, root }: RootProps) => {
     return undefined;
   }, [root]);
 
+  useDeferredCleanup(() => {
+    ownCtx.unmount();
+  });
+
   return <rootContext.Provider value={existingRootCtx ?? ownCtx}>{children}</rootContext.Provider>;
 };
 

packages/typegpu-react/tests/root-context.test.tsx

@@ -0,0 +1,139 @@
+import { act, render } from '@testing-library/react';
+import { afterEach, beforeEach, describe, expect, vi } from 'vitest';
+import tgpu from 'typegpu';
+import type { TgpuRoot } from 'typegpu';
+import { Root, useRootWithStatus } from '@typegpu/react';
+import { useEffect } from 'react';
+import { it } from './utils/extended-test.tsx';
+
+describe('Root unmount cleanup', () => {
+  beforeEach(() => {
+    vi.useFakeTimers();
+  });
+
+  afterEach(() => {
+    vi.useRealTimers();
+  });
+
+  it('should destroy own root on unmount (deferred cleanup)', async () => {
+    let capturedRoot: TgpuRoot | undefined;
+
+    function TestConsumer() {
+      const result = useRootWithStatus();
+
+      useEffect(() => {
+        if (result.status === 'resolved') {
+          capturedRoot = result.value;
+        }
+      });
+
+      return null;
+    }
+
+    const { unmount } = render(
+      <Root>
+        <TestConsumer />
+      </Root>,
+    );
+
+    // Flush microtasks so tgpu.init() resolves and React re-renders
+    await act(async () => {
+      await Promise.resolve();
+    });
+
+    expect(capturedRoot).toBeDefined();
+    const destroySpy = vi.spyOn(capturedRoot!, 'destroy');
+
+    unmount();
+    vi.runAllTimers();
+
+    expect(destroySpy).toHaveBeenCalledTimes(1);
+  });
+
+  it('should not throw when existing root unmounts', async ({ RootWrapper }) => {
+    const { unmount } = render(<div />, { wrapper: RootWrapper });
+    expect(() => unmount()).not.toThrow();
+  });
+
+  it('should destroy root when init promise resolves after unmount', async ({
+    root: fixtureRoot,
+  }) => {
+    let resolveInit!: (root: TgpuRoot) => void;
+    const initPromise = new Promise<TgpuRoot>((resolve) => {
+      resolveInit = resolve;
+    });
+
+    using _initSpy = vi.spyOn(tgpu, 'init').mockReturnValue(initPromise);
+    const destroySpy = vi.spyOn(fixtureRoot, 'destroy');
+
+    function TestConsumer() {
+      useRootWithStatus();
+      return null;
+    }
+
+    const { unmount } = render(
+      <Root>
+        <TestConsumer />
+      </Root>,
+    );
+
+    // Unmount and run deferred cleanup before init resolves
+    unmount();
+    vi.runAllTimers();
+
+    // Resolve init after context has been destroyed
+    resolveInit(fixtureRoot);
+
+    // Flush microtasks so the .then() callback runs
+    await act(async () => {
+      await Promise.resolve();
+    });
+
+    expect(destroySpy).toHaveBeenCalledTimes(1);
+  });
+
+  describe('React StrictMode compatibility', () => {
+    it('should survive the double mount/unmount cycle', async () => {
+      let capturedRoot: TgpuRoot | undefined;
+
+      function TestConsumer() {
+        const result = useRootWithStatus();
+
+        useEffect(() => {
+          if (result.status === 'resolved') {
+            capturedRoot = result.value;
+          }
+        });
+
+        return null;
+      }
+
+      const { unmount } = render(
+        <Root>
+          <TestConsumer />
+        </Root>,
+        { reactStrictMode: true },
+      );
+
+      // Flush microtasks so tgpu.init() resolves and React re-renders
+      await act(async () => {
+        await Promise.resolve();
+      });
+
+      expect(capturedRoot).toBeDefined();
+
+      // Run any pending deferred cleanup timeouts (from StrictMode's first mount)
+      vi.runAllTimers();
+
+      // Root should still be alive (deferred cleanup was cancelled by remount)
+      const destroySpy = vi.spyOn(capturedRoot!, 'destroy');
+      expect(destroySpy).not.toHaveBeenCalled();
+
+      // Now actually unmount
+      unmount();
+      vi.runAllTimers();
+
+      expect(destroySpy).toHaveBeenCalledTimes(1);
+    });
+  });
+});