WebSockets, properly: real-time systems without connection chaos

GET /ws HTTP/1.1
Host: api.example.com
Upgrade: websocket
Connection: Upgrade
Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==
Sec-WebSocket-Version: 13
Sec-WebSocket-Protocol: v1.chat

HTTP/1.1 101 Switching Protocols
Upgrade: websocket
Connection: Upgrade
Sec-WebSocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo=
Sec-WebSocket-Protocol: v1.chat

const ws = new WebSocket('wss://api.example.com/ws', ['v1.chat'])

ws.readyState
// 0 CONNECTING
// 1 OPEN        ← does NOT mean the peer is reachable
// 2 CLOSING
// 3 CLOSED

ws.onopen    = () => { /* handshake complete */ }
ws.onmessage = (e) => { /* a frame arrived */ }
ws.onerror   = (e) => { /* deliberately uninformative, for security */ }
ws.onclose   = (e) => { console.log(e.code, e.reason, e.wasClean) }

// Server: ping every client, drop anyone who missed the last pong
const HEARTBEAT_MS = 30_000

wss.on('connection', (ws) => {
  ws.isAlive = true
  ws.on('pong', () => { ws.isAlive = true })
})

const interval = setInterval(() => {
  for (const ws of wss.clients) {
    if (!ws.isAlive) {
      ws.terminate()          // not close() - the peer is gone, do not wait
      continue
    }
    ws.isAlive = false
    ws.ping()
  }
}, HEARTBEAT_MS)

wss.on('close', () => clearInterval(interval))

// 1. Client asks an ordinary authenticated endpoint for a ticket
const { ticket } = await fetch('/api/ws-ticket', {
  method: 'POST',
  credentials: 'include',
}).then((r) => r.json())

// 2. Connect with it. Valid for ~30 seconds, single use.
const ws = new WebSocket(`wss://api.example.com/ws?ticket=${ticket}`)

// Server: issue
app.post('/api/ws-ticket', requireAuth, async (c) => {
  const ticket = crypto.randomUUID()
  await redis.set(
    `ws:ticket:${ticket}`,
    JSON.stringify({ userId: c.var.user.id }),
    'EX', 30,
  )
  return c.json({ ticket })
})

// Server: redeem, during the upgrade, before accepting the socket
server.on('upgrade', async (req, socket, head) => {
  const url = new URL(req.url!, 'http://localhost')

  // CSRF: cookies are sent automatically, so the origin must be checked
  if (!ALLOWED_ORIGINS.has(req.headers.origin ?? '')) {
    socket.write('HTTP/1.1 403 Forbidden\r\n\r\n')
    return socket.destroy()
  }

  // GETDEL: redeem atomically so a captured ticket cannot be replayed
  const raw = await redis.getdel(`ws:ticket:${url.searchParams.get('ticket')}`)
  if (!raw) {
    socket.write('HTTP/1.1 401 Unauthorized\r\n\r\n')
    return socket.destroy()
  }

  wss.handleUpgrade(req, socket, head, (ws) => {
    wss.emit('connection', ws, req, JSON.parse(raw))
  })
})

// Connections outlive tokens. Handle expiry explicitly.
const EXPIRY_GRACE_MS = 60_000

function scheduleTokenExpiry(ws: Socket, expiresAt: number) {
  const ms = expiresAt - Date.now() - EXPIRY_GRACE_MS
  ws.expiryTimer = setTimeout(() => {
    ws.send(JSON.stringify({ type: 'auth.refresh_required' }))
    // Give the client a window to send a fresh token before closing
    ws.graceTimer = setTimeout(
      () => ws.close(4002, 'token expired'),
      EXPIRY_GRACE_MS,
    )
  }, Math.max(0, ms))
}

class ResilientSocket {
  private attempt = 0
  private ws?: WebSocket
  private closedByUs = false

  connect() {
    this.ws = new WebSocket(this.url())

    this.ws.onopen = () => {
      this.attempt = 0                    // reset only on a real open
      this.resubscribe()
    }

    this.ws.onclose = (e) => {
      if (this.closedByUs || e.code === 1000) return
      if (e.code === 4001) return this.reauthenticate()
      this.scheduleReconnect()
    }
  }

  private scheduleReconnect() {
    const base = Math.min(1000 * 2 ** this.attempt, 30_000)
    // Full jitter: pick uniformly in [0, base), not base ± a bit
    const delay = Math.random() * base
    this.attempt++
    setTimeout(() => this.connect(), delay)
  }
}

