rivetkit 2.0.22-rc.1 → 2.0.22

package/src/actor/instance.ts

@@ -15,7 +15,9 @@ import { PERSISTED_ACTOR_VERSIONED } from "@/schemas/actor-persist/versioned";
 import type * as protocol from "@/schemas/client-protocol/mod";
 import { TO_CLIENT_VERSIONED } from "@/schemas/client-protocol/versioned";
 import {
+	arrayBuffersEqual,
 	bufferToArrayBuffer,
+	EXTRA_ERROR_LOG,
 	getEnvUniversal,
 	promiseWithResolvers,
 	SinglePromiseQueue,
@@ -26,7 +28,7 @@ import {
 	Conn,
 	type ConnId,
 	generateConnId,
-	generateConnSocketId,
+	generateConnRequestId,
 	generateConnToken,
 } from "./conn";
 import {
@@ -45,6 +47,7 @@ import { serializeActorKey } from "./keys";
 import type {
 	PersistedActor,
 	PersistedConn,
+	PersistedHibernatableWebSocket,
 	PersistedScheduleEvent,
 } from "./persisted";
 import { processMessage } from "./protocol/old";
@@ -52,6 +55,8 @@ import { CachedSerializer } from "./protocol/serde";
 import { Schedule } from "./schedule";
 import { DeadlineError, deadline } from "./utils";
 
+export const PERSIST_SYMBOL = Symbol("persist");
+
 /**
  * Options for the `_saveState` method.
  */
@@ -131,6 +136,14 @@ export type ExtractActorConnState<A extends AnyActorInstance> =
 		? ConnState
 		: never;
 
+enum CanSleep {
+	Yes,
+	NotReady,
+	ActiveConns,
+	ActiveHonoHttpRequests,
+	ActiveRawWebSockets,
+}
+
 export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 	// Shared actor context for this instance
 	actorContext: ActorContext<S, CP, CS, V, I, DB>;
@@ -145,7 +158,7 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 	#stopCalled = false;
 
 	get isStopping() {
-		return this.#stopCalled || this.#sleepCalled;
+		return this.#stopCalled;
 	}
 
 	#persistChanged = false;
@@ -158,6 +171,10 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 	 */
 	#persist!: PersistedActor<S, CP, CS, I>;
 
+	get [PERSIST_SYMBOL](): PersistedActor<S, CP, CS, I> {
+		return this.#persist;
+	}
+
 	/** Raw state without the proxy wrapper */
 	#persistRaw!: PersistedActor<S, CP, CS, I>;
 
@@ -186,8 +203,11 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 
 	#sleepTimeout?: NodeJS.Timeout;
 
-	// Track active raw requests so sleep logic can account for them
-	#activeRawFetchCount = 0;
+	/**
+	 * Track active HTTP requests through Hono router so sleep logic can
+	 * account for them. Does not include WebSockets.
+	 **/
+	#activeHonoHttpRequests = 0;
 	#activeRawWebSockets = new Set<UniversalWebSocket>();
 
 	#schedule!: Schedule;
@@ -217,9 +237,17 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 				return Array.from(this.#connections.entries()).map(
 					([id, conn]) => ({
 						id,
-						stateEnabled: conn.__stateEnabled,
 						params: conn.params as any,
 						state: conn.__stateEnabled ? conn.state : undefined,
+						status: conn.status,
+						subscriptions: conn.subscriptions.size,
+						lastSeen: conn.lastSeen,
+						stateEnabled: conn.__stateEnabled,
+						isHibernatable: conn.isHibernatable,
+						requestId: conn.__socket?.requestId,
+						driver: conn.__driverState
+							? getConnDriverKindFromState(conn.__driverState)
+							: undefined,
 					}),
 				);
 			},
@@ -235,10 +263,10 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 				await this.saveState({ immediate: true });
 			},
 			executeAction: async (name, params) => {
-				const socketId = generateConnSocketId();
+				const requestId = generateConnRequestId();
 				const conn = await this.createConn(
 					{
-						socketId,
+						requestId: requestId,
 						driverState: { [ConnDriverKind.HTTP]: {} },
 					},
 					undefined,
@@ -252,7 +280,7 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 						params || [],
 					);
 				} finally {
-					this.__connDisconnected(conn, true, socketId);
+					this.__connDisconnected(conn, true, requestId);
 				}
 			},
 		};
@@ -271,7 +299,7 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 	}
 
 	get #sleepingSupported(): boolean {
