rivetkit 2.3.0-rc.5 → 2.3.0-rc.7

package/src/registry/index.ts

@@ -1,6 +1,8 @@
+import { Hono } from "hono";
 import { ENGINE_ENDPOINT } from "@/common/engine";
 import { configureServerlessPool } from "@/serverless/configure";
-import { VERSION } from "@/utils";
+import { detectRuntime, VERSION } from "@/utils";
+import { crossPlatformServe, loadRuntimeServeStatic } from "@/utils/serve";
 import {
 	type RegistryActors,
 	type RegistryConfig,
@@ -13,6 +15,22 @@ import type { RuntimeServerlessResponseHead } from "./runtime";
 
 type ShutdownSignal = "SIGINT" | "SIGTERM";
 
+function signalExitCode(signal: ShutdownSignal): number {
+	switch (signal) {
+		case "SIGINT":
+			return 130;
+		case "SIGTERM":
+			return 143;
+	}
+}
+
+function finishShutdownSignal(signal: ShutdownSignal): void {
+	if (process.pid === 1) {
+		process.exit(signalExitCode(signal));
+	}
+	process.kill(process.pid, signal);
+}
+
 export type FetchHandler = (
 	request: Request,
 	...args: any
@@ -279,6 +297,36 @@ export class Registry<A extends RegistryActors> {
 		};
 	}
 
+	/**
+	 * Starts an HTTP server that dispatches every request through the
+	 * serverless handler. Uses `crossPlatformServe` to pick the right
+	 * runtime (Node, Bun, Deno).
+	 *
+	 * @param opts.port      Port to listen on. Defaults to 3000.
+	 * @param opts.publicDir If set, serves static files from this directory
+	 *                       before falling through to the registry handler.
+	 *
+	 * @example
+	 * ```ts
+	 * await registry.listen();
+	 * await registry.listen({ port: 8080, publicDir: "./public" });
+	 * ```
+	 */
+	public async listen(
+		opts: { port?: number; publicDir?: string } = {},
+	): Promise<void> {
+		const port = opts.port ?? 3000;
+		const config = this.parseConfig();
+		const runtime = detectRuntime();
+		const app = new Hono();
+		if (opts.publicDir) {
+			const serveStatic = await loadRuntimeServeStatic(runtime);
+			app.use("*", serveStatic({ root: opts.publicDir }));
+		}
+		app.all("*", (c) => this.handler(c.req.raw));
+		await crossPlatformServe(config, port, app, runtime);
+	}
+
 	public async diagnostics(): Promise<RegistryDiagnostics> {
 		const candidates = [
 			this.#runtimeServerlessPromise,
@@ -307,12 +355,12 @@ export class Registry<A extends RegistryActors> {
 				.then(async ({ runtime, registry, serveConfig }) => {
 					await runtime.serveRegistry(registry, serveConfig);
 				})
-				.catch((err) => {
+				.catch((error) => {
 					// Always-attached catch so the stored promise never leaves a
 					// rejection unhandled. Downstream awaits (e.g. #runShutdown's
 					// Promise.race) attach their own catches and still observe
 					// resolution via the race.
-					logger().warn({ err }, "runtime registry serve errored");
+					logger().warn({ error }, "runtime registry serve errored");
 				});
 			// Install signal handlers once an envoy lifecycle has begun. Only
 			// Mode A ever reaches here. Mode B (handler(request)) intentionally
@@ -364,18 +412,19 @@ export class Registry<A extends RegistryActors> {
 	): void {
 		if (this.#shutdownInFlight !== null) {
 			// Second delivery of the same (or another) shutdown signal.
-			// Remove our handler only (preserving any user-installed listeners)
-			// and re-raise so Node proceeds with its default exit path.
+			// Remove our handler only, preserving any user-installed listeners.
+			// PID 1 must exit directly because re-raised default signals can be
+			// swallowed by the container signal path.
 			this.#removeSignalHandlers();
-			process.kill(process.pid, signal);
+			finishShutdownSignal(signal);
 			return;
 		}
 		this.#shutdownInFlight = this.#runShutdown(
 			signal,
 			config,
 			configuredRegistryPromise,
-		).catch((err) => {
-			logger().warn({ err }, "shutdown error");
+		).catch((error) => {
+			logger().warn({ error }, "shutdown error");
 		});
 	}
 
@@ -384,11 +433,13 @@ export class Registry<A extends RegistryActors> {
 		config: RegistryConfig,
 		configuredRegistryPromise: ReturnType<typeof buildConfiguredRegistry>,
 	): Promise<void> {
-		const gracePeriodMs = config.shutdown?.gracePeriodMs ?? 30_000;
+		const gracePeriodMs =
+			config.shutdown?.gracePeriodMs ??
+			(await this.#actorStopThresholdMs(configuredRegistryPromise)) ??
+			30 * 60 * 1000;
 		// Race the entire drain sequence (both modes + serve promise) against
-		// a single grace ceiling. Without this, each mode's Rust-side drain
-		// (20s) could stack sequentially and blow past gracePeriodMs before
-		// we re-raise the signal.
+		// a single grace ceiling. By default, this uses the engine-provided
+		// actor stop threshold, matching Pegboard's hard cutoff for actors.
 		const drain = async () => {
 			// Shut down every live `CoreRegistry` we know about. Mode A
 			// (`start()`) and Mode B (`handler()`) each build a separate
@@ -400,9 +451,9 @@ export class Registry<A extends RegistryActors> {
 						const { runtime, registry } =
 							await configuredRegistryPromise;
 						await runtime.shutdownRegistry(registry);
-					} catch (err) {
+					} catch (error) {
 						logger().warn(
-							{ err },
+							{ error },
 							"runtime registry shutdown errored (mode A)",
 						);
 					}
@@ -418,7 +469,7 @@ export class Registry<A extends RegistryActors> {
 							await runtime.shutdownRegistry(registry);
 						} catch (err) {
 							logger().warn(
-								{ err },
+								{ error: err },
 								"runtime registry shutdown errored (mode B)",
 							);
 						}
@@ -442,7 +493,30 @@ export class Registry<A extends RegistryActors> {
 			),
 		]);
 		this.#removeSignalHandlers();
-		process.kill(process.pid, signal);
+		finishShutdownSignal(signal);
+	}
+
+	async #actorStopThresholdMs(
+		configuredRegistryPromise: ReturnType<typeof buildConfiguredRegistry>,
+	): Promise<number | undefined> {
+		try {
+			const { runtime, registry } = await configuredRegistryPromise;
+			const thresholdMs =
+				await runtime.registryActorStopThresholdMs?.(registry);
+			if (
+				thresholdMs !== undefined &&
+				Number.isFinite(thresholdMs) &&
+				thresholdMs > 0
+			) {
+				return thresholdMs;
+			}
+		} catch (err) {
+			logger().warn(
+				{ err },
+				"failed to read actor stop threshold for shutdown grace",
+			);
+		}
+		return undefined;
 	}
 
 	#removeSignalHandlers(): void {

package/src/registry/napi-runtime.ts

@@ -212,6 +212,15 @@ export class NapiCoreRuntime implements CoreRuntime {
 		};
 	}
 
+	async registryActorStopThresholdMs(
+		registry: RegistryHandle,
+	): Promise<number | undefined> {
+		return (
+			(await asNativeRegistry(registry).actorStopThresholdMs()) ??
+			undefined
+		);
+	}
+
 	async handleServerlessRequest(
 		registry: RegistryHandle,
 		req: RuntimeServerlessRequest,
@@ -414,6 +423,12 @@ export class NapiCoreRuntime implements CoreRuntime {
 		asNativeActorContext(ctx).waitUntil(promise);
 	}
 
+	async actorWaitForTrackedShutdownWork(
+		ctx: ActorContextHandle,
+	): Promise<boolean> {
+		return await asNativeActorContext(ctx).waitForTrackedShutdownWork();
+	}
+
 	actorKeepAwake(ctx: ActorContextHandle, promise: Promise<unknown>): void {
 		asNativeActorContext(ctx).keepAwake(promise);
 	}