feat(shutdown): drain WebSocket clients on SIGTERM

.changeset/ws-drain-on-shutdown.md

@@ -0,0 +1,7 @@
+---
+"nostream": minor
+---
+
+feat(shutdown): drain WebSocket clients on SIGTERM
+
+On SIGTERM, `/readyz` returns 503, new WebSocket connections are rejected, and existing clients receive Nostr CLOSED messages before the socket closes. Drain is bounded by `WS_DRAIN_TIMEOUT_MS` (default 30s).

deploy/README.md

@@ -116,8 +116,12 @@ Each dependency ping uses the default 3s timeout (`ADMIN_DEPENDENCY_PING_TIMEOUT
 Set your load balancer check timeout above that (for example HAProxy
 `timeout check 5s`) so slow-but-healthy backends do not flap during probes.
 Responses are cached in-process for 1s to absorb polling without hammering the DB pool.
-Use readiness before routing traffic to a new instance during deploys; graceful
-WebSocket draining on shutdown is planned as a follow-up.
+Use readiness before routing traffic to a new instance during deploys. On
+SIGTERM the relay sets `/readyz` to `503` with `"status":"draining"` while
+the HTTP listener remains up, rejects new WebSocket connections, drains
+existing clients, then closes (`WS_DRAIN_TIMEOUT_MS`, default 30s). Set
+`stop_grace_period` above that timeout (reference compose uses 45s) so Docker
+does not SIGKILL the container mid-drain.
 
 ## Image delivery on restricted networks
 

deploy/docker-compose.prod.yml

@@ -34,6 +34,7 @@ services:
       nostream-migrate:
         condition: service_completed_successfully
     restart: on-failure
+    stop_grace_period: 45s
 
   nostream-db:
     image: postgres:15

docker-compose.yml

@@ -81,6 +81,7 @@ services:
       otel-collector:
         condition: service_started
     restart: on-failure
+    stop_grace_period: 45s
     networks:
       default:
         ipv4_address: 10.10.10.2

src/@types/adapters.ts

@@ -19,6 +19,7 @@ export type IWebSocketAdapter = EventEmitter & {
   getAuthenticatedPubkeys(): ReadonlySet<string>
   /** Returns false if this AUTH event id was already accepted on this socket. */
   addAuthenticatedPubkey(pubkey: string, authEventId: string): boolean
+  drainAndClose(reason?: string): void
 }
 
 export interface ICacheAdapter {

src/adapters/web-socket-adapter.ts

@@ -5,7 +5,12 @@ import { WebSocket } from 'ws'
 import { ZodError } from 'zod'
 
 import { ContextMetadata, Factory } from '../@types/base'
-import { createAuthChallengeMessage, createNoticeMessage, createOutgoingEventMessage } from '../utils/messages'
+import {
+  createAuthChallengeMessage,
+  createClosedMessage,
+  createNoticeMessage,
+  createOutgoingEventMessage,
+} from '../utils/messages'
 import { IAbortable, IMessageHandler } from '../@types/message-handlers'
 import { IncomingMessage, OutgoingMessage } from '../@types/messages'
 import { IWebSocketAdapter, IWebSocketServerAdapter } from '../@types/adapters'
@@ -156,6 +161,16 @@ export class WebSocketAdapter extends EventEmitter implements IWebSocketAdapter
     return new Map(this.subscriptions)
   }
 
+  public drainAndClose(reason = 'relay shutting down'): void {
+    this.subscriptions.forEach((_filters, subscriptionId) => {
+      this.sendMessage(createClosedMessage(subscriptionId, `error: ${reason}`))
+    })
+
+    if (this.client.readyState === WebSocket.OPEN) {
+      this.client.close(1001, reason)
+    }
+  }
+
   // NIP-42
   public getChallenge(): string {
     return this.session.getChallenge()

src/adapters/web-socket-server-adapter.ts

@@ -1,5 +1,5 @@
 import { IncomingMessage, Server } from 'http'
-import WebSocket, { OPEN, WebSocketServer } from 'ws'
+import WebSocket, { CLOSED, CLOSING, OPEN, WebSocketServer } from 'ws'
 import { propEq } from 'ramda'
 
 import { IWebSocketAdapter, IWebSocketServerAdapter } from '../@types/adapters'
@@ -10,6 +10,7 @@ import { Factory } from '../@types/base'
 import { getRemoteAddress } from '../utils/http'
 import { isRateLimited } from '../handlers/request-handlers/rate-limiter-middleware'
 import { Settings } from '../@types/settings'
+import { getWsDrainTimeoutMs, isDraining } from '../utils/shutdown-state'
 import { WebServerAdapter } from './web-server-adapter'
 
 const logger = createLogger('web-socket-server-adapter')
@@ -46,26 +47,62 @@ export class WebSocketServerAdapter extends WebServerAdapter implements IWebSock
   }
 
   public close(callback?: () => void): void {
-    super.close(() => {
-      logger('closing')
-      clearInterval(this.heartbeatInterval)
-      this.webSocketServer.clients.forEach((webSocket: WebSocket) => {
-        const webSocketAdapter = this.webSocketsAdapters.get(webSocket)
-        if (webSocketAdapter) {
-          logger('terminating client %s: %s', webSocketAdapter.getClientId(), webSocketAdapter.getClientAddress())
-        }
-        webSocket.terminate()
-      })
-      logger('closing web socket server')
-      this.webSocketServer.close(() => {
-        this.webSocketServer.removeAllListeners()
-        if (typeof callback !== 'undefined') {
-          callback()
-        }
-        logger('closed')
+    logger('closing')
+    clearInterval(this.heartbeatInterval)
+    void this.drainClients(getWsDrainTimeoutMs()).finally(() => {
+      super.close(() => {
+        logger('closing web socket server')
+        this.webSocketServer.close(() => {
+          this.webSocketServer.removeAllListeners()
+          if (typeof callback !== 'undefined') {
+            callback()
+          }
+          logger('closed')
+        })
       })
+      this.removeAllListeners()
+    })
+  }
+
+  private async drainClients(timeoutMs: number): Promise<void> {
+    const clients = [...this.webSocketServer.clients] as WebSocket[]
+    if (clients.length === 0) {
+      return
+    }
+
+    logger('draining %d websocket client(s)', clients.length)
+
+    for (const webSocket of clients) {
+      const webSocketAdapter = this.webSocketsAdapters.get(webSocket)
+      if (webSocketAdapter) {
+        logger('closing client %s: %s', webSocketAdapter.getClientId(), webSocketAdapter.getClientAddress())
+        webSocketAdapter.drainAndClose()
+      } else if (webSocket.readyState === OPEN) {
+        webSocket.close(1001, 'relay shutting down')
+      }
+    }
+
+    await Promise.race([
+      Promise.all(clients.map((webSocket) => this.waitForWebSocketClose(webSocket))),
+      new Promise<void>((resolve) => setTimeout(resolve, timeoutMs)),
+    ])
+
+    for (const webSocket of this.webSocketServer.clients) {
+      if (webSocket.readyState === OPEN || webSocket.readyState === CLOSING) {
+        logger('terminating client after drain timeout')
+        webSocket.terminate()
+      }
+    }
+  }
+
+  private waitForWebSocketClose(webSocket: WebSocket): Promise<void> {
+    if (webSocket.readyState === CLOSED) {
+      return Promise.resolve()
+    }
+
+    return new Promise((resolve) => {
+      webSocket.once('close', () => resolve())
     })
-    this.removeAllListeners()
   }
 
   private onBroadcast(event: Event) {
@@ -86,6 +123,12 @@ export class WebSocketServerAdapter extends WebServerAdapter implements IWebSock
   }
 
   private async onConnection(client: WebSocket, req: IncomingMessage) {
+    if (isDraining()) {
+      logger('client rejected: draining')
+      client.close(1001, 'relay shutting down')
+      return
+    }
+
     const currentSettings = this.settings()
     const remoteAddress = getRemoteAddress(req, currentSettings)
 

src/app/app.ts

@@ -11,12 +11,14 @@ import { Serializable } from 'child_process'
 import { Settings } from '../@types/settings'
 import { SettingsStatic } from '../utils/settings'
 import { shutdownMetricsTelemetry } from '../telemetry/metrics'
+import { getPrimaryShutdownDeadlineMs } from '../utils/shutdown-state'
 
 const logger = createLogger('app-primary')
 
 export class App implements IRunnable {
   private workers: WeakMap<Worker, Record<string, string>>
   private watchers: FSWatcher[] | undefined
+  private shuttingDown = false
 
   public constructor(
     private readonly process: NodeJS.Process,
@@ -155,7 +157,7 @@ export class App implements IRunnable {
   private onClusterExit(deadWorker: Worker, code: number, signal: string) {
     logger('worker %s died', deadWorker.process.pid)
 
-    if (code === 0 || signal === 'SIGINT') {
+    if (this.shuttingDown || code === 0 || signal === 'SIGINT') {
       return
     }
     setTimeout(() => {
@@ -172,7 +174,48 @@ export class App implements IRunnable {
   }
 
   private onExit() {
+    if (this.shuttingDown) {
+      return
+    }
+    this.shuttingDown = true
     logger.info('exiting')
+
+    const workers = Object.values(this.cluster.workers ?? {}) as Worker[]
+    if (workers.length === 0) {
+      this.finishExit()
+      return
+    }
+
+    let remaining = workers.length
+    let finished = false
+    const finishOnce = () => {
+      if (finished) {
+        return
+      }
+      finished = true
+      clearTimeout(deadline)
+      this.finishExit()
+    }
+
+    const onWorkerDone = () => {
+      remaining -= 1
+      if (remaining <= 0) {
+        finishOnce()
+      }
+    }
+
+    const deadline = setTimeout(() => {
+      logger.warn('shutdown deadline exceeded, exiting primary')
+      finishOnce()
+    }, getPrimaryShutdownDeadlineMs())
+
+    for (const worker of workers) {
+      worker.once('exit', onWorkerDone)
+      worker.kill()
+    }
+  }
+
+  private finishExit() {
     void shutdownMetricsTelemetry().finally(() => {
       this.close(() => {
         this.process.exit(0)

src/app/worker.ts

@@ -6,9 +6,11 @@ import { createLogger } from '../factories/logger-factory'
 import { FSWatcher } from 'fs'
 import { SettingsStatic } from '../utils/settings'
 import { shutdownMetricsTelemetry } from '../telemetry/metrics'
+import { beginDraining } from '../utils/shutdown-state'
 
 const logger = createLogger('app-worker')
 export class AppWorker implements IRunnable {
+  private exiting = false
   private watchers: FSWatcher[] | undefined
 
   public constructor(
@@ -46,6 +48,11 @@ export class AppWorker implements IRunnable {
   }
 
   private onExit() {
+    if (this.exiting) {
+      return
+    }
+    this.exiting = true
+    beginDraining()
     logger('exiting')
     void shutdownMetricsTelemetry().finally(() => {
       this.close(() => {

src/handlers/request-handlers/get-readyz-request-handler.ts

@@ -1,13 +1,14 @@
 import { NextFunction, Request, Response } from 'express'
 
 import { AdminDependencyHealth, collectAdminHealthSnapshot } from '../../utils/admin-health'
+import { isDraining } from '../../utils/shutdown-state'
 
 // Public readiness probe for load balancers (e.g. HAProxy blue/green). Unlike /healthz
 // (liveness), /readyz returns non-200 when Postgres or Redis is unavailable.
 const READY_SNAPSHOT_CACHE_TTL_MS = 1000
 
 export interface ReadyzSnapshot {
-  status: 'ok' | 'unavailable'
+  status: 'ok' | 'unavailable' | 'draining'
   database: AdminDependencyHealth
   redis: AdminDependencyHealth
 }
@@ -72,6 +73,25 @@ const sendReadyzResponse = (res: Response, statusCode: number, snapshot: ReadyzS
 }
 
 export const getReadyzRequestHandler = async (_req: Request, res: Response, next: NextFunction) => {
+  if (isDraining()) {
+    try {
+      const health = await collectAdminHealthSnapshot()
+      sendReadyzResponse(res, 503, {
+        status: 'draining',
+        database: health.database,
+        redis: health.redis,
+      })
+    } catch {
+      sendReadyzResponse(res, 503, {
+        status: 'draining',
+        database: { ok: false },
+        redis: { ok: false },
+      })
+    }
+    next()
+    return
+  }
+
   try {
     const snapshot = await collectReadyzSnapshot()
     const statusCode = snapshot.status === 'ok' ? 200 : 503

src/utils/shutdown-state.ts

@@ -0,0 +1,33 @@
+const DEFAULT_WS_DRAIN_TIMEOUT_MS = 30_000
+const PRIMARY_SHUTDOWN_DEADLINE_BUFFER_MS = 5_000
+
+let draining = false
+
+export const beginDraining = (): void => {
+  draining = true
+}
+
+export const isDraining = (): boolean => draining
+
+/** Resets in-process drain state. Used by unit tests only. */
+export const resetDrainingState = (): void => {
+  draining = false
+}
+
+export const getWsDrainTimeoutMs = (): number => {
+  const raw = process.env.WS_DRAIN_TIMEOUT_MS
+  if (raw === undefined || raw === '') {
+    return DEFAULT_WS_DRAIN_TIMEOUT_MS
+  }
+
+  const parsed = Number(raw)
+  if (!Number.isFinite(parsed) || parsed < 0) {
+    return DEFAULT_WS_DRAIN_TIMEOUT_MS
+  }
+
+  return parsed
+}
+
+export const getPrimaryShutdownDeadlineMs = (): number => {
+  return getWsDrainTimeoutMs() + PRIMARY_SHUTDOWN_DEADLINE_BUFFER_MS
+}