feat(core): disconnect stale connections (#1154)

Closes FRONT-742

### TL;DR

Added connection liveness checking to detect and clean up stale WebSocket connections.

### What changed?

- Implemented a connection liveness system that periodically checks if WebSocket and SSE connections are still active
- Added `connectionLivenessInterval` and `connectionLivenessTimeout` configuration options to the actor lifecycle settings
- Enhanced the `Conn` class with liveness tracking and checking capabilities
- Added `getConnectionReadyState` method to connection drivers to report connection status
- Created a test actor (`connLivenessActor`) and test suite to verify liveness functionality
- Implemented internal event scheduling for periodic liveness checks
- Added timestamp tracking for last seen activity on connections

### How to test?

The PR includes a comprehensive test suite in `actor-conn.ts` that verifies:
1. Connections report correct liveness status
2. Dead connections are properly detected
3. Stale connections are automatically cleaned up after the timeout period
4. Active connections continue to function normally

Run the tests with:
```
pnpm test
```

### Why make this change?

WebSocket connections can become stale or disconnected without proper notification, leading to resource leaks and potential issues with connection management. This change implements a robust mechanism to detect and clean up these stale connections automatically, improving system reliability and resource utilization.

The liveness check system ensures that the actor framework maintains an accurate view of active connections and can properly manage resources by removing connections that are no longer active.

Author: jog1t

packages/core/fixtures/driver-test-suite/conn-liveness.ts

@@ -0,0 +1,52 @@
+import { actor, CONNECTION_DRIVER_WEBSOCKET } from "@rivetkit/core";
+
+export const connLivenessActor = actor({
+	onAuth: () => {},
+	state: {
+		counter: 0,
+		acceptingConnections: true,
+	},
+	options: {
+		lifecycle: {
+			connectionLivenessInterval: 5_000,
+			connectionLivenessTimeout: 2_500,
+		},
+	},
+	onConnect: (c, conn) => {
+		if (!c.state.acceptingConnections) {
+			conn.disconnect();
+			throw new Error("Actor is not accepting connections");
+		}
+	},
+	actions: {
+		getWsConnectionsLiveness: (c) => {
+			return Array.from(c.conns.values())
+				.filter((conn) => conn.driver === CONNECTION_DRIVER_WEBSOCKET)
+				.map((conn) => ({
+					id: conn.id,
+					status: conn.status,
+					lastSeen: conn.lastSeen,
+				}));
+		},
+		getConnectionId: (c) => {
+			return c.conn.id;
+		},
+		kill: (c, connId: string) => {
+			c.state.acceptingConnections = false;
+			// Disconnect the connection with the given ID
+			// This simulates a network failure or a manual disconnection
+			// The connection will be cleaned up by the actor manager after the timeout
+			const conn = c.conns.get(connId);
+			if (conn) {
+				conn.disconnect();
+			}
+		},
+		getCounter: (c) => {
+			return c.state.counter;
+		},
+		increment: (c, amount: number) => {
+			c.state.counter += amount;
+			return c.state.counter;
+		},
+	},
+});

packages/core/fixtures/driver-test-suite/registry.ts

@@ -19,6 +19,7 @@ import {
 	noAuthActor,
 	publicActor,
 } from "./auth";
+import { connLivenessActor } from "./conn-liveness";
 import { counterWithParams } from "./conn-params";
 import { connStateActor } from "./conn-state";
 // Import actors from individual files
@@ -87,6 +88,8 @@ export const registry = setup({
 		counterWithParams,
 		// From conn-state.ts
 		connStateActor,
+		// From actor-conn.ts
+		connLivenessActor,
 		// From metadata.ts
 		metadataActor,
 		// From vars.ts

packages/core/scripts/dump-openapi.ts

@@ -1,25 +1,10 @@
 import * as fs from "node:fs/promises";
 import { resolve } from "node:path";
-import { Context } from "hono";
-import WebSocket from "ws";
 import type { ClientDriver } from "@/client/client";
 import { createFileSystemOrMemoryDriver } from "@/drivers/file-system/mod";
-import type {
-	ActorOutput,
-	CreateInput,
-	GetForIdInput,
-	GetOrCreateWithKeyInput,
-	GetWithKeyInput,
-	ManagerDriver,
-} from "@/manager/driver";
-import { ActorQuery } from "@/manager/protocol/query";
+import type { ManagerDriver } from "@/manager/driver";
 import { createManagerRouter } from "@/manager/router";
-import {
-	Encoding,
-	type RegistryConfig,
-	RegistryConfigSchema,
-	setup,
-} from "@/mod";
+import { type RegistryConfig, RegistryConfigSchema, setup } from "@/mod";
 import { type RunConfig, RunConfigSchema } from "@/registry/run-config";
 import { VERSION } from "@/utils";
 

packages/core/src/actor/config.ts

@@ -65,6 +65,8 @@ export const ActorConfigSchema = z
 						createVarsTimeout: z.number().positive().default(5000),
 						createConnStateTimeout: z.number().positive().default(5000),
 						onConnectTimeout: z.number().positive().default(5000),
+						connectionLivenessTimeout: z.number().positive().default(2500),
+						connectionLivenessInterval: z.number().positive().default(5000),
 					})
 					.strict()
 					.default({}),

packages/core/src/actor/connection.ts

@@ -1,9 +1,10 @@
 import type * as messageToClient from "@/actor/protocol/message/to-client";
 import type * as wsToClient from "@/actor/protocol/message/to-client";
 import type { AnyDatabaseProvider } from "./database";
-import type { ConnDriver } from "./driver";
+import { type ConnDriver, ConnectionReadyState } from "./driver";
 import * as errors from "./errors";
 import type { ActorInstance } from "./instance";
+import { logger } from "./log";
 import type { PersistedConn } from "./persisted";
 import { CachedSerializer } from "./protocol/serde";
 import { generateSecureToken } from "./utils";
@@ -20,6 +21,19 @@ export type ConnId = string;
 
 export type AnyConn = Conn<any, 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.
  *
@@ -35,6 +49,8 @@ export class Conn<S, CP, CS, V, I, AD, DB extends AnyDatabaseProvider> {
 	// TODO: Remove this cyclical reference
 	#actor: ActorInstance<S, CP, CS, V, I, AD, DB>;
 
+	#status: ConnectionStatus = "connected";
+
 	/**
 	 * The proxied state that notifies of changes automatically.
 	 *
@@ -57,6 +73,10 @@ export class Conn<S, CP, CS, V, I, AD, DB extends AnyDatabaseProvider> {
 		return this.__persist.a as AD;
 	}
 
+	public get driver(): ConnectionDriver {
+		return this.__persist.d as ConnectionDriver;
+	}
+
 	public get _stateEnabled() {
 		return this.#stateEnabled;
 	}
@@ -96,6 +116,20 @@ export class Conn<S, CP, CS, V, I, AD, DB extends AnyDatabaseProvider> {
 		return this.__persist.t;
 	}
 
+	/**
+	 * Status of the connection.
+	 */
+	public get status(): ConnectionStatus {
+		return this.#status;
+	}
+
+	/**
+	 * Timestamp of the last time the connection was seen, i.e. the last time the connection was active and checked for liveness.
+	 */
+	public get lastSeen(): number {
+		return this.__persist.l;
+	}
+
 	/**
 	 * Initializes a new instance of the Connection class.
 	 *
@@ -164,6 +198,50 @@ export class Conn<S, CP, CS, V, I, AD, 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.ds, 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";
+
+		logger().debug("liveness probe for connection", {
+			connId: this.id,
+			actorId: this.#actor.id,
+			readyState,
+
+			status: this.#status,
+			newStatus,
+
+			lastSeen: this.__persist.l,
+			currentTs: newLastSeen,
+		});
+
+		if (!isConnectionClosed) {
+			this.__persist.l = newLastSeen;
+		}
+
+		this.#status = newStatus;
+		return {
+			status: this.#status,
+			lastSeen: this.__persist.l,
+		};
+	}
 }

packages/core/src/actor/driver.ts

@@ -1,15 +1,14 @@
 import type * as messageToClient from "@/actor/protocol/message/to-client";
 import type { CachedSerializer } from "@/actor/protocol/serde";
 import type { AnyClient } from "@/client/client";
-import type { ActorInspector } from "@/inspector/actor";
 import type { ManagerDriver } from "@/manager/driver";
 import type { RegistryConfig } from "@/registry/config";
 import type { RunConfig } from "@/registry/run-config";
-import type { AnyConn } from "./connection";
+import type { AnyConn, ConnectionDriver } from "./connection";
 import type { GenericConnGlobalState } from "./generic-conn-driver";
 import type { AnyActorInstance } from "./instance";
 
-export type ConnDrivers = Record<string, ConnDriver>;
+export type ConnectionDriversMap = Record<ConnectionDriver, ConnDriver>;
 
 export type ActorDriverBuilder = (
 	registryConfig: RegistryConfig,
@@ -45,6 +44,14 @@ export interface ActorDriver {
 	//readState(): void;
 }
 
+export enum ConnectionReadyState {
+	UNKNOWN = -1,
+	CONNECTING = 0,
+	OPEN = 1,
+	CLOSING = 2,
+	CLOSED = 3,
+}
+
 export interface ConnDriver<ConnDriverState = unknown> {
 	sendMessage?(
 		actor: AnyActorInstance,
@@ -62,4 +69,13 @@ export interface ConnDriver<ConnDriverState = unknown> {
 		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;
 }

packages/core/src/actor/generic-conn-driver.ts

@@ -1,13 +1,21 @@
 import type { SSEStreamingApi } from "hono/streaming";
 import type { WSContext } from "hono/ws";
 import type { WebSocket } from "ws";
-import type { AnyConn } from "@/actor/connection";
-import type { ConnDriver } from "@/actor/driver";
+import {
+	type AnyConn,
+	CONNECTION_DRIVER_HTTP,
+	CONNECTION_DRIVER_SSE,
+	CONNECTION_DRIVER_WEBSOCKET,
+} from "@/actor/connection";
+import {
+	type ConnDriver,
+	type ConnectionDriversMap,
+	ConnectionReadyState,
+} from "@/actor/driver";
 import type { AnyActorInstance } from "@/actor/instance";
 import type * as messageToClient from "@/actor/protocol/message/to-client";
 import type { CachedSerializer, Encoding } from "@/actor/protocol/serde";
 import { encodeDataToString } from "@/actor/protocol/serde";
-import { dbg } from "@/utils";
 import { logger } from "./log";
 
 // This state is different than `PersistedConn` state since the connection-specific state is persisted & must be serializable. This is also part of the connection driver, not part of the core actor.
@@ -25,17 +33,15 @@ export class GenericConnGlobalState {
  */
 export function createGenericConnDrivers(
 	globalState: GenericConnGlobalState,
-): Record<string, ConnDriver> {
+): ConnectionDriversMap {
 	return {
-		[CONN_DRIVER_GENERIC_WEBSOCKET]: createGenericWebSocketDriver(globalState),
-		[CONN_DRIVER_GENERIC_SSE]: createGenericSseDriver(globalState),
-		[CONN_DRIVER_GENERIC_HTTP]: createGeneircHttpDriver(),
+		[CONNECTION_DRIVER_WEBSOCKET]: createGenericWebSocketDriver(globalState),
+		[CONNECTION_DRIVER_SSE]: createGenericSseDriver(globalState),
+		[CONNECTION_DRIVER_HTTP]: createGenericHttpDriver(),
 	};
 }
 
 // MARK: WebSocket
-export const CONN_DRIVER_GENERIC_WEBSOCKET = "genericWebSocket";
-
 export interface GenericWebSocketDriverState {
 	encoding: Encoding;
 }
@@ -114,7 +120,6 @@ export function createGenericWebSocketDriver(
 				});
 				return;
 			}