-		return this.#actorDriver.sleep !== undefined;
+		return this.#actorDriver.startSleep !== undefined;
 	}
 
 	/**
@@ -300,6 +328,9 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 			actorId,
 		};
 
+		const extraLogParams = actorDriver.getExtraActorLogParams?.();
+		if (extraLogParams) Object.assign(logParams, extraLogParams);
+
 		this.#log = getBaseLogger().child(
 			Object.assign(
 				getIncludeTarget() ? { target: "actor" } : {},
@@ -861,16 +892,20 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 	__connDisconnected(
 		conn: Conn<S, CP, CS, V, I, DB>,
 		wasClean: boolean,
-		socketId: string,
+		requestId: string,
 	) {
 		// If socket ID is provided, check if it matches the current socket ID
 		// If it doesn't match, this is a stale disconnect event from an old socket
-		if (socketId && conn.__socket && socketId !== conn.__socket.socketId) {
+		if (
+			requestId &&
+			conn.__socket &&
+			requestId !== conn.__socket.requestId
+		) {
 			this.#rLog.debug({
 				msg: "ignoring stale disconnect event",
 				connId: conn.id,
-				eventSocketId: socketId,
-				currentSocketId: conn.__socket.socketId,
+				eventRequestId: requestId,
+				currentRequestId: conn.__socket.requestId,
 			});
 			return;
 		}
@@ -1270,8 +1305,6 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 
 	#assertReady(allowStoppingState: boolean = false) {
 		if (!this.#ready) throw new errors.InternalError("Actor not ready");
-		if (!allowStoppingState && this.#sleepCalled)
-			throw new errors.InternalError("Actor is going to sleep");
 		if (!allowStoppingState && this.#stopCalled)
 			throw new errors.InternalError("Actor is stopping");
 	}
@@ -1283,13 +1316,19 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 	#checkConnectionsLiveness() {
 		this.#rLog.debug({ msg: "checking connections liveness" });
 
+		let connected = 0;
+		let reconnecting = 0;
+		let removed = 0;
 		for (const conn of this.#connections.values()) {
 			if (conn.__status === "connected") {
+				connected += 1;
 				this.#rLog.debug({
 					msg: "connection is alive",
 					connId: conn.id,
 				});
 			} else {
+				reconnecting += 1;
+
 				const lastSeen = conn.__persist.lastSeen;
 				const sinceLastSeen = Date.now() - lastSeen;
 				if (
@@ -1313,9 +1352,18 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 				});
 
 				// Assume that the connection is dead here, no need to disconnect anything
+				removed += 1;
 				this.#removeConn(conn);
 			}
 		}
+
+		this.#rLog.debug({
+			msg: "checked connection liveness",
+			total: connected + reconnecting,
+			connected,
+			reconnecting,
+			removed,
+		});
 	}
 
 	/**
@@ -1482,10 +1530,6 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 			throw new errors.FetchHandlerNotDefined();
 		}
 
-		// Track active raw fetch while handler runs
-		this.#activeRawFetchCount++;
-		this.#resetSleepTimer();
-
 		try {
 			const response = await this.#config.onFetch(
 				this.actorContext,
@@ -1503,12 +1547,6 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 			});
 			throw error;
 		} finally {
-			// Decrement active raw fetch counter and re-evaluate sleep
-			this.#activeRawFetchCount = Math.max(
-				0,
-				this.#activeRawFetchCount - 1,
-			);
-			this.#resetSleepTimer();
 			this.#savePersistThrottled();
 		}
 	}
@@ -1534,17 +1572,116 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 			this.#activeRawWebSockets.add(websocket);
 			this.#resetSleepTimer();
 
-			// Track socket close
-			const onSocketClosed = () => {
+			// Track hibernatable WebSockets
+			let rivetRequestId: ArrayBuffer | undefined;
+			let persistedHibernatableWebSocket:
+				| PersistedHibernatableWebSocket
+				| undefined;
+
+			const onSocketOpened = (event: any) => {
+				rivetRequestId = event?.rivetRequestId;
+
+				// Find hibernatable WS
+				if (rivetRequestId) {
+					const rivetRequestIdLocal = rivetRequestId;
+					persistedHibernatableWebSocket =
+						this.#persist.hibernatableWebSocket.find((ws) =>
+							arrayBuffersEqual(
+								ws.requestId,
+								rivetRequestIdLocal,
+							),
+						);
+
+					if (persistedHibernatableWebSocket) {
+						persistedHibernatableWebSocket.lastSeenTimestamp =
+							BigInt(Date.now());
+					}
+				}
+
+				this.#rLog.debug({
+					msg: "actor instance onSocketOpened",
+					rivetRequestId,
+					isHibernatable: !!persistedHibernatableWebSocket,
+					hibernationMsgIndex:
+						persistedHibernatableWebSocket?.msgIndex,
+				});
+			};
+
+			const onSocketMessage = (event: any) => {
+				// Update state of hibernatable WS
+				if (persistedHibernatableWebSocket) {
+					persistedHibernatableWebSocket.lastSeenTimestamp = BigInt(
+						Date.now(),
+					);
+					persistedHibernatableWebSocket.msgIndex = BigInt(
+						event.rivetMessageIndex,
+					);
+				}
+
+				this.#rLog.debug({
+					msg: "actor instance onSocketMessage",
+					rivetRequestId,
+					isHibernatable: !!persistedHibernatableWebSocket,
+					hibernationMsgIndex:
+						persistedHibernatableWebSocket?.msgIndex,
+				});
+			};
+
+			const onSocketClosed = (_event: any) => {
+				// Remove hibernatable WS
+				if (rivetRequestId) {
+					const rivetRequestIdLocal = rivetRequestId;
+					const wsIndex =
+						this.#persist.hibernatableWebSocket.findIndex((ws) =>
+							arrayBuffersEqual(
+								ws.requestId,
+								rivetRequestIdLocal,
+							),
+						);
+
+					const removed = this.#persist.hibernatableWebSocket.splice(
+						wsIndex,
+						1,
+					);
+					if (removed.length > 0) {
+						this.#rLog.debug({
+							msg: "removed hibernatable websocket",
+							rivetRequestId,
+							hibernationMsgIndex:
+								persistedHibernatableWebSocket?.msgIndex,
+						});
+					} else {
+						this.#rLog.warn({
+							msg: "could not find hibernatable websocket to remove",
+							rivetRequestId,
+							hibernationMsgIndex:
+								persistedHibernatableWebSocket?.msgIndex,
+						});
+					}
+				}
+
+				this.#rLog.debug({
+					msg: "actor instance onSocketMessage",
+					rivetRequestId,
+					isHibernatable: !!persistedHibernatableWebSocket,
+					hibernatableWebSocketCount:
+						this.#persist.hibernatableWebSocket.length,
+				});
+
 				// Remove listener and socket from tracking
 				try {
+					websocket.removeEventListener("open", onSocketOpened);
+					websocket.removeEventListener("message", onSocketMessage);
 					websocket.removeEventListener("close", onSocketClosed);
 					websocket.removeEventListener("error", onSocketClosed);
 				} catch {}
 				this.#activeRawWebSockets.delete(websocket);
 				this.#resetSleepTimer();
 			};
+
 			try {
+				websocket.addEventListener("open", onSocketOpened);
+				websocket.addEventListener("message", onSocketMessage);
 				websocket.addEventListener("close", onSocketClosed);
 				websocket.addEventListener("error", onSocketClosed);
 			} catch {}
@@ -1746,6 +1883,29 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 		}
 	}
 
+	/**
+	 * Called by router middleware when an HTTP request begins.
+	 */
+	__beginHonoHttpRequest() {
+		this.#activeHonoHttpRequests++;
+		this.#resetSleepTimer();
+	}
+
+	/**
+	 * Called by router middleware when an HTTP request ends.
+	 */
+	__endHonoHttpRequest() {
+		this.#activeHonoHttpRequests--;
+		if (this.#activeHonoHttpRequests < 0) {
+			this.#activeHonoHttpRequests = 0;
+			this.#rLog.warn({
+				msg: "active hono requests went below 0, this is a RivetKit bug",
+				...EXTRA_ERROR_LOG,
+			});
+		}
+		this.#resetSleepTimer();
+	}
+
 	// MARK: Sleep
 	/**
 	 * Reset timer from the last actor interaction that allows it to be put to sleep.
@@ -1764,8 +1924,9 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 
 		this.#rLog.debug({
 			msg: "resetting sleep timer",
-			canSleep,
+			canSleep: CanSleep[canSleep],
 			existingTimeout: !!this.#sleepTimeout,
+			timeout: this.#config.options.sleepTimeout,
 		});
 
 		if (this.#sleepTimeout) {
@@ -1776,64 +1937,84 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 		// Don't set a new timer if already sleeping
 		if (this.#sleepCalled) return;
 
-		if (canSleep) {
+		if (canSleep === CanSleep.Yes) {
 			this.#sleepTimeout = setTimeout(() => {
-				this._sleep().catch((error) => {
-					this.#rLog.error({
-						msg: "error during sleep",
-						error: stringifyError(error),
-					});
-				});
+				this._startSleep();
 			}, this.#config.options.sleepTimeout);
 		}
 	}
 
 	/** If this actor can be put in a sleeping state. */
