From 6ead0c03931a080eabcc6f9d53d8bff0ab0bc0a7 Mon Sep 17 00:00:00 2001 From: syntaxbullet Date: Tue, 6 Jan 2026 17:21:50 +0100 Subject: [PATCH] feat: implement graceful shutdown handling --- .gitignore | 1 + src/index.ts | 6 ++++- src/lib/BotClient.ts | 23 +++++++++++++++++ src/lib/DrizzleClient.ts | 8 ++++-- src/lib/db.test.ts | 56 ++++++++++++++++++++++++++++++++++++++++ src/lib/db.ts | 16 +++++++++--- src/lib/shutdown.test.ts | 56 ++++++++++++++++++++++++++++++++++++++++ src/lib/shutdown.ts | 30 +++++++++++++++++++++ 8 files changed, 190 insertions(+), 6 deletions(-) create mode 100644 src/lib/db.test.ts create mode 100644 src/lib/shutdown.test.ts create mode 100644 src/lib/shutdown.ts diff --git a/.gitignore b/.gitignore index 4af5e53..6d2792d 100644 --- a/.gitignore +++ b/.gitignore @@ -44,3 +44,4 @@ report.[0-9]_.[0-9]_.[0-9]_.[0-9]_.json src/db/data src/db/log scratchpad/ +tickets/ diff --git a/src/index.ts b/src/index.ts index 0797a0b..1ded559 100644 --- a/src/index.ts +++ b/src/index.ts @@ -11,4 +11,8 @@ await AuroraClient.deployCommands(); if (!env.DISCORD_BOT_TOKEN) { throw new Error("❌ DISCORD_BOT_TOKEN is not set in environment variables."); } -AuroraClient.login(env.DISCORD_BOT_TOKEN); \ No newline at end of file +AuroraClient.login(env.DISCORD_BOT_TOKEN); + +// Handle graceful shutdown +process.on("SIGINT", () => AuroraClient.shutdown()); +process.on("SIGTERM", () => AuroraClient.shutdown()); \ No newline at end of file diff --git a/src/lib/BotClient.ts b/src/lib/BotClient.ts index b97404d..b883b54 100644 --- a/src/lib/BotClient.ts +++ b/src/lib/BotClient.ts @@ -93,6 +93,29 @@ export class Client extends DiscordClient { } } } + + async shutdown() { + const { setShuttingDown, waitForTransactions } = await import("./shutdown"); + const { closeDatabase } = await import("./DrizzleClient"); + + logger.info("🛑 Shutdown signal received. Starting graceful shutdown..."); + setShuttingDown(true); + + // Wait for transactions to complete + logger.info("⏳ Waiting for active transactions to complete..."); + await waitForTransactions(10000); + + // Destroy Discord client + logger.info("🔌 Disconnecting from Discord..."); + this.destroy(); + + // Close database + logger.info("🗄️ Closing database connection..."); + await closeDatabase(); + + logger.success("👋 Graceful shutdown complete. Exiting."); + process.exit(0); + } } export const AuroraClient = new Client({ intents: [GatewayIntentBits.Guilds, GatewayIntentBits.MessageContent, GatewayIntentBits.GuildMessages, GatewayIntentBits.GuildMembers] }); \ No newline at end of file diff --git a/src/lib/DrizzleClient.ts b/src/lib/DrizzleClient.ts index 0f86be7..7b193bd 100644 --- a/src/lib/DrizzleClient.ts +++ b/src/lib/DrizzleClient.ts @@ -4,6 +4,10 @@ import * as schema from "@db/schema"; import { env } from "@lib/env"; const connectionString = env.DATABASE_URL; -const postgres = new SQL(connectionString); +export const postgres = new SQL(connectionString); -export const DrizzleClient = drizzle(postgres, { schema }); \ No newline at end of file +export const DrizzleClient = drizzle(postgres, { schema }); + +export const closeDatabase = async () => { + await postgres.close(); +}; \ No newline at end of file diff --git a/src/lib/db.test.ts b/src/lib/db.test.ts new file mode 100644 index 0000000..242d0c1 --- /dev/null +++ b/src/lib/db.test.ts @@ -0,0 +1,56 @@ +import { describe, it, expect, mock, beforeEach } from "bun:test"; + +// Mock shutdown module before importing db +mock.module("./shutdown", () => { + let shuttingDown = false; + let transactions = 0; + return { + isShuttingDown: () => shuttingDown, + setShuttingDown: (v: boolean) => { shuttingDown = v; }, + incrementTransactions: () => { transactions++; }, + decrementTransactions: () => { transactions--; }, + getActiveTransactions: () => transactions, + }; +}); + +// Mock DrizzleClient +mock.module("./DrizzleClient", () => ({ + DrizzleClient: { + transaction: async (cb: any) => cb("MOCK_TX") + } +})); + +import { withTransaction } from "./db"; +import { setShuttingDown, getActiveTransactions } from "./shutdown"; + +describe("db withTransaction", () => { + beforeEach(() => { + setShuttingDown(false); + }); + + it("should allow transactions when not shutting down", async () => { + const result = await withTransaction(async (tx) => { + return "success"; + }); + expect(result).toBe("success"); + expect(getActiveTransactions()).toBe(0); + }); + + it("should throw error when shutting down", async () => { + setShuttingDown(true); + expect(withTransaction(async (tx) => { + return "success"; + })).rejects.toThrow("System is shutting down, no new transactions allowed."); + expect(getActiveTransactions()).toBe(0); + }); + + it("should increment and decrement transaction count", async () => { + let countDuring = 0; + await withTransaction(async (tx) => { + countDuring = getActiveTransactions(); + return "ok"; + }); + expect(countDuring).toBe(1); + expect(getActiveTransactions()).toBe(0); + }); +}); diff --git a/src/lib/db.ts b/src/lib/db.ts index e3b856b..a4c9922 100644 --- a/src/lib/db.ts +++ b/src/lib/db.ts @@ -1,5 +1,6 @@ import { DrizzleClient } from "./DrizzleClient"; import type { Transaction } from "./types"; +import { isShuttingDown, incrementTransactions, decrementTransactions } from "./shutdown"; export const withTransaction = async ( callback: (tx: Transaction) => Promise, @@ -8,8 +9,17 @@ export const withTransaction = async ( if (tx) { return await callback(tx); } else { - return await DrizzleClient.transaction(async (newTx) => { - return await callback(newTx); - }); + if (isShuttingDown()) { + throw new Error("System is shutting down, no new transactions allowed."); + } + + incrementTransactions(); + try { + return await DrizzleClient.transaction(async (newTx) => { + return await callback(newTx); + }); + } finally { + decrementTransactions(); + } } }; diff --git a/src/lib/shutdown.test.ts b/src/lib/shutdown.test.ts new file mode 100644 index 0000000..99cad41 --- /dev/null +++ b/src/lib/shutdown.test.ts @@ -0,0 +1,56 @@ +import { describe, it, expect, beforeEach } from "bun:test"; +import { isShuttingDown, setShuttingDown, incrementTransactions, decrementTransactions, getActiveTransactions, waitForTransactions } from "./shutdown"; + +describe("shutdown logic", () => { + beforeEach(() => { + setShuttingDown(false); + while (getActiveTransactions() > 0) { + decrementTransactions(); + } + }); + + it("should initialize with shuttingDown as false", () => { + expect(isShuttingDown()).toBe(false); + }); + + it("should update shuttingDown state", () => { + setShuttingDown(true); + expect(isShuttingDown()).toBe(true); + }); + + it("should track active transactions", () => { + expect(getActiveTransactions()).toBe(0); + incrementTransactions(); + expect(getActiveTransactions()).toBe(1); + decrementTransactions(); + expect(getActiveTransactions()).toBe(0); + }); + + it("should wait for transactions to complete", async () => { + incrementTransactions(); + + const start = Date.now(); + const waitPromise = waitForTransactions(1000); + + // Simulate completion after 200ms + setTimeout(() => { + decrementTransactions(); + }, 200); + + await waitPromise; + const duration = Date.now() - start; + + expect(duration).toBeGreaterThanOrEqual(200); + expect(getActiveTransactions()).toBe(0); + }); + + it("should timeout if transactions never complete", async () => { + incrementTransactions(); + const start = Date.now(); + await waitForTransactions(500); + const duration = Date.now() - start; + + expect(duration).toBeGreaterThanOrEqual(500); + expect(getActiveTransactions()).toBe(1); // Still 1 because we didn't decrement + }); +}); diff --git a/src/lib/shutdown.ts b/src/lib/shutdown.ts new file mode 100644 index 0000000..eddc1cf --- /dev/null +++ b/src/lib/shutdown.ts @@ -0,0 +1,30 @@ +import { logger } from "@lib/logger"; + +let shuttingDown = false; +let activeTransactions = 0; + +export const isShuttingDown = () => shuttingDown; +export const setShuttingDown = (value: boolean) => { + shuttingDown = value; +}; + +export const incrementTransactions = () => { + activeTransactions++; +}; + +export const decrementTransactions = () => { + activeTransactions--; +}; + +export const getActiveTransactions = () => activeTransactions; + +export const waitForTransactions = async (timeoutMs: number = 10000) => { + const start = Date.now(); + while (activeTransactions > 0) { + if (Date.now() - start > timeoutMs) { + logger.warn(`Shutdown timed out waiting for ${activeTransactions} transactions after ${timeoutMs}ms`); + break; + } + await new Promise(resolve => setTimeout(resolve, 100)); + } +};