rivetkit 2.0.5 → 2.0.7-rc.1

package/src/actor/{connection.ts → conn.ts}

@@ -1,12 +1,19 @@
 import * as cbor from "cbor-x";
+import invariant from "invariant";
 import type * as protocol from "@/schemas/client-protocol/mod";
 import { TO_CLIENT_VERSIONED } from "@/schemas/client-protocol/versioned";
 import { bufferToArrayBuffer } from "@/utils";
+import {
+	CONN_DRIVERS,
+	ConnDriverKind,
+	type ConnDriverState,
+	ConnReadyState,
+	getConnDriverKindFromState,
+} from "./conn-drivers";
+import type { ConnSocket } from "./conn-socket";
 import type { AnyDatabaseProvider } from "./database";
-import { type ConnDriver, ConnectionReadyState } from "./driver";
 import * as errors from "./errors";
 import type { ActorInstance } from "./instance";
-import { loggerWithoutContext } from "./log";
 import type { PersistedConn } from "./persisted";
 import { CachedSerializer } from "./protocol/serde";
 import { generateSecureToken } from "./utils";
@@ -19,23 +26,16 @@ export function generateConnToken(): string {
 	return generateSecureToken(32);
 }
 
+export function generateConnSocketId(): string {
+	return crypto.randomUUID();
+}
+
 export type ConnId = string;
 
 export type AnyConn = Conn<any, any, any, any, any, any>;
 
-export const CONNECTION_DRIVER_WEBSOCKET = "webSocket";
-export const CONNECTION_DRIVER_SSE = "sse";
-export const CONNECTION_DRIVER_HTTP = "http";
-
-export type ConnectionDriver =
-	| typeof CONNECTION_DRIVER_WEBSOCKET
-	| typeof CONNECTION_DRIVER_SSE
-	| typeof CONNECTION_DRIVER_HTTP;
-
 export type ConnectionStatus = "connected" | "reconnecting";
 
-export const CONNECTION_CHECK_LIVENESS_SYMBOL = Symbol("checkLiveness");
-
 /**
  * Represents a client connection to a actor.
  *
@@ -46,13 +46,9 @@ export const CONNECTION_CHECK_LIVENESS_SYMBOL = Symbol("checkLiveness");
 export class Conn<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 	subscriptions: Set<string> = new Set<string>();
 
-	#stateEnabled: boolean;
-
 	// TODO: Remove this cyclical reference
 	#actor: ActorInstance<S, CP, CS, V, I, DB>;
 
-	#status: ConnectionStatus = "connected";
-
 	/**
 	 * The proxied state that notifies of changes automatically.
 	 *
@@ -60,23 +56,31 @@ export class Conn<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 	 */
 	__persist: PersistedConn<CP, CS>;
 
+	get __driverState(): ConnDriverState | undefined {
+		return this.__socket?.driverState;
+	}
+
 	/**
-	 * Driver used to manage realtime connection communication.
+	 * Socket connected to this connection.
 	 *
-	 * @protected
+	 * If undefined, then nothing is connected to this.
 	 */
-	#driver: ConnDriver;
+	__socket?: ConnSocket;
 
-	public get params(): CP {
-		return this.__persist.params;
+	get __status(): ConnectionStatus {
+		if (this.__socket) {
+			return "connected";
+		} else {
+			return "reconnecting";
+		}
 	}
 
-	public get driver(): ConnectionDriver {
-		return this.__persist.connDriver as ConnectionDriver;
+	public get params(): CP {
+		return this.__persist.params;
 	}
 
-	public get _stateEnabled() {
-		return this.#stateEnabled;
+	public get __stateEnabled() {
+		return this.#actor.connStateEnabled;
 	}
 
 	/**
@@ -118,7 +122,7 @@ export class Conn<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 	 * Status of the connection.
 	 */
 	public get status(): ConnectionStatus {
-		return this.#status;
+		return this.__status;
 	}
 
 	/**
@@ -138,17 +142,13 @@ export class Conn<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 	public constructor(
 		actor: ActorInstance<S, CP, CS, V, I, DB>,
 		persist: PersistedConn<CP, CS>,
-		driver: ConnDriver,
-		stateEnabled: boolean,
 	) {
 		this.#actor = actor;
 		this.__persist = persist;
-		this.#driver = driver;
-		this.#stateEnabled = stateEnabled;
 	}
 
 	#validateStateEnabled() {
-		if (!this.#stateEnabled) {
+		if (!this.__stateEnabled) {
 			throw new errors.ConnStateNotEnabled();
 		}
 	}
@@ -161,12 +161,28 @@ export class Conn<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 	 * @protected
 	 */
 	public _sendMessage(message: CachedSerializer<protocol.ToClient>) {
-		this.#driver.sendMessage?.(
-			this.#actor,
-			this,
-			this.__persist.connDriverState,
-			message,
-		);
+		if (this.__driverState) {
+			const driverKind = getConnDriverKindFromState(this.__driverState);
+			const driver = CONN_DRIVERS[driverKind];
+			if (driver.sendMessage) {
+				driver.sendMessage(
+					this.#actor,
+					this,
+					(this.__driverState as any)[driverKind],
+					message,
+				);
+			} else {
+				this.#actor.rLog.debug({
+					msg: "conn driver does not support sending messages",
+					conn: this.id,
+				});
+			}
+		} else {
+			this.#actor.rLog.warn({
+				msg: "missing connection driver state for send message",
+				conn: this.id,
+			});
+		}
 	}
 
 	/**
@@ -205,53 +221,31 @@ export class Conn<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 	 * @param reason - The reason for disconnection.
 	 */
 	public async disconnect(reason?: string) {
-		this.#status = "reconnecting";
-		await this.#driver.disconnect(
-			this.#actor,
-			this,
-			this.__persist.connDriverState,
-			reason,
-		);
-	}
-
-	/**
-	 * This method checks the connection's liveness by querying the driver for its ready state.
-	 * If the connection is not closed, it updates the last liveness timestamp and returns `true`.
-	 * Otherwise, it returns `false`.
-	 * @internal
-	 */
-	[CONNECTION_CHECK_LIVENESS_SYMBOL]() {
-		const readyState = this.#driver.getConnectionReadyState(this.#actor, this);
-
-		const isConnectionClosed =
-			readyState === ConnectionReadyState.CLOSED ||
-			readyState === ConnectionReadyState.CLOSING ||
-			readyState === undefined;
-
-		const newLastSeen = Date.now();
-		const newStatus = isConnectionClosed ? "reconnecting" : "connected";
-
-		this.#actor.rLog.debug({
-			msg: "liveness probe for connection",
-			connId: this.id,
-			actorId: this.#actor.id,
-			readyState,
-
-			status: this.#status,
-			newStatus,
-
-			lastSeen: this.__persist.lastSeen,
-			currentTs: newLastSeen,
-		});
-
-		if (!isConnectionClosed) {
-			this.__persist.lastSeen = newLastSeen;
+		if (this.__socket && this.__driverState) {
+			const driverKind = getConnDriverKindFromState(this.__driverState);
+			const driver = CONN_DRIVERS[driverKind];
+			if (driver.disconnect) {
+				driver.disconnect(
+					this.#actor,
+					this,
+					(this.__driverState as any)[driverKind],
+					reason,
+				);
+			} else {
+				this.#actor.rLog.debug({
+					msg: "no disconnect handler for conn driver",
+					conn: this.id,
+				});
+			}
+
+			this.#actor.__connDisconnected(this, true, this.__socket.socketId);
+		} else {
+			this.#actor.rLog.warn({
+				msg: "missing connection driver state for disconnect",
+				conn: this.id,
+			});
 		}
 
