rivetkit 2.0.24-rc.1 → 2.0.24

package/src/drivers/file-system/global-state.ts

@@ -1,18 +1,11 @@
-import * as crypto from "node:crypto";
-import * as fsSync from "node:fs";
-import * as fs from "node:fs/promises";
-import * as path from "node:path";
 import invariant from "invariant";
 import { lookupInRegistry } from "@/actor/definition";
-import { ActorAlreadyExists } from "@/actor/errors";
-import type { AnyActorInstance } from "@/actor/instance";
+import { ActorDuplicateKey } from "@/actor/errors";
+import type { AnyActorInstance } from "@/actor/instance/mod";
 import type { ActorKey } from "@/actor/mod";
 import { generateRandomString } from "@/actor/utils";
 import type { AnyClient } from "@/client/client";
-import {
-	type ActorDriver,
-	serializeEmptyPersistData,
-} from "@/driver-helpers/mod";
+import { type ActorDriver, getInitialActorKvState } from "@/driver-helpers/mod";
 import type { RegistryConfig } from "@/registry/config";
 import type { RunnerConfig } from "@/registry/run-config";
 import type * as schema from "@/schemas/file-system-driver/mod";
@@ -21,12 +14,19 @@ import {
 	ACTOR_STATE_VERSIONED,
 } from "@/schemas/file-system-driver/versioned";
 import {
+	arrayBuffersEqual,
 	bufferToArrayBuffer,
 	type LongTimeoutHandle,
 	promiseWithResolvers,
 	setLongTimeout,
 	stringifyError,
 } from "@/utils";
