rivetkit 2.3.11-rc.9 → 2.3.11

package/src/registry/run-handler-coordinator.test.ts

@@ -0,0 +1,252 @@
+import { describe, expect, test, vi } from "vitest";
+import { defineRunHandler, type RunControl } from "@/actor/config";
+import { RunHandlerCoordinator } from "./run-handler-coordinator";
+
+class Deferred<T = void> {
+	readonly promise: Promise<T>;
+	resolve!: (value: T | PromiseLike<T>) => void;
+	reject!: (reason?: unknown) => void;
+
+	constructor() {
+		this.promise = new Promise<T>((resolve, reject) => {
+			this.resolve = resolve;
+			this.reject = reject;
+		});
+	}
+}
+
+function createCoordinator() {
+	let control: RunControl | undefined;
+	const dispose = vi.fn();
+	const createInspector = vi.fn((context: { control: RunControl }) => {
+		control = context.control;
+		return {
+			inspector: {
+				workflow: {
+					getHistory: () => null,
+					getState: async () => null,
+					onHistoryUpdated: () => () => {},
+					replayFromStep: async () => null,
+				},
+			},
+			dispose,
+		};
+	});
+	const run = defineRunHandler(async () => {}, {
+		inspectorKind: "workflow",
+		createInspector,
+	});
+	const coordinator = new RunHandlerCoordinator(run);
+	const restart = vi.fn();
+	const inspector = coordinator.getInspector("actor-1", restart);
+
+	return {
+		control: control as unknown as RunControl,
+		coordinator,
+		createInspector,
+		dispose,
+		inspector,
+		restart,
+	};
+}
+
+describe("RunHandlerCoordinator", () => {
+	test("fails explicitly when a JavaScript factory omits required workflow controls", async () => {
+		const coordinator = new RunHandlerCoordinator(
+			defineRunHandler(async () => {}, {
+				inspectorKind: "workflow",
+				createInspector: (() => ({
+					inspector: { workflow: { getHistory: () => null } },
+				})) as never,
+			}),
+		);
+
+		await expect(
+			coordinator.run(
+				"actor-1",
+				() => {},
+				async () => {},
+			),
+		).rejects.toThrow(
+			"createInspector returned an invalid workflow adapter for actor actor-1",
+		);
+	});
+
+	test("initializes and disposes Inspector state even when it is never queried", async () => {
+		const dispose = vi.fn();
+		const createInspector = vi.fn(() => ({
+			inspector: {
+				workflow: {
+					getHistory: () => null,
+					getState: async () => null,
+					onHistoryUpdated: () => () => {},
+					replayFromStep: async () => null,
+				},
+			},
+			dispose,
+		}));
+		const coordinator = new RunHandlerCoordinator(
+			defineRunHandler(async () => {}, {
+				inspectorKind: "workflow",
+				createInspector,
+			}),
+		);
+
+		await coordinator.run(
+			"actor-1",
+			() => {},
+			async () => {},
+		);
+		expect(createInspector).toHaveBeenCalledOnce();
+		coordinator.destroy("actor-1");
+		expect(dispose).toHaveBeenCalledOnce();
+	});
+
+	test("creates one inspector per live actor and disposes it exactly once", () => {
+		const subject = createCoordinator();
+		expect(
+			subject.coordinator.getInspector("actor-1", subject.restart),
+		).toBe(subject.inspector);
+		expect(subject.createInspector).toHaveBeenCalledOnce();
+
+		subject.coordinator.destroy("actor-1");
+		subject.coordinator.destroy("actor-1");
+		expect(subject.dispose).toHaveBeenCalledOnce();
+
+		const recreated = subject.coordinator.getInspector(
+			"actor-1",
+			subject.restart,
+		);
+		expect(recreated).not.toBe(subject.inspector);
+		expect(subject.createInspector).toHaveBeenCalledTimes(2);
+	});
+
+	test("rejects replay while a run is active", async () => {
+		const subject = createCoordinator();
+		const active = new Deferred();
+		const started = new Deferred();
+		const running = subject.coordinator.run(
+			"actor-1",
+			subject.restart,
+			async () => {
+				started.resolve();
+				await active.promise;
+			},
+		);
+		await started.promise;
+
+		await expect(
+			subject.control.run.withInactive(
+				{ restartOnSuccess: true },
+				async () => {},
+			),
+		).rejects.toMatchObject({
+			group: "actor",
+			code: "run_handler_unavailable",
+		});
+
+		active.resolve();
+		await running;
+	});
+
+	test("uses a start queued during replay as the single restart", async () => {
+		const subject = createCoordinator();
+		const replay = new Deferred();
+		const replayStarted = new Deferred();
+		const events: string[] = [];
+		subject.restart.mockImplementation(() => {
+			events.push("restart");
+		});
+
+		const exclusive = subject.control.run.withInactive(
+			{ restartOnSuccess: true },
+			async () => {
+				events.push("replay");
+				replayStarted.resolve();
+				await replay.promise;
+			},
+		);
+		await replayStarted.promise;
+
+		const queued = subject.coordinator.run(
+			"actor-1",
+			subject.restart,
+			async () => {
+				events.push("run");
+			},
+		);
+		await Promise.resolve();
+		expect(events).toEqual(["replay"]);
+
+		replay.resolve();
+		await exclusive;
+		await queued;
+		expect(events).toEqual(["replay", "run"]);
+		expect(subject.restart).not.toHaveBeenCalled();
+	});
+
+	test("drops starts queued behind a failed replay and leaves the run inactive", async () => {
+		const subject = createCoordinator();
+		const replay = new Deferred();
+		const replayStarted = new Deferred();
+		const run = vi.fn();
+		const expected = new Error("rewrite failed");
+
+		const exclusive = subject.control.run.withInactive(
+			{ restartOnSuccess: true },
+			async () => {
+				replayStarted.resolve();
+				await replay.promise;
+				throw expected;
+			},
+		);
+		await replayStarted.promise;
+		const queued = subject.coordinator.run("actor-1", subject.restart, run);
+
+		replay.resolve();
+		await expect(exclusive).rejects.toBe(expected);
+		await expect(queued).resolves.toBe("suppressed");
+		expect(run).not.toHaveBeenCalled();
+		expect(subject.restart).not.toHaveBeenCalled();
+	});
+
+	test("restores a durable wake suppressed by a failed replay", async () => {
+		const subject = createCoordinator();
+		const replay = new Deferred();
+		const replayStarted = new Deferred();
+		const run = vi.fn();
+
+		const exclusive = subject.control.run.withInactive({}, async () => {
+			replayStarted.resolve();
+			await replay.promise;
+			throw new Error("rewrite failed");
+		});
+		await replayStarted.promise;
+		const queued = subject.coordinator.run("actor-1", subject.restart, run);
+
+		replay.resolve();
+		await expect(exclusive).rejects.toThrow("rewrite failed");
+		await expect(queued).resolves.toBe("suppressed");
+		expect(run).not.toHaveBeenCalled();
+	});
+
+	test("allows only one concurrent replay and closes old controls on destroy", async () => {
+		const subject = createCoordinator();
+		const replay = new Deferred();
+		const first = subject.control.run.withInactive({}, async () => {
+			await replay.promise;
+		});
+
+		await expect(
+			subject.control.run.withInactive({}, async () => {}),
+		).rejects.toMatchObject({ code: "run_handler_unavailable" });
+		subject.coordinator.destroy("actor-1");
+		replay.resolve();
+		await expect(first).rejects.toMatchObject({
+			code: "run_handler_unavailable",
+		});
+		await expect(
+			subject.control.run.withInactive({}, async () => {}),
+		).rejects.toMatchObject({ code: "run_handler_unavailable" });
+	});
+});

