fix: scope dev defer registrations per HTML render to bound registry growth

packages/static/src/rsc/defer.tsx

@@ -1,15 +1,11 @@
 import type { ReactElement, ReactNode } from "react";
-import { renderToReadableStream } from "@vitejs/plugin-rsc/react/rsc";
 import { DeferredComponent } from "#rsc-client";
-import { drainStream } from "../util/drainStream";
 import { getPayloadIDFor } from "./rscModule";
+import { DeferRegistry } from "./deferRegistry";
 import { rscPayloadDir } from "virtual:funstack/config";
 
-export interface DeferEntry {
-  state: DeferEntryState;
-  name?: string;
-  drainPromise?: Promise<string>;
-}
+export { DeferRegistry } from "./deferRegistry";
+export type { DeferEntry, LoadedDeferEntry } from "./deferRegistry";
 
 /**
  * Options for the defer function.
@@ -23,28 +19,6 @@ export interface DeferOptions {
   name?: string;
 }
 
-export interface LoadedDeferEntry extends DeferEntry {
-  state: Exclude<DeferEntryState, { state: "pending" }>;
-  drainPromise: Promise<string>;
-}
-
-type DeferEntryState =
-  | {
-      state: "pending";
-      element: ReactElement;
-    }
-  | {
-      state: "streaming";
-      stream: ReadableStream<Uint8Array>;
-    }
-  | {
-      state: "ready";
-    }
-  | {
-      state: "error";
-      error: unknown;
-    };
-
 /**
  * Sanitizes a name for use in file paths.
  * Replaces non-alphanumeric characters with underscores and limits length.
@@ -57,124 +31,6 @@ function sanitizeName(name: string): string {
     .slice(0, 50);
 }
 
-export class DeferRegistry {
-  #registry = new Map<string, DeferEntry>();
-
-  register(element: ReactElement, id: string, name?: string) {
-    this.#registry.set(id, { state: { element, state: "pending" }, name });
-  }
-
-  load(id: string): LoadedDeferEntry | undefined {
-    const entry = this.#registry.get(id);
-    if (!entry) {
-      return undefined;
-    }
-    return this.#loadEntry(entry);
-  }
-
-  #loadEntry(entry: DeferEntry): LoadedDeferEntry {
-    const { state } = entry;
-    switch (state.state) {
-      case "pending": {
-        const stream = renderToReadableStream<ReactNode>(state.element);
-        const [stream1, stream2] = stream.tee();
-        entry.state = { state: "streaming", stream: stream1 };
-        const drainPromise = drainStream(stream2);
-        entry.drainPromise = drainPromise;
-        drainPromise.then(
-          () => {
-            entry.state = { state: "ready" };
-          },
-          (error) => {
-            entry.state = { state: "error", error };
-          },
-        );
-        return entry as LoadedDeferEntry;
-      }
-      case "streaming":
-      case "ready":
-      case "error":
-        return entry as LoadedDeferEntry;
-    }
-  }
-
-  has(id: string): boolean {
-    return this.#registry.has(id);
-  }
-
-  /**
-   * Iterates over all entries in parallel.
-   * Yields results as each stream completes.
-   *
-   * Rendering a deferred element may itself call `defer()` (nested defer),
-   * registering new entries while earlier ones are still draining. Entries
-   * registered mid-iteration are picked up too, until none are left.
-   */
-  async *loadAll() {
-    const errors: unknown[] = [];
-
-    type Result = { id: string; data: string; name?: string };
-
-    // Completion queue
-    const completed: Array<Result | { error: unknown }> = [];
-    let waiting: (() => void) | undefined;
-    let remainingCount = 0;
-    const started = new Set<string>();
-
-    // Start loading every entry not started yet and track its drain promise.
-    // We use drain promises (which drain stream2 from tee) instead of
-    // draining stream1 directly, because stream1 may have been locked
-    // by createFromReadableStream during SSR.
-    const startPending = () => {
-      for (const [id, entry] of this.#registry) {
-        if (started.has(id)) continue;
-        started.add(id);
-        const loaded = this.#loadEntry(entry);
-        remainingCount++;
-        loaded.drainPromise.then(
-          (data) => {
-            completed.push({ id, data, name: entry.name });
-            remainingCount--;
-            waiting?.();
-          },
-          (error) => {
-            completed.push({ error });
-            remainingCount--;
-            waiting?.();
-          },
-        );
-      }
-    };
-
-    startPending();
-
-    // Yield from queue as results arrive
-    while (remainingCount > 0 || completed.length > 0) {
-      if (completed.length === 0) {
-        await new Promise<void>((r) => {
-          waiting = r;
-        });
-        waiting = undefined;
-      }
-      for (const result of completed.splice(0)) {
-        if ("error" in result) {
-          errors.push(result.error);
-        } else {
-          yield result;
-        }
-      }
-      // A drained entry may have registered nested entries during its
-      // render; any registration happens before its parent's drain promise
-      // resolves, so once remainingCount hits 0 no new entries can appear.
-      startPending();
-    }
-
-    if (errors.length > 0) {
-      throw new AggregateError(errors);
-    }
-  }
-}
-
 export const deferRegistry = new DeferRegistry();
 
 /**

packages/static/src/rsc/deferRegistry.test.ts

@@ -0,0 +1,176 @@
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import type { ReactElement } from "react";
+import { DeferRegistry, renderEvictionGracePeriodMs } from "./deferRegistry";
+
+// The registry only passes elements through to renderToReadableStream, so
+// mock rendering: emit `data` as the payload and invoke `onRender` (used by
+// tests to simulate nested defer() calls made during rendering).
+vi.mock("@vitejs/plugin-rsc/react/rsc", () => ({
+  renderToReadableStream: (element: unknown) => {
+    const el = element as { data: string; onRender?: () => void };
+    el.onRender?.();
+    return new ReadableStream<Uint8Array>({
+      start(controller) {
+        controller.enqueue(new TextEncoder().encode(el.data));
+        controller.close();
+      },
+    });
+  },
+}));
+
+function el(data: string, onRender?: () => void): ReactElement {
+  return { data, onRender } as unknown as ReactElement;
+}
+
+describe("DeferRegistry", () => {
+  beforeEach(() => {
+    vi.useFakeTimers();
+  });
+  afterEach(() => {
+    vi.useRealTimers();
+  });
+
+  it("loads a registered entry and drains its payload", async () => {
+    const registry = new DeferRegistry();
+    registry.register(el("payload-a"), "id-a");
+
+    const entry = registry.load("id-a");
+    expect(entry).toBeDefined();
+    expect(entry?.state.state).toBe("streaming");
+    await expect(entry?.drainPromise).resolves.toBe("payload-a");
+    expect(entry?.state.state).toBe("ready");
+  });
+
+  it("returns the render function's result from startRender", () => {
+    const registry = new DeferRegistry();
+    const result = registry.startRender("index.html", () => "rendered");
+    expect(result).toBe("rendered");
+  });
+
+  it("evicts the previous render's entries after the grace period when a new render for the same key starts", async () => {
+    const registry = new DeferRegistry();
+    registry.startRender("index.html", () => {
+      registry.register(el("a"), "id-a");
+    });
+    registry.startRender("index.html", () => {
+      registry.register(el("b"), "id-b");
+    });
+
+    // Within the grace period, the previous render's entries survive.
+    expect(registry.has("id-a")).toBe(true);
+    expect(registry.has("id-b")).toBe(true);
+
+    await vi.advanceTimersByTimeAsync(renderEvictionGracePeriodMs);
+    expect(registry.has("id-a")).toBe(false);
+    expect(registry.has("id-b")).toBe(true);
+  });
+
+  it("keeps the previous render's entries loadable during the grace period", async () => {
+    const registry = new DeferRegistry();
+    registry.startRender("index.html", () => {
+      registry.register(el("stale"), "id-stale");
+    });
+    registry.startRender("index.html", () => {});
+
+    const entry = registry.load("id-stale");
+    await expect(entry?.drainPromise).resolves.toBe("stale");
+  });
+
+  it("does not evict entries belonging to a different key", async () => {
+    const registry = new DeferRegistry();
+    registry.startRender("a.html", () => {
+      registry.register(el("a"), "id-a");
+    });
+    registry.startRender("b.html", () => {
+      registry.register(el("b"), "id-b");
+    });
+
+    await vi.advanceTimersByTimeAsync(renderEvictionGracePeriodMs);
+    expect(registry.has("id-a")).toBe(true);
+    expect(registry.has("id-b")).toBe(true);
+  });
+
+  it("never evicts entries registered outside of a scoped render (build)", async () => {
+    const registry = new DeferRegistry();
+    registry.register(el("build-time"), "id-build");
+
+    registry.startRender("index.html", () => {});
+    registry.startRender("index.html", () => {});
+
+    await vi.advanceTimersByTimeAsync(renderEvictionGracePeriodMs * 2);
+    expect(registry.has("id-build")).toBe(true);
+  });
+
+  it("attributes nested defer registrations to the parent's render and evicts the whole tree", async () => {
+    const registry = new DeferRegistry();
+    registry.startRender("index.html", () => {
+      registry.register(
+        el("parent", () => {
+          // Nested defer() during the parent's render; registered through
+          // an async continuation to exercise AsyncLocalStorage propagation.
+          queueMicrotask(() => {
+            registry.register(el("child"), "id-child");
+          });
+        }),
+        "id-parent",
+      );
+    });
+
+    // Loading the parent (e.g. during SSR or a client fetch) triggers its
+    // render, which registers the nested child under the same render.
+    const parent = registry.load("id-parent");
+    await parent?.drainPromise;
+    expect(registry.has("id-child")).toBe(true);
+
+    registry.startRender("index.html", () => {});
+    await vi.advanceTimersByTimeAsync(renderEvictionGracePeriodMs);
+    expect(registry.has("id-parent")).toBe(false);
+    expect(registry.has("id-child")).toBe(false);
+  });
+
+  it("evicts nested registrations made after the render was superseded", async () => {
+    const registry = new DeferRegistry();
+    registry.startRender("index.html", () => {
+      registry.register(
+        el("parent", () => {
+          registry.register(el("child"), "id-child");
+        }),
+        "id-parent",
+      );
+    });
+
+    // The render is superseded before the parent is ever loaded.
+    registry.startRender("index.html", () => {});
+
+    // A client fetch during the grace period loads the parent, registering
+    // the nested child into the already-invalidated render.
+    const parent = registry.load("id-parent");
+    await expect(parent?.drainPromise).resolves.toBe("parent");
+    expect(registry.has("id-child")).toBe(true);
+
+    await vi.advanceTimersByTimeAsync(renderEvictionGracePeriodMs);
+    expect(registry.has("id-parent")).toBe(false);
+    expect(registry.has("id-child")).toBe(false);
+  });
+
+  it("loadAll picks up nested entries registered mid-iteration", async () => {
+    const registry = new DeferRegistry();
+    registry.register(
+      el("parent", () => {
+        registry.register(el("child"), "id-child");
+      }),
+      "id-parent",
+    );
+
+    const results = new Map<string, string>();
+    for await (const { id, data } of registry.loadAll()) {
+      results.set(id, data);
+    }
+    expect(results).toEqual(
+      new Map([
+        ["id-parent", "parent"],
+        ["id-child", "child"],
+      ]),
+    );
+  });
+});