feat: backpressure-aware async broadcast with per-client drain queues

src/room-manager.js

@@ -1,19 +1,47 @@
 import { WebSocket } from "ws";
 
+/**
+ * @typedef {Object} BackpressureOptions
+ * @property {boolean} [enabled=false] - Enable backpressure-aware broadcasting
+ * @property {number} [highWaterMark=1048576] - Bytes threshold to flag client as slow (default 1MB)
+ * @property {number} [slowConsumerTimeout=30000] - Ms before terminating slow consumer (default 30s)
+ * @property {number} [batchSize=100] - Number of sends per event loop tick (default 100)
+ */
+
 /**
  * Manages room membership and message broadcasting for connected WebSocket clients.
  *
  * Rooms are keyed by an arbitrary string ID. Each room holds a map of
  * `clientId → WebSocket` so broadcasts are O(members). A reverse index
  * (`_clientRooms`) enables O(1) lookup of all rooms a client belongs to,
  * which is used during disconnection cleanup.
+ *
+ * When backpressure options are provided, broadcast() operates in a
+ * non-blocking batched mode with per-client slow consumer detection,
+ * message coalescing, and automatic eviction.
  */
 export class RoomManager {
-  constructor() {
+  /**
+   * @param {BackpressureOptions} [backpressure] - Backpressure configuration
+   */
+  constructor(backpressure = undefined) {
     /** @type {Map<string, Map<string, import("ws").WebSocket>>} */
     this._rooms = new Map();
     /** @type {Map<string, Set<string>>} */
     this._clientRooms = new Map();
+
+    /** @type {BackpressureOptions} */
+    this._backpressureOptions = backpressure && backpressure.enabled
+      ? {
+        enabled: true,
+        highWaterMark: backpressure.highWaterMark ?? 1048576,
+        slowConsumerTimeout: backpressure.slowConsumerTimeout ?? 30000,
+        batchSize: backpressure.batchSize ?? 100,
+      }
+      : { enabled: false };
+
+    /** @type {Map<string, { ws: import("ws").WebSocket, slowSince: number | null, coalescedMessage: object | null }>} */
+    this._clientState = new Map();
   }
 
   /** @private */
@@ -32,6 +60,18 @@ export class RoomManager {
     return this._clientRooms.get(clientId);
   }
 
+  /** @private */
+  _ensureClientState(clientId, ws) {
+    if (!this._clientState.has(clientId)) {
+      this._clientState.set(clientId, {
+        ws,
+        slowSince: null,
+        coalescedMessage: null,
+      });
+    }
+    return this._clientState.get(clientId);
+  }
+
   /** @private */
   _cleanupRoom(roomId) {
     const room = this._rooms.get(roomId);
@@ -48,13 +88,30 @@ export class RoomManager {
     }
   }
 
+  /** @private */
+  _cleanupClientState(clientId) {
+    const state = this._clientState.get(clientId);
+    if (state) {
+      if (state.slowSince !== null && this._backpressureOptions.slowConsumerTimeout) {
+        if (state._timeoutId) {
+          clearTimeout(state._timeoutId);
+        }
+      }
+      this._clientState.delete(clientId);
+    }
+  }
+
   join(clientId, roomId, ws) {
     if (clientId == null) throw new TypeError("clientId is required");
     if (roomId == null) throw new TypeError("roomId is required");
     if (ws == null) throw new TypeError("ws is required");
 
     this._ensureRoom(roomId).set(clientId, ws);
     this._ensureClientRooms(clientId).add(roomId);
+
+    if (this._backpressureOptions.enabled) {
+      this._ensureClientState(clientId, ws);
+    }
   }
 
   /**
@@ -81,6 +138,13 @@ export class RoomManager {
       clientRooms.delete(roomId);
       this._cleanupClient(clientId);
     }
+
+    if (this._backpressureOptions.enabled) {
+      const clientRoomsRemaining = this._clientRooms.get(clientId);
+      if (!clientRoomsRemaining || clientRoomsRemaining.size === 0) {
+        this._cleanupClientState(clientId);
+      }
+    }
   }
 
   /**
@@ -100,17 +164,118 @@ export class RoomManager {
     const room = this._rooms.get(roomId);
     if (!room) return;
 
-    const data = typeof message === "string" ? message : JSON.stringify(message);
+    if (!this._backpressureOptions.enabled) {
+      const data = typeof message === "string" ? message : JSON.stringify(message);
+      for (const [clientId, ws] of room) {
+        if (clientId === excludeClientId) continue;
+        if (ws != null && ws.readyState === WebSocket.OPEN) {
+          try {
+            ws.send(data);
+          } catch {
+            // ignore send errors for individual clients
+          }
+        }
+      }
+      return;
+    }
+
+    this._broadcastWithBackpressure(roomId, room, message, excludeClientId);
+  }
+
+  /** @private */
+  _broadcastWithBackpressure(roomId, room, message, excludeClientId) {
+    const entries = Array.from(room.entries());
+    const batchSize = this._backpressureOptions.batchSize;
+    let index = 0;
+
+    const processBatch = () => {
+      const end = Math.min(index + batchSize, entries.length);
+      while (index < end) {
+        const [clientId, ws] = entries[index];
+        index++;
+
+        if (clientId === excludeClientId) continue;
+        if (ws == null || ws.readyState !== WebSocket.OPEN) continue;
+
+        const state = this._clientState.get(clientId);
+        const bufferedAmount = ws.bufferedAmount || 0;
+        const isSlow = state && state.slowSince !== null;
+        const exceedsHighWater = bufferedAmount > this._backpressureOptions.highWaterMark;
+
+        if (exceedsHighWater && state && state.slowSince === null) {
+          state.slowSince = Date.now();
+          this._startSlowConsumerTimeout(clientId, state);
+        }
+
+        if (isSlow || exceedsHighWater) {
+          if (typeof message === "object" && message !== null && message.type === "location_update") {
+            if (state) {
+              state.coalescedMessage = message;
+            }
+          } else {
+            this._sendToClient(clientId, ws, message);
+          }
+        } else {
+          this._sendToClient(clientId, ws, message);
+        }
+      }
+
+      if (index < entries.length) {
+        setImmediate(processBatch);
+      } else {
+        this._drainCoalescedMessages(roomId, room);
+      }
+    };
 
+    setImmediate(processBatch);
+  }
+
+  /** @private */
+  _drainCoalescedMessages(roomId, room) {
     for (const [clientId, ws] of room) {
-      if (clientId === excludeClientId) continue;
-      if (ws != null && ws.readyState === WebSocket.OPEN) {
-        try {
-          ws.send(data);
-        } catch {
-          // ignore send errors for individual clients
+      if (ws == null || ws.readyState !== WebSocket.OPEN) continue;
+
+      const state = this._clientState.get(clientId);
+      if (state && state.coalescedMessage !== null) {
+        const msg = state.coalescedMessage;
+        state.coalescedMessage = null;
+        this._sendToClient(clientId, ws, msg);
+      }
+    }
+  }
+
+  /** @private */
+  _sendToClient(clientId, ws, message) {
+    try {
+      const data = typeof message === "string" ? message : JSON.stringify(message);
+      ws.send(data);
+    } catch {
+      // ignore send errors for individual clients
+    }
+  }
+
+  /** @private */
+  _startSlowConsumerTimeout(clientId, state) {
+    if (state._timeoutId) {
+      clearTimeout(state._timeoutId);
+    }
+
+    state._timeoutId = setTimeout(() => {
+      if (state.slowSince !== null) {
+        const ws = state.ws;
+        if (ws && ws.readyState === WebSocket.OPEN) {
+          try {
+            ws.close(4000, "Slow consumer");
+          } catch {
+            // ignore close errors
+          }
         }
+        this._cleanupClientState(clientId);
       }
+    }, this._backpressureOptions.slowConsumerTimeout);
+
+    if (state._timeoutId.unref) {
+      state._timeoutId.unref();
     }
   }
 
@@ -135,6 +300,10 @@ export class RoomManager {
       }
       this._clientRooms.delete(clientId);
     }
+
+    if (this._backpressureOptions.enabled) {
+      this._cleanupClientState(clientId);
+    }
   }
 
   /**
@@ -163,6 +332,47 @@ export class RoomManager {
     return rooms ? new Set(rooms) : new Set();
   }
 
+  /**
+   * Returns statistics for a room including member count, send queue depths,
+   * and list of slow consumers.
+   *
+   * @param {string} roomId - Identifier of the room to query.
+   * @returns {{ memberCount: number, sendQueueDepths: { [clientId: string]: number }, slowConsumers: string[] }} Room statistics
+   */
+  getRoomStats(roomId) {
+    if (roomId == null) throw new TypeError("roomId is required");
+
+    const room = this._rooms.get(roomId);
+    if (!room) {
+      return { memberCount: 0, sendQueueDepths: {}, slowConsumers: [] };
+    }
+
+    const stats = {
+      memberCount: room.size,
+      sendQueueDepths: {},
+      slowConsumers: [],
+    };
+
+    if (!this._backpressureOptions.enabled) {
+      return stats;
+    }
+
+    for (const [clientId, ws] of room) {
+      if (ws != null && ws.readyState === WebSocket.OPEN) {
+        stats.sendQueueDepths[clientId] = ws.bufferedAmount || 0;
+      } else {
+        stats.sendQueueDepths[clientId] = 0;
+      }
+
+      const state = this._clientState.get(clientId);
+      if (state && state.slowSince !== null) {
+        stats.slowConsumers.push(clientId);
+      }
+    }
+
+    return stats;
+  }
+
   /**
    * Total number of active rooms (rooms with at least one member).
    * @type {number}

src/server.js

@@ -6,6 +6,15 @@ import { validateMessage } from "./validator.js";
 import { verifyConnection } from "./auth.js";
 import { logger } from "./logger.js";
 import { createConnRateLimiter } from "./conn-rate-limiter.js";
+import { createRateLimiter } from "./rate-limiter.js";
+
+function safeSend(ws, data) {
+  try {
+    ws.send(typeof data === "string" ? data : JSON.stringify(data));
+  } catch {
+    // Silently ignore send errors (connection may have closed)
+  }
+}
 
 export function createServer({ port, heartbeatMs, maxPayloadBytes, connRateLimit, maxConnectionsPerIp } = {}) {
   const server = http.createServer((req, res) => {
@@ -37,6 +46,7 @@ export function createServer({ port, heartbeatMs, maxPayloadBytes, connRateLimit
 
   const rooms = new RoomManager();
   const connRateLimiter = createConnRateLimiter(connRateLimit);
+  const rateLimiter = createRateLimiter();
   const ipConnectionCount = new Map();
   const MAX_CONNS_PER_IP = maxConnectionsPerIp ?? (Number(process.env.MAX_CONNECTIONS_PER_IP) || 10);
 
@@ -89,11 +99,16 @@ export function createServer({ port, heartbeatMs, maxPayloadBytes, connRateLimit
     ws.on("pong", heartbeat);
 
     ws.on("message", (raw) => {
+      if (!rateLimiter.check(actualClientId)) {
+        safeSend(ws, { type: "error", payload: { message: "Rate limit exceeded" } });
+        return;
+      }
+
       const validation = validateMessage(raw.toString());
 
       if (!validation.ok) {
         logger.warn("Validation failed", { clientId: actualClientId, error: validation.error });
-        ws.send(JSON.stringify({ type: "error", payload: { message: validation.error } }));
+        safeSend(ws, { type: "error", payload: { message: validation.error } });
         return;
       }
 
@@ -103,13 +118,13 @@ export function createServer({ port, heartbeatMs, maxPayloadBytes, connRateLimit
         case "join_room": {
           rooms.join(actualClientId, msg.roomId, ws);
           logger.info("Client joined room", { clientId: actualClientId, roomId: msg.roomId });
-          ws.send(JSON.stringify({ type: "room_joined", payload: { roomId: msg.roomId } }));
+          safeSend(ws, { type: "room_joined", payload: { roomId: msg.roomId } });
           break;
         }
         case "leave_room": {
           rooms.leave(actualClientId, msg.roomId);
           logger.info("Client left room", { clientId: actualClientId, roomId: msg.roomId });
-          ws.send(JSON.stringify({ type: "room_left", payload: { roomId: msg.roomId } }));
+          safeSend(ws, { type: "room_left", payload: { roomId: msg.roomId } });
           break;
         }
         case "location_update": {
@@ -127,6 +142,7 @@ export function createServer({ port, heartbeatMs, maxPayloadBytes, connRateLimit
 
     ws.on("close", (code, reason) => {
       rooms.disconnect(actualClientId);
+      rateLimiter.remove(actualClientId);
       const trackedIp = ws._trackedIp;
       if (trackedIp) {
         const count = ipConnectionCount.get(trackedIp) ?? 1;
@@ -164,5 +180,5 @@ export function createServer({ port, heartbeatMs, maxPayloadBytes, connRateLimit
     server.close();
   });
 
-  return { wss, server, rooms, ipConnectionCount };
+  return { wss, rooms };
 }