package/src/registry/run-handler-coordinator.ts

@@ -0,0 +1,219 @@
+import {
+	createRunInspector,
+	getRunInspectorKind,
+	type RunControl,
+	type RunInspectorConfig,
+	type RunInspectorFactoryResult,
+	type RunInspectorKind,
+} from "@/actor/config";
+import { RivetError } from "@/actor/errors";
+
+type RunHandler = ((...args: any[]) => any) | { run: (...args: any[]) => any };
+
+interface ActorRunState {
+	active: boolean;
+	closed: boolean;
+	exclusive: boolean;
+	exclusiveGeneration: number;
+	exclusiveOutcome?: "success" | "failure";
+	queued: number;
+	restart?: () => void | Promise<void>;
+	inspectorInitialized: boolean;
+	inspector?: RunInspectorFactoryResult;
+	waiters: Set<() => void>;
+}
+
+type RunHandlerOutcome = "ran" | "suppressed" | "closed";
+
+function runUnavailable(actorId: string, reason: string): RivetError {
+	return new RivetError(
+		"actor",
+		"run_handler_unavailable",
+		`Run handler control is unavailable for actor ${actorId}: ${reason}.`,
+		{
+			public: true,
+			statusCode: 409,
+			metadata: { actorId, reason },
+		},
+	);
+}
+
+function notifyStateChanged(state: ActorRunState): void {
+	const waiters = [...state.waiters];
+	state.waiters.clear();
+	for (const waiter of waiters) {
+		waiter();
+	}
+}
+
+function waitForStateChange(state: ActorRunState): Promise<void> {
+	return new Promise((resolve) => {
+		state.waiters.add(resolve);
+	});
+}
+
+/** Coordinates a run handler and its Inspector controls within one registry. */
+export class RunHandlerCoordinator {
+	readonly #run: RunHandler | undefined;
+	readonly #states = new Map<string, ActorRunState>();
+
+	constructor(run: RunHandler | undefined) {
+		this.#run = run;
+	}
+
+	get inspectorKind(): RunInspectorKind | undefined {
+		return getRunInspectorKind(this.#run);
+	}
+
+	async run(
+		actorId: string,
+		restart: () => void | Promise<void>,
+		callback: () => void | Promise<void>,
+	): Promise<RunHandlerOutcome> {
+		const state = this.#getOrCreate(actorId);
+		state.restart = restart;
+		this.#initializeInspector(actorId, state);
+		state.queued += 1;
+		const blockedGeneration = state.exclusive
+			? state.exclusiveGeneration
+			: undefined;
+
+		try {
+			while (!state.closed && (state.exclusive || state.active)) {
+				await waitForStateChange(state);
+			}
+		} finally {
+			state.queued -= 1;
+		}
+
+		if (state.closed) {
+			return "closed";
+		}
+		if (
+			blockedGeneration !== undefined &&
+			state.exclusiveGeneration === blockedGeneration &&
+			state.exclusiveOutcome === "failure"
+		) {
+			return "suppressed";
+		}
+
+		state.active = true;
+		try {
+			await callback();
+			return "ran";
+		} finally {
+			state.active = false;
+			notifyStateChanged(state);
+		}
+	}
+
+	getInspector(
+		actorId: string,
+		restart: () => void | Promise<void>,
+	): RunInspectorConfig | undefined {
+		const state = this.#getOrCreate(actorId);
+		state.restart = restart;
+		this.#initializeInspector(actorId, state);
+		return state.inspector?.inspector;
+	}
+
+	destroy(actorId: string): void {
+		const state = this.#states.get(actorId);
+		if (!state) return;
+
+		state.closed = true;
+		state.inspector?.dispose?.();
+		state.inspector = undefined;
+		notifyStateChanged(state);
+		this.#states.delete(actorId);
+	}
+
+	#control(actorId: string, state: ActorRunState): RunControl {
+		return {
+			run: {
+				withInactive: async (options, callback) => {
+					if (state.closed) {
+						throw runUnavailable(actorId, "actor is destroyed");
+					}
+					if (state.active || state.queued > 0 || state.exclusive) {
+						throw runUnavailable(
+							actorId,
+							"the run handler is active or waiting to start",
+						);
+					}
+
+					state.exclusive = true;
+					state.exclusiveGeneration += 1;
+					state.exclusiveOutcome = undefined;
+					try {
+						const result = await callback();
+						if (state.closed) {
+							throw runUnavailable(actorId, "actor is destroyed");
+						}
+						state.exclusiveOutcome = "success";
+						if (options.restartOnSuccess && state.queued === 0) {
+							if (!state.restart) {
+								throw runUnavailable(
+									actorId,
+									"the runtime restart callback is not available",
+								);
+							}
+							await state.restart();
+						}
+						return result;
+					} catch (error) {
+						state.exclusiveOutcome = "failure";
+						throw error;
+					} finally {
+						state.exclusive = false;
+						notifyStateChanged(state);
+					}
+				},
+			},
+		};
+	}
+
+	#getOrCreate(actorId: string): ActorRunState {
+		let state = this.#states.get(actorId);
+		if (!state) {
+			state = {
+				active: false,
+				closed: false,
+				exclusive: false,
+				exclusiveGeneration: 0,
+				inspectorInitialized: false,
+				queued: 0,
+				waiters: new Set(),
+			};
+			this.#states.set(actorId, state);
+		}
+		return state;
+	}
+
+	#initializeInspector(actorId: string, state: ActorRunState): void {
+		if (state.inspectorInitialized) return;
+		const inspector = createRunInspector(this.#run, {
+			actorId,
+			control: this.#control(actorId, state),
+		});
+		if (this.inspectorKind === "workflow") {
+			const workflow =
+				inspector === undefined
+					? undefined
+					: inspector.inspector.workflow;
+			if (
+				!workflow ||
+				typeof workflow.getHistory !== "function" ||
+				typeof workflow.getState !== "function" ||
+				typeof workflow.onHistoryUpdated !== "function" ||
+				typeof workflow.replayFromStep !== "function"
+			) {
+				throw new TypeError(
+					`defineRunHandler createInspector returned an invalid workflow adapter for actor ${actorId}`,
+				);
+			}
+		}
+		state.inspector = inspector;
+		state.inspectorInitialized = true;
+	}
+}

