Scaling WebSockets

// Chatterbox subset (hello, join, message, ack) on ws, with a Redis pub/sub backplane and a draining shutdown.
import http from "node:http";
import { randomBytes } from "node:crypto";
import { setTimeout as sleep } from "node:timers/promises";
import { WebSocketServer } from "ws";
import Redis from "ioredis";

const NODE = process.env.NODE_ID ?? randomBytes(3).toString("hex");
const PORT = Number(process.env.PORT ?? 8080);
const REDIS_URL = process.env.REDIS_URL ?? "redis://localhost:6379";
const READY_DELAY_MS = Number(process.env.READY_DELAY_MS ?? 500); // time for the load balancer to notice /readyz
const DRAIN_WINDOW_MS = Number(process.env.DRAIN_WINDOW_MS ?? 5000); // spread of the 4503 closes; 0 = all at once
const CHANNEL = "cb:fanout";

const rooms = new Map(); // room -> Set<WebSocket> (this node only)
const accepted = []; // accept times, for the demo
let draining = false;

const data = new Redis(REDIS_URL); // commands
const sub = new Redis(REDIS_URL); // a subscribed connection cannot run other commands
await sub.subscribe(CHANNEL); // ioredis re-subscribes by itself after a reconnect

const envelope = (type, room, payload) =>
  JSON.stringify({ v: 1, type, id: "s_" + randomBytes(8).toString("hex"), room, ts: Date.now(), payload });

function deliverLocal(room, frame) {
  for (const ws of rooms.get(room) ?? []) if (ws.readyState === ws.OPEN) ws.send(frame);
}

sub.on("message", (_channel, raw) => {
  const { room, frame } = JSON.parse(raw);
  deliverLocal(room, frame);
});

const server = http.createServer((req, res) => {
  if (req.url === "/readyz") {
    res.statusCode = draining ? 503 : 200;
    return res.end(draining ? "draining" : "ok");
  }
  if (req.url === "/stats") {
    res.setHeader("content-type", "application/json");
    return res.end(JSON.stringify({ node: NODE, connections: wss.clients.size, accepted }));
  }
  res.statusCode = 404;
  res.end();
});

const wss = new WebSocketServer({ noServer: true, maxPayload: 16384 });

server.on("upgrade", (req, socket, head) => {
  if (draining) {
    socket.end("HTTP/1.1 503 Service Unavailable\r\nRetry-After: 1\r\nConnection: close\r\n\r\n");
    return;
  }
  wss.handleUpgrade(req, socket, head, (ws) => wss.emit("connection", ws, req));
});

wss.on("connection", (ws) => {
  accepted.push(Date.now());
  ws.rooms = new Set();
  ws.send(envelope("hello", null, { node: NODE }));
  ws.on("message", async (buf) => {
    let m;
    try {
      m = JSON.parse(buf.toString());
    } catch {
      return ws.close(4400, "bad json");
    }
    const ack = (result, seq = null) =>
      ws.send(envelope("ack", m.room ?? null, { ref: m.id, ok: true, seq, dup: false, result }));
    if (m.type === "join") {
      if (!rooms.has(m.room)) rooms.set(m.room, new Set());
      rooms.get(m.room).add(ws);
      ws.rooms.add(m.room);
      ack({ node: NODE });
    } else if (m.type === "message") {
      const seq = await data.incr(`cb:seq:${m.room}`); // one counter for the whole cluster
      const frame = envelope("message", m.room, { text: m.payload.text, seq, via: NODE });
      await data.publish(CHANNEL, JSON.stringify({ room: m.room, frame })); // every node, this one included
      ack({}, seq);
    }
  });
  ws.on("close", () => {
    for (const r of ws.rooms) rooms.get(r)?.delete(ws);
  });
});

server.listen(PORT, () => console.log(`${NODE} listening on ${PORT}`));

process.on("SIGTERM", async () => {
  console.log(`${NODE} draining: ${wss.clients.size} connections, window ${DRAIN_WINDOW_MS} ms`);
  draining = true; // /readyz now answers 503 and new upgrades are refused
  await sleep(READY_DELAY_MS);
  for (const ws of wss.clients) setTimeout(() => ws.close(4503, "restarting"), Math.random() * DRAIN_WINDOW_MS);
  await sleep(DRAIN_WINDOW_MS + 1000);
  for (const ws of wss.clients) ws.terminate();
  await Promise.all([data.quit(), sub.quit()]);
  server.close();
  process.exit(0);
});

