akanjs 2.2.12 → 2.2.13-rc.0

package/server/rscWorkerHost.ts

@@ -5,14 +5,18 @@ import { type AkanI18nConfig, DEFAULT_AKAN_I18N, Logger } from "akanjs/common";
 import type { AkanTheme } from "akanjs/fetch";
 import type { AkanMetricsReport } from "akanjs/service";
 import type { ClientManifest } from "./artifact";
+import type { SsrLateRedirect } from "./ssrTypes";
 import type { BaseBuildArtifact, CssAsset } from "./types";
 
-interface RscPending {
+const DEFAULT_RSC_HOST_MAX_PENDING_CHUNKS = 256;
+
+export interface RscPending {
   onChunk: (data: Uint8Array) => void;
   onEnd: () => void;
   onError: (message: string) => void;
   onMeta?: (meta: { theme?: AkanTheme; status?: number }) => void;
   onRedirect?: (location: string, method: RscRedirectMethod, status: RscRedirectStatus) => void;
+  onLateRedirect?: (location: string, method: RscRedirectMethod, status: RscRedirectStatus) => void;
   onNotFound?: () => void;
 }
 
@@ -20,10 +24,182 @@ export type RscRedirectMethod = "replace" | "push";
 export type RscRedirectStatus = 303 | 307 | 308;
 
 export type RscRenderResult =
-  | { type: "stream"; stream: ReadableStream<Uint8Array>; theme?: AkanTheme; status?: number }
+  | {
+      type: "stream";
+      stream: ReadableStream<Uint8Array>;
+      theme?: AkanTheme;
+      status?: number;
+      lateControl: Promise<SsrLateRedirect | null>;
+      cancel: (reason?: unknown) => void;
+    }
   | { type: "redirect"; location: string; method: RscRedirectMethod; status: RscRedirectStatus }
   | { type: "not-found" };
 
+export function getRscHostMaxPendingChunks(value = process.env.AKAN_RSC_HOST_MAX_PENDING_CHUNKS): number {
+  const parsed = Number.parseInt(value ?? "", 10);
+  return Number.isFinite(parsed) && parsed > 0 ? parsed : DEFAULT_RSC_HOST_MAX_PENDING_CHUNKS;
+}
+
+export function nextRscHostPendingChunkCount(currentPendingChunks: number, desiredSize: number | null): number {
+  return desiredSize !== null && desiredSize <= 0 ? currentPendingChunks + 1 : 0;
+}
+
+export function isRscHostPendingChunkOverflow(pendingChunks: number, maxPendingChunks: number): boolean {
+  return pendingChunks > maxPendingChunks;
+}
+
+function createRscRenderAbortError(reason?: unknown): Error {
+  if (reason instanceof Error) return reason;
+  const error = new Error(reason === undefined ? "rsc render aborted" : String(reason));
+  error.name = "AbortError";
+  return error;
+}
+
+export function createIdempotentRscRenderCancel(onCancel: (reason?: unknown) => void): (reason?: unknown) => void {
+  let cancelled = false;
+  return (reason?: unknown) => {
+    if (cancelled) return;
+    cancelled = true;
+    onCancel(reason);
+  };
+}
+
+export function createRscHostRenderStream(input: {
+  setPending: (pending: RscPending) => void;
+  deletePending: () => void;
+  sendRenderOrQueue: () => void;
+  cancelRender: (reason?: unknown) => void;
+  maxPendingChunks?: number;
+  signal?: AbortSignal;
+  onPendingChunkOverflow?: () => void;
+}): Promise<RscRenderResult> {
+  let settled = false;
+  let stream!: ReadableStream<Uint8Array>;
+  let theme: AkanTheme | undefined;
+  let status: number | undefined;
+  let resolveLateControl!: (control: SsrLateRedirect | null) => void;
+  const lateControl = new Promise<SsrLateRedirect | null>((resolve) => {
+    resolveLateControl = resolve;
+  });
+  let lateControlSettled = false;
+  const settleLateControl = (control: SsrLateRedirect | null) => {
+    if (lateControlSettled) return;
+    lateControlSettled = true;
+    resolveLateControl(control);
+  };
+  const maxPendingChunks = input.maxPendingChunks ?? getRscHostMaxPendingChunks();
+  let pendingChunks = 0;
+  let removeAbortListener: (() => void) | undefined;
+  const cleanupAbortListener = () => {
+    removeAbortListener?.();
+    removeAbortListener = undefined;
+  };
+  const cancelRender = createIdempotentRscRenderCancel((reason) => {
+    input.deletePending();
+    settleLateControl(null);
+    input.cancelRender(reason);
+    cleanupAbortListener();
+  });
+
+  return new Promise<RscRenderResult>((resolve, reject) => {
+    stream = new ReadableStream<Uint8Array>({
+      start: (controller) => {
+        const abortRender = () => {
+          const error = createRscRenderAbortError(input.signal?.reason);
+          cancelRender(error);
+          if (!settled) {
+            settled = true;
+            reject(error);
+            return;
+          }
+          controller.error(error);
+        };
+        const settleStream = () => {
+          if (settled) return;
+          settled = true;
+          resolve({ type: "stream", stream, theme, status, lateControl, cancel: cancelRender });
+        };
+        input.setPending({
+          onMeta: (meta) => {
+            theme = meta.theme;
+            status = meta.status;
+            settleStream();
+          },
+          onChunk: (data) => {
+            settleStream();
+            pendingChunks = nextRscHostPendingChunkCount(pendingChunks, controller.desiredSize);
+            if (isRscHostPendingChunkOverflow(pendingChunks, maxPendingChunks)) {
+              const msg = `rsc worker host queue exceeded ${maxPendingChunks} pending chunks`;
+              const error = new Error(msg);
+              input.onPendingChunkOverflow?.();
+              cancelRender(error);
+              controller.error(error);
+              return;
+            }
+            controller.enqueue(data);
+          },
+          onEnd: () => {
+            settleLateControl(null);
+            cleanupAbortListener();
+            settleStream();
+            controller.close();
+          },
+          onError: (msg) => {
+            settleLateControl(null);
+            cleanupAbortListener();
+            if (!settled) {
+              settled = true;
+              reject(new Error(msg));
+              return;
+            }
+            controller.error(new Error(msg));
+          },
+          onRedirect: (location, method, status) => {
+            settleLateControl(null);
+            cleanupAbortListener();
+            if (!settled) {
+              settled = true;
+              resolve({ type: "redirect", location, method, status });
+              controller.close();
+              return;
+            }
+            controller.error(new Error(`redirect after stream started: ${location}`));
+          },
+          onLateRedirect: (location, method, status) => {
+            settleLateControl({ type: "redirect", location, method, status });
+          },
+          onNotFound: () => {
+            settleLateControl(null);
+            cleanupAbortListener();
+            if (!settled) {
+              settled = true;
+              resolve({ type: "not-found" });
+              controller.close();
+              return;
+            }
+            controller.error(new Error("not-found after stream started"));
+          },
+        });
+        if (input.signal) {
+          if (input.signal.aborted) {
+            abortRender();
+            return;
+          }
+          input.signal.addEventListener("abort", abortRender, { once: true });
+          removeAbortListener = () => input.signal?.removeEventListener("abort", abortRender);
+        }
+        input.sendRenderOrQueue();
+      },
+      cancel: (reason) => {
+        cancelRender(reason);
+      },
+      pull: () => {
+        pendingChunks = Math.max(0, pendingChunks - 1);
+      },
+    });
+  });
+}
+
 type RscInMsg =
   | { type: "hello" }
   | { type: "ready" }
@@ -32,6 +208,13 @@ type RscInMsg =
   | { type: "chunk"; requestId: string; data: Uint8Array }
   | { type: "end"; requestId: string }
   | { type: "redirect"; requestId: string; location: string; method?: RscRedirectMethod; status?: RscRedirectStatus }
+  | {
+      type: "late-redirect";
+      requestId: string;
+      location: string;
+      method?: RscRedirectMethod;
+      status?: RscRedirectStatus;
+    }
   | { type: "not-found"; requestId: string }
   | { type: "metrics"; metrics: AkanMetricsReport }
   | { type: "error"; requestId: string; message: string; buildId?: number };
@@ -106,6 +289,7 @@ export class RscWorker {
   #recycleCount = 0;
   #lastRecycleReason: string | undefined;
   #lastWorkerMetrics: AkanMetricsReport = {};
+  #hostPendingChunkOverflowCount = 0;
   #restartTimer: ReturnType<typeof setTimeout> | null = null;
   #recycleTimer: ReturnType<typeof setTimeout> | null = null;
   #rollingRecycle: { oldProc: Bun.Subprocess<"ignore", "inherit", "inherit">; reason: string } | null = null;
@@ -172,73 +356,34 @@ export class RscWorker {
       },
       cancel: () => {
         this.#pending.delete(requestId);
+        this.#cancelRender(requestId);
       },
     });
   }
 
