From 8f1d90defb87f86dcc3795cf3bfcab02fcdc1daf Mon Sep 17 00:00:00 2001 From: DevNetlife Date: Sat, 29 Aug 2026 01:15:14 +0100 Subject: [PATCH] feat: Implement Graceful Shutdown Drain Protocol (#1187) --- .../middleware/drainProtection.middleware.ts | 44 ++ .../__tests__/drainProtocol.routes.test.ts | 123 ++++++ .../src/api/routes/drainProtocol.routes.ts | 66 +++ .../api/routes/route-groups/admin-routes.ts | 6 + .../20260829000000_shutdown_drain_protocol.ts | 45 ++ backend/src/index.ts | 8 +- .../__tests__/drainProtocol.service.test.ts | 102 +++++ backend/src/services/drainProtocol.service.ts | 408 ++++++++++++++++++ docs/SHUTDOWN_DRAIN_PROTOCOL.md | 95 ++++ 9 files changed, 896 insertions(+), 1 deletion(-) create mode 100644 backend/src/api/middleware/drainProtection.middleware.ts create mode 100644 backend/src/api/routes/__tests__/drainProtocol.routes.test.ts create mode 100644 backend/src/api/routes/drainProtocol.routes.ts create mode 100644 backend/src/database/migrations/20260829000000_shutdown_drain_protocol.ts create mode 100644 backend/src/services/__tests__/drainProtocol.service.test.ts create mode 100644 backend/src/services/drainProtocol.service.ts create mode 100644 docs/SHUTDOWN_DRAIN_PROTOCOL.md diff --git a/backend/src/api/middleware/drainProtection.middleware.ts b/backend/src/api/middleware/drainProtection.middleware.ts new file mode 100644 index 00000000..a4f29e4a --- /dev/null +++ b/backend/src/api/middleware/drainProtection.middleware.ts @@ -0,0 +1,44 @@ +import type { FastifyInstance, FastifyRequest, FastifyReply } from "fastify"; +import { drainProtocolService } from "../../services/drainProtocol.service.js"; + +const EXEMPT_PATHS = [ + "/health", + "/healthz", + "/readyz", + "/api/v1/health", + "/api/v1/admin/shutdown/drain/start", + "/api/v1/admin/shutdown/drain/status", + "/api/v1/admin/shutdown/drain/cancel", + "/api/v1/admin/shutdown/drain/force", + "/api/v1/admin/shutdown/drain/history", +]; + +export async function registerDrainProtectionMiddleware(server: FastifyInstance): Promise { + // Track in-flight request lifecycle + server.addHook("onRequest", async (request: FastifyRequest, reply: FastifyReply) => { + drainProtocolService.incrementInFlight(); + + reply.raw.on("finish", () => { + drainProtocolService.decrementInFlight(); + }); + + const isExempt = EXEMPT_PATHS.some((path) => request.url.startsWith(path)); + if (isExempt) { + return; + } + + if (drainProtocolService.isDraining()) { + const isMutatingMethod = ["POST", "PUT", "DELETE", "PATCH"].includes(request.method.toUpperCase()); + + if (isMutatingMethod || drainProtocolService.getMode() === "force") { + reply.header("Retry-After", "30"); + return reply.status(503).send({ + error: "Service Unavailable", + message: "Server is currently undergoing graceful shutdown drain", + state: drainProtocolService.getState(), + retryAfterSeconds: 30, + }); + } + } + }); +} diff --git a/backend/src/api/routes/__tests__/drainProtocol.routes.test.ts b/backend/src/api/routes/__tests__/drainProtocol.routes.test.ts new file mode 100644 index 00000000..85260aec --- /dev/null +++ b/backend/src/api/routes/__tests__/drainProtocol.routes.test.ts @@ -0,0 +1,123 @@ +import { describe, it, expect, beforeEach, vi } from "vitest"; +import Fastify from "fastify"; +import { drainProtocolRoutes } from "../drainProtocol.routes.js"; +import { drainProtocolService } from "../../../services/drainProtocol.service.js"; + +vi.mock("../../../services/drainProtocol.service.js", () => { + const mockStatus = { + sessionId: "mock-session-id", + nodeId: "node-1", + state: "ACTIVE", + drainMode: "graceful", + inFlightRequests: 0, + activeConnections: 0, + activeStreams: 0, + startedAt: null, + drainedAt: null, + reason: null, + initiatedBy: null, + timeoutSeconds: 30, + }; + + return { + drainProtocolService: { + getStatus: vi.fn().mockImplementation(() => mockStatus), + startDrain: vi.fn().mockImplementation(async (opts) => ({ + ...mockStatus, + state: "DRAINED", + reason: opts?.reason || "Graceful drain", + initiatedBy: opts?.initiatedBy || "admin", + })), + cancelDrain: vi.fn().mockImplementation(async (by) => ({ + ...mockStatus, + state: "ACTIVE", + initiatedBy: by, + })), + forceShutdown: vi.fn().mockImplementation(async (reason) => ({ + ...mockStatus, + state: "FAILED", + reason, + })), + getDrainHistory: vi.fn().mockResolvedValue([mockStatus]), + }, + }; +}); + +vi.mock("../../middleware/auth.js", () => ({ + authMiddleware: () => async () => {}, +})); + +describe("drainProtocolRoutes", () => { + let app: ReturnType; + + beforeEach(async () => { + app = Fastify(); + await app.register(drainProtocolRoutes); + vi.clearAllMocks(); + }); + + it("GET /status returns current drain status", async () => { + const res = await app.inject({ + method: "GET", + url: "/status", + }); + + expect(res.statusCode).toBe(200); + const body = res.json(); + expect(body.nodeId).toBe("node-1"); + expect(body.state).toBe("ACTIVE"); + }); + + it("POST /start initiates drain and returns HTTP 202 Accepted", async () => { + const res = await app.inject({ + method: "POST", + url: "/start", + payload: { + timeoutSeconds: 15, + reason: "Maintenance shutdown", + }, + }); + + expect(res.statusCode).toBe(202); + expect(drainProtocolService.startDrain).toHaveBeenCalledWith( + expect.objectContaining({ + timeoutSeconds: 15, + reason: "Maintenance shutdown", + }) + ); + }); + + it("POST /cancel cancels active drain protocol", async () => { + const res = await app.inject({ + method: "POST", + url: "/cancel", + payload: { cancelledBy: "operator-admin" }, + }); + + expect(res.statusCode).toBe(200); + expect(drainProtocolService.cancelDrain).toHaveBeenCalledWith("operator-admin"); + }); + + it("POST /force triggers force shutdown", async () => { + const res = await app.inject({ + method: "POST", + url: "/force", + payload: { reason: "Urgent abort" }, + }); + + expect(res.statusCode).toBe(200); + expect(drainProtocolService.forceShutdown).toHaveBeenCalledWith("Urgent abort"); + }); + + it("GET /history returns past drain sessions", async () => { + const res = await app.inject({ + method: "GET", + url: "/history?limit=5", + }); + + expect(res.statusCode).toBe(200); + const body = res.json(); + expect(body.history).toHaveLength(1); + expect(body.count).toBe(1); + }); +}); diff --git a/backend/src/api/routes/drainProtocol.routes.ts b/backend/src/api/routes/drainProtocol.routes.ts new file mode 100644 index 00000000..2eb5508f --- /dev/null +++ b/backend/src/api/routes/drainProtocol.routes.ts @@ -0,0 +1,66 @@ +import type { FastifyInstance } from "fastify"; +import { z } from "zod"; +import { drainProtocolService } from "../../services/drainProtocol.service.js"; +import { authMiddleware } from "../middleware/auth.js"; +import { sendApiError } from "../utils/response.js"; + +const startDrainSchema = z.object({ + nodeId: z.string().optional(), + timeoutSeconds: z.number().int().min(1).max(300).optional(), + reason: z.string().optional(), + initiatedBy: z.string().optional(), + mode: z.enum(["graceful", "force", "read_only"]).optional(), + metadata: z.record(z.unknown()).optional(), +}); + +const cancelDrainSchema = z.object({ + cancelledBy: z.string().optional(), +}); + +const forceDrainSchema = z.object({ + reason: z.string().optional(), +}); + +export async function drainProtocolRoutes(server: FastifyInstance): Promise { + const adminAuth = authMiddleware({ requiredScopes: ["admin:write"] }); + + // Start graceful shutdown drain protocol + server.post("/start", { preHandler: [adminAuth] }, async (request, reply) => { + const parsed = startDrainSchema.safeParse(request.body || {}); + if (!parsed.success) { + return sendApiError(reply, 400, "Invalid drain options", { issues: parsed.error.errors }); + } + + const status = await drainProtocolService.startDrain(parsed.data); + return reply.status(202).send(status); + }); + + // Get current drain status + server.get("/status", async (_request, reply) => { + const status = drainProtocolService.getStatus(); + return reply.status(200).send(status); + }); + + // Cancel drain and resume normal operation + server.post("/cancel", { preHandler: [adminAuth] }, async (request, reply) => { + const parsed = cancelDrainSchema.safeParse(request.body || {}); + const cancelledBy = parsed.success ? parsed.data.cancelledBy || "admin" : "admin"; + const status = await drainProtocolService.cancelDrain(cancelledBy); + return reply.status(200).send(status); + }); + + // Force immediate shutdown + server.post("/force", { preHandler: [adminAuth] }, async (request, reply) => { + const parsed = forceDrainSchema.safeParse(request.body || {}); + const reason = parsed.success ? parsed.data.reason || "Force shutdown requested" : "Force shutdown requested"; + const status = await drainProtocolService.forceShutdown(reason); + return reply.status(200).send(status); + }); + + // Get drain history + server.get("/history", { preHandler: [adminAuth] }, async (request, reply) => { + const limit = Number((request.query as { limit?: string })?.limit) || 20; + const history = await drainProtocolService.getDrainHistory(limit); + return reply.status(200).send({ history, count: history.length }); + }); +} diff --git a/backend/src/api/routes/route-groups/admin-routes.ts b/backend/src/api/routes/route-groups/admin-routes.ts index 03e9acfe..8b15f353 100644 --- a/backend/src/api/routes/route-groups/admin-routes.ts +++ b/backend/src/api/routes/route-groups/admin-routes.ts @@ -121,5 +121,11 @@ export async function registerAdminRoutes(server: FastifyInstance): Promise { + await knex.schema.createTable("shutdown_drain_sessions", (t) => { + t.uuid("id").primary().defaultTo(knex.raw("gen_random_uuid()")); + t.string("node_id", 128).notNullable(); + t.string("state", 32).notNullable().defaultTo("ACTIVE"); // ACTIVE, DRAINING, DRAINED, CANCELLED, FAILED + t.string("drain_mode", 32).notNullable().defaultTo("graceful"); // graceful, force, read_only + t.string("reason", 500).nullable(); + t.string("initiated_by", 128).notNullable().defaultTo("system"); + t.integer("timeout_seconds").notNullable().defaultTo(30); + t.integer("pending_jobs_count").notNullable().defaultTo(0); + t.integer("active_connections_count").notNullable().defaultTo(0); + t.integer("active_streams_count").notNullable().defaultTo(0); + t.timestamp("started_at", { useTz: true }).notNullable().defaultTo(knex.fn.now()); + t.timestamp("drained_at", { useTz: true }).nullable(); + t.timestamp("cancelled_at", { useTz: true }).nullable(); + t.jsonb("metadata").notNullable().defaultTo("{}"); + t.timestamp("created_at", { useTz: true }).notNullable().defaultTo(knex.fn.now()); + t.timestamp("updated_at", { useTz: true }).notNullable().defaultTo(knex.fn.now()); + + t.index(["node_id", "state"], "idx_shutdown_drain_node_state"); + t.index(["created_at"], "idx_shutdown_drain_created"); + }); + + await knex.schema.createTable("shutdown_drain_logs", (t) => { + t.uuid("id").primary().defaultTo(knex.raw("gen_random_uuid()")); + t.uuid("session_id").notNullable().references("id").inTable("shutdown_drain_sessions").onDelete("CASCADE"); + t.string("event_type", 64).notNullable(); // DRAIN_INITIATED, JOBS_PAUSED, WS_DRAINED, STREAMS_STOPPED, DRAIN_COMPLETED, DRAIN_CANCELLED, DRAIN_FAILED, FORCE_SHUTDOWN + t.string("message", 500).notNullable(); + t.jsonb("details").notNullable().defaultTo("{}"); + t.timestamp("timestamp", { useTz: true }).notNullable().defaultTo(knex.fn.now()); + + t.index(["session_id", "timestamp"], "idx_shutdown_drain_logs_session"); + }); +} + +export async function down(knex: Knex): Promise { + await knex.schema.dropTableIfExists("shutdown_drain_logs"); + await knex.schema.dropTableIfExists("shutdown_drain_sessions"); +} diff --git a/backend/src/index.ts b/backend/src/index.ts index c150f7c2..729a502a 100644 --- a/backend/src/index.ts +++ b/backend/src/index.ts @@ -25,6 +25,8 @@ import { registerRequestLoggingMiddleware } from "./api/middleware/logging.middl import { registerTracing } from "./api/middleware/tracing.js"; import { getTelegramBotService } from "./services/telegram.bot.service.js"; import { registerCompatibilityMiddleware } from "./api/compatibility/middleware.js"; +import { registerDrainProtectionMiddleware } from "./api/middleware/drainProtection.middleware.js"; +import { drainProtocolService } from "./services/drainProtocol.service.js"; export async function buildServer() { const server = Fastify({ @@ -75,6 +77,9 @@ export async function buildServer() { // Register metrics middleware (to capture all requests) await registerMetrics(server as any); + // Register graceful shutdown drain protection middleware + await registerDrainProtectionMiddleware(server as any); + // Register plugins await server.register(cors, { origin: (origin, callback) => { @@ -181,7 +186,8 @@ async function start() { // ─── Graceful shutdown ────────────────────────────────────────────────────── const shutdown = async (signal: string) => { - logger.info({ signal }, "Shutdown signal received"); + logger.info({ signal }, "Shutdown signal received; initiating drain protocol"); + await drainProtocolService.startDrain({ reason: `Received signal ${signal}`, initiatedBy: "system" }); // Stop Telegram bot service const telegramService = getTelegramBotService(); diff --git a/backend/src/services/__tests__/drainProtocol.service.test.ts b/backend/src/services/__tests__/drainProtocol.service.test.ts new file mode 100644 index 00000000..3a6a95fa --- /dev/null +++ b/backend/src/services/__tests__/drainProtocol.service.test.ts @@ -0,0 +1,102 @@ +import { describe, it, expect, beforeEach, vi } from "vitest"; +import { DrainProtocolService } from "../drainProtocol.service.js"; + +vi.mock("../database/connection.js", () => { + const mockDb: any = vi.fn().mockImplementation(() => mockDb); + mockDb.schema = { + hasTable: vi.fn().mockResolvedValue(true), + }; + mockDb.where = vi.fn().mockReturnValue(mockDb); + mockDb.update = vi.fn().mockResolvedValue([1]); + mockDb.insert = vi.fn().mockReturnValue({ + returning: vi.fn().mockResolvedValue([{ id: "test-session-uuid-1234" }]), + }); + mockDb.select = vi.fn().mockReturnValue(mockDb); + mockDb.orderBy = vi.fn().mockReturnValue(mockDb); + mockDb.limit = vi.fn().mockResolvedValue([]); + return { getDatabase: () => mockDb }; +}); + +vi.mock("../api/websocket/websocket.server.js", () => ({ + wsServer: { + shutdown: vi.fn().mockResolvedValue(undefined), + }, +})); + +vi.mock("../workers/queue.js", () => ({ + JobQueue: { + getInstance: () => ({ + stop: vi.fn().mockResolvedValue(undefined), + }), + }, +})); + +vi.mock("../jobs/supplyVerification.job.js", () => ({ + getSupplyVerificationQueue: () => ({ + stop: vi.fn().mockResolvedValue(undefined), + }), +})); + +vi.mock("../workers/webhookDelivery.worker.js", () => ({ + stopWebhookWorker: vi.fn().mockResolvedValue(undefined), +})); + +describe("DrainProtocolService", () => { + let service: DrainProtocolService; + + beforeEach(() => { + service = new DrainProtocolService(); + vi.clearAllMocks(); + }); + + it("initializes with ACTIVE state and 0 in-flight requests", () => { + expect(service.getState()).toBe("ACTIVE"); + expect(service.isDraining()).toBe(false); + expect(service.getInFlightCount()).toBe(0); + }); + + it("increments and decrements in-flight requests accurately", () => { + service.incrementInFlight(); + service.incrementInFlight(); + expect(service.getInFlightCount()).toBe(2); + + service.decrementInFlight(); + expect(service.getInFlightCount()).toBe(1); + + service.decrementInFlight(); + expect(service.getInFlightCount()).toBe(0); + + // Prevents negative counter + service.decrementInFlight(); + expect(service.getInFlightCount()).toBe(0); + }); + + it("transitions state to DRAINING and completes immediately when no in-flight requests exist", async () => { + const status = await service.startDrain({ + reason: "Maintenance test", + initiatedBy: "operator", + timeoutSeconds: 10, + }); + + expect(status.state).toBe("DRAINED"); + expect(status.reason).toBe("Maintenance test"); + expect(status.initiatedBy).toBe("operator"); + expect(service.isDraining()).toBe(true); + }); + + it("allows cancelling active or completed drain session and resets to ACTIVE", async () => { + await service.startDrain({ reason: "Cancel test" }); + const cancelledStatus = await service.cancelDrain("operator-admin"); + + expect(service.getState()).toBe("ACTIVE"); + expect(service.isDraining()).toBe(false); + expect(cancelledStatus.state).toBe("ACTIVE"); + }); + + it("supports force shutdown execution", async () => { + await service.startDrain({ reason: "Force test" }); + const forcedStatus = await service.forceShutdown("Emergency stop"); + + expect(forcedStatus.state).toBe("FAILED"); + }); +}); diff --git a/backend/src/services/drainProtocol.service.ts b/backend/src/services/drainProtocol.service.ts new file mode 100644 index 00000000..1884a9af --- /dev/null +++ b/backend/src/services/drainProtocol.service.ts @@ -0,0 +1,408 @@ +import { getDatabase } from "../database/connection.js"; +import { logger } from "../utils/logger.js"; +import { Gauge, Counter } from "prom-client"; +import { wsServer } from "../api/websocket/websocket.server.js"; +import { JobQueue } from "../workers/queue.js"; +import { getSupplyVerificationQueue } from "../jobs/supplyVerification.job.js"; +import { stopWebhookWorker } from "../workers/webhookDelivery.worker.js"; + +export type DrainState = "ACTIVE" | "DRAINING" | "DRAINED" | "CANCELLED" | "FAILED"; +export type DrainMode = "graceful" | "force" | "read_only"; + +export interface DrainOptions { + nodeId?: string; + timeoutSeconds?: number; + reason?: string; + initiatedBy?: string; + mode?: DrainMode; + metadata?: Record; +} + +export interface DrainStatus { + sessionId: string | null; + nodeId: string; + state: DrainState; + drainMode: DrainMode; + inFlightRequests: number; + activeConnections: number; + activeStreams: number; + startedAt: string | null; + drainedAt: string | null; + reason: string | null; + initiatedBy: string | null; + timeoutSeconds: number; +} + +// Prometheus metrics for Graceful Shutdown Drain Protocol +export const drainStatusMetric = new Gauge({ + name: "bridge_watch_drain_status", + help: "System drain status (0 = ACTIVE, 1 = DRAINING, 2 = DRAINED, 3 = CANCELLED, 4 = FAILED)", +}); + +export const inFlightRequestsMetric = new Gauge({ + name: "bridge_watch_drain_in_flight_requests", + help: "Current number of in-flight HTTP requests during drain", +}); + +export const drainEventsCounter = new Counter({ + name: "bridge_watch_drain_events_total", + help: "Total count of drain protocol lifecycle events", + labelNames: ["event_type", "state"], +}); + +export class DrainProtocolService { + private readonly db = getDatabase(); + private currentState: DrainState = "ACTIVE"; + private currentMode: DrainMode = "graceful"; + private currentSessionId: string | null = null; + private nodeId: string = process.env.NODE_ID || "node-1"; + private inFlightRequests = 0; + private activeStreams = 0; + private drainStartedAt: Date | null = null; + private drainDrainedAt: Date | null = null; + private drainReason: string | null = null; + private drainInitiatedBy: string | null = null; + private timeoutSeconds = 30; + private drainTimer: NodeJS.Timeout | null = null; + + constructor() { + drainStatusMetric.set(0); // 0 = ACTIVE + } + + public getNodeId(): string { + return this.nodeId; + } + + public getState(): DrainState { + return this.currentState; + } + + public getMode(): DrainMode { + return this.currentMode; + } + + public isDraining(): boolean { + return this.currentState === "DRAINING" || this.currentState === "DRAINED"; + } + + public incrementInFlight(): void { + this.inFlightRequests++; + inFlightRequestsMetric.set(this.inFlightRequests); + } + + public decrementInFlight(): void { + if (this.inFlightRequests > 0) { + this.inFlightRequests--; + } + inFlightRequestsMetric.set(this.inFlightRequests); + } + + public getInFlightCount(): number { + return this.inFlightRequests; + } + + /** + * Initiates the Graceful Shutdown Drain Protocol + */ + public async startDrain(options: DrainOptions = {}): Promise { + if (this.currentState === "DRAINING") { + logger.warn({ nodeId: this.nodeId, sessionId: this.currentSessionId }, "Drain protocol is already in progress"); + return this.getStatus(); + } + + this.nodeId = options.nodeId || this.nodeId; + this.currentState = "DRAINING"; + this.currentMode = options.mode || "graceful"; + this.timeoutSeconds = options.timeoutSeconds || 30; + this.drainStartedAt = new Date(); + this.drainDrainedAt = null; + this.drainReason = options.reason || "Scheduled shutdown / maintenance"; + this.drainInitiatedBy = options.initiatedBy || "admin"; + + drainStatusMetric.set(1); // 1 = DRAINING + drainEventsCounter.inc({ event_type: "DRAIN_INITIATED", state: "DRAINING" }); + + logger.info( + { + nodeId: this.nodeId, + mode: this.currentMode, + timeoutSeconds: this.timeoutSeconds, + reason: this.drainReason, + }, + "Initiating Graceful Shutdown Drain Protocol" + ); + + // Persist drain session to database + try { + const hasTable = await this.db.schema.hasTable("shutdown_drain_sessions"); + if (hasTable) { + const [inserted] = await this.db("shutdown_drain_sessions") + .insert({ + node_id: this.nodeId, + state: "DRAINING", + drain_mode: this.currentMode, + reason: this.drainReason, + initiated_by: this.drainInitiatedBy, + timeout_seconds: this.timeoutSeconds, + pending_jobs_count: 0, + active_connections_count: 0, + active_streams_count: this.activeStreams, + started_at: this.drainStartedAt, + metadata: JSON.stringify(options.metadata || {}), + }) + .returning("*"); + + this.currentSessionId = inserted?.id || null; + + if (this.currentSessionId) { + await this.logDrainEvent("DRAIN_INITIATED", "Drain protocol initiated", { + mode: this.currentMode, + timeoutSeconds: this.timeoutSeconds, + reason: this.drainReason, + }); + } + } + } catch (err) { + logger.error({ err }, "Failed to persist shutdown drain session to DB"); + } + + // Step 1: Pause background job queues & workers + try { + await JobQueue.getInstance().stop().catch(() => {}); + await getSupplyVerificationQueue().stop().catch(() => {}); + await stopWebhookWorker().catch(() => {}); + await this.logDrainEvent("JOBS_PAUSED", "Background queues and workers stopped"); + } catch (err) { + logger.error({ err }, "Error stopping background workers during drain"); + } + + // Step 2: Gracefully shutdown WebSocket server (broadcast disconnect frame) + try { + await wsServer.shutdown().catch(() => {}); + await this.logDrainEvent("WS_DRAINED", "WebSocket connections gracefully drained"); + } catch (err) { + logger.error({ err }, "Error shutting down WebSockets during drain"); + } + + // Set up timeout timer for forceful completion if graceful timeout is reached + this.drainTimer = setTimeout(() => { + this.handleDrainTimeout().catch((err) => { + logger.error({ err }, "Error during drain timeout resolution"); + }); + }, this.timeoutSeconds * 1000); + + // Check if system is immediately drained + if (this.inFlightRequests === 0) { + await this.completeDrain(); + } + + return this.getStatus(); + } + + /** + * Completes the drain protocol when all tasks/connections are finished + */ + public async completeDrain(): Promise { + if (this.currentState !== "DRAINING") { + return this.getStatus(); + } + + if (this.drainTimer) { + clearTimeout(this.drainTimer); + this.drainTimer = null; + } + + this.currentState = "DRAINED"; + this.drainDrainedAt = new Date(); + drainStatusMetric.set(2); // 2 = DRAINED + drainEventsCounter.inc({ event_type: "DRAIN_COMPLETED", state: "DRAINED" }); + + logger.info({ nodeId: this.nodeId, sessionId: this.currentSessionId }, "Graceful Shutdown Drain Protocol completed successfully"); + + if (this.currentSessionId) { + try { + const hasTable = await this.db.schema.hasTable("shutdown_drain_sessions"); + if (hasTable) { + await this.db("shutdown_drain_sessions") + .where({ id: this.currentSessionId }) + .update({ + state: "DRAINED", + drained_at: this.drainDrainedAt, + pending_jobs_count: this.inFlightRequests, + updated_at: new Date(), + }); + + await this.logDrainEvent("DRAIN_COMPLETED", "Graceful shutdown drain completed successfully"); + } + } catch (err) { + logger.error({ err }, "Failed to update DB on drain completion"); + } + } + + return this.getStatus(); + } + + /** + * Cancels an active drain session and resumes normal operation + */ + public async cancelDrain(cancelledBy = "admin"): Promise { + if (this.currentState !== "DRAINING" && this.currentState !== "DRAINED") { + logger.warn("No active drain session to cancel"); + return this.getStatus(); + } + + if (this.drainTimer) { + clearTimeout(this.drainTimer); + this.drainTimer = null; + } + + const previousState = this.currentState; + this.currentState = "CANCELLED"; + drainStatusMetric.set(3); // 3 = CANCELLED + drainEventsCounter.inc({ event_type: "DRAIN_CANCELLED", state: "CANCELLED" }); + + logger.info({ nodeId: this.nodeId, cancelledBy }, "Graceful shutdown drain protocol cancelled by operator"); + + if (this.currentSessionId) { + try { + const hasTable = await this.db.schema.hasTable("shutdown_drain_sessions"); + if (hasTable) { + await this.db("shutdown_drain_sessions") + .where({ id: this.currentSessionId }) + .update({ + state: "CANCELLED", + cancelled_at: new Date(), + updated_at: new Date(), + }); + + await this.logDrainEvent("DRAIN_CANCELLED", `Drain protocol cancelled by ${cancelledBy}`, { + previousState, + }); + } + } catch (err) { + logger.error({ err }, "Failed to update DB on drain cancellation"); + } + } + + // Reset back to ACTIVE state + this.currentState = "ACTIVE"; + this.currentSessionId = null; + drainStatusMetric.set(0); // 0 = ACTIVE + + return this.getStatus(); + } + + /** + * Forcefully completes the drain session + */ + public async forceShutdown(reason = "Force shutdown requested"): Promise { + if (this.drainTimer) { + clearTimeout(this.drainTimer); + this.drainTimer = null; + } + + this.currentState = "FAILED"; + drainStatusMetric.set(4); // 4 = FAILED + drainEventsCounter.inc({ event_type: "FORCE_SHUTDOWN", state: "FAILED" }); + + logger.warn({ nodeId: this.nodeId, reason }, "Forceful shutdown drain executed"); + + if (this.currentSessionId) { + try { + const hasTable = await this.db.schema.hasTable("shutdown_drain_sessions"); + if (hasTable) { + await this.db("shutdown_drain_sessions") + .where({ id: this.currentSessionId }) + .update({ + state: "FAILED", + updated_at: new Date(), + }); + + await this.logDrainEvent("FORCE_SHUTDOWN", `Force shutdown executed: ${reason}`); + } + } catch (err) { + logger.error({ err }, "Failed to log force shutdown in DB"); + } + } + + return this.getStatus(); + } + + /** + * Returns current drain status + */ + public getStatus(): DrainStatus { + return { + sessionId: this.currentSessionId, + nodeId: this.nodeId, + state: this.currentState, + drainMode: this.currentMode, + inFlightRequests: this.inFlightRequests, + activeConnections: 0, + activeStreams: this.activeStreams, + startedAt: this.drainStartedAt ? this.drainStartedAt.toISOString() : null, + drainedAt: this.drainDrainedAt ? this.drainDrainedAt.toISOString() : null, + reason: this.drainReason, + initiatedBy: this.drainInitiatedBy, + timeoutSeconds: this.timeoutSeconds, + }; + } + + /** + * Fetches drain session history + */ + public async getDrainHistory(limit = 20): Promise { + try { + const hasTable = await this.db.schema.hasTable("shutdown_drain_sessions"); + if (!hasTable) return []; + + return await this.db("shutdown_drain_sessions") + .select("*") + .orderBy("created_at", "desc") + .limit(limit); + } catch (err) { + logger.error({ err }, "Failed to fetch shutdown drain history"); + return []; + } + } + + private async handleDrainTimeout(): Promise { + if (this.currentState === "DRAINING") { + logger.warn( + { + nodeId: this.nodeId, + remainingInFlight: this.inFlightRequests, + timeoutSeconds: this.timeoutSeconds, + }, + "Drain timeout reached before all in-flight requests completed; executing completion fallback" + ); + + await this.logDrainEvent("DRAIN_TIMEOUT", "Timeout reached during drain", { + remainingInFlight: this.inFlightRequests, + }); + + await this.completeDrain(); + } + } + + private async logDrainEvent(eventType: string, message: string, details: Record = {}): Promise { + if (!this.currentSessionId) return; + + try { + const hasTable = await this.db.schema.hasTable("shutdown_drain_logs"); + if (hasTable) { + await this.db("shutdown_drain_logs").insert({ + session_id: this.currentSessionId, + event_type: eventType, + message, + details: JSON.stringify(details), + timestamp: new Date(), + }); + } + } catch (err) { + logger.warn({ err, eventType }, "Failed to log drain event to database"); + } + } +} + +export const drainProtocolService = new DrainProtocolService(); diff --git a/docs/SHUTDOWN_DRAIN_PROTOCOL.md b/docs/SHUTDOWN_DRAIN_PROTOCOL.md new file mode 100644 index 00000000..8c6a2805 --- /dev/null +++ b/docs/SHUTDOWN_DRAIN_PROTOCOL.md @@ -0,0 +1,95 @@ +# Graceful Shutdown Drain Protocol + +## Overview + +The Graceful Shutdown Drain Protocol guarantees safe node teardown and maintenance transition across Bridge Watch instances. When a node receives an OS signal (`SIGTERM`, `SIGINT`) or an administrative trigger, the drain protocol halts incoming mutating API calls, gracefully drains WebSocket frames, pauses background workers (BullMQ queues, Horizon streams, webhook dispatchers), and waits for in-flight tasks to complete before exiting. + +--- + +## Drain Lifecycle States + +``` + +--------+ Trigger Drain +----------+ Tasks Completed +----------+ + | ACTIVE | -------------------------> | DRAINING | -------------------------> | DRAINED | + +--------+ +----------+ +----------+ + ^ | | + | Cancel Drain | Timeout / Force Expiry v + +--------------------------------------+ <-------------------------------- + FAILED | + +----------+ +``` + +1. **`ACTIVE`**: Node is operating normally. +2. **`DRAINING`**: Ingestion queues paused, WS server draining, mutating HTTP requests rejected with HTTP 503 (`Retry-After: 30`). +3. **`DRAINED`**: All in-flight requests and connections have cleanly drained. Ready for process exit. +4. **`CANCELLED`**: Active drain was aborted by an operator; normal traffic resumes. +5. **`FAILED`**: Force shutdown was executed or timeout expired before all tasks finished. + +--- + +## Data Model & Migration + +Database table: `shutdown_drain_sessions` +- `id` (UUID, Primary Key) +- `node_id` (String) +- `state` (`ACTIVE` | `DRAINING` | `DRAINED` | `CANCELLED` | `FAILED`) +- `drain_mode` (`graceful` | `force` | `read_only`) +- `reason` (String) +- `initiated_by` (String) +- `timeout_seconds` (Integer) +- `pending_jobs_count` (Integer) +- `active_connections_count` (Integer) +- `active_streams_count` (Integer) +- `started_at` / `drained_at` / `cancelled_at` (Timestamps) + +Database table: `shutdown_drain_logs` +- Audit trail for drain protocol events (`DRAIN_INITIATED`, `JOBS_PAUSED`, `WS_DRAINED`, `DRAIN_COMPLETED`, `DRAIN_CANCELLED`, `FORCE_SHUTDOWN`). + +--- + +## Operational Control API + +### Initiate Drain Session +```http +POST /api/v1/admin/shutdown/drain/start +Content-Type: application/json +Authorization: Bearer + +{ + "timeoutSeconds": 30, + "reason": "Scheduled node maintenance", + "mode": "graceful" +} +``` + +### Check Drain Status +```http +GET /api/v1/admin/shutdown/drain/status +``` + +### Cancel Drain Session +```http +POST /api/v1/admin/shutdown/drain/cancel +Authorization: Bearer + +{ + "cancelledBy": "operator-alice" +} +``` + +### Force Immediate Shutdown +```http +POST /api/v1/admin/shutdown/drain/force +Authorization: Bearer + +{ + "reason": "Emergency node replacement" +} +``` + +--- + +## Observability + +- **`bridge_watch_drain_status`**: Prometheus gauge (0 = ACTIVE, 1 = DRAINING, 2 = DRAINED, 3 = CANCELLED, 4 = FAILED). +- **`bridge_watch_drain_in_flight_requests`**: Current in-flight HTTP request counter. +- **`bridge_watch_drain_events_total`**: Counter metric tracking drain lifecycle events.