-		this.#status = newStatus;
-		return {
-			status: this.#status,
-			lastSeen: this.__persist.lastSeen,
-		};
+		this.__socket = undefined;
 	}
 }

package/src/actor/context.ts

@@ -2,7 +2,7 @@ import type { ActorKey } from "@/actor/mod";
 import type { Client } from "@/client/client";
 import type { Logger } from "@/common/log";
 import type { Registry } from "@/registry/mod";
-import type { Conn, ConnId } from "./connection";
+import type { Conn, ConnId } from "./conn";
 import type { AnyDatabaseProvider, InferDatabaseClient } from "./database";
 import type { ActorInstance, SaveStateOptions } from "./instance";
 import type { Schedule } from "./schedule";

package/src/actor/driver.ts

@@ -1,15 +1,10 @@
-import type { CachedSerializer } from "@/actor/protocol/serde";
+import type { Context as HonoContext } from "hono";
 import type { AnyClient } from "@/client/client";
 import type { ManagerDriver } from "@/manager/driver";
 import type { RegistryConfig } from "@/registry/config";
 import type { RunConfig } from "@/registry/run-config";
-import type * as protocol from "@/schemas/client-protocol/mod";
-import type { AnyConn, ConnectionDriver } from "./connection";
-import type { GenericConnGlobalState } from "./generic-conn-driver";
 import type { AnyActorInstance } from "./instance";
 
-export type ConnectionDriversMap = Record<ConnectionDriver, ConnDriver>;
-
 export type ActorDriverBuilder = (
 	registryConfig: RegistryConfig,
 	runConfig: RunConfig,
@@ -22,8 +17,6 @@ export interface ActorDriver {
 
 	loadActor(actorId: string): Promise<AnyActorInstance>;
 
-	getGenericConnGlobalState(actorId: string): GenericConnGlobalState;
-
 	getContext(actorId: string): unknown;
 
 	readPersistedData(actorId: string): Promise<Uint8Array | undefined>;
@@ -45,40 +38,8 @@ export interface ActorDriver {
 	sleep?(actorId: string): Promise<void>;
 
 	shutdown?(immediate: boolean): Promise<void>;
-}
-
-export enum ConnectionReadyState {
-	UNKNOWN = -1,
-	CONNECTING = 0,
-	OPEN = 1,
-	CLOSING = 2,
-	CLOSED = 3,
-}
-
-export interface ConnDriver<ConnDriverState = unknown> {
-	sendMessage?(
-		actor: AnyActorInstance,
-		conn: AnyConn,
-		state: ConnDriverState,
-		message: CachedSerializer<protocol.ToClient>,
-	): void;
 
-	/**
-	 * This returns a promise since we commonly disconnect at the end of a program, and not waiting will cause the socket to not close cleanly.
-	 */
-	disconnect(
-		actor: AnyActorInstance,
-		conn: AnyConn,
-		state: ConnDriverState,
-		reason?: string,
-	): Promise<void>;
-
-	/**
-	 * Returns the ready state of the connection.
-	 * This is used to determine if the connection is ready to send messages, or if the connection is stale.
-	 */
-	getConnectionReadyState(
-		actor: AnyActorInstance,
-		conn: AnyConn,
-	): ConnectionReadyState | undefined;
+	// Serverless
+	/** This handles the serverless start request. This should manage the lifecycle of the runner tied to the request lifecycle. */
+	serverlessHandleStart?(c: HonoContext): Promise<Response>;
 }