-	#canSleep(): boolean {
-		if (!this.#ready) return false;
+	#canSleep(): CanSleep {
+		if (!this.#ready) return CanSleep.NotReady;
+
+		// Do not sleep if Hono HTTP requests are in-flight
+		if (this.#activeHonoHttpRequests > 0)
+			return CanSleep.ActiveHonoHttpRequests;
+
+		// TODO: When WS hibernation is ready, update this to only count non-hibernatable websockets
+		// Do not sleep if there are raw websockets open
+		if (this.#activeRawWebSockets.size > 0)
+			return CanSleep.ActiveRawWebSockets;
 
 		// Check for active conns. This will also cover active actions, since all actions have a connection.
 		for (const conn of this.#connections.values()) {
-			if (conn.status === "connected") return false;
-		}
-
-		// Do not sleep if raw fetches are in-flight
-		if (this.#activeRawFetchCount > 0) return false;
+			// TODO: Enable this when hibernation is implemented. We're waiting on support for Guard to not auto-wake the actor if it sleeps.
+			// if (conn.status === "connected" && !conn.isHibernatable)
+			// 	return false;
 
-		// Do not sleep if there are raw websockets open
-		if (this.#activeRawWebSockets.size > 0) return false;
+			if (conn.status === "connected") return CanSleep.ActiveConns;
+		}
 
-		return true;
+		return CanSleep.Yes;
 	}
 
-	/** Puts an actor to sleep. This should just start the sleep sequence, most shutdown logic should be in _stop (which is called by the ActorDriver when sleeping). */
-	async _sleep() {
-		const sleep = this.#actorDriver.sleep?.bind(
-			this.#actorDriver,
-			this.#actorId,
-		);
-		invariant(this.#sleepingSupported, "sleeping not supported");
-		invariant(sleep, "no sleep on driver");
-
+	/**
+	 * Puts an actor to sleep. This should just start the sleep sequence, most shutdown logic should be in _stop (which is called by the ActorDriver when sleeping).
+	 *
+	 * For the engine, this will:
+	 * 1. Publish EventActorIntent with ActorIntentSleep (via driver.startSleep)
+	 * 2. Engine runner will wait for CommandStopActor
+	 * 3. Engine runner will call _onStop and wait for it to finish
+	 * 4. Engine runner will publish EventActorStateUpdate with ActorStateSTop
+	 **/
+	_startSleep() {
+		// IMPORTANT: #sleepCalled should have no effect on the actor's
+		// behavior aside from preventing calling _startSleep twice. Wait for
+		// `_onStop` before putting in a stopping state.
 		if (this.#sleepCalled) {
 			this.#rLog.warn({ msg: "already sleeping actor" });
 			return;
 		}
 		this.#sleepCalled = true;
 
+		// NOTE: Publishes ActorIntentSleep
+		const sleep = this.#actorDriver.startSleep?.bind(
+			this.#actorDriver,
+			this.#actorId,
+		);
+		invariant(this.#sleepingSupported, "sleeping not supported");
+		invariant(sleep, "no sleep on driver");
+
 		this.#rLog.info({ msg: "actor sleeping" });
 
 		// Schedule sleep to happen on the next tick. This allows for any action that calls _sleep to complete.
-		setImmediate(async () => {
+		setImmediate(() => {
 			// The actor driver should call stop when ready to stop
 			//
 			// This will call _stop once Pegboard responds with the new status
-			await sleep();
+			sleep();
 		});
 	}
 
 	// MARK: Stop
-	async _stop() {
+	/**
+	 * For the engine:
+	 * 1. Engine runner receives CommandStopActor
+	 * 2. Engine runner calls _onStop and waits for it to finish
+	 * 3. Engine runner publishes EventActorStateUpdate with ActorStateSTop
+	 */
+	async _onStop() {
 		if (this.#stopCalled) {
 			this.#rLog.warn({ msg: "already stopping actor" });
 			return;
@@ -1980,6 +2161,11 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 					},
 				},
 			})),
+			hibernatableWebSocket: persist.hibernatableWebSocket.map((ws) => ({
+				requestId: ws.requestId,
+				lastSeenTimestamp: ws.lastSeenTimestamp,
+				msgIndex: ws.msgIndex,
+			})),
 		};
 	}
 
@@ -2012,6 +2198,11 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 					},
 				},
 			})),
+			hibernatableWebSocket: bareData.hibernatableWebSocket.map((ws) => ({
+				requestId: ws.requestId,
+				lastSeenTimestamp: ws.lastSeenTimestamp,
+				msgIndex: ws.msgIndex,
+			})),
 		};
 	}
 }

package/src/actor/persisted.ts

@@ -5,6 +5,7 @@ export interface PersistedActor<S, CP, CS, I> {
 	state: S;
 	connections: PersistedConn<CP, CS>[];
 	scheduledEvents: PersistedScheduleEvent[];
+	hibernatableWebSocket: PersistedHibernatableWebSocket[];
 }
 
 /** Object representing connection that gets persisted to storage. */
@@ -37,3 +38,9 @@ export interface PersistedScheduleEvent {
 	timestamp: number;
 	kind: PersistedScheduleEventKind;
 }
+
+export interface PersistedHibernatableWebSocket {
+	requestId: ArrayBuffer;
+	lastSeenTimestamp: bigint;
+	msgIndex: bigint;
+}