+import {
+	getNodeCrypto,
+	getNodeFs,
+	getNodeFsSync,
+	getNodePath,
+} from "@/utils/node";
 import { logger } from "./log";
 import {
 	ensureDirectoryExists,
@@ -36,10 +36,19 @@ import {
 
 // Actor handler to track running instances
 
+enum ActorLifecycleState {
+	NONEXISTENT, // Entry exists but actor not yet created
+	AWAKE, // Actor is running normally
+	STARTING_SLEEP, // Actor is being put to sleep
+	STARTING_DESTROY, // Actor is being destroyed
+	DESTROYED, // Actor was destroyed, should not be recreated
+}
+
 interface ActorEntry {
 	id: string;
 
 	state?: schema.ActorState;
+
 	/** Promise for loading the actor state. */
 	loadPromise?: Promise<ActorEntry>;
 
@@ -54,8 +63,12 @@ interface ActorEntry {
 	/** Resolver for pending write operations that need to be notified when any write completes */
 	pendingWriteResolver?: PromiseWithResolvers<void>;
 
-	/** If the actor has been removed by destroy or sleep. */
-	removed: boolean;
+	lifecycleState: ActorLifecycleState;
+
+	// TODO: This might make sense to move in to actorstate, but we have a
+	// single reader/writer so it's not an issue
+	/** Generation of this actor when creating/destroying. */
+	generation: string;
 }
 
 /**
@@ -68,7 +81,12 @@ export class FileSystemGlobalState {
 	#alarmsDir: string;
 
 	#persist: boolean;
+
+	// IMPORTANT: Never delete from this map. Doing so will result in race
+	// conditions since the actor generation will cease to be tracked
+	// correctly. Always increment generation if a new actor is created.
 	#actors = new Map<string, ActorEntry>();
+
 	#actorCountOnStartup: number = 0;
 
 	#runnerParams?: {
@@ -92,7 +110,8 @@ export class FileSystemGlobalState {
 
 	constructor(persist: boolean = true, customPath?: string) {
 		this.#persist = persist;
-		this.#storagePath = persist ? getStoragePath(customPath) : "/tmp";
+		this.#storagePath = persist ? (customPath ?? getStoragePath()) : "/tmp";
+		const path = getNodePath();
 		this.#stateDir = path.join(this.#storagePath, "state");
 		this.#dbsDir = path.join(this.#storagePath, "databases");
 		this.#alarmsDir = path.join(this.#storagePath, "alarms");
@@ -104,6 +123,7 @@ export class FileSystemGlobalState {
 			ensureDirectoryExistsSync(this.#alarmsDir);
 
 			try {
+				const fsSync = getNodeFsSync();
 				const actorIds = fsSync.readdirSync(this.#stateDir);
 				this.#actorCountOnStartup = actorIds.length;
 			} catch (error) {
@@ -131,15 +151,15 @@ export class FileSystemGlobalState {
 	}
 
 	getActorStatePath(actorId: string): string {
-		return path.join(this.#stateDir, actorId);
+		return getNodePath().join(this.#stateDir, actorId);
 	}
 
 	getActorDbPath(actorId: string): string {
-		return path.join(this.#dbsDir, `${actorId}.db`);
+		return getNodePath().join(this.#dbsDir, `${actorId}.db`);
 	}
 
 	getActorAlarmPath(actorId: string): string {
-		return path.join(this.#alarmsDir, actorId);
+		return getNodePath().join(this.#alarmsDir, actorId);
 	}
 
 	async *getActorsIterator(params: {
@@ -148,6 +168,7 @@ export class FileSystemGlobalState {
 		let actorIds = Array.from(this.#actors.keys()).sort();
 
 		// Check if state directory exists first
+		const fsSync = getNodeFsSync();
 		if (fsSync.existsSync(this.#stateDir)) {
 			actorIds = fsSync
 				.readdirSync(this.#stateDir)
@@ -191,7 +212,8 @@ export class FileSystemGlobalState {
 
 		entry = {
 			id: actorId,
-			removed: false,
+			lifecycleState: ActorLifecycleState.NONEXISTENT,
+			generation: crypto.randomUUID(),
 		};
 		this.#actors.set(actorId, entry);
 		return entry;
@@ -208,21 +230,44 @@ export class FileSystemGlobalState {
 	): Promise<ActorEntry> {
 		// TODO: Does not check if actor already exists on fs
 
-		if (this.#actors.has(actorId)) {
-			throw new ActorAlreadyExists(name, key);
+		const entry = this.#upsertEntry(actorId);
+
+		// Check if actor already exists (has state or is being stopped)
+		if (entry.state) {
+			throw new ActorDuplicateKey(name, key);
+		}
+		if (this.isActorStopping(actorId)) {
+			throw new Error(`Actor ${actorId} is stopping`);
 		}
 
-		const entry = this.#upsertEntry(actorId);
+		// If actor was destroyed, reset to NONEXISTENT and increment generation
+		if (entry.lifecycleState === ActorLifecycleState.DESTROYED) {
+			entry.lifecycleState = ActorLifecycleState.NONEXISTENT;
+			entry.generation = crypto.randomUUID();
+		}
+
+		// Initialize storage
+		const kvStorage: schema.ActorKvEntry[] = [];
+		const initialKvState = getInitialActorKvState(input);
+		for (const [key, value] of initialKvState) {
+			kvStorage.push({
+				key: bufferToArrayBuffer(key),
+				value: bufferToArrayBuffer(value),
+			});
+		}
+
+		// Initialize metadata
 		entry.state = {
 			actorId,
 			name,
 			key,
 			createdAt: BigInt(Date.now()),
-			persistedData: bufferToArrayBuffer(
-				serializeEmptyPersistData(input),
-			),
+			kvStorage,
 		};
-		await this.writeActor(actorId, entry.state);
+		entry.lifecycleState = ActorLifecycleState.AWAKE;
+
+		await this.writeActor(actorId, entry.generation, entry.state);
+
 		return entry;
 	}
 
@@ -232,6 +277,11 @@ export class FileSystemGlobalState {
 	async loadActor(actorId: string): Promise<ActorEntry> {
 		const entry = this.#upsertEntry(actorId);
 
+		// Check if destroyed - don't load from disk
+		if (entry.lifecycleState === ActorLifecycleState.DESTROYED) {
+			return entry;
+		}
+
 		// Check if already loaded
 		if (entry.state) {
 			return entry;
@@ -258,6 +308,7 @@ export class FileSystemGlobalState {
 
 		// Read & parse file
 		try {
+			const fs = getNodeFs();
 			const stateData = await fs.readFile(stateFilePath);
 
 			// Cache the loaded state in handler
@@ -292,16 +343,34 @@ export class FileSystemGlobalState {
 
 		// If no state for this actor, then create & write state
 		if (!entry.state) {
+			if (this.isActorStopping(actorId)) {
+				throw new Error(`Actor ${actorId} stopping`);
+			}
+
+			// If actor was destroyed, reset to NONEXISTENT and increment generation
+			if (entry.lifecycleState === ActorLifecycleState.DESTROYED) {
+				entry.lifecycleState = ActorLifecycleState.NONEXISTENT;
+				entry.generation = crypto.randomUUID();
+			}
+
+			// Initialize kvStorage with the initial persist data
+			const kvStorage: schema.ActorKvEntry[] = [];
+			const initialKvState = getInitialActorKvState(input);
+			for (const [key, value] of initialKvState) {
+				kvStorage.push({
+					key: bufferToArrayBuffer(key),
+					value: bufferToArrayBuffer(value),
+				});
+			}
+
 			entry.state = {
 				actorId,
 				name,
 				key: key as readonly string[],
 				createdAt: BigInt(Date.now()),
-				persistedData: bufferToArrayBuffer(
-					serializeEmptyPersistData(input),
-				),
+				kvStorage,
 			};
-			await this.writeActor(actorId, entry.state);
+			await this.writeActor(actorId, entry.generation, entry.state);
 		}
 		return entry;
 	}
@@ -312,29 +381,128 @@ export class FileSystemGlobalState {
 			"cannot sleep actor with memory driver, must use file system driver",
 		);
 
-		const actor = this.#actors.get(actorId);
+		// Get the actor. We upsert it even though we're about to destroy it so we have a lock on flagging `destroying` as true.
+		const actor = this.#upsertEntry(actorId);
 		invariant(actor, `tried to sleep ${actorId}, does not exist`);
 
+		// Check if already destroying
+		if (this.isActorStopping(actorId)) {
+			return;
+		}
+		actor.lifecycleState = ActorLifecycleState.STARTING_SLEEP;
+
 		// Wait for actor to fully start before stopping it to avoid race conditions
 		if (actor.loadPromise) await actor.loadPromise.catch();
 		if (actor.startPromise?.promise)
 			await actor.startPromise.promise.catch();
 
-		// Mark as removed
-		actor.removed = true;
-
 		// Stop actor
 		invariant(actor.actor, "actor should be loaded");
-		await actor.actor._onStop();
+		await actor.actor.onStop("sleep");
 
 		// Remove from map after stop is complete
 		this.#actors.delete(actorId);
 	}
 
+	async destroyActor(actorId: string) {
+		// Get the actor. We upsert it even though we're about to destroy it so we have a lock on flagging `destroying` as true.
+		const actor = this.#upsertEntry(actorId);
+
+		// If actor is loaded, stop it first
+		// Check if already destroying
+		if (this.isActorStopping(actorId)) {
+			return;
+		}
+		actor.lifecycleState = ActorLifecycleState.STARTING_DESTROY;
+
+		// Wait for actor to fully start before stopping it to avoid race conditions
+		if (actor.loadPromise) await actor.loadPromise.catch();
+		if (actor.startPromise?.promise)
+			await actor.startPromise.promise.catch();
+
+		// Stop actor if it's running
+		if (actor.actor) {
+			await actor.actor.onStop("destroy");
+		}
+
+		// Clear alarm timeout if exists
+		if (actor.alarmTimeout) {
+			actor.alarmTimeout.abort();
+		}
+
+		// Delete persisted files if using file system driver
+		if (this.#persist) {
+			const fs = getNodeFs();
+
+			// Delete all actor files in parallel
+			await Promise.all([
+				// Delete actor state file
+				(async () => {
+					try {
+						await fs.unlink(this.getActorStatePath(actorId));
+					} catch (err: any) {
+						if (err?.code !== "ENOENT") {
+							logger().error({
+								msg: "failed to delete actor state file",
+								actorId,
+								error: stringifyError(err),
+							});
+						}
+					}
+				})(),
+				// Delete actor database file
+				(async () => {
+					try {
+						await fs.unlink(this.getActorDbPath(actorId));
+					} catch (err: any) {
+						if (err?.code !== "ENOENT") {
+							logger().error({
+								msg: "failed to delete actor database file",
+								actorId,
+								error: stringifyError(err),
+							});
+						}
+					}
+				})(),
+				// Delete actor alarm file
+				(async () => {
+					try {
+						await fs.unlink(this.getActorAlarmPath(actorId));
+					} catch (err: any) {
+						if (err?.code !== "ENOENT") {
+							logger().error({
+								msg: "failed to delete actor alarm file",
+								actorId,
+								error: stringifyError(err),
+							});
+						}
+					}
+				})(),
+			]);
+		}
+
+		// Reset the entry
+		//
+		// Do not remove entry in order to avoid race condition with
+		// destroying. Next actor creation will increment the generation.
+		actor.state = undefined;
+		actor.loadPromise = undefined;
+		actor.actor = undefined;
+		actor.startPromise = undefined;
+		actor.alarmTimeout = undefined;
+		actor.alarmTimeout = undefined;
+		actor.pendingWriteResolver = undefined;
+		actor.lifecycleState = ActorLifecycleState.DESTROYED;
+	}
+
 	/**
 	 * Save actor state to disk.
 	 */
-	async writeActor(actorId: string, state: schema.ActorState): Promise<void> {
+	async writeActor(
+		actorId: string,
+		generation: string,
+		state: schema.ActorState,
+	): Promise<void> {
 		if (!this.#persist) {
 			return;
 		}
@@ -342,18 +510,49 @@ export class FileSystemGlobalState {
 		const entry = this.#actors.get(actorId);
 		invariant(entry, "actor entry does not exist");
 
-		await this.#performWrite(actorId, state);
+		await this.#performWrite(actorId, generation, state);
+	}
+
+	isGenerationCurrentAndNotDestroyed(
+		actorId: string,
+		generation: string,
+	): boolean {
+		const entry = this.#upsertEntry(actorId);
+		if (!entry) return false;
+		return (
+			entry.generation === generation &&
+			entry.lifecycleState !== ActorLifecycleState.STARTING_DESTROY
+		);
+	}
+
+	isActorStopping(actorId: string) {
+		const entry = this.#upsertEntry(actorId);
+		if (!entry) return false;
+		return (
+			entry.lifecycleState === ActorLifecycleState.STARTING_SLEEP ||
+			entry.lifecycleState === ActorLifecycleState.STARTING_DESTROY
+		);
 	}
 
 	async setActorAlarm(actorId: string, timestamp: number) {
 		const entry = this.#actors.get(actorId);
 		invariant(entry, "actor entry does not exist");
 
+		// Track generation of the actor when the write started to detect
+		// destroy/create race condition
+		const writeGeneration = entry.generation;
+		if (this.isActorStopping(actorId)) {
+			logger().info("skipping set alarm since actor stopping");
+			return;
+		}
+
 		// Persist alarm to disk
 		if (this.#persist) {
 			const alarmPath = this.getActorAlarmPath(actorId);
+			const crypto = getNodeCrypto();
 			const tempPath = `${alarmPath}.tmp.${crypto.randomUUID()}`;
 			try {
+				const path = getNodePath();
 				await ensureDirectoryExists(path.dirname(alarmPath));
 				const alarmData: schema.ActorAlarm = {
 					actorId,
@@ -363,10 +562,25 @@ export class FileSystemGlobalState {
 					ACTOR_ALARM_VERSIONED.serializeWithEmbeddedVersion(
 						alarmData,
 					);
+				const fs = getNodeFs();
 				await fs.writeFile(tempPath, data);
+
+				if (
+					!this.isGenerationCurrentAndNotDestroyed(
+						actorId,
+						writeGeneration,
+					)
+				) {
+					logger().debug(
+						"skipping writing alarm since actor destroying or new generation",
+					);
+					return;
+				}
+
 				await fs.rename(tempPath, alarmPath);
 			} catch (error) {
 				try {
+					const fs = getNodeFs();
 					await fs.unlink(tempPath);
 				} catch {}
 				logger().error({
@@ -387,14 +601,17 @@ export class FileSystemGlobalState {
 	 */
 	async #performWrite(
 		actorId: string,
+		generation: string,
 		state: schema.ActorState,
 	): Promise<void> {
 		const dataPath = this.getActorStatePath(actorId);
 		// Generate unique temp filename to prevent any race conditions
+		const crypto = getNodeCrypto();
 		const tempPath = `${dataPath}.tmp.${crypto.randomUUID()}`;
 
 		try {
 			// Create directory if needed
+			const path = getNodePath();
 			await ensureDirectoryExists(path.dirname(dataPath));
 
 			// Convert to BARE types for serialization
@@ -403,17 +620,27 @@ export class FileSystemGlobalState {
 				name: state.name,
 				key: state.key,
 				createdAt: state.createdAt,
-				persistedData: state.persistedData,
+				kvStorage: state.kvStorage,
 			};
 
 			// Perform atomic write
 			const serializedState =
 				ACTOR_STATE_VERSIONED.serializeWithEmbeddedVersion(bareState);
+			const fs = getNodeFs();
 			await fs.writeFile(tempPath, serializedState);
+
+			if (!this.isGenerationCurrentAndNotDestroyed(actorId, generation)) {
+				logger().debug(
+					"skipping writing alarm since actor destroying or new generation",
+				);
+				return;
+			}
+
 			await fs.rename(tempPath, dataPath);
 		} catch (error) {
 			// Cleanup temp file on error
 			try {
+				const fs = getNodeFs();
 				await fs.unlink(tempPath);
 			} catch {
 				// Ignore cleanup errors
@@ -548,12 +775,14 @@ export class FileSystemGlobalState {
 	 */
 	#loadAlarmsSync(): void {
 		try {
+			const fsSync = getNodeFsSync();
 			const files = fsSync.existsSync(this.#alarmsDir)
 				? fsSync.readdirSync(this.#alarmsDir)
 				: [];
 			for (const file of files) {
 				// Skip temp files
 				if (file.includes(".tmp.")) continue;
+				const path = getNodePath();
 				const fullPath = path.join(this.#alarmsDir, file);
 				try {
 					const buf = fsSync.readFileSync(fullPath);
@@ -622,6 +851,7 @@ export class FileSystemGlobalState {
 			// On trigger: remove persisted alarm file
 			if (this.#persist) {
 				try {
+					const fs = getNodeFs();
 					await fs.unlink(this.getActorAlarmPath(actorId));
 				} catch (err: any) {
 					if (err?.code !== "ENOENT") {
@@ -656,7 +886,7 @@ export class FileSystemGlobalState {
 				}
 
 				invariant(loaded.actor, "actor should be loaded after wake");
-				await loaded.actor._onAlarm();
+				await loaded.actor.onAlarm();
 			} catch (err) {
 				logger().error({
 					msg: "failed to handle alarm",
@@ -668,6 +898,8 @@ export class FileSystemGlobalState {
 	}
 
 	getOrCreateInspectorAccessToken(): string {
+		const path = getNodePath();
+		const fsSync = getNodeFsSync();
 		const tokenPath = path.join(this.#storagePath, "inspector-token");
 		if (fsSync.existsSync(tokenPath)) {
 			return fsSync.readFileSync(tokenPath, "utf-8");
@@ -683,6 +915,7 @@ export class FileSystemGlobalState {
 	 */
 	#cleanupTempFilesSync(): void {
 		try {
+			const fsSync = getNodeFsSync();
 			const files = fsSync.readdirSync(this.#stateDir);
 			const tempFiles = files.filter((f) => f.includes(".tmp."));
 
@@ -690,6 +923,7 @@ export class FileSystemGlobalState {
 
 			for (const tempFile of tempFiles) {
 				try {
+					const path = getNodePath();
 					const fullPath = path.join(this.#stateDir, tempFile);
 					const stat = fsSync.statSync(fullPath);
 
@@ -716,4 +950,160 @@ export class FileSystemGlobalState {
 			});
 		}
 	}
+
+	/**
+	 * Batch put KV entries for an actor.
+	 */
+	async kvBatchPut(
+		actorId: string,
+		entries: [Uint8Array, Uint8Array][],
+	): Promise<void> {
+		const entry = await this.loadActor(actorId);
+		if (!entry.state) {
+			if (this.isActorStopping(actorId)) {
+				return;
+			} else {
+				throw new Error(`Actor ${actorId} state not loaded`);
+			}
+		}
+
+		// Create a mutable copy of kvStorage
+		const newKvStorage = [...entry.state.kvStorage];
+
+		// Update kvStorage with new entries
+		for (const [key, value] of entries) {
+			// Find existing entry with the same key
+			const existingIndex = newKvStorage.findIndex((e) =>
+				arrayBuffersEqual(e.key, bufferToArrayBuffer(key)),
+			);
+
+			if (existingIndex >= 0) {
+				// Replace existing entry with new one
+				newKvStorage[existingIndex] = {
+					key: bufferToArrayBuffer(key),
+					value: bufferToArrayBuffer(value),
+				};
+			} else {
+				// Add new entry
+				newKvStorage.push({
+					key: bufferToArrayBuffer(key),
+					value: bufferToArrayBuffer(value),
+				});
+			}
+		}
+
+		// Update state with new kvStorage
+		entry.state = {
+			...entry.state,
+			kvStorage: newKvStorage,
+		};
+
+		// Save state to disk
+		await this.writeActor(actorId, entry.generation, entry.state);
+	}
+
+	/**
+	 * Batch get KV entries for an actor.
+	 */
+	async kvBatchGet(
+		actorId: string,
+		keys: Uint8Array[],
+	): Promise<(Uint8Array | null)[]> {
+		const entry = await this.loadActor(actorId);
+		if (!entry.state) {
+			if (this.isActorStopping(actorId)) {
+				throw new Error(`Actor ${actorId} is stopping`);
+			} else {
+				throw new Error(`Actor ${actorId} state not loaded`);
+			}
+		}
+
+		const results: (Uint8Array | null)[] = [];
+		for (const key of keys) {
+			// Find entry with the same key
+			const foundEntry = entry.state.kvStorage.find((e) =>
+				arrayBuffersEqual(e.key, bufferToArrayBuffer(key)),
+			);
+
+			if (foundEntry) {
+				results.push(new Uint8Array(foundEntry.value));
+			} else {
+				results.push(null);
+			}
+		}
+		return results;
+	}
+
+	/**
+	 * Batch delete KV entries for an actor.
+	 */
+	async kvBatchDelete(actorId: string, keys: Uint8Array[]): Promise<void> {
+		const entry = await this.loadActor(actorId);
+		if (!entry.state) {
+			if (this.isActorStopping(actorId)) {
+				return;
+			} else {
+				throw new Error(`Actor ${actorId} state not loaded`);
+			}
+		}
+
+		// Create a mutable copy of kvStorage
+		const newKvStorage = [...entry.state.kvStorage];
+
+		// Delete entries from kvStorage
+		for (const key of keys) {
+			const indexToDelete = newKvStorage.findIndex((e) =>
+				arrayBuffersEqual(e.key, bufferToArrayBuffer(key)),
+			);
+
+			if (indexToDelete >= 0) {
+				newKvStorage.splice(indexToDelete, 1);
+			}
+		}
+
+		// Update state with new kvStorage
+		entry.state = {
+			...entry.state,
+			kvStorage: newKvStorage,
+		};
+
+		// Save state to disk
+		await this.writeActor(actorId, entry.generation, entry.state);
+	}
+
+	/**
+	 * List KV entries with a given prefix for an actor.
+	 */
+	async kvListPrefix(
+		actorId: string,
+		prefix: Uint8Array,
+	): Promise<[Uint8Array, Uint8Array][]> {
+		const entry = await this.loadActor(actorId);
+		if (!entry.state) {
+			if (this.isActorStopping(actorId)) {
+				throw new Error(`Actor ${actorId} is destroying`);
+			} else {
+				throw new Error(`Actor ${actorId} state not loaded`);
+			}
+		}
+
+		const results: [Uint8Array, Uint8Array][] = [];
+		for (const kvEntry of entry.state.kvStorage) {
+			const keyBytes = new Uint8Array(kvEntry.key);
+			// Check if key starts with prefix
+			if (keyBytes.length >= prefix.length) {
+				let hasPrefix = true;
+				for (let i = 0; i < prefix.length; i++) {
+					if (keyBytes[i] !== prefix[i]) {
+						hasPrefix = false;
+						break;
+					}
+				}
+				if (hasPrefix) {
+					results.push([keyBytes, new Uint8Array(kvEntry.value)]);
+				}
+			}
+		}
+		return results;
+	}
 }