FROM node:24-alpine
WORKDIR /app
COPY package.json ./
RUN npm install --omit=dev
COPY server.mjs ./
USER node
CMD ["node", "server.mjs"]

name: g7-scaling-websockets
services:
  redis:
    image: redis:8.8.3-alpine
    container_name: g7-scaling-websockets-redis
    ports: ["18300:6379"]
  chat-a:
    build: ./app
    image: g7-scaling-websockets-app
    container_name: g7-scaling-websockets-a
    environment:
      NODE_ID: a
      REDIS_URL: redis://redis:6379
      DRAIN_WINDOW_MS: ${DRAIN_WINDOW_MS:-5000}
    ports: ["18301:8080"]
    depends_on: [redis]
    stop_grace_period: 30s
  chat-b:
    image: g7-scaling-websockets-app
    container_name: g7-scaling-websockets-b
    environment:
      NODE_ID: b
      REDIS_URL: redis://redis:6379
    ports: ["18302:8080"]
    depends_on: [redis, chat-a]

// Ada is on instance a (18301), Bob on instance b (18302): one room, one Redis.
import assert from "node:assert/strict";
import WebSocket from "ws";

function open(port, name) {
  const ws = new WebSocket(`ws://localhost:${port}/ws`, "chatterbox.v1");
  const inbox = [];
  ws.on("message", (b) => inbox.push(JSON.parse(b.toString())));
  const send = (type, room, payload) =>
    ws.send(JSON.stringify({ v: 1, type, id: `${name}-${inbox.length}-${Math.random().toString(16).slice(2)}`, room, ts: 0, payload }));
  return new Promise((resolve) => ws.once("open", () => resolve({ ws, inbox, send })));
}
const until = async (fn) => { for (let i = 0; i < 100; i++) { if (fn()) return; await new Promise((r) => setTimeout(r, 50)); } throw new Error("timeout"); };

const ada = await open(18301, "ada");
const bob = await open(18302, "bob");
await until(() => ada.inbox.length && bob.inbox.length);
console.log("Ada is on node", ada.inbox[0].payload.node, "- Bob is on node", bob.inbox[0].payload.node);
ada.send("join", "general", {});
bob.send("join", "general", {});
await until(() => ada.inbox.length >= 2 && bob.inbox.length >= 2);

ada.send("message", "general", { text: "hello from a" });
bob.send("message", "general", { text: "hello from b" });
await until(() => [ada, bob].every((c) => c.inbox.filter((e) => e.type === "message").length === 2));

for (const [who, c] of [["ada", ada], ["bob", bob]]) {
  const msgs = c.inbox.filter((e) => e.type === "message").map((e) => `${e.payload.seq}:${e.payload.text}(via ${e.payload.via})`);
  console.log(who, "received", msgs.join(" | "));
}
const seqs = (c) => c.inbox.filter((e) => e.type === "message").map((e) => e.payload.seq);
assert.deepEqual(seqs(ada), seqs(bob), "same order on both nodes");
assert.deepEqual([...seqs(ada)].sort(), [1, 2], "the cluster-wide counter gave 1 and 2");
console.log("OK: one room, two nodes, same seq order");
ada.ws.close(); bob.ws.close();

Ada is on node a - Bob is on node b
ada received 1:hello from b(via b) | 2:hello from a(via a)
bob received 1:hello from b(via b) | 2:hello from a(via a)
OK: one room, two nodes, same seq order

// A node that was away resumes from the last id it delivered; pub/sub would have dropped these entries.
import Redis from "ioredis";
const STREAM = "cb:events:replay-demo";
const producer = new Redis("redis://localhost:18300");
await producer.del(STREAM);

let lastId = await producer.xadd(STREAM, "MAXLEN", "~", 10000, "*", "room", "general", "text", "delivered before the outage");

// the consumer node is "down": three more entries are appended meanwhile
for (const text of ["missed 1", "missed 2", "missed 3"])
  await producer.xadd(STREAM, "MAXLEN", "~", 10000, "*", "room", "general", "text", text);

const reader = new Redis("redis://localhost:18300"); // XREAD BLOCK needs a connection of its own
const rows = await reader.xread("COUNT", 100, "BLOCK", 1000, "STREAMS", STREAM, lastId);
for (const [id, fields] of rows[0][1]) {
  console.log("replayed", id, fields[3]);
  lastId = id;
}
const empty = await reader.xread("BLOCK", 300, "STREAMS", STREAM, lastId);
console.log("then nothing new:", empty);
await producer.del(STREAM);
producer.disconnect(); reader.disconnect();

