Implement alarm methods for other drivers

Author: zephraph

packages/actor-core/src/test/driver/actor.ts

@@ -7,6 +7,7 @@ export interface ActorDriverContext {
 
 export class TestActorDriver implements ActorDriver {
 	#state: TestGlobalState;
+	#alarms: Map<string, { timeout: NodeJS.Timeout; timestamp: number }> = new Map();
 
 	constructor(state: TestGlobalState) {
 		this.#state = state;
@@ -64,9 +65,31 @@ export class TestActorDriver implements ActorDriver {
 	}
 
 	async setAlarm(actor: AnyActorInstance, timestamp: number): Promise<void> {
-		setTimeout(() => {
+		// Clear any existing alarm for this actor
+		await this.deleteAlarm(actor);
+
+		const expires = Math.max(timestamp - Date.now(), 0);
+		// Set new alarm
+		const timeout = setTimeout(() => {
 			actor.onAlarm();
-		}, timestamp - Date.now());
+			this.#alarms.delete(actor.id);
+		}, expires);
+
+		this.#alarms.set(actor.id, { timeout, timestamp });
+	}
+
+	getAlarm(actor: AnyActorInstance): Promise<number | null> {
+		const alarm = this.#alarms.get(actor.id);
+		return Promise.resolve(alarm ? alarm.timestamp : null);
+	}
+
+	deleteAlarm(actor: AnyActorInstance): Promise<void> {
+		const alarm = this.#alarms.get(actor.id);
+		if (alarm) {
+			clearTimeout(alarm.timeout);
+			this.#alarms.delete(actor.id);
+		}
+		return Promise.resolve();
 	}
 
 	// Simple key serialization without depending on keys.ts

packages/drivers/memory/src/actor.ts

@@ -7,6 +7,7 @@ export interface ActorDriverContext {
 
 export class MemoryActorDriver implements ActorDriver {
 	#state: MemoryGlobalState;
+	#alarms = new Map<string, { timeout: NodeJS.Timeout; timestamp: number }>();
 
 	constructor(state: MemoryGlobalState) {
 		this.#state = state;
@@ -64,9 +65,29 @@ export class MemoryActorDriver implements ActorDriver {
 	}
 
 	async setAlarm(actor: AnyActorInstance, timestamp: number): Promise<void> {
-		setTimeout(() => {
+		// Clear any existing alarm
+		await this.deleteAlarm(actor);
+
+		// Set new alarm
+		const timeout = setTimeout(() => {
 			actor.onAlarm();
+			this.#alarms.delete(actor.id);
 		}, timestamp - Date.now());
+
+		this.#alarms.set(actor.id, { timeout, timestamp });
+	}
+
+	async getAlarm(actor: AnyActorInstance): Promise<number | null> {
+		const alarm = this.#alarms.get(actor.id);
+		return alarm?.timestamp ?? null;
+	}
+
+	async deleteAlarm(actor: AnyActorInstance): Promise<void> {
+		const alarm = this.#alarms.get(actor.id);
+		if (alarm) {
+			clearTimeout(alarm.timeout);
+			this.#alarms.delete(actor.id);
+		}
 	}
 
 	// Simple key serialization without depending on keys.ts

packages/platforms/cloudflare-workers/src/actor_driver.ts

@@ -1,4 +1,4 @@
-import { ActorDriver, KvKey, KvValue, AnyActorInstance } from "actor-core/driver-helpers";
+import type { ActorDriver, KvKey, KvValue, AnyActorInstance } from "actor-core/driver-helpers";
 
 export interface ActorDriverContext {
 	ctx: DurableObjectState;
@@ -57,6 +57,14 @@ export class CloudflareWorkersActorDriver implements ActorDriver {
 		await this.#doCtx.storage.setAlarm(timestamp);
 	}
 
+	async getAlarm(_actor: AnyActorInstance): Promise<number | null> {
+		return this.#doCtx.storage.getAlarm();
+	}
+
+	async deleteAlarm(_actor: AnyActorInstance): Promise<void> {
+		await this.#doCtx.storage.deleteAlarm();
+	}
+
 	#serializeKey(key: KvKey): string {
 		return JSON.stringify(key);
 	}

packages/platforms/rivet/src/actor_driver.ts

@@ -7,6 +7,7 @@ export interface ActorDriverContext {
 
 export class RivetActorDriver implements ActorDriver {
 	#ctx: ActorContext;
+	#alarm: { timeout: ReturnType<typeof setTimeout>; timestamp: number } | null = null;
 
 	constructor(ctx: ActorContext) {
 		this.#ctx = ctx;
@@ -48,9 +49,25 @@ export class RivetActorDriver implements ActorDriver {
 	}
 
 	async setAlarm(actor: AnyActorInstance, timestamp: number): Promise<void> {
+		await this.deleteAlarm(actor);
+
 		const timeout = Math.max(0, timestamp - Date.now());
-		setTimeout(() => {
+		const timeoutId = setTimeout(() => {
 			actor.onAlarm();
+			this.#alarm = null;
 		}, timeout);
+
+		this.#alarm = { timeout: timeoutId, timestamp };
+	}
+
+	async getAlarm(_actor: AnyActorInstance): Promise<number | null> {
+		return this.#alarm?.timestamp ?? null;
+	}
+
+	async deleteAlarm(_actor: AnyActorInstance): Promise<void> {
+		if (this.#alarm) {
+			clearTimeout(this.#alarm.timeout);
+			this.#alarm = null;
+		}
 	}
 }