diff --git a/packages/static/src/rsc/defer.tsx b/packages/static/src/rsc/defer.tsx index 82e56a5..4b5d071 100644 --- a/packages/static/src/rsc/defer.tsx +++ b/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; -} +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; - drainPromise: Promise; -} - -type DeferEntryState = - | { - state: "pending"; - element: ReactElement; - } - | { - state: "streaming"; - stream: ReadableStream; - } - | { - 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(); - - 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(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 = []; - let waiting: (() => void) | undefined; - let remainingCount = 0; - const started = new Set(); - - // 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((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(); /** diff --git a/packages/static/src/rsc/deferRegistry.test.ts b/packages/static/src/rsc/deferRegistry.test.ts new file mode 100644 index 0000000..44837ae --- /dev/null +++ b/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({ + 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(); + for await (const { id, data } of registry.loadAll()) { + results.set(id, data); + } + expect(results).toEqual( + new Map([ + ["id-parent", "parent"], + ["id-child", "child"], + ]), + ); + }); +}); diff --git a/packages/static/src/rsc/deferRegistry.ts b/packages/static/src/rsc/deferRegistry.ts new file mode 100644 index 0000000..d82f230 --- /dev/null +++ b/packages/static/src/rsc/deferRegistry.ts @@ -0,0 +1,243 @@ +import type { ReactElement, ReactNode } from "react"; +import { AsyncLocalStorage } from "node:async_hooks"; +import { renderToReadableStream } from "@vitejs/plugin-rsc/react/rsc"; +import { drainStream } from "../util/drainStream"; + +export interface DeferEntry { + state: DeferEntryState; + name?: string; + drainPromise?: Promise; + /** + * The top-level render this entry belongs to. + * Only set in dev, where renders are scoped via `startRender`. + */ + render?: RenderHandle; +} + +export interface LoadedDeferEntry extends DeferEntry { + state: Exclude; + drainPromise: Promise; +} + +type DeferEntryState = + | { + state: "pending"; + element: ReactElement; + } + | { + state: "streaming"; + stream: ReadableStream; + } + | { + state: "ready"; + } + | { + state: "error"; + error: unknown; + }; + +/** + * Represents one top-level HTML render in the dev server. + * + * Every `defer()` registration made while rendering a page — including + * nested `defer()` calls made when a deferred element itself is rendered + * later — is attributed to the same RenderHandle, forming a "rendered-by" + * tree rooted at the top-level render. When a new top-level render for the + * same key begins, the previous render's whole tree is evicted after a + * grace period, so the registry does not grow unboundedly over a dev + * session (see issue #144). + */ +export interface RenderHandle { + /** Identifies what is being rendered (the entry path in dev). */ + readonly key: string; + /** IDs of all entries registered during this render, incl. nested ones. */ + readonly ids: Set; + /** True once this render has been superseded by a newer one. */ + invalidated: boolean; +} + +/** + * How long entries of a superseded render are kept before eviction. + * The grace period lets in-flight client fetches for the previous + * page load still resolve. + */ +export const renderEvictionGracePeriodMs = 30_000; + +export class DeferRegistry { + #registry = new Map(); + #renderContext = new AsyncLocalStorage(); + #currentRenders = new Map(); + + /** + * Marks the start of a new top-level render identified by `key` and runs + * `fn` so that all `defer()` registrations made during it are attributed + * to this render. Attribution propagates through async continuations, so + * registrations made while React renders the tree started inside `fn` + * are captured too. + * + * If a previous render exists for the same key, it is invalidated and + * its entries (the whole rendered-by tree) are evicted after a grace + * period. + * + * Only the dev server scopes renders; build-time rendering never calls + * this, so build entries are kept until the process exits. + */ + startRender(key: string, fn: () => T): T { + const previous = this.#currentRenders.get(key); + if (previous) { + previous.invalidated = true; + this.#scheduleEviction(previous); + } + const handle: RenderHandle = { key, ids: new Set(), invalidated: false }; + this.#currentRenders.set(key, handle); + return this.#renderContext.run(handle, fn); + } + + #scheduleEviction(handle: RenderHandle) { + const timer = setTimeout(() => { + for (const id of handle.ids) { + const entry = this.#registry.get(id); + if (entry?.render === handle) { + this.#registry.delete(id); + } + } + handle.ids.clear(); + }, renderEvictionGracePeriodMs); + // In Node, don't keep the process alive just for eviction. + (timer as unknown as { unref?: () => void }).unref?.(); + } + + register(element: ReactElement, id: string, name?: string) { + const render = this.#renderContext.getStore(); + this.#registry.set(id, { + state: { element, state: "pending" }, + name, + render, + }); + if (render) { + render.ids.add(id); + if (render.invalidated) { + // Late registration from an already-superseded render (e.g. the + // previous render is still streaming, or a nested defer loads + // after invalidation). Its eviction timer may have already fired, + // so re-arm eviction to make sure this straggler is dropped too. + this.#scheduleEviction(render); + } + } + } + + 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": { + // Render within the owning render's context so that nested + // `defer()` registrations attach to the same top-level render. + const doRender = () => renderToReadableStream(state.element); + const stream = entry.render + ? this.#renderContext.run(entry.render, doRender) + : doRender(); + 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 = []; + let waiting: (() => void) | undefined; + let remainingCount = 0; + const started = new Set(); + + // 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((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); + } + } +} diff --git a/packages/static/src/rsc/entry.tsx b/packages/static/src/rsc/entry.tsx index 24d6203..4d0b938 100644 --- a/packages/static/src/rsc/entry.tsx +++ b/packages/static/src/rsc/entry.tsx @@ -85,9 +85,13 @@ async function renderEntryToResponse( if (ssrEnabled) { // SSR on: single RSC stream with full tree const rscStart = performance.now(); - const rootRscStream = renderToReadableStream({ - root: {appNode}, - }); + // Scope defer() registrations to this render so the previous render's + // entries for this entry can be evicted (see DeferRegistry#startRender). + const rootRscStream = deferRegistry.startRender(entry.path, () => + renderToReadableStream({ + root: {appNode}, + }), + ); timings.push(`rsc;dur=${performance.now() - rscStart}`); const ssrStart = performance.now(); @@ -109,16 +113,23 @@ async function renderEntryToResponse( } else { // SSR off: shell RSC for SSR, full RSC for client const rscStart = performance.now(); - const shellRscStream = renderToReadableStream({ - root: ( - - - - ), - }); - const clientRscStream = renderToReadableStream({ - root: {appNode}, - }); + // Scope defer() registrations to this render so the previous render's + // entries for this entry can be evicted (see DeferRegistry#startRender). + const [shellRscStream, clientRscStream] = deferRegistry.startRender( + entry.path, + () => [ + renderToReadableStream({ + root: ( + + + + ), + }), + renderToReadableStream({ + root: {appNode}, + }), + ], + ); timings.push(`rsc;dur=${performance.now() - rscStart}`); const ssrStart = performance.now(); @@ -209,9 +220,13 @@ export async function serveRSC(request: Request): Promise { timings.push(`resolve;dur=${performance.now() - resolveStart}`); const rscStart = performance.now(); - const rootRscStream = renderToReadableStream({ - root: {appNode}, - }); + // An HMR re-fetch is a new top-level render of the entry: scope it so + // the previous render's defer() registrations can be evicted. + const rootRscStream = deferRegistry.startRender(entry.path, () => + renderToReadableStream({ + root: {appNode}, + }), + ); timings.push(`rsc;dur=${performance.now() - rscStart}`); return new Response(rootRscStream, {