replayed 1792885712584-1 missed 1
replayed 1792885712585-0 missed 2
replayed 1792885712585-1 missed 3
then nothing new: null

worker_processes auto;
worker_rlimit_nofile 40000;          # each proxied WebSocket uses two descriptors

events {
  worker_connections 20000;
}

http {
  map $http_upgrade $connection_upgrade {
    default upgrade;
    ""      close;
  }

  upstream chatterbox {
    least_conn;                      # long-lived connections: balance by open connections, not by requests
    server chat-a:8080;
    server chat-b:8080;
  }

  server {
    listen 80;

    location /ws {
      proxy_pass http://chatterbox;
      proxy_http_version 1.1;
      proxy_set_header Upgrade $http_upgrade;
      proxy_set_header Connection $connection_upgrade;
      proxy_set_header Host $host;
      proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
      proxy_read_timeout 70s;        # more than 2 x the 20 s application heartbeat
      proxy_send_timeout 70s;
    }

    location = /readyz {
      proxy_pass http://chatterbox;  # only reports ready while an instance answers
    }
  }
}

apiVersion: apps/v1
kind: Deployment
metadata:
  name: chatterbox
spec:
  replicas: 4
  strategy:
    rollingUpdate:
      maxUnavailable: 1      # replace one pod at a time: about a quarter of the clients reconnect per step
      maxSurge: 1
  selector:
    matchLabels: { app: chatterbox }
  template:
    metadata:
      labels: { app: chatterbox }
    spec:
      terminationGracePeriodSeconds: 60   # preStop + READY_DELAY_MS + DRAIN_WINDOW_MS + margin
      containers:
        - name: app
          image: registry.example.test/chatterbox:1.0.0
          ports: [{ containerPort: 8080 }]
          env:
            - { name: READY_DELAY_MS, value: "3000" }
            - { name: DRAIN_WINDOW_MS, value: "30000" }
          readinessProbe:
            httpGet: { path: /readyz, port: 8080 }
            periodSeconds: 2
            failureThreshold: 1
          resources:
            requests: { cpu: 500m, memory: 512Mi }
            limits: { memory: 512Mi }
---
apiVersion: policy/v1
kind: PodDisruptionBudget
metadata:
  name: chatterbox
spec:
  maxUnavailable: 1
  selector:
    matchLabels: { app: chatterbox }

// N clients on node a; a is told to stop; clients reconnect (to a, then b). MODE=jitter|none
import WebSocket from "ws";
const N = Number(process.env.N ?? 300);
const MODE = process.env.MODE ?? "jitter";
const TARGETS = [18301, 18302];
const BASE = 1000, CAP = 10000;
let attempts = 0, failures = 0, online = 0;
const reachedB = [];

function client(i) {
  let attempt = 0, t = 0;
  const connect = () => {
    const port = TARGETS[t % TARGETS.length];
    attempts++;
    const ws = new WebSocket(`ws://localhost:${port}/ws`);
    let opened = false;
    ws.on("error", () => {});
    ws.on("open", () => { opened = true; online++; attempt = 0; if (port === 18302) reachedB.push(Date.now()); });
    ws.on("close", (code) => {
      if (opened) online--; else { failures++; t++; }
      if (opened && port === 18301) t = 1; // after an unexpected close, prefer the other node
      const ceiling = Math.min(CAP, BASE * 2 ** attempt++);
      const delay = MODE === "jitter" ? Math.random() * ceiling : 0;
      setTimeout(connect, delay);
    });
  };
  connect();
}
for (let i = 0; i < N; i++) client(i);
await new Promise((r) => setTimeout(r, 1500));
console.log(`before: ${online}/${N} online on a`);
console.log("DRAIN_NOW");
const t0 = Date.now();
await new Promise((r) => setTimeout(r, Number(process.env.WAIT ?? 20000)));
const buckets = new Map();
for (const ts of reachedB) { const k = Math.floor((ts - t0) / 250); buckets.set(k, (buckets.get(k) ?? 0) + 1); }
const peak = Math.max(0, ...buckets.values());
const last = reachedB.length ? Math.max(...reachedB) - t0 : 0;
console.log(`mode=${MODE}: ${reachedB.length}/${N} reached b, peak ${peak} connections per 250 ms, last at ${last} ms, ${attempts} attempts, ${failures} failed attempts`);
process.exit(0);