Implement alarm methods for other drivers

Author: zephraph

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

@@ -10,6 +10,7 @@ export type ActorDriverContext = Record<never, never>;
 
 export class TestActorDriver implements ActorDriver {
 	#state: TestGlobalState;
+	#alarms: Map<string, { timeout: NodeJS.Timeout; timestamp: number }> = new Map();
 
 	constructor(state: TestGlobalState) {
 		this.#state = state;
@@ -67,9 +68,30 @@ 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 });
+	}
+
+	async getAlarm(actor: AnyActorInstance): Promise<number | null> {
+		const alarm = this.#alarms.get(actor.id);
+		return alarm ? 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/drivers/memory/src/actor.ts

@@ -5,6 +5,7 @@ export type ActorDriverContext = Record<never, never>;
 
 export class MemoryActorDriver implements ActorDriver {
 	#state: MemoryGlobalState;
+	#alarms = new Map<string, { timeout: NodeJS.Timeout; timestamp: number }>();
 
 	constructor(state: MemoryGlobalState) {
 		this.#state = state;
@@ -62,9 +63,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 {
+import type {
 	ActorDriver,
 	KvKey,
 	KvValue,
@@ -91,6 +91,14 @@ export class CloudflareWorkersActorDriver implements ActorDriver {
 		await this.#getDOCtx(actor.id).storage.setAlarm(timestamp);
 	}
 
+	async getAlarm(actor: AnyActorInstance): Promise<number | null> {
+		return this.#getDOCtx(actor.id).storage.getAlarm();
+	}
+
+	async deleteAlarm(actor: AnyActorInstance): Promise<void> {
+		await this.#getDOCtx(actor.id).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;
+		}
 	}
 }