rivetkit 2.0.24-rc.1 → 2.0.24

package/src/actor/instance/persisted.ts

@@ -0,0 +1,67 @@
+/**
+ * Persisted data structures for actors.
+ *
+ * Keep this file in sync with the Connection section of rivetkit-typescript/packages/rivetkit/schemas/actor-persist/
+ */
+
+import * as cbor from "cbor-x";
+import type * as persistSchema from "@/schemas/actor-persist/mod";
+import { bufferToArrayBuffer } from "@/utils";
+
+export type Cbor = ArrayBuffer;
+
+// MARK: Schedule Event
+/** Scheduled event to be executed at a specific timestamp */
+export interface PersistedScheduleEvent {
+	eventId: string;
+	timestamp: number;
+	action: string;
+	args?: Cbor;
+}
+
+// MARK: Actor
+/** State object that gets automatically persisted to storage */
+export interface PersistedActor<S, I> {
+	/** Input data passed to the actor on initialization */
+	input?: I;
+	hasInitialized: boolean;
+	state: S;
+	scheduledEvents: PersistedScheduleEvent[];
+}
+
+export function convertActorToBarePersisted<S, I>(
+	persist: PersistedActor<S, I>,
+): persistSchema.Actor {
+	return {
+		input:
+			persist.input !== undefined
+				? bufferToArrayBuffer(cbor.encode(persist.input))
+				: null,
+		hasInitialized: persist.hasInitialized,
+		state: bufferToArrayBuffer(cbor.encode(persist.state)),
+		scheduledEvents: persist.scheduledEvents.map((event) => ({
+			eventId: event.eventId,
+			timestamp: BigInt(event.timestamp),
+			action: event.action,
+			args: event.args ?? null,
+		})),
+	};
+}
+
+export function convertActorFromBarePersisted<S, I>(
+	bareData: persistSchema.Actor,
+): PersistedActor<S, I> {
+	return {
+		input: bareData.input
+			? cbor.decode(new Uint8Array(bareData.input))
+			: undefined,
+		hasInitialized: bareData.hasInitialized,
+		state: cbor.decode(new Uint8Array(bareData.state)),
+		scheduledEvents: bareData.scheduledEvents.map((event) => ({
+			eventId: event.eventId,
+			timestamp: Number(event.timestamp),
+			action: event.action,
+			args: event.args ?? undefined,
+		})),
+	};
+}

package/src/actor/instance/schedule-manager.ts

