From bbdf7abcff42835c07be4158b7e09226276200a9 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Sun, 23 Aug 2026 07:13:33 -0700 Subject: [PATCH] perf(gateway): index targeted connections (#128198) Amp-Thread-ID: https://ampcode.com/threads/T-01a021f5-984a-7628-a30c-491c166ff247 Co-authored-by: Amp --- src/gateway/server-broadcast.ts | 14 ++- src/gateway/server-connection-state.test.ts | 99 +++++++++++++++++++++ src/gateway/server-connection-state.ts | 12 +-- src/gateway/server-lifecycle.ts | 2 +- src/gateway/server/client-registry.ts | 60 +++++++++++++ 5 files changed, 176 insertions(+), 11 deletions(-) create mode 100644 src/gateway/server-connection-state.test.ts create mode 100644 src/gateway/server/client-registry.ts diff --git a/src/gateway/server-broadcast.ts b/src/gateway/server-broadcast.ts index 773412aae888..de4b6951520e 100644 --- a/src/gateway/server-broadcast.ts +++ b/src/gateway/server-broadcast.ts @@ -32,6 +32,7 @@ import type { } from "./server-broadcast-types.js"; import type { SessionMessageSubscriberRegistry } from "./server-chat-state.js"; import { MAX_BUFFERED_BYTES, WEBSOCKET_OPEN_READY_STATE } from "./server-constants.js"; +import { GatewayClientRegistry } from "./server/client-registry.js"; import type { GatewayWsClient } from "./server/ws-types.js"; import { logWs, summarizeAgentEventForWsLog } from "./ws-log.js"; @@ -218,6 +219,8 @@ export function createGatewayBroadcaster(params: { }) { const clientSeq = new WeakMap(); const reportedSlowPayloadClients = new WeakSet(); + const indexedClients = + params.clients instanceof GatewayClientRegistry ? params.clients : undefined; const broadcastInternal = ( event: string, @@ -268,12 +271,16 @@ export function createGatewayBroadcaster(params: { const isSessionSubscriptionEvent = SESSION_SUBSCRIPTION_EVENTS.has(event); const sessionMessageSubscribers = params.sessionMessageSubscribers; let sessionSubscriberConnIdsByKey: Array | undefined> | undefined; - for (const c of params.clients) { + const recipients = + targetConnIds && indexedClients + ? indexedClients.getByConnectionIds(targetConnIds) + : params.clients; + for (const c of recipients) { // Closing nodes remain discoverable until their owner drains admitted lifecycle work. if (c.invalidated === true || c.socket.readyState !== WEBSOCKET_OPEN_READY_STATE) { continue; } - if (targetConnIds && !targetConnIds.has(c.connId)) { + if (targetConnIds && !indexedClients && !targetConnIds.has(c.connId)) { continue; } if (!hasEventScope(c, event, explicitPluginScope)) { @@ -393,6 +400,9 @@ export function createGatewayBroadcaster(params: { }; const getBufferedAmount: GatewayBufferedAmountFn = (connId) => { + if (indexedClients) { + return indexedClients.getByConnectionId(connId)?.socket.bufferedAmount; + } for (const client of params.clients) { if (client.connId === connId) { return client.socket.bufferedAmount; diff --git a/src/gateway/server-connection-state.test.ts b/src/gateway/server-connection-state.test.ts new file mode 100644 index 000000000000..dab157a55b83 --- /dev/null +++ b/src/gateway/server-connection-state.test.ts @@ -0,0 +1,99 @@ +import { describe, expect, it, vi } from "vitest"; +import { WebSocket } from "ws"; +import type { OpenClawConfig } from "../config/types.openclaw.js"; +import { createGatewayConnectionState } from "./server-connection-state.js"; +import type { GatewayWsClient } from "./server/ws-types.js"; + +type ConnectionIdReads = { count: number }; + +function makeClient( + connId: string, + reads: ConnectionIdReads, + sendOrder?: string[], +): { + client: GatewayWsClient; + socket: { readyState: number }; + send: ReturnType; +} { + const send = vi.fn(() => sendOrder?.push(connId)); + const socket = { + readyState: WebSocket.OPEN, + bufferedAmount: 0, + close: vi.fn(), + send, + }; + const client = { + socket: socket as unknown as GatewayWsClient["socket"], + connect: { + role: "operator", + scopes: ["operator.read"], + } as GatewayWsClient["connect"], + usesSharedGatewayAuth: false, + } as GatewayWsClient; + Object.defineProperty(client, "connId", { + enumerable: true, + get: () => { + reads.count += 1; + return connId; + }, + }); + return { client, socket, send }; +} + +describe("gateway connection state", () => { + it("bounds targeted delivery and connection lookups to the requested connection", () => { + const state = createGatewayConnectionState({ cfg: {} as OpenClawConfig }); + const reads = { count: 0 }; + for (let index = 0; index < 256; index += 1) { + state.clients.add(makeClient(`other-${index}`, reads).client); + } + const target = makeClient("target", reads); + state.clients.add(target.client); + reads.count = 0; + + state.broadcastToConnIds("tick", { ts: 1 }, new Set(["target"])); + + expect(target.send).toHaveBeenCalledTimes(1); + expect(reads.count).toBe(0); + + target.socket.readyState = WebSocket.CLOSING; + state.broadcastToConnIds("tick", { ts: 2 }, new Set(["target"])); + + expect(target.send).toHaveBeenCalledTimes(1); + + reads.count = 0; + expect(state.getBufferedAmount("target")).toBe(0); + expect(state.isConnectionActive("target")).toBe(true); + expect(reads.count).toBe(0); + + state.clients.delete(target.client); + reads.count = 0; + state.broadcastToConnIds("tick", { ts: 3 }, new Set(["target"])); + + expect(target.send).toHaveBeenCalledTimes(1); + expect(state.getBufferedAmount("target")).toBeUndefined(); + expect(state.isConnectionActive("target")).toBe(false); + expect(reads.count).toBe(0); + + state.clients.add(target.client); + state.clients.clear(); + reads.count = 0; + + expect(state.getBufferedAmount("target")).toBeUndefined(); + expect(state.isConnectionActive("target")).toBe(false); + expect(reads.count).toBe(0); + }); + + it("preserves connection insertion order for targeted fanout", () => { + const state = createGatewayConnectionState({ cfg: {} as OpenClawConfig }); + const reads = { count: 0 }; + const sendOrder: string[] = []; + state.clients.add(makeClient("first", reads, sendOrder).client); + state.clients.add(makeClient("unrelated", reads, sendOrder).client); + state.clients.add(makeClient("last", reads, sendOrder).client); + + state.broadcastToConnIds("tick", { ts: 1 }, new Set(["last", "first"])); + + expect(sendOrder).toEqual(["first", "last"]); + }); +}); diff --git a/src/gateway/server-connection-state.ts b/src/gateway/server-connection-state.ts index 604d15bd16ba..158428c93fbe 100644 --- a/src/gateway/server-connection-state.ts +++ b/src/gateway/server-connection-state.ts @@ -7,7 +7,7 @@ import { createSessionEventSubscriberRegistry, createSessionMessageSubscriberRegistry, } from "./server-chat-state.js"; -import type { GatewayWsClient } from "./server/ws-types.js"; +import { GatewayClientRegistry } from "./server/client-registry.js"; import { canReceiveSessionEvent } from "./session-sharing.js"; /** Creates transport-independent connection, subscription, and run state. */ @@ -16,16 +16,12 @@ export function createGatewayConnectionState(params: { getRuntimeConfig?: () => import("../config/config.js").OpenClawConfig; }) { const loadRuntimeConfig = params.getRuntimeConfig ?? (() => params.cfg); - const clients = new Set(); + const clients = new GatewayClientRegistry(); // Detached RPC dispatch can resume after close cleanup, so connection-owned // producers must validate against the live transport owner before mutation. const isConnectionActive = (connId: string) => { - for (const client of clients) { - if (client.connId === connId && !client.invalidated) { - return true; - } - } - return false; + const client = clients.getByConnectionId(connId); + return Boolean(client && !client.invalidated); }; const sessionEventSubscribers = createSessionEventSubscriberRegistry(isConnectionActive); const sessionMessageSubscribers = createSessionMessageSubscriberRegistry(isConnectionActive); diff --git a/src/gateway/server-lifecycle.ts b/src/gateway/server-lifecycle.ts index 288255b882ba..bdbf573e9f26 100644 --- a/src/gateway/server-lifecycle.ts +++ b/src/gateway/server-lifecycle.ts @@ -348,7 +348,7 @@ export async function prepareGatewayLifecycle(params: { runtimeState.sessionViewerPresence = createSessionViewerPresenceDeclarations({ isConnectionActive, onReplace: (connId, sessionKeys) => { - const client = [...clients].find((candidate) => candidate.connId === connId); + const client = clients.getByConnectionId(connId); if (!client?.presenceKey) { return; } diff --git a/src/gateway/server/client-registry.ts b/src/gateway/server/client-registry.ts new file mode 100644 index 000000000000..0a2cd40accf0 --- /dev/null +++ b/src/gateway/server/client-registry.ts @@ -0,0 +1,60 @@ +import type { GatewayWsClient } from "./ws-types.js"; + +type IndexedClient = { + client: GatewayWsClient; + order: number; +}; + +export class GatewayClientRegistry extends Set { + readonly #byConnectionId = new Map(); + #nextOrder = 0; + + constructor(clients?: Iterable) { + super(); + for (const client of clients ?? []) { + this.add(client); + } + } + + override add(client: GatewayWsClient): this { + if (!this.has(client)) { + this.#byConnectionId.set(client.connId, { client, order: this.#nextOrder++ }); + } + return super.add(client); + } + + override delete(client: GatewayWsClient): boolean { + if (!super.delete(client)) { + return false; + } + if (this.#byConnectionId.get(client.connId)?.client === client) { + this.#byConnectionId.delete(client.connId); + } + return true; + } + + override clear(): void { + super.clear(); + this.#byConnectionId.clear(); + } + + getByConnectionId(connId: string): GatewayWsClient | undefined { + return this.#byConnectionId.get(connId)?.client; + } + + getByConnectionIds(connIds: ReadonlySet): GatewayWsClient[] { + const indexed: IndexedClient[] = []; + for (const connId of connIds) { + const entry = this.#byConnectionId.get(connId); + if (entry) { + indexed.push(entry); + } + } + // Targeted fanout keeps authenticated-client insertion order without + // walking unrelated sockets. + if (indexed.length > 1) { + indexed.sort((a, b) => a.order - b.order); + } + return indexed.map((entry) => entry.client); + } +}