package/src/registry/runtime.ts

@@ -1,4 +1,7 @@
-import type { SqliteNativeMetrics } from "@/common/database/config";
+import type {
+	SqliteNativeMetrics,
+	SqliteProfilingOptions,
+} from "@/common/database/config";
 import { stringifyError } from "@/common/utils";
 import type { RegistryConfig } from "./config";
 import { logger } from "./log";
@@ -15,6 +18,7 @@ export type ActorContextHandle = OpaqueHandle<"actorContext">;
 export type ConnHandle = OpaqueHandle<"conn">;
 export type WebSocketHandle = OpaqueHandle<"webSocket">;
 export type CancellationTokenHandle = OpaqueHandle<"cancellationToken">;
+export type ActorStateTransactionHandle = OpaqueHandle<"actorStateTransaction">;
 export type SqliteTransactionHandle = OpaqueHandle<"sqliteTransaction">;
 
 export type RuntimeBytes = Uint8Array;
@@ -255,6 +259,7 @@ export interface RuntimeSqlDatabase {
 	): Promise<RuntimeSqlRunResult>;
 	beginTransaction(
 		timeoutMs?: number,
+		name?: string,
 	): Promise<RuntimeSqlTransactionDatabase>;
 	metrics?(): SqliteNativeMetrics | null;
 	takeLastKvError?(): string | null;
