feat(rivetkit): add ConnectionMap readonly Map wrapper for actor connections

Author: abcxff

rivetkit-rust/packages/rivetkit-core/src/actor/context.rs

@@ -13,6 +13,7 @@ use parking_lot::{Mutex, RwLock};
 use rivet_error::ActorSpecifier;
 use rivet_envoy_client::handle::EnvoyHandle;
 use rivet_envoy_client::tunnel::HibernatingWebSocketMetadata;
+use rivet_error::ActorSpecifier;
 use scc::HashMap as SccHashMap;
 use tokio::runtime::Handle;
 use tokio::sync::{Mutex as AsyncMutex, Notify, OnceCell, broadcast, mpsc, oneshot};

rivetkit-typescript/packages/rivetkit/src/registry/connection-map.ts

@@ -0,0 +1,89 @@
+import type { ActorContextHandle, ConnHandle, CoreRuntime } from "./runtime";
+import type { NativeValidationConfig } from "./native-validation";
+import { callNativeSync, NativeConnAdapter } from "./native";
+
+export class ConnectionMap implements ReadonlyMap<string, NativeConnAdapter> {
+	#runtime: CoreRuntime;
+	#ctx: ActorContextHandle;
+	#schemas: NativeValidationConfig;
+
+	constructor(
+		runtime: CoreRuntime,
+		ctx: ActorContextHandle,
+		schemas: NativeValidationConfig,
+	) {
+		this.#runtime = runtime;
+		this.#ctx = ctx;
+		this.#schemas = schemas;
+	}
+
+	#connToAdapter(conn: ConnHandle): NativeConnAdapter {
+		return new NativeConnAdapter(
+			this.#runtime,
+			conn,
+			this.#schemas,
+			this.#ctx,
+			(connId) =>
+				callNativeSync(() =>
+					this.#runtime.actorQueueHibernationRemoval(
+						this.#ctx,
+						connId,
+					),
+				),
+		);
+	}
+
+	get size(): number {
+		return callNativeSync(() => this.#runtime.actorConns(this.#ctx)).length;
+	}
+
+	get(key: string): NativeConnAdapter | undefined {
+		const conns = callNativeSync(() => this.#runtime.actorConns(this.#ctx));
+		const conn = conns.find(
+			(c) => this.#runtime.connId(c) === key,
+		);
+		if (!conn) return undefined;
+		return this.#connToAdapter(conn);
+	}
+
+	has(key: string): boolean {
+		const conns = callNativeSync(() => this.#runtime.actorConns(this.#ctx));
+		return conns.some((c) => this.#runtime.connId(c) === key);
+	}
+
+	keys(): MapIterator<string> {
+		const conns = callNativeSync(() => this.#runtime.actorConns(this.#ctx));
+		return conns.map((c) => this.#runtime.connId(c))[Symbol.iterator]() as MapIterator<string>;
+	}
+
+	values(): MapIterator<NativeConnAdapter> {
+		const conns = callNativeSync(() => this.#runtime.actorConns(this.#ctx));
+		return conns.map((c) => this.#connToAdapter(c))[Symbol.iterator]() as MapIterator<NativeConnAdapter>;
+	}
+
+	entries(): MapIterator<[string, NativeConnAdapter]> {
+		const conns = callNativeSync(() => this.#runtime.actorConns(this.#ctx));
+		return conns.map(
+			(c) => [this.#runtime.connId(c), this.#connToAdapter(c)] as [string, NativeConnAdapter],
+		)[Symbol.iterator]() as MapIterator<[string, NativeConnAdapter]>;
+	}
+
+	forEach(
+		callback: (value: NativeConnAdapter, key: string, map: ReadonlyMap<string, NativeConnAdapter>) => void,
+		thisArg?: unknown,
+	): void {
+		const conns = callNativeSync(() => this.#runtime.actorConns(this.#ctx));
+		for (const conn of conns) {
+			const id = this.#runtime.connId(conn);
+			callback.call(thisArg, this.#connToAdapter(conn), id, this);
+		}
+	}
+
+	[Symbol.iterator](): MapIterator<[string, NativeConnAdapter]> {
+		return this.entries();
+	}
+
+	get [Symbol.toStringTag](): string {
+		return "ConnectionMap";
+	}
+}

rivetkit-typescript/packages/rivetkit/src/registry/native.ts

@@ -1,4 +1,5 @@
 import { VirtualWebSocket } from "@rivetkit/virtual-websocket";
+import { ConnectionMap } from "./connection-map";
 import {
 	ACTOR_CONTEXT_INTERNAL_SYMBOL,
 	CONN_STATE_MANAGER_SYMBOL,
@@ -658,7 +659,7 @@ async function callNative<T>(invoke: () => Promise<T>): Promise<T> {
 	}
 }
 
-function callNativeSync<T>(invoke: () => T): T {
+export function callNativeSync<T>(invoke: () => T): T {
 	try {
 		return invoke();
 	} catch (error) {
@@ -1153,7 +1154,7 @@ function toActorKey(
 	);
 }
 
-class NativeConnAdapter {
+export class NativeConnAdapter {
 	#runtime: CoreRuntime;
 	#conn: ConnHandle;
 	#schemas: NativeValidationConfig;
@@ -2503,27 +2504,8 @@ export class ActorContextHandleAdapter {
 		return callNativeSync(() => this.#runtime.actorRegion(this.#ctx));
 	}
 
-	get conns(): Map<string, NativeConnAdapter> {
-		return new Map(
-			callNativeSync(() => this.#runtime.actorConns(this.#ctx)).map(
-				(conn) => [
-					this.#runtime.connId(conn),
-					new NativeConnAdapter(
-						this.#runtime,
-						conn,
-						this.#schemas,
-						this.#ctx,
-						(connId) =>
-							callNativeSync(() =>
-								this.#runtime.actorQueueHibernationRemoval(
-									this.#ctx,
-									connId,
-								),
-							),
-					),
-				],
-			),
-		);
+	get conns(): ConnectionMap {
+		return new ConnectionMap(this.#runtime, this.#ctx, this.#schemas);
 	}
 
 	get log() {