fix: evict stale defer entries in dev server to bound registry memory

packages/static/src/rsc-client/clientWrapper.tsx

@@ -2,7 +2,7 @@ import React from "react";
 import { createFromFetch } from "@vitejs/plugin-rsc/browser";
 import { getModulePathFor } from "../rsc/rscModule";
 import { createContext, use } from "react";
-import type { LoadedDeferEntry, DeferRegistry } from "../rsc/defer";
+import type { LoadedDeferEntry, DeferRegistry } from "../rsc/deferRegistry";
 import { withBasePath } from "../util/basePath";
 
 interface DeferContextValue {

packages/static/src/rsc/defer.tsx

@@ -1,16 +1,10 @@
 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 { DeferRegistry } from "./deferRegistry";
 import { getPayloadIDFor } from "./rscModule";
 import { rscPayloadDir } from "virtual:funstack/config";
 
-export interface DeferEntry {
-  state: DeferEntryState;
-  name?: string;
-  drainPromise?: Promise<string>;
-}
-
 /**
  * Options for the defer function.
  */
@@ -23,28 +17,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,125 +29,9 @@ 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();
+export const deferRegistry = new DeferRegistry((element) =>
+  renderToReadableStream<ReactNode>(element),
+);
 
 /**
  * Renders given Server Component into a separate RSC payload.

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

@@ -0,0 +1,189 @@
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import { createElement } from "react";
+import type { ReactElement } from "react";
+import { DeferRegistry, devDeferEvictionOptions } from "./deferRegistry";
+
+const { ttlMs } = devDeferEvictionOptions;
+
+const encoder = new TextEncoder();
+
+/**
+ * Fake renderer that emits the element's `data-payload` prop as the stream
+ * content, so tests can verify what was rendered without the RSC runtime.
+ */
+function fakeRender(element: ReactElement): ReadableStream<Uint8Array> {
+  const payload = String(
+    (element.props as { "data-payload"?: string })["data-payload"] ?? "",
+  );
+  return new ReadableStream({
+    start(controller) {
+      controller.enqueue(encoder.encode(payload));
+      controller.close();
+    },
+  });
+}
+
+function element(payload: string): ReactElement {
+  return createElement("div", { "data-payload": payload });
+}
+
+describe("DeferRegistry", () => {
+  let registry: DeferRegistry;
+
+  beforeEach(() => {
+    vi.useFakeTimers();
+    registry = new DeferRegistry(fakeRender);
+  });
+
+  afterEach(() => {
+    vi.useRealTimers();
+  });
+
+  it("returns undefined for unknown ids", () => {
+    expect(registry.load("nope")).toBeUndefined();
+  });
+
+  it("loads a registered entry and drains its payload", async () => {
+    registry.register(element("hello"), "id1");
+    const entry = registry.load("id1");
+    expect(entry).toBeDefined();
+    expect(entry?.state.state).toBe("streaming");
+    await expect(entry?.drainPromise).resolves.toBe("hello");
+    expect(entry?.state.state).toBe("ready");
+  });
+
+  describe("evictStale", () => {
+    it("never evicts pending entries by time", () => {
+      // A pending entry may be fetched arbitrarily late, e.g. deferred
+      // content inside an accordion that is rarely opened.
+      registry.register(element("accordion"), "a");
+      vi.advanceTimersByTime(ttlMs * 100);
+
+      registry.evictStale(devDeferEvictionOptions);
+
+      expect(registry.has("a")).toBe(true);
+    });
+
+    it("evicts settled entries older than the TTL and keeps fresh ones", async () => {
+      registry.register(element("old"), "old");
+      await registry.load("old")?.drainPromise;
+      vi.advanceTimersByTime(ttlMs + 1);
+      registry.register(element("fresh"), "fresh");
+      await registry.load("fresh")?.drainPromise;
+
+      registry.evictStale(devDeferEvictionOptions);
+
+      expect(registry.has("old")).toBe(false);
+      expect(registry.has("fresh")).toBe(true);
+    });
+
+    it("keeps settled entries within the TTL", async () => {
+      registry.register(element("a"), "a");
+      await registry.load("a")?.drainPromise;
+      vi.advanceTimersByTime(ttlMs);
+
+      registry.evictStale(devDeferEvictionOptions);
+
+      expect(registry.has("a")).toBe(true);
+    });
+
+    it("treats a load as an access that refreshes the TTL", async () => {
+      registry.register(element("a"), "a");
+      await registry.load("a")?.drainPromise;
+      vi.advanceTimersByTime(ttlMs - 1);
+      registry.load("a");
+      vi.advanceTimersByTime(ttlMs - 1);
+
+      registry.evictStale(devDeferEvictionOptions);
+      expect(registry.has("a")).toBe(true);
+
+      vi.advanceTimersByTime(2);
+      registry.evictStale(devDeferEvictionOptions);
+      expect(registry.has("a")).toBe(false);
+    });
+
+    it("does not break an already-loaded entry when it is evicted", async () => {
+      registry.register(element("in-flight"), "a");
+      const entry = registry.load("a");
+
+      vi.advanceTimersByTime(ttlMs + 1);
+      registry.evictStale(devDeferEvictionOptions);
+      expect(registry.has("a")).toBe(false);
+
+      // The response holding the entry can still drain it.
+      await expect(entry?.drainPromise).resolves.toBe("in-flight");
+    });
+
+    it("caps pending entries, evicting the oldest first", () => {
+      registry.register(element("p1"), "p1");
+      registry.register(element("p2"), "p2");
+      registry.register(element("p3"), "p3");
+
+      registry.evictStale({ ttlMs, maxPending: 2 });
+
+      expect(registry.has("p1")).toBe(false);
+      expect(registry.has("p2")).toBe(true);
+      expect(registry.has("p3")).toBe(true);
+    });
+
+    it("does not count settled entries toward the pending cap", async () => {
+      registry.register(element("s1"), "s1");
+      await registry.load("s1")?.drainPromise;
+      registry.register(element("p1"), "p1");
+      registry.register(element("p2"), "p2");
+
+      registry.evictStale({ ttlMs, maxPending: 2 });
+
+      expect(registry.has("s1")).toBe(true);
+      expect(registry.has("p1")).toBe(true);
+      expect(registry.has("p2")).toBe(true);
+    });
+  });
+
+  describe("loadAll", () => {
+    it("yields all registered entries", async () => {
+      registry.register(element("one"), "id1", "first");
+      registry.register(element("two"), "id2");
+
+      const results = [];
+      for await (const result of registry.loadAll()) {
+        results.push(result);
+      }
+
+      results.sort((a, b) => a.id.localeCompare(b.id));
+      expect(results).toEqual([
+        { id: "id1", data: "one", name: "first" },
+        { id: "id2", data: "two", name: undefined },
+      ]);
+    });
+
+    it("picks up entries registered while draining (nested defer)", async () => {
+      const nested = element("nested");
+      const parent: ReactElement = createElement("div", {
+        "data-payload": "parent",
+      });
+      const renderWithNested = (el: ReactElement) => {
+        if (el === parent) {
+          return new ReadableStream<Uint8Array>({
+            start(controller) {
+              // Simulates a nested defer() call during the parent's render.
+              registry.register(nested, "nested-id");
+              controller.enqueue(encoder.encode("parent"));
+              controller.close();
+            },
+          });
+        }
+        return fakeRender(el);
+      };
+      registry = new DeferRegistry(renderWithNested);
+      registry.register(parent, "parent-id");
+
+      const ids = [];
+      for await (const result of registry.loadAll()) {
+        ids.push(result.id);
+      }
+
+      expect(ids.sort()).toEqual(["nested-id", "parent-id"]);
+    });
+  });
+});