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,10 @@ 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 returns `503` on `/readyz`, stops accepting new WebSocket
+connections, and drains existing clients before exit (`WS_DRAIN_TIMEOUT_MS`,
+default 30s).
 
 ## Image delivery on restricted networks
 

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, `closed: ${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')
@@ -49,25 +50,61 @@ export class WebSocketServerAdapter extends WebServerAdapter implements IWebSock
     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')
+      void this.drainClients(getWsDrainTimeoutMs()).finally(() => {
+        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())
+    })
+  }
+
   private onBroadcast(event: Event) {
     this.webSocketServer.clients.forEach((webSocket: WebSocket) => {
       if (!propEq('readyState', OPEN)(webSocket)) {
@@ -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

@@ -17,6 +17,7 @@ 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 +156,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 +173,33 @@ 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
+    const onWorkerDone = () => {
+      remaining -= 1
+      if (remaining <= 0) {
+        this.finishExit()
+      }
+    }
+
+    for (const worker of workers) {
+      worker.once('exit', onWorkerDone)
+      worker.process.kill('SIGTERM')
+    }
+  }
+
+  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,6 +1,7 @@
 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.
@@ -72,6 +73,16 @@ const sendReadyzResponse = (res: Response, statusCode: number, snapshot: ReadyzS
 }
 
 export const getReadyzRequestHandler = async (_req: Request, res: Response, next: NextFunction) => {
+  if (isDraining()) {
+    sendReadyzResponse(res, 503, {
+      status: 'unavailable',
+      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,27 @@
+const DEFAULT_WS_DRAIN_TIMEOUT_MS = 30_000
+
+let draining = false
+
+export const beginDraining = (): void => {
+  draining = true
+}
+
+export const isDraining = (): boolean => draining
+
+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
+}