rivetkit 2.3.0-rc.5 → 2.3.0-rc.6

package/src/drivers/engine/actor-driver.ts

@@ -78,7 +78,7 @@ import {
 import { logger } from "./log";
 
 const ENVOY_SSE_PING_INTERVAL = 1000;
-const ENVOY_STOP_WAIT_MS = 15_000;
+const FALLBACK_ENVOY_STOP_WAIT_MS = 15_000;
 const INITIAL_SLEEP_TIMEOUT_MS = 250;
 const REMOTE_ACK_HOOK_QUERY_PARAM = "__rivetkitAckHook";
 
@@ -847,6 +847,7 @@ export class EngineActorDriver implements ActorDriver {
 			return;
 		}
 		this.#isShuttingDown = true;
+		const envoyStopWaitMs = this.#envoyStopWaitMs();
 
 		logger().info({ msg: "stopping engine actor driver", immediate });
 		if (!immediate) {
@@ -861,7 +862,7 @@ export class EngineActorDriver implements ActorDriver {
 				this.startSleep(actorId);
 			}
 
-			const actorSleepDeadline = Date.now() + ENVOY_STOP_WAIT_MS;
+			const actorSleepDeadline = Date.now() + envoyStopWaitMs;
 			while (this.#actors.size > 0 && Date.now() < actorSleepDeadline) {
 				await new Promise((resolve) => setTimeout(resolve, 50));
 			}
@@ -870,7 +871,7 @@ export class EngineActorDriver implements ActorDriver {
 				logger().warn({
 					msg: "timed out waiting for actors to stop before envoy drain",
 					remainingActors: this.#actors.size,
-					waitMs: ENVOY_STOP_WAIT_MS,
+					waitMs: envoyStopWaitMs,
 				});
 				// Snapshot so concurrent removals from `stopActor` do not
 				// invalidate the iterator.
@@ -912,13 +913,13 @@ export class EngineActorDriver implements ActorDriver {
 		const stopped = await Promise.race([
 			this.#envoyStopped.promise.then(() => true),
 			new Promise<false>((resolve) =>
-				setTimeout(() => resolve(false), ENVOY_STOP_WAIT_MS),
+				setTimeout(() => resolve(false), envoyStopWaitMs),
 			),
 		]);
 		if (!stopped) {
 			logger().warn({
 				msg: "timed out waiting for envoy shutdown",
-				waitMs: ENVOY_STOP_WAIT_MS,
+				waitMs: envoyStopWaitMs,
 			});
 		}
 
@@ -929,6 +930,16 @@ export class EngineActorDriver implements ActorDriver {
 		await this.#envoy.started();
 	}
 
+	#envoyStopWaitMs(): number {
+		const actorStopThreshold = Number(
+			this.#envoy.getProtocolMetadata()?.actorStopThreshold,
+		);
+		if (Number.isFinite(actorStopThreshold) && actorStopThreshold > 0) {
+			return actorStopThreshold;
+		}
+		return FALLBACK_ENVOY_STOP_WAIT_MS;
+	}
+
 	async #bindHibernatableConnectSocket(
 		binding: HibernatableConnectBinding,
 		isRestoringHibernatable: boolean,
@@ -2190,12 +2201,7 @@ export class EngineActorDriver implements ActorDriver {
 					isRestoringHibernatable,
 				);
 			} catch (error) {
-				const { group, code } = deconstructError(
-					error,
-					logger(),
-					{},
-					false,
-				);
+				const { group, code } = deconstructError(error, false);
 				logger().error({
 					msg: "failed to bind dynamic hibernatable websocket",
 					actorId,
@@ -2233,12 +2239,7 @@ export class EngineActorDriver implements ActorDriver {
 				},
 			);
 		} catch (error) {
-			const { group, code } = deconstructError(
-				error,
-				logger(),
-				{},
-				false,
-			);
+			const { group, code } = deconstructError(error, false);
 			logger().error({
 				msg: "failed to open dynamic websocket",
 				actorId,

package/src/registry/config/index.ts

@@ -264,8 +264,8 @@ export const RegistryConfigSchema = z
 			.object({
 				/**
 				 * Wait this many milliseconds for the serve promise to resolve
-				 * after calling `CoreRegistry::shutdown()`. Defaults to 30s,
-				 * matching Kubernetes `terminationGracePeriodSeconds`.
+				 * after calling `CoreRegistry::shutdown()`. Defaults to the
+				 * engine-provided actor stop threshold once the envoy connects.
 				 *
 				 * Must be >= rivetkit-core's drain timeout (20s) + margin.
 				 */
@@ -273,8 +273,7 @@ export const RegistryConfigSchema = z
 					.number()
 					.int()
 					.min(1_000)
-					.optional()
-					.default(30_000),
+					.optional(),
 				/**
 				 * If true, rivetkit will not install SIGINT/SIGTERM handlers.
 				 * Use when the host application owns signal policy and will
@@ -284,7 +283,6 @@ export const RegistryConfigSchema = z
 			})
 			.optional()
 			.default(() => ({
-				gracePeriodMs: 30_000,
 				disableSignalHandlers: false,
 			})),
 	})

package/src/registry/index.ts

@@ -13,6 +13,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
@@ -22,13 +38,15 @@ export interface ServerlessHandler {
 	fetch: FetchHandler;
 }
 
-export interface RegistryDiagnostics {
-	mode: string;
-	envoyActiveActorCount?: number | null;
+export interface RegistryRoutes {
+	health(): Promise<Response>;
+	metadata(): Promise<Response>;
+	prometheusMetrics(request?: Request): Promise<Response>;
 }
 
 export class Registry<A extends RegistryActors> {
 	#config: RegistryConfigInput<A>;
+	public readonly routes: RegistryRoutes;
 
 	get config(): RegistryConfigInput<A> {
 		return this.#config;
@@ -49,6 +67,12 @@ export class Registry<A extends RegistryActors> {
 
 	constructor(config: RegistryConfigInput<A>) {
 		this.#config = config;
+		this.routes = {
+			health: () => this.#healthRoute(),
+			metadata: () => this.#metadataRoute(),
+			prometheusMetrics: (request?: Request) =>
+				this.#prometheusMetricsRoute(request),
+		};
 	}
 
 	#ensureServerlessPoolConfigured(config: RegistryConfig): Promise<void> | undefined {
@@ -279,7 +303,92 @@ export class Registry<A extends RegistryActors> {
 		};
 	}
 
-	public async diagnostics(): Promise<RegistryDiagnostics> {
+	/**
+	 * Returns a health response suitable for mounting in a user-owned router.
+	 */
+	async #healthRoute(): Promise<Response> {
+		const configured = await this.#activeConfiguredRegistry();
+		if (!configured) {
+			return jsonRouteResponse(503, {
+				status: "not_started",
+				runtime: "rivetkit",
+				version: VERSION,
+			});
+		}
+
+		const { runtime, registry } = configured;
+		if (!runtime.registryHealth) {
+			return jsonRouteResponse(501, {
+				status: "unsupported",
+				runtime: "rivetkit",
+				version: VERSION,
+			});
+		}
+
+		const response = await runtime.registryHealth(registry);
+		return new Response(new Uint8Array(response.body), {
+			status: response.status,
+			headers: response.headers,
+		});
+	}
+
+	/**
+	 * Returns serverless metadata suitable for mounting in a user-owned router.
+	 */
+	async #metadataRoute(): Promise<Response> {
+		const configured = await this.#activeConfiguredRegistry();
+		if (!configured) {
+			return new Response("registry not started\n", {
+				status: 503,
+				headers: { "content-type": "text/plain; charset=utf-8" },
+			});
+		}
+
+		const { runtime, registry } = configured;
+		if (!runtime.registryMetadata) {
+			return new Response("metadata is not supported by this runtime\n", {
+				status: 501,
+				headers: { "content-type": "text/plain; charset=utf-8" },
+			});
+		}
+
+		const response = await runtime.registryMetadata(registry);
+		return new Response(new Uint8Array(response.body), {
+			status: response.status,
+			headers: response.headers,
+		});
+	}
+
+	/**
+	 * Returns a Prometheus metrics response suitable for mounting in a user-owned router.
+	 */
+	async #prometheusMetricsRoute(_request?: Request): Promise<Response> {
+		const configured = await this.#activeConfiguredRegistry();
+		if (!configured) {
+			return new Response("registry not started\n", {
+				status: 503,
+				headers: { "content-type": "text/plain; charset=utf-8" },
+			});
+		}
+
+		const { runtime, registry } = configured;
+		if (!runtime.registryMetrics) {
+			return new Response("metrics are not supported by this runtime\n", {
+				status: 501,
+				headers: { "content-type": "text/plain; charset=utf-8" },
+			});
+		}
+
+		const response = await runtime.registryMetrics(registry);
+		return new Response(new Uint8Array(response.body), {
+			status: response.status,
+			headers: response.headers,
+		});
+	}
+
+	async #activeConfiguredRegistry(): Promise<
+		Awaited<ReturnType<typeof buildConfiguredRegistry>> | undefined
+	> {
 		const candidates = [
 			this.#runtimeServerlessPromise,
 			this.#runtimeServeConfiguredPromise,
@@ -287,13 +396,8 @@ export class Registry<A extends RegistryActors> {
 			candidate !== undefined
 		);
 
-		for (const candidate of candidates) {
-			const { runtime, registry } = await candidate;
-			const diagnostics = await runtime.registryDiagnostics?.(registry);
-			if (diagnostics) return diagnostics;
-		}
-
-		return { mode: "not_started", envoyActiveActorCount: null };
+		if (candidates.length === 0) return undefined;
+		return await candidates[0]!;
 	}
 
 	/**
@@ -364,10 +468,11 @@ 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(
@@ -384,11 +489,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
@@ -442,7 +549,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 {
@@ -524,6 +654,13 @@ function isServerlessMetadataRequest(request: Request, basePath: string): boolea
 	return parsed.pathname === `${normalizedBase}/metadata`;
 }
 
+function jsonRouteResponse(status: number, body: unknown): Response {
+	return new Response(JSON.stringify(body), {
+		status,
+		headers: { "content-type": "application/json" },
+	});
+}
+
 export function setup<A extends RegistryActors>(
 	input: RegistryConfigInput<A>,
 ): Registry<A> {

package/src/registry/napi-runtime.ts

@@ -24,7 +24,7 @@ import type {
 	RuntimeQueueTryNextBatchOptions,
 	RuntimeQueueWaitOptions,
 	RuntimeRequestSaveOpts,
-	RuntimeRegistryDiagnostics,
+	RuntimeRegistryRouteResponse,
 	RuntimeServeConfig,
 	RuntimeServerlessRequest,
 	RuntimeServerlessResponseHead,
@@ -202,13 +202,45 @@ export class NapiCoreRuntime implements CoreRuntime {
 		await asNativeRegistry(registry).shutdown();
 	}
 
-	async registryDiagnostics(
+	async registryActorStopThresholdMs(
 		registry: RegistryHandle,
-	): Promise<RuntimeRegistryDiagnostics> {
-		const diagnostics = await asNativeRegistry(registry).diagnostics();
+	): Promise<number | undefined> {
+		return (
+			(await asNativeRegistry(registry).actorStopThresholdMs()) ??
+			undefined
+		);
+	}
+
+	async registryHealth(
+		registry: RegistryHandle,
+	): Promise<RuntimeRegistryRouteResponse> {
+		const response = await asNativeRegistry(registry).health();
+		return {
+			status: response.status,
+			headers: response.headers,
+			body: response.body,
+		};
+	}
+
+	async registryMetadata(
+		registry: RegistryHandle,
+	): Promise<RuntimeRegistryRouteResponse> {
+		const response = asNativeRegistry(registry).metadata();
 		return {
-			mode: diagnostics.mode,
-			envoyActiveActorCount: diagnostics.envoyActiveActorCount,
+			status: response.status,
+			headers: response.headers,
+			body: response.body,
+		};
+	}
+
+	async registryMetrics(
+		registry: RegistryHandle,
+	): Promise<RuntimeRegistryRouteResponse> {
+		const response = asNativeRegistry(registry).metrics();
+		return {
+			status: response.status,
+			headers: response.headers,
+			body: response.body,
 		};
 	}
 
@@ -414,6 +446,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);
 	}

package/src/registry/native.ts

@@ -29,22 +29,9 @@ import {
 	createClientWithDriver,
 } from "@/client/client";
 import { convertRegistryConfigToClientConfig } from "@/client/config";
-import {
-	HEADER_CONN_PARAMS,
-	HEADER_ENCODING,
-} from "@/common/actor-router-consts";
-import type * as protocol from "@/common/client-protocol";
-import {
-	CURRENT_VERSION as CLIENT_PROTOCOL_CURRENT_VERSION,
-	HTTP_RESPONSE_ERROR_VERSIONED,
-} from "@/common/client-protocol-versioned";
-import {
-	type HttpResponseError as HttpResponseErrorJson,
-	HttpResponseErrorSchema,
-} from "@/common/client-protocol-zod";
+import { HEADER_CONN_PARAMS } from "@/common/actor-router-consts";
 import type { AnyDatabaseProvider } from "@/common/database/config";
 import { wrapJsNativeDatabase } from "@/common/database/native-database";
-import type { Encoding } from "@/common/encoding";
 import { decodeWorkflowHistoryTransport } from "@/common/inspector-transport";
 import { deconstructError, stringifyError } from "@/common/utils";
 import type {
@@ -61,11 +48,9 @@ import type {
 	SqliteBackend,
 } from "@/registry/config";
 import {
-	contentTypeForEncoding,
 	decodeCborCompat,
 	decodeCborJsonCompat,
 	encodeCborCompat,
-	serializeWithEncoding,
 } from "@/serde";
 import { getEnvUniversal, VERSION } from "@/utils";
 import { logger } from "./log";
@@ -269,9 +254,6 @@ type NativeDatabaseClientState = {
 type NativeActorRuntimeState = {
 	sql?: ReturnType<typeof wrapJsNativeDatabase>;
 	databaseClient?: NativeDatabaseClientState;
-	keepAwakeCount?: number;
-	deferSleepCleanupUntilKeepAwakeIdle?: boolean;
-	deferredSleepCleanupActorCtx?: ActorContextHandleAdapter;
 	varsInitialized?: boolean;
 	vars?: unknown;
 	destroyGate?: NativeDestroyGate;
@@ -424,26 +406,12 @@ async function cleanupNativeSleepRuntimeState(
 	runtime: CoreRuntime,
 	ctx: ActorContextHandle,
 ): Promise<void> {
+	await runtime.actorWaitForTrackedShutdownWork(ctx);
 	await closeNativeDatabaseClient(runtime, ctx);
 	await closeNativeSqlDatabase(runtime, ctx);
 	clearNativeRuntimeState(runtime, ctx);
 }
 
-async function cleanupDeferredNativeSleepRuntimeState(
-	runtime: CoreRuntime,
-	ctx: ActorContextHandle,
-	runtimeState: NativeActorRuntimeState,
-): Promise<void> {
-	runtimeState.deferSleepCleanupUntilKeepAwakeIdle = false;
-	const actorCtx = runtimeState.deferredSleepCleanupActorCtx;
-	runtimeState.deferredSleepCleanupActorCtx = undefined;
-	try {
-		await cleanupNativeSleepRuntimeState(runtime, ctx);
-	} finally {
-		await actorCtx?.dispose();
-	}
-}
-
 function closeNativeSqlDatabase(
 	runtime: CoreRuntime,
 	ctx: ActorContextHandle,
@@ -670,28 +638,7 @@ function isStructuredBridgeError(
 function encodeNativeCallbackError(error: unknown): Error {
 	const structuredError = isStructuredBridgeError(error)
 		? error
-		: deconstructError(error, logger(), {
-				bridge: "native_callback",
-			});
-	let stack: string | undefined;
-	if (error instanceof Error) {
-		try {
-			stack = error.stack;
-		} catch {
-			stack = undefined;
-		}
-	}
-
-	logger().warn({
-		msg: "native callback error encoded for bridge",
-		group: structuredError.group,
-		code: structuredError.code,
-		message: structuredError.message,
-		metadata: structuredError.metadata,
-		originalError: stringifyError(error),
-		stack,
-		bridge: "native_callback",
-	});
+		: deconstructError(error, true);
 
 	const bridgeError = new Error(encodeBridgeRivetError(structuredError), {
 		cause: error instanceof Error ? error : undefined,
@@ -2807,11 +2754,6 @@ export class ActorContextHandleAdapter {
 	}
 
 	keepAwake<T>(promise: Promise<T>): Promise<T> {
-		const runtimeState = getNativeRuntimeState(this.#runtime, this.#ctx);
-		// Increment before native registration so sleep cleanup observes JS work
-		// even if the promise settles immediately.
-		runtimeState.keepAwakeCount = (runtimeState.keepAwakeCount ?? 0) + 1;
-		let registered = false;
 		const trackedPromise = Promise.resolve(promise)
 			.catch((error) => {
 				logger().warn({
@@ -2819,36 +2761,12 @@ export class ActorContextHandleAdapter {
 					error: stringifyError(error),
 				});
 			})
-			.finally(async () => {
-				if (!registered) {
-					return;
-				}
-				runtimeState.keepAwakeCount = Math.max(
-					(runtimeState.keepAwakeCount ?? 1) - 1,
-					0,
-				);
-				if (
-					runtimeState.keepAwakeCount === 0 &&
-					runtimeState.deferSleepCleanupUntilKeepAwakeIdle
-				) {
-					await cleanupDeferredNativeSleepRuntimeState(
-						this.#runtime,
-						this.#ctx,
-						runtimeState,
-					);
-				}
-			})
 			.then(() => null);
 		try {
 			callNativeSync(() =>
 				this.#runtime.actorKeepAwake(this.#ctx, trackedPromise),
 			);
-			registered = true;
 		} catch (error) {
-			runtimeState.keepAwakeCount = Math.max(
-				(runtimeState.keepAwakeCount ?? 1) - 1,
-				0,
-			);
 			if (!isClosedTaskRegistrationError(error)) {
 				throw error;
 			}
@@ -3265,52 +3183,6 @@ function withConnContext(
 	);
 }
 
-function buildNativeRequestErrorResponse(
-	encoding: Encoding,
-	path: string,
-	error: unknown,
-): Response {
-	const { statusCode, group, code, message, metadata } = deconstructError(
-		error,
-		logger(),
-		{
-			path,
-			runtime: "native",
-		},
-		false,
-	);
-	const body = serializeWithEncoding<
-		protocol.HttpResponseError,
-		HttpResponseErrorJson,
-		{ group: string; code: string; message: string; metadata?: unknown }
-	>(
-		encoding,
-		{ group, code, message, metadata },
-		HTTP_RESPONSE_ERROR_VERSIONED,
-		CLIENT_PROTOCOL_CURRENT_VERSION,
-		HttpResponseErrorSchema,
-		(value) => value,
-		(value) => ({
-			group: value.group,
-			code: value.code,
-			message: value.message,
-			metadata:
-				value.metadata === undefined
-					? null
-					: runtimeBytesToArrayBuffer(
-							encodeCborCompat(value.metadata),
-						),
-		}),
-	);
-
-	return new Response(body, {
-		status: statusCode,
-		headers: {
-			"Content-Type": contentTypeForEncoding(encoding),
-		},
-	});
-}
-
 function buildActorConfig(
 	definition: AnyActorDefinition,
 	registryConfig: RegistryConfig,
@@ -4013,12 +3885,9 @@ export function buildNativeFactory(
 						}
 					}
 				} finally {
-					const runtimeState = getNativeRuntimeState(runtime, ctx);
-					if ((runtimeState.keepAwakeCount ?? 0) > 0) {
-						runtimeState.deferSleepCleanupUntilKeepAwakeIdle = true;
-						runtimeState.deferredSleepCleanupActorCtx = actorCtx;
-					} else {
+					try {
 						await cleanupNativeSleepRuntimeState(runtime, ctx);
+					} finally {
 						await actorCtx.dispose();
 					}
 				}
@@ -4366,22 +4235,6 @@ export function buildNativeFactory(
 							);
 						}
 						return await toRuntimeHttpResponse(response);
-					} catch (error) {
-						const encodingHeader =
-							jsRequest.headers.get(HEADER_ENCODING);
-						const encoding: Encoding =
-							encodingHeader === "cbor" ||
-							encodingHeader === "bare"
-								? encodingHeader
-								: "json";
-						const path = new URL(jsRequest.url).pathname;
-						return await toRuntimeHttpResponse(
-							buildNativeRequestErrorResponse(
-								encoding,
-								path,
-								error,
-							),
-						);
 					} finally {
 						await requestCtx?.dispose();
 						if (conn) {