-
 			const raw = ws.raw as WebSocket;
 			if (!raw) {
 				logger().warn("ws.raw does not exist");
@@ -130,12 +135,27 @@ export function createGenericWebSocketDriver(
 
 			await promise;
 		},
+
+		getConnectionReadyState: (
+			_actor: AnyActorInstance,
+			conn: AnyConn,
+		): ConnectionReadyState | undefined => {
+			const ws = globalState.websockets.get(conn.id);
+			if (!ws) {
+				logger().warn("missing ws for getConnectionReadyState", {
+					connId: conn.id,
+				});
+				return undefined;
+			}
+
+			const raw = ws.raw as WebSocket;
+
+			return raw.readyState as ConnectionReadyState;
+		},
 	};
 }
 
 // MARK: SSE
-export const CONN_DRIVER_GENERIC_SSE = "genericSse";
-
 export interface GenericSseDriverState {
 	encoding: Encoding;
 }
@@ -174,15 +194,32 @@ export function createGenericSseDriver(globalState: GenericConnGlobalState) {
 
 			stream.close();
 		},
+
+		getConnectionReadyState: (
+			_actor: AnyActorInstance,
+			conn: AnyConn,
+		): ConnectionReadyState | undefined => {
+			const stream = globalState.sseStreams.get(conn.id);
+			if (!stream) {
+				logger().warn("missing sse stream for getConnectionReadyState", {
+					connId: conn.id,
+				});
+				return undefined;
+			}
+
+			if (stream.aborted || stream.closed) {
+				return ConnectionReadyState.CLOSED;
+			}
+
+			return ConnectionReadyState.OPEN;
+		},
 	};
 }
 
 // MARK: HTTP
-export const CONN_DRIVER_GENERIC_HTTP = "genericHttp";
-
 export type GenericHttpDriverState = Record<never, never>;
 
-export function createGeneircHttpDriver() {
+export function createGenericHttpDriver() {
 	return {
 		disconnect: async () => {
 			// Noop