-  renderWithMeta(req: Request, options: { clientManifest?: ClientManifest } = {}): Promise<RscRenderResult> {
+  renderWithMeta(
+    req: Request,
+    options: { clientManifest?: ClientManifest; signal?: AbortSignal } = {},
+  ): Promise<RscRenderResult> {
     const requestId = crypto.randomUUID();
-    let settled = false;
-    let stream!: ReadableStream<Uint8Array>;
-    let theme: AkanTheme | undefined;
-    let status: number | undefined;
-    const result = new Promise<RscRenderResult>((resolve, reject) => {
-      stream = new ReadableStream<Uint8Array>({
-        start: (controller) => {
-          const settleStream = () => {
-            if (settled) return;
-            settled = true;
-            resolve({ type: "stream", stream, theme, status });
-          };
-          this.#pending.set(requestId, {
-            onMeta: (meta) => {
-              theme = meta.theme;
-              status = meta.status;
-              settleStream();
-            },
-            onChunk: (data) => {
-              settleStream();
-              controller.enqueue(data);
-            },
-            onEnd: () => {
-              settleStream();
-              controller.close();
-            },
-            onError: (msg) => {
-              if (!settled) {
-                settled = true;
-                reject(new Error(msg));
-                return;
-              }
-              controller.error(new Error(msg));
-            },
-            onRedirect: (location, method, status) => {
-              if (!settled) {
-                settled = true;
-                resolve({ type: "redirect", location, method, status });
-                controller.close();
-                return;
-              }
-              controller.error(new Error(`redirect after stream started: ${location}`));
-            },
-            onNotFound: () => {
-              if (!settled) {
-                settled = true;
-                resolve({ type: "not-found" });
-                controller.close();
-                return;
-              }
-              controller.error(new Error("not-found after stream started"));
-            },
-          });
-          this.#sendRenderOrQueue(requestId, req, options.clientManifest);
-        },
-        cancel: () => {
-          this.#pending.delete(requestId);
-        },
-      });
+    return createRscHostRenderStream({
+      setPending: (pending) => {
+        this.#pending.set(requestId, pending);
+      },
+      deletePending: () => {
+        this.#pending.delete(requestId);
+      },
+      sendRenderOrQueue: () => {
+        this.#sendRenderOrQueue(requestId, req, options.clientManifest);
+      },
+      cancelRender: () => {
+        this.#cancelRender(requestId);
+      },
+      signal: options.signal,
+      onPendingChunkOverflow: () => {
+        this.#hostPendingChunkOverflowCount += 1;
+      },
     });
-    return result;
   }
 
   kill(): void {
@@ -258,8 +403,18 @@ export class RscWorker {
   }
 
   getMetrics(): AkanMetricsReport {
-    return {
-      ...this.#lastWorkerMetrics,
+    const {
+      rscWorkerPid: _rscWorkerPid,
+      rscWorkerStatus: _rscWorkerStatus,
+      rscWorkerRestartCount: _rscWorkerRestartCount,
+      rscWorkerRecycleCount: _rscWorkerRecycleCount,
+      rscWorkerLastRecycleReason: _rscWorkerLastRecycleReason,
+      rscPendingRenderCount: _rscPendingRenderCount,
+      rscQueuedSendCount: _rscQueuedSendCount,
+      rscHostPendingChunkOverflowCount: _rscHostPendingChunkOverflowCount,
+      ...workerMetrics
+    } = this.#lastWorkerMetrics;
+    return Object.assign(workerMetrics, {
       rscWorkerPid: this.#proc.pid,
       rscWorkerStatus: this.#status,
       rscWorkerRestartCount: this.#restartCount,
@@ -267,7 +422,8 @@ export class RscWorker {
       rscWorkerLastRecycleReason: this.#lastRecycleReason,
       rscPendingRenderCount: this.#pending.size,
       rscQueuedSendCount: this.#queuedSends.length,
-    };
+      rscHostPendingChunkOverflowCount: this.#hostPendingChunkOverflowCount,
+    });
   }
 
   restartWhenIdle(reason: string): boolean {
@@ -411,6 +567,11 @@ export class RscWorker {
           p.onRedirect?.(message.location, message.method ?? "replace", message.status ?? 307),
         );
         return;
+      case "late-redirect":
+        this.#pending
+          .get(message.requestId)
+          ?.onLateRedirect?.(message.location, message.method ?? "replace", message.status ?? 307);
+        return;
       case "not-found":
         this.#resolvePending(message.requestId, (p) => p.onNotFound?.());
         return;
@@ -470,6 +631,14 @@ export class RscWorker {
     }
   }
 
+  #cancelRender(requestId: string): void {
+    if (this.#status !== "ready") return;
+    try {
+      this.#proc.send({ type: "cancel", requestId });
+    } catch {
+    }
+  }
+
   #flushQueuedSends(): void {
     const queue = this.#queuedSends;
     this.#queuedSends = [];

package/server/rscWorkerReplay.ts

@@ -0,0 +1,35 @@
+export type CachedRscReplayMessage =
+  | { type: "meta"; requestId: string; theme?: string; status?: number }
+  | { type: "chunk"; requestId: string; data: Uint8Array }
+  | { type: "end"; requestId: string };
+
+function yieldToHostEventLoop(): Promise<void> {
+  return new Promise((resolve) => setTimeout(resolve, 0));
+}
+
+export async function replayCachedRscResult(input: {
+  requestId: string;
+  chunks: readonly Uint8Array[];
+  theme?: string;
+  status?: number;
+  send: (message: CachedRscReplayMessage) => void;
+  isCancelled: () => boolean;
+  yieldEveryChunks?: number;
+  yieldToHost?: () => Promise<void>;
+}): Promise<boolean> {
+  const yieldEveryChunks =
+    input.yieldEveryChunks !== undefined && Number.isFinite(input.yieldEveryChunks) && input.yieldEveryChunks > 0
+      ? Math.floor(input.yieldEveryChunks)
+      : 1;
+  const yieldToHost = input.yieldToHost ?? yieldToHostEventLoop;
+  if (input.isCancelled()) return false;
+  input.send({ type: "meta", requestId: input.requestId, theme: input.theme, status: input.status });
+  for (let index = 0; index < input.chunks.length; index += 1) {
+    if (input.isCancelled()) return false;
+    input.send({ type: "chunk", requestId: input.requestId, data: input.chunks[index] });
+    if ((index + 1) % yieldEveryChunks === 0) await yieldToHost();
+  }
+  if (input.isCancelled()) return false;
+  input.send({ type: "end", requestId: input.requestId });
+  return true;
+}