wip: cf workers fix

Author: NathanFlurry

rivetkit-typescript/packages/cloudflare-workers/src/actor-driver.ts

@@ -1,5 +1,7 @@
 import invariant from "invariant";
 import type {
+	ActorKey,
+	ActorRouter,
 	AnyActorInstance as CoreAnyActorInstance,
 	RegistryConfig,
 	RunConfig,
@@ -12,10 +14,10 @@ import type {
 	ManagerDriver,
 } from "rivetkit/driver-helpers";
 import { promiseWithResolvers } from "rivetkit/utils";
-import { buildActorId, parseActorId } from "./actor-id";
-import { GLOBAL_KV_KEYS } from "./global_kv";
+import { parseActorId } from "./actor-id";
+import { kvDelete, kvGet, kvListPrefix, kvPut } from "./actor-kv";
+import { GLOBAL_KV_KEYS } from "./global-kv";
 import { getCloudflareAmbientEnv } from "./handler";
-import { kvDelete, kvGet, kvListPrefix, kvPut } from "./kv_query";
 
 interface DurableObjectGlobalState {
 	ctx: DurableObjectState;
@@ -31,8 +33,8 @@ export class CloudflareDurableObjectGlobalState {
 	// Map of actor ID -> DO state
 	#dos: Map<string, DurableObjectGlobalState> = new Map();
 
-	// Map of DO ID -> ActorHandler
-	#actors: Map<string, ActorHandler> = new Map();
+	// WeakMap of DO state -> ActorGlobalState for proper GC
+	#actors: WeakMap<DurableObjectState, ActorGlobalState> = new WeakMap();
 
 	getDOState(doId: string): DurableObjectGlobalState {
 		const state = this.#dos.get(doId);
@@ -47,20 +49,40 @@ export class CloudflareDurableObjectGlobalState {
 		this.#dos.set(doId, state);
 	}
 
-	get actors() {
-		return this.#actors;
+	getActorState(ctx: DurableObjectState): ActorGlobalState | undefined {
+		return this.#actors.get(ctx);
+	}
+
+	setActorState(ctx: DurableObjectState, actorState: ActorGlobalState): void {
+		this.#actors.set(ctx, actorState);
 	}
 }
 
 export interface DriverContext {
 	state: DurableObjectState;
 }
 
-// Actor handler to track running instances
-class ActorHandler {
-	actor?: AnyActorInstance;
-	actorPromise?: ReturnType<typeof promiseWithResolvers<void>> =
-		promiseWithResolvers();
+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.
@@ -70,6 +92,14 @@ class ActorHandler {
 	 * 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 {
@@ -103,23 +133,36 @@ export class CloudflareActorsActorDriver implements ActorDriver {
 		// 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.#globalState.actors.get(doId);
-		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.#globalState.actors.set(doId, handler);
+		// Create new actor state if it doesn't exist
+		if (!actorState) {
+			actorState = new ActorGlobalState();
+			actorState.actorPromise = promiseWithResolvers();
+			this.#globalState.setActorState(doState.ctx, actorState);
+		}
 
-		// Get the actor metadata from Durable Object storage
-		const doState = this.#globalState.getDOState(doId);
-		const sql = doState.ctx.storage.sql;
+		// Another request is already loading this actor, wait for it
+		if (actorState.actorPromise) {
+			await actorState.actorPromise.promise;
+			if (!actorState.actorInstance) {
+				throw new Error(
+					`Actor ${actorId} failed to load in concurrent request`,
+				);
+			}
+			return actorState.actorInstance;
+		}
 
-		// Load actor metadata from SQL table
+		// Load actor metadata
+		const sql = doState.ctx.storage.sql;
 		const cursor = sql.exec(
 			"SELECT name, key, destroyed, generation FROM _rivetkit_metadata LIMIT 1",
 		);
@@ -150,10 +193,10 @@ export class CloudflareActorsActorDriver implements ActorDriver {
 
 		// Create actor instance
 		const definition = lookupInRegistry(this.#registryConfig, name);
-		handler.actor = definition.instantiate();
+		actorState.actorInstance = definition.instantiate();
 
 		// Start actor
-		await handler.actor.start(
+		await actorState.actorInstance.start(
 			this,
 			this.#inlineClient,
 			actorId,
@@ -163,10 +206,10 @@ export class CloudflareActorsActorDriver implements ActorDriver {
 		);
 
 		// Finish
-		handler.actorPromise?.resolve();
-		handler.actorPromise = undefined;
+		actorState.actorPromise?.resolve();
+		actorState.actorPromise = undefined;
 
-		return handler.actor;
+		return actorState.actorInstance;
 	}
 
 	getContext(actorId: string): DriverContext {
@@ -231,21 +274,32 @@ export class CloudflareActorsActorDriver implements ActorDriver {
 		// Parse actor ID to get DO ID and generation
 		const [doId, generation] = parseActorId(actorId);
 
-		const handler = this.#globalState.actors.get(doId);
+		// Get the DO state
+		const doState = this.#globalState.getDOState(doId);
+		const actorState = this.#globalState.getActorState(doState.ctx);
 
 		// Actor not loaded, nothing to destroy
-		if (!handler || !handler.actor) {
+		if (!actorState?.actorInstance) {
 			return;
 		}
 
 		// Check if already destroying
-		if (handler.destroying) {
+		if (actorState.destroying) {
 			return;
 		}
-		handler.destroying = true;
+		actorState.destroying = true;
+
+		// Spawn onStop in background
+		this.#callOnStopAsync(actorId, doId, actorState.actorInstance);
+	}
 
+	async #callOnStopAsync(
+		actorId: string,
+		doId: string,
+		actor: CoreAnyActorInstance,
+	) {
 		// Stop
-		handler.actor.onStop("destroy");
+		await actor.onStop("destroy");
 
 		// Remove state
 		const doState = this.#globalState.getDOState(doId);
@@ -254,15 +308,17 @@ export class CloudflareActorsActorDriver implements ActorDriver {
 		sql.exec("DELETE FROM _rivetkit_kv_storage");
 
 		// Clear any scheduled alarms
-		doState.ctx.storage.deleteAlarm();
+		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)),
 		);
 
-		this.#globalState.actors.delete(doId);
+		// Reset global state using the DO context
+		const actorHandle = this.#globalState.getActorState(doState.ctx);
+		actorHandle?.reset();
 	}
 }