@@ -0,0 +1,349 @@
+import * as cbor from "cbor-x";
+import {
+	bufferToArrayBuffer,
+	SinglePromiseQueue,
+	stringifyError,
+} from "@/utils";
+import type { AnyDatabaseProvider } from "../database";
+import type { ActorDriver } from "../driver";
+import type { ActorInstance } from "./mod";
+import type { PersistedScheduleEvent } from "./persisted";
+
+/**
+ * Manages scheduled events and alarms for actor instances.
+ * Handles event scheduling, alarm triggers, and automatic event execution.
+ */
+export class ScheduleManager<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
+	#actor: ActorInstance<S, CP, CS, V, I, DB>;
+	#actorDriver: ActorDriver;
+	#alarmWriteQueue = new SinglePromiseQueue();
+	#config: any; // ActorConfig type
+	#persist: any; // Reference to PersistedActor
+
+	constructor(
+		actor: ActorInstance<S, CP, CS, V, I, DB>,
+		actorDriver: ActorDriver,
+		config: any,
+	) {
+		this.#actor = actor;
+		this.#actorDriver = actorDriver;
+		this.#config = config;
+	}
+
+	// MARK: - Public API
+
+	/**
+	 * Sets the persist object reference.
+	 * Called after StateManager initializes the persist proxy.
+	 */
+	setPersist(persist: any) {
+		this.#persist = persist;
+	}
+
+	/**
+	 * Schedules an event to be executed at a specific timestamp.
+	 *
+	 * @param timestamp - Unix timestamp in milliseconds when the event should fire
+	 * @param action - The name of the action to execute
+	 * @param args - Arguments to pass to the action
+	 */
+	async scheduleEvent(
+		timestamp: number,
+		action: string,
+		args: unknown[],
+	): Promise<void> {
+		const newEvent: PersistedScheduleEvent = {
+			eventId: crypto.randomUUID(),
+			timestamp,
+			action,
+			args: bufferToArrayBuffer(cbor.encode(args)),
+		};
+
+		await this.#scheduleEventInner(newEvent);
+	}
+
+	/**
+	 * Triggers any pending alarms that are due.
+	 * This method is idempotent and safe to call multiple times.
+	 */
+	async onAlarm(): Promise<void> {
+		const now = Date.now();
+		this.#actor.log.debug({
+			msg: "alarm triggered",
+			now,
+			events: this.#persist?.scheduledEvents?.length || 0,
+		});
+
+		if (!this.#persist?.scheduledEvents) {
+			this.#actor.rLog.debug({ msg: "no scheduled events" });
+			return;
+		}
+
+		// Find events that are due
+		const dueIndex = this.#persist.scheduledEvents.findIndex(
+			(x: PersistedScheduleEvent) => x.timestamp <= now,
+		);
+
+		if (dueIndex === -1) {
+			// No events are due yet
+			this.#actor.rLog.debug({ msg: "no events are due yet" });
+
+			// Reschedule alarm for next event if any exist
+			if (this.#persist.scheduledEvents.length > 0) {
+				const nextTs = this.#persist.scheduledEvents[0].timestamp;
+				this.#actor.log.debug({
+					msg: "alarm fired early, rescheduling for next event",
+					now,
+					nextTs,
+					delta: nextTs - now,
+				});
+				await this.#queueSetAlarm(nextTs);
+			}
+			return;
+		}
+
+		// Remove and process due events
+		const dueEvents = this.#persist.scheduledEvents.splice(0, dueIndex + 1);
+		this.#actor.log.debug({
+			msg: "running events",
+			count: dueEvents.length,
+		});
+
+		// Schedule next alarm if more events remain
+		if (this.#persist.scheduledEvents.length > 0) {
+			const nextTs = this.#persist.scheduledEvents[0].timestamp;
+			this.#actor.log.info({
+				msg: "setting next alarm",
+				nextTs,
+				remainingEvents: this.#persist.scheduledEvents.length,
+			});
+			await this.#queueSetAlarm(nextTs);
+		}
+
+		// Execute due events
+		await this.#executeDueEvents(dueEvents);
+	}
+
+	/**
+	 * Initializes alarms on actor startup.
+	 * Sets the alarm for the next scheduled event if any exist.
+	 */
+	async initializeAlarms(): Promise<void> {
+		if (this.#persist?.scheduledEvents?.length > 0) {
+			await this.#queueSetAlarm(
+				this.#persist.scheduledEvents[0].timestamp,
+			);
+		}
+	}
+
+	/**
+	 * Waits for any pending alarm write operations to complete.
+	 */
+	async waitForPendingAlarmWrites(): Promise<void> {
+		if (this.#alarmWriteQueue.runningDrainLoop) {
+			await this.#alarmWriteQueue.runningDrainLoop;
+		}
+	}
+
+	/**
+	 * Gets statistics about scheduled events.
+	 */
+	getScheduleStats(): {
+		totalEvents: number;
+		nextEventTime: number | null;
+		overdueCount: number;
+	} {
+		if (!this.#persist?.scheduledEvents) {
+			return {
+				totalEvents: 0,
+				nextEventTime: null,
+				overdueCount: 0,
+			};
+		}
+
+		const now = Date.now();
+		const events = this.#persist.scheduledEvents;
+
+		return {
+			totalEvents: events.length,
+			nextEventTime: events.length > 0 ? events[0].timestamp : null,
+			overdueCount: events.filter(
+				(e: PersistedScheduleEvent) => e.timestamp <= now,
+			).length,
+		};
+	}
+
+	/**
+	 * Cancels a scheduled event by its ID.
+	 *
+	 * @param eventId - The ID of the event to cancel
+	 * @returns True if the event was found and cancelled
+	 */
+	async cancelEvent(eventId: string): Promise<boolean> {
+		if (!this.#persist?.scheduledEvents) {
+			return false;
+		}
+
+		const index = this.#persist.scheduledEvents.findIndex(
+			(e: PersistedScheduleEvent) => e.eventId === eventId,
+		);
+
+		if (index === -1) {
+			return false;
+		}
+
+		// Remove the event
+		const wasFirst = index === 0;
+		this.#persist.scheduledEvents.splice(index, 1);
+
+		// If we removed the first event, update the alarm
+		if (wasFirst && this.#persist.scheduledEvents.length > 0) {
+			await this.#queueSetAlarm(
+				this.#persist.scheduledEvents[0].timestamp,
+			);
+		}
+
+		this.#actor.log.info({
+			msg: "cancelled scheduled event",
+			eventId,
+			remainingEvents: this.#persist.scheduledEvents.length,
+		});
+
+		return true;
+	}
+
+	// MARK: - Private Helpers
+
+	async #scheduleEventInner(newEvent: PersistedScheduleEvent): Promise<void> {
+		this.#actor.log.info({
+			msg: "scheduling event",
+			eventId: newEvent.eventId,
+			timestamp: newEvent.timestamp,
+			action: newEvent.action,
+		});
+
+		if (!this.#persist?.scheduledEvents) {
+			throw new Error("Persist not initialized");
+		}
+
+		// Find insertion point (events are sorted by timestamp)
+		const insertIndex = this.#persist.scheduledEvents.findIndex(
+			(x: PersistedScheduleEvent) => x.timestamp > newEvent.timestamp,
+		);
+
+		if (insertIndex === -1) {
+			// Add to end
+			this.#persist.scheduledEvents.push(newEvent);
+		} else {
+			// Insert at correct position
+			this.#persist.scheduledEvents.splice(insertIndex, 0, newEvent);
+		}
+
+		// Update alarm if this is the newest event
+		if (insertIndex === 0 || this.#persist.scheduledEvents.length === 1) {
+			this.#actor.log.info({
+				msg: "setting alarm for new event",
+				timestamp: newEvent.timestamp,
+				eventCount: this.#persist.scheduledEvents.length,
+			});
+			await this.#queueSetAlarm(newEvent.timestamp);
+		}
+	}
+
+	async #executeDueEvents(events: PersistedScheduleEvent[]): Promise<void> {
+		for (const event of events) {
+			try {
+				this.#actor.log.info({
+					msg: "executing scheduled event",
+					eventId: event.eventId,
+					timestamp: event.timestamp,
+					action: event.action,
+				});
+
+				// Look up the action function
+				const fn = this.#config.actions[event.action];
+
+				if (!fn) {
+					throw new Error(
+						`Missing action for scheduled event: ${event.action}`,
+					);
+				}
+
+				if (typeof fn !== "function") {
+					throw new Error(
+						`Scheduled event action ${event.action} is not a function (got ${typeof fn})`,
+					);
+				}
+
+				// Decode arguments and execute
+				const args = event.args
+					? cbor.decode(new Uint8Array(event.args))
+					: [];
+
+				const result = fn.call(
+					undefined,
+					this.#actor.actorContext,
+					...args,
+				);
+
+				// Handle async actions
+				if (result instanceof Promise) {
+					await result;
+				}
+
+				this.#actor.log.debug({
+					msg: "scheduled event completed",
+					eventId: event.eventId,
+					action: event.action,
+				});
+			} catch (error) {
+				this.#actor.log.error({
+					msg: "error executing scheduled event",
+					error: stringifyError(error),
+					eventId: event.eventId,
+					timestamp: event.timestamp,
+					action: event.action,
+				});
+
+				// Continue processing other events even if one fails
+			}
+		}
+	}
+
+	async #queueSetAlarm(timestamp: number): Promise<void> {
+		await this.#alarmWriteQueue.enqueue(async () => {
+			await this.#actorDriver.setAlarm(this.#actor, timestamp);
+		});
+	}
+
+	/**
+	 * Gets the next scheduled event, if any.
+	 */
+	getNextEvent(): PersistedScheduleEvent | null {
+		if (
+			!this.#persist?.scheduledEvents ||
+			this.#persist.scheduledEvents.length === 0
+		) {
+			return null;
+		}
+		return this.#persist.scheduledEvents[0];
+	}
+
+	/**
+	 * Gets all scheduled events.
+	 */
+	getAllEvents(): PersistedScheduleEvent[] {
+		return this.#persist?.scheduledEvents || [];
+	}
+
+	/**
+	 * Clears all scheduled events.
+	 * Use with caution - this removes all pending scheduled events.
+	 */
+	clearAllEvents(): void {
+		if (this.#persist?.scheduledEvents) {
+			this.#persist.scheduledEvents = [];
+			this.#actor.log.warn({ msg: "cleared all scheduled events" });
+		}
+	}
+}