@@ -276,6 +281,7 @@ export interface RuntimeActorConfig {
 	icon?: string;
 	hasDatabase?: boolean;
 	remoteSqlite?: boolean;
+	sqliteProfiling?: SqliteProfilingOptions;
 	enableActorRuntimeSocket?: boolean;
 	hasState?: boolean;
 	canHibernateWebsocket?: boolean;
@@ -479,6 +485,10 @@ export interface CoreRuntime {
 		ctx: ActorContextHandle,
 		timestampMs?: number | undefined | null,
 	): void;
+	actorSetRunWakeAt(
+		ctx: ActorContextHandle,
+		timestampMs?: number | undefined | null,
+	): Promise<void>;
 	actorRequestSave(
 		ctx: ActorContextHandle,
 		opts?: RuntimeRequestSaveOpts | undefined | null,
@@ -601,6 +611,7 @@ export interface CoreRuntime {
 	actorSqlBeginTransaction(
 		ctx: ActorContextHandle,
 		timeoutMs?: number,
+		name?: string,
 	): Promise<SqliteTransactionHandle>;
 	actorSqlTransactionExec(
 		transaction: SqliteTransactionHandle,
@@ -617,6 +628,22 @@ export interface CoreRuntime {
 	actorSqlTransactionRollback(
 		transaction: SqliteTransactionHandle,
 	): Promise<void>;
+	actorBeginStateTransaction(
+		ctx: ActorContextHandle,
+		timeoutMs?: number,
+	): Promise<ActorStateTransactionHandle>;
+	actorStateTransactionExecute(
+		transaction: ActorStateTransactionHandle,
+		sql: string,
+		params?: RuntimeSqlBindParams,
+	): Promise<RuntimeSqlExecuteResult>;
+	actorStateTransactionCommit(
+		transaction: ActorStateTransactionHandle,
+		payload: RuntimeStateDeltaPayload,
+	): Promise<void>;
+	actorStateTransactionRollback(
+		transaction: ActorStateTransactionHandle,
+	): Promise<void>;
 	actorSqlQuery(
 		ctx: ActorContextHandle,
 		sql: string,
@@ -656,6 +683,12 @@ export interface CoreRuntime {
 		options?: RuntimeQueueWaitOptions | undefined | null,
 		signal?: CancellationTokenHandle | undefined | null,
 	): Promise<void>;
+	actorQueueCompletePersisted(
+		ctx: ActorContextHandle,
+		messageId: bigint,
+		expectedName: string,
+		response?: RuntimeBytes | undefined | null,
+	): Promise<boolean>;
 	actorQueueEnqueueAndWait(
 		ctx: ActorContextHandle,
 		name: string,