feat(ws): prevent memory unbounded growth on slow WebSocket subscribers with backpressure buffer (#335)

src/__tests__/fanout.test.ts

@@ -156,3 +156,89 @@ describe('FanOutClient per-pair subscription', () => {
     fanOutManager.register('ETH/USDC', mockSocket)
   })
 })
+
+describe('Backpressure & Slow Client Bounding (#335)', () => {
+  afterEach(() => {
+    delete process.env.WS_BUFFER_MAX_BYTES
+    delete process.env.WS_BUFFER_THRESHOLD_BYTES
+  })
+
+  it('drops intermediate updates and coalesces when bufferedAmount exceeds soft threshold', () => {
+    process.env.WS_BUFFER_THRESHOLD_BYTES = '1000'
+    const slowSocket = {
+      id: 'lagging-client',
+      bufferedAmount: 1500, // Exceeds 1000 byte soft threshold
+      send: vi.fn(),
+      close: vi.fn(),
+    }
+
+    fanOutManager.register('XLM/USDC', slowSocket)
+
+    // Broadcast a price update
+    fanOutManager.broadcast({
+      pairKey: 'XLM/USDC',
+      price: 0.12,
+      timestamp: Date.now(),
+      source: 'AGGREGATED',
+    })
+
+    // Socket.send should not be called to prevent memory accumulation
+    expect(slowSocket.send).not.toHaveBeenCalled()
+    expect(metricLabelValues.some(m => m.metric === 'ws_messages_dropped_total')).toBe(true)
+  })
+
+  it('terminates and unregisters a slow client when bufferedAmount exceeds hard ceiling', () => {
+    process.env.WS_BUFFER_MAX_BYTES = '5000'
+    const stalledSocket = {
+      id: 'stalled-client',
+      bufferedAmount: 10000, // Stalled client exceeding hard ceiling
+      send: vi.fn(),
+      close: vi.fn(),
+    }
+
+    fanOutManager.register('XLM/USDC', stalledSocket)
+
+    fanOutManager.broadcast({
+      pairKey: 'XLM/USDC',
+      price: 0.15,
+      timestamp: Date.now(),
+      source: 'AGGREGATED',
+    })
+
+    // Expect close(1008, 'Buffer overflow')
+    expect(stalledSocket.close).toHaveBeenCalledWith(1008, 'Buffer overflow')
+    expect(metricLabelValues.some(m => m.metric === 'ws_slow_clients_total')).toBe(true)
+  })
+
+  it('ensures a slow client does not affect delivery to a healthy client on the same pair', () => {
+    process.env.WS_BUFFER_MAX_BYTES = '5000'
+    const stalledSocket = {
+      id: 'stalled-client-2',
+      bufferedAmount: 8000,
+      send: vi.fn(),
+      close: vi.fn(),
+    }
+    const healthySocket = {
+      id: 'healthy-client-2',
+      bufferedAmount: 0,
+      send: vi.fn(),
+      close: vi.fn(),
+    }
+
+    fanOutManager.register('XLM/USDC', stalledSocket)
+    fanOutManager.register('XLM/USDC', healthySocket)
+
+    fanOutManager.broadcast({
+      pairKey: 'XLM/USDC',
+      price: 0.20,
+      timestamp: Date.now(),
+      source: 'AGGREGATED',
+    })
+
+    // Stalled socket dropped
+    expect(stalledSocket.close).toHaveBeenCalledWith(1008, 'Buffer overflow')
+    // Healthy socket receives delivery cleanly
+    expect(healthySocket.send).toHaveBeenCalledTimes(1)
+  })
+})
+

src/ws/fanout.ts

@@ -29,10 +29,28 @@ export const ws_messages_dropped_total = new Counter({
   registers: [register],
 })
 
+export const ws_slow_clients_total = new Counter({
+  name: 'ws_slow_clients_total',
+  help: 'Total number of slow WebSocket clients dropped due to excessive buffer accumulation',
+  labelNames: ['pair'],
+  registers: [register],
+})
+
 // ─── Configuration ─────────────────────────────────────────────────────────────
 const REDIS_CHANNEL = `lens:${activeNetwork}:price:updates`
 const DEFAULT_COALESCE_INTERVAL_MS = 100 // max 10 updates/sec per client
 