// Client stores the last message id it processed
const ws = new WebSocket(`${base}?ticket=${ticket}&since=${lastSeenId ?? ''}`)

// Server, on connect: replay from a bounded buffer before live delivery
async function onConnect(ws: Socket, rooms: string[], since?: string) {
  if (since) {
    for (const room of rooms) {
      // A Redis stream per room gives you exactly this, with a cap
      const missed = await redis.xrange(`room:${room}`, `(${since}`, '+', 'COUNT', 500)
      for (const [id, fields] of missed) {
        ws.send(JSON.stringify({ id, ...decode(fields) }))
      }
    }
  }
  subscribeLive(ws, rooms)
}

// One key per connection, not per user. TTL slightly over 2 heartbeats.
const PRESENCE_TTL = 75

async function register(userId: string, connId: string, room: string) {
  await redis.set(`presence:${room}:${userId}:${connId}`, nodeId, 'EX', PRESENCE_TTL)
}

// Called from the same heartbeat loop that pings the socket
async function refresh(userId: string, connId: string, room: string) {
  await redis.expire(`presence:${room}:${userId}:${connId}`, PRESENCE_TTL)
}

// A user is present if ANY of their connections is
async function whoIsHere(room: string): Promise<string[]> {
  const users = new Set<string>()
  for await (const key of redis.scanIterator({ MATCH: `presence:${room}:*` })) {
    users.add(key.split(':')[2])
  }
  return [...users]
}

        Alice                     Bob
          │                        │
     ┌────▼────┐              ┌────▼────┐
     │ node A  │              │ node B  │
     └────┬────┘              └────┬────┘
          │                        │
          └──────► Redis ◄──────────┘
                  pub/sub

  Node B publishes to `room:42`. Every node subscribed to that
  channel receives it and writes to its own local sockets.

import Redis from 'ioredis'

// Two connections: a subscriber cannot issue other commands
const sub = new Redis(url)
const pub = new Redis(url)

const localRooms = new Map<string, Set<Socket>>()

sub.on('messageBuffer', (channel, payload) => {
  const room = channel.toString().slice('room:'.length)
  const sockets = localRooms.get(room)
  if (!sockets) return

  const frame = payload.toString()
  for (const ws of sockets) {
    if (ws.readyState === WebSocket.OPEN) ws.send(frame)
  }
})

export async function joinRoom(ws: Socket, room: string) {
  let sockets = localRooms.get(room)
  if (!sockets) {
    sockets = new Set()
    localRooms.set(room, sockets)
    await sub.subscribe(`room:${room}`)   // subscribe once per node, not per socket
  }
  sockets.add(ws)
}

export async function leaveRoom(ws: Socket, room: string) {
  const sockets = localRooms.get(room)
  if (!sockets) return
  sockets.delete(ws)
  if (sockets.size === 0) {
    localRooms.delete(room)
    await sub.unsubscribe(`room:${room}`)  // stop paying for rooms with nobody here
  }
}

const MAX_BUFFERED = 1024 * 1024      // 1MB of unsent data per socket

function safeSend(ws: Socket, frame: string): boolean {
  if (ws.readyState !== WebSocket.OPEN) return false

  if (ws.bufferedAmount > MAX_BUFFERED) {
    // This client cannot keep up. Do not keep queueing for it.
    metrics.increment('ws.slow_consumer_dropped')
    ws.close(4008, 'too slow')
    return false
  }

  ws.send(frame)
  return true
}

// Coalescing: one frame per tick instead of one per event
class Batcher {
  private pending = new Map<string, unknown>()   // keyed: newest wins
  private timer?: NodeJS.Timeout

  push(key: string, value: unknown) {
    this.pending.set(key, value)
    this.timer ??= setTimeout(() => this.flush(), 50)
  }

  private flush() {
    this.timer = undefined
    if (this.pending.size === 0) return

    const batch = [...this.pending.entries()]
    this.pending.clear()
    safeSend(this.ws, JSON.stringify({ type: 'batch', updates: batch }))
  }
}

// Graceful shutdown: stagger the disconnect, tell clients why
async function shutdown() {
  server.close()                       // stop accepting new upgrades

  const clients = [...wss.clients]
  const windowMs = 20_000
  const step = windowMs / Math.max(clients.length, 1)

  clients.forEach((ws, i) => {
    setTimeout(() => {
      // 1012 = service restarting. Clients can reconnect promptly.
      ws.close(1012, 'restarting')
    }, i * step)
  })

  await sleep(windowMs + 5_000)
  process.exit(0)
}

process.on('SIGTERM', shutdown)