perf(gateway): index targeted connections (#128198)

Amp-Thread-ID: https://ampcode.com/threads/T-01a021f5-984a-7628-a30c-491c166ff247

Co-authored-by: Amp <amp@ampcode.com>

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<GatewayWsClient, number>();
   const reportedSlowPayloadClients = new WeakSet<GatewayWsClient>();
+  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<ReadonlySet<string> | 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;

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<typeof vi.fn>;
+} {
+  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"]);
+  });
+});

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<GatewayWsClient>();
+  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);

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;
       }

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<GatewayWsClient> {
+  readonly #byConnectionId = new Map<string, IndexedClient>();
+  #nextOrder = 0;
+
+  constructor(clients?: Iterable<GatewayWsClient>) {
+    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<string>): 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);
+  }
+}