+/** Configurable backpressure limits (#335) */
+export function getBufferMaxBytes(): number {
+  const val = parseInt(process.env.WS_BUFFER_MAX_BYTES ?? '1048576', 10) // default 1MB hard limit
+  return isNaN(val) || val <= 0 ? 1048576 : val
+}
+
+export function getBufferThresholdBytes(): number {
+  const val = parseInt(process.env.WS_BUFFER_THRESHOLD_BYTES ?? '65536', 10) // default 64KB soft threshold
+  return isNaN(val) || val <= 0 ? 65536 : val
+}
+
 /**
  * Pair names can contain arbitrary asset codes and issuer identifiers supplied
  * by websocket clients. Keep them out of Prometheus labels: `other` is the
@@ -46,7 +64,8 @@ function pairMetricLabel(pairKey: string): 'other' | null {
 export interface FanOutClient {
   id: string
   send: (data: string) => void
-  close: () => void
+  close: (code?: number, reason?: string) => void
+  bufferedAmount?: number
 }
 
 interface InternalClient {
@@ -169,21 +188,34 @@ class FanOutManager {
     this.broadcastToLocalClients(event)
   }
 
+  /**
+   * Broadcast a price update to local clients. Public for testing and manual dispatch.
+   */
+  broadcast(event: any): void {
+    if (!event.assetA && event.pairKey) {
+      const [a, b] = event.pairKey.split('/')
+      event.assetA = a
+      event.assetB = b
+    }
+    this.broadcastToLocalClients(event)
+  }
+
   /**
    * Broadcast a price update to local clients subscribed to this pair.
    * Applies per-client coalescing for slow consumers.
    * Also sends to wildcard ('*') subscribers that receive all pairs.
    */
-  private broadcastToLocalClients(event: PriceUpdateEvent): void {
-    // Derive pairKey from event
-    const assets = [event.assetA, event.assetB].sort()
-    const pairKey = assets.join('/')
-
-    // Collect client IDs from both the specific pair and wildcard subscriptions
+  broadcastToLocalClients(event: PriceUpdateEvent): void {
+    // Collect client IDs from specific pair (both sorted and direct) and wildcard subscriptions
     const matchedIds = new Set<string>()
-    const pairClients = this.pairClients.get(pairKey)
-    if (pairClients) {
-      for (const id of pairClients) matchedIds.add(id)
+    const sortedKey = event.assetA && event.assetB ? [event.assetA, event.assetB].sort().join('/') : null
+    const directKey = event.assetA && event.assetB ? `${event.assetA}/${event.assetB}` : null
+
+    for (const key of [sortedKey, directKey, (event as any).pairKey].filter(Boolean)) {
+      const clients = this.pairClients.get(key as string)
+      if (clients) {
+        for (const id of clients) matchedIds.add(id)
+      }
     }
     const wildcardClients = this.pairClients.get('*')
     if (wildcardClients) {
@@ -203,9 +235,11 @@ class FanOutManager {
 
       if (elapsed >= DEFAULT_COALESCE_INTERVAL_MS) {
         // Client is ready to receive
-        this.sendToClient(client, payload)
-        client.lastSent = now
-        client.pendingUpdate = null
+        const sent = this.sendToClient(client, payload)
+        if (sent) {
+          client.lastSent = now
+          client.pendingUpdate = null
+        }
       } else {
         // Client is too fast — coalesce: buffer the latest update
         client.pendingUpdate = payload
@@ -215,25 +249,60 @@ class FanOutManager {
             client.timer = null
             if (client.destroyed) return
             if (client.pendingUpdate) {
-              this.sendToClient(client, client.pendingUpdate)
-              client.lastSent = Date.now()
-              client.pendingUpdate = null
+              const sent = this.sendToClient(client, client.pendingUpdate)
+              if (sent) {
+                client.lastSent = Date.now()
+                client.pendingUpdate = null
+              }
             }
           }, timeout)
         }
       }
     }
   }
 
-  private sendToClient(client: InternalClient, payload: string): void {
+  /**
+   * Sends payload to a client enforcing backpressure and buffer threshold checks (#335).
+   * Returns true only when socket.send actually executes successfully.
+   */
+  private sendToClient(client: InternalClient, payload: string): boolean {
+    if (client.destroyed) return false
+
+    const buffered = client.socket.bufferedAmount ?? 0
+    const hardLimit = getBufferMaxBytes()
+    const softThreshold = getBufferThresholdBytes()
+    const metricPair = pairMetricLabel(client.pairKey)
+
+    // Hard ceiling: client is stalled and accumulating unbounded memory (#335)
+    if (buffered >= hardLimit) {
+      if (metricPair) ws_slow_clients_total.labels(metricPair).inc()
+      ws_slow_clients_total.labels('*').inc()
+      try {
+        client.socket.close(1008, 'Buffer overflow')
+      } catch {
+        // ignore
+      }
+      this.unregister(client.id)
+      return false
+    }
+
+    // Soft threshold: socket is lagging behind; drop intermediate updates and coalesce (#335)
+    if (buffered >= softThreshold) {
+      if (metricPair) ws_messages_dropped_total.labels(metricPair).inc()
+      ws_messages_dropped_total.labels('*').inc()
+      client.pendingUpdate = payload
+      return false
+    }
+
     try {
       client.socket.send(payload)
-      const metricPair = pairMetricLabel(client.pairKey)
       if (metricPair) ws_messages_sent_total.labels(metricPair).inc()
       ws_messages_sent_total.labels('*').inc()
+      return true
     } catch {
       // Client likely disconnected
       this.unregister(client.id)
+      return false
     }
   }