@rivetkit/cloudflare-workers 2.0.2 → 2.0.4-rc.1

package/src/actor-driver.ts

@@ -1,21 +1,23 @@
+import invariant from "invariant";
 import type {
+	ActorKey,
+	ActorRouter,
 	AnyActorInstance as CoreAnyActorInstance,
 	RegistryConfig,
-	RunConfig,
-} from "@rivetkit/core";
-import {
-	createGenericConnDrivers,
-	GenericConnGlobalState,
-	lookupInRegistry,
-} from "@rivetkit/core";
-import type { Client } from "@rivetkit/core/client";
+} from "rivetkit";
+import { lookupInRegistry } from "rivetkit";
+import type { Client } from "rivetkit/client";
 import type {
 	ActorDriver,
 	AnyActorInstance,
 	ManagerDriver,
-} from "@rivetkit/core/driver-helpers";
-import invariant from "invariant";
-import { KEYS } from "./actor-handler-do";
+} from "rivetkit/driver-helpers";
+import { promiseWithResolvers } from "rivetkit/utils";
+import { logger } from "./log";
+import { kvDelete, kvGet, kvListPrefix, kvPut } from "./actor-kv";
+import { GLOBAL_KV_KEYS } from "./global-kv";
+import { getCloudflareAmbientEnv } from "./handler";
+import { parseActorId } from "./actor-id";
 
 interface DurableObjectGlobalState {
 	ctx: DurableObjectState;
@@ -28,17 +30,31 @@ interface DurableObjectGlobalState {
  * This allows for storing the actor context globally and looking it up by ID in `CloudflareActorsActorDriver`.
  */
 export class CloudflareDurableObjectGlobalState {
-	// Single map for all actor state
+	// Map of actor ID -> DO state
 	#dos: Map<string, DurableObjectGlobalState> = new Map();
 
-	getDOState(actorId: string): DurableObjectGlobalState {
-		const state = this.#dos.get(actorId);
-		invariant(state !== undefined, "durable object state not in global state");
+	// WeakMap of DO state -> ActorGlobalState for proper GC
+	#actors: WeakMap<DurableObjectState, ActorGlobalState> = new WeakMap();
+
+	getDOState(doId: string): DurableObjectGlobalState {
+		const state = this.#dos.get(doId);
+		invariant(
+			state !== undefined,
+			"durable object state not in global state",
+		);
 		return state;
 	}
 
-	setDOState(actorId: string, state: DurableObjectGlobalState) {
-		this.#dos.set(actorId, state);
+	setDOState(doId: string, state: DurableObjectGlobalState) {
+		this.#dos.set(doId, state);
+	}
+
+	getActorState(ctx: DurableObjectState): ActorGlobalState | undefined {
+		return this.#actors.get(ctx);
+	}
+
+	setActorState(ctx: DurableObjectState, actorState: ActorGlobalState): void {
+		this.#actors.set(ctx, actorState);
 	}
 }
 
@@ -46,79 +62,136 @@ export interface DriverContext {
 	state: DurableObjectState;
 }
 
-// Actor handler to track running instances
-class ActorHandler {
-	actor?: AnyActorInstance;
-	actorPromise?: PromiseWithResolvers<void> = Promise.withResolvers();
-	genericConnGlobalState = new GenericConnGlobalState();
+interface InitializedData {
+	name: string;
+	key: ActorKey;
+	generation: number;
+}
+
+interface LoadedActor {
+	actorRouter: ActorRouter;
+	actorDriver: ActorDriver;
+	generation: number;
+}
+
+// Actor global state to track running instances
+export class ActorGlobalState {
+	// Initialization state
+	initialized?: InitializedData;
+
+	// Loaded actor state
+	actor?: LoadedActor;
+	actorInstance?: AnyActorInstance;
+	actorPromise?: ReturnType<typeof promiseWithResolvers<void>>;
+
+	/**
+	 * Indicates if `startDestroy` has been called.
+	 *
+	 * This is stored in memory instead of SQLite since the destroy may be cancelled.
+	 *
+	 * See the corresponding `destroyed` property in SQLite metadata.
+	 */
+	destroying: boolean = false;
+
+	reset() {
+		this.initialized = undefined;
+		this.actor = undefined;
+		this.actorInstance = undefined;
+		this.actorPromise = undefined;
+		this.destroying = false;
+	}
 }
 
 export class CloudflareActorsActorDriver implements ActorDriver {
 	#registryConfig: RegistryConfig;
-	#runConfig: RunConfig;
 	#managerDriver: ManagerDriver;
 	#inlineClient: Client<any>;
 	#globalState: CloudflareDurableObjectGlobalState;
-	#actors: Map<string, ActorHandler> = new Map();
 
 	constructor(
 		registryConfig: RegistryConfig,
-		runConfig: RunConfig,
 		managerDriver: ManagerDriver,
 		inlineClient: Client<any>,
 		globalState: CloudflareDurableObjectGlobalState,
 	) {
 		this.#registryConfig = registryConfig;
-		this.#runConfig = runConfig;
 		this.#managerDriver = managerDriver;
 		this.#inlineClient = inlineClient;
 		this.#globalState = globalState;
 	}
 
 	#getDOCtx(actorId: string) {
-		return this.#globalState.getDOState(actorId).ctx;
+		// Parse actor ID to get DO ID
+		const [doId] = parseActorId(actorId);
+		return this.#globalState.getDOState(doId).ctx;
 	}
 
 	async loadActor(actorId: string): Promise<AnyActorInstance> {
+		// Parse actor ID to get DO ID and generation
+		const [doId, expectedGeneration] = parseActorId(actorId);
+
+		// Get the DO state
+		const doState = this.#globalState.getDOState(doId);
+
 		// Check if actor is already loaded
-		let handler = this.#actors.get(actorId);
-		if (handler) {
-			if (handler.actorPromise) await handler.actorPromise.promise;
-			if (!handler.actor) throw new Error("Actor should be loaded");
-			return handler.actor;
+		let actorState = this.#globalState.getActorState(doState.ctx);
+		if (actorState?.actorInstance) {
+			// Actor is already loaded, return it
+			return actorState.actorInstance;
 		}
 
-		// Create new actor handler
-		handler = new ActorHandler();
-		this.#actors.set(actorId, handler);
-
-		// Get the actor metadata from Durable Object storage
-		const doState = this.#globalState.getDOState(actorId);
-		const storage = doState.ctx.storage;
+		// Create new actor state if it doesn't exist
+		if (!actorState) {
+			actorState = new ActorGlobalState();
+			actorState.actorPromise = promiseWithResolvers((reason) => logger().warn({ msg: "unhandled actor promise rejection", reason }));
+			this.#globalState.setActorState(doState.ctx, actorState);
+		} else if (actorState.actorPromise) {
+			// Another request is already loading this actor, wait for it
+			await actorState.actorPromise.promise;
+			if (!actorState.actorInstance) {
+				throw new Error(
+					`Actor ${actorId} failed to load in concurrent request`,
+				);
+			}
+			return actorState.actorInstance;
+		}
 
 		// Load actor metadata
-		const [name, key] = await Promise.all([
-			storage.get<string>(KEYS.NAME),
-			storage.get<string[]>(KEYS.KEY),
-		]);
+		const sql = doState.ctx.storage.sql;
+		const cursor = sql.exec(
+			"SELECT name, key, destroyed, generation FROM _rivetkit_metadata LIMIT 1",
+		);
+		const result = cursor.raw().next();
+
+		if (result.done || !result.value) {
+			throw new Error(
+				`Actor ${actorId} is not initialized - missing metadata`,
+			);
+		}
+
+		const name = result.value[0] as string;
+		const key = JSON.parse(result.value[1] as string) as string[];
+		const destroyed = result.value[2] as number;
+		const generation = result.value[3] as number;
 
-		if (!name) {
-			throw new Error(`Actor ${actorId} is not initialized - missing name`);
+		// Check if actor is destroyed
+		if (destroyed) {
+			throw new Error(`Actor ${actorId} is destroyed`);
 		}
-		if (!key) {
-			throw new Error(`Actor ${actorId} is not initialized - missing key`);
+
+		// Check if generation matches
+		if (generation !== expectedGeneration) {
+			throw new Error(
+				`Actor ${actorId} generation mismatch - expected ${expectedGeneration}, got ${generation}`,
+			);
 		}
 
 		// Create actor instance
 		const definition = lookupInRegistry(this.#registryConfig, name);
-		handler.actor = definition.instantiate();
+		actorState.actorInstance = definition.instantiate();
 
 		// Start actor
-		const connDrivers = createGenericConnDrivers(
-			handler.genericConnGlobalState,
-		);
-		await handler.actor.start(
-			connDrivers,
+		await actorState.actorInstance.start(
 			this,
 			this.#inlineClient,
 			actorId,
@@ -128,39 +201,119 @@ export class CloudflareActorsActorDriver implements ActorDriver {
 		);
 
 		// Finish
-		handler.actorPromise?.resolve();
-		handler.actorPromise = undefined;
+		actorState.actorPromise?.resolve();
+		actorState.actorPromise = undefined;
+
+		return actorState.actorInstance;
+	}
+
+	getContext(actorId: string): DriverContext {
+		// Parse actor ID to get DO ID
+		const [doId] = parseActorId(actorId);
+		const state = this.#globalState.getDOState(doId);
+		return { state: state.ctx };
+	}
 
-		return handler.actor;
+	async setAlarm(actor: AnyActorInstance, timestamp: number): Promise<void> {
+		await this.#getDOCtx(actor.id).storage.setAlarm(timestamp);
+	}
+
+	async getDatabase(actorId: string): Promise<unknown | undefined> {
+		return this.#getDOCtx(actorId).storage.sql;
 	}
 
-	getGenericConnGlobalState(actorId: string): GenericConnGlobalState {
-		const handler = this.#actors.get(actorId);
-		if (!handler) {
-			throw new Error(`Actor ${actorId} not loaded`);
+	// Batch KV operations
+	async kvBatchPut(
+		actorId: string,
+		entries: [Uint8Array, Uint8Array][],
+	): Promise<void> {
+		const sql = this.#getDOCtx(actorId).storage.sql;
+
+		for (const [key, value] of entries) {
+			kvPut(sql, key, value);
 		}
-		return handler.genericConnGlobalState;
 	}
 
-	getContext(actorId: string): DriverContext {
-		const state = this.#globalState.getDOState(actorId);
-		return { state: state.ctx };
+	async kvBatchGet(
+		actorId: string,
+		keys: Uint8Array[],
+	): Promise<(Uint8Array | null)[]> {
+		const sql = this.#getDOCtx(actorId).storage.sql;
+
+		const results: (Uint8Array | null)[] = [];
+		for (const key of keys) {
+			results.push(kvGet(sql, key));
+		}
+
+		return results;
 	}
 
-	async readPersistedData(actorId: string): Promise<Uint8Array | undefined> {
-		return await this.#getDOCtx(actorId).storage.get(KEYS.PERSIST_DATA);
+	async kvBatchDelete(actorId: string, keys: Uint8Array[]): Promise<void> {
+		const sql = this.#getDOCtx(actorId).storage.sql;
+
+		for (const key of keys) {
+			kvDelete(sql, key);
+		}
 	}
 
-	async writePersistedData(actorId: string, data: Uint8Array): Promise<void> {
-		await this.#getDOCtx(actorId).storage.put(KEYS.PERSIST_DATA, data);
+	async kvListPrefix(
+		actorId: string,
+		prefix: Uint8Array,
+	): Promise<[Uint8Array, Uint8Array][]> {
+		const sql = this.#getDOCtx(actorId).storage.sql;
+
+		return kvListPrefix(sql, prefix);
 	}
 
-	async setAlarm(actor: AnyActorInstance, timestamp: number): Promise<void> {
-		await this.#getDOCtx(actor.id).storage.setAlarm(timestamp);
+	startDestroy(actorId: string): void {
+		// Parse actor ID to get DO ID and generation
+		const [doId, generation] = parseActorId(actorId);
+
+		// Get the DO state
+		const doState = this.#globalState.getDOState(doId);
+		const actorState = this.#globalState.getActorState(doState.ctx);
+
+		// Actor not loaded, nothing to destroy
+		if (!actorState?.actorInstance) {
+			return;
+		}
+
+		// Check if already destroying
+		if (actorState.destroying) {
+			return;
+		}
+		actorState.destroying = true;
+
+		// Spawn onStop in background
+		this.#callOnStopAsync(actorId, doId, actorState.actorInstance);
 	}
 
-	async getDatabase(actorId: string): Promise<unknown | undefined> {
-		return this.#getDOCtx(actorId).storage.sql;
+	async #callOnStopAsync(
+		actorId: string,
+		doId: string,
+		actor: CoreAnyActorInstance,
+	) {
+		// Stop
+		await actor.onStop("destroy");
+
+		// Remove state
+		const doState = this.#globalState.getDOState(doId);
+		const sql = doState.ctx.storage.sql;
+		sql.exec("UPDATE _rivetkit_metadata SET destroyed = 1 WHERE 1=1");
+		sql.exec("DELETE FROM _rivetkit_kv_storage");
+
+		// Clear any scheduled alarms
+		await doState.ctx.storage.deleteAlarm();
+
+		// Delete from ACTOR_KV in the background - use full actorId including generation
+		const env = getCloudflareAmbientEnv();
+		doState.ctx.waitUntil(
+			env.ACTOR_KV.delete(GLOBAL_KV_KEYS.actorMetadata(actorId)),
+		);
+
+		// Reset global state using the DO context
+		const actorHandle = this.#globalState.getActorState(doState.ctx);
+		actorHandle?.reset();
 	}
 }
 
@@ -168,14 +321,12 @@ export function createCloudflareActorsActorDriverBuilder(
 	globalState: CloudflareDurableObjectGlobalState,
 ) {
 	return (
-		registryConfig: RegistryConfig,
-		runConfig: RunConfig,
+		config: RegistryConfig,
 		managerDriver: ManagerDriver,
 		inlineClient: Client<any>,
 	) => {
 		return new CloudflareActorsActorDriver(
-			registryConfig,
-			runConfig,
+			config,
 			managerDriver,
 			inlineClient,
 			globalState,