feat: implement graceful shutdown handling

.gitignore

@@ -44,3 +44,4 @@ report.[0-9]_.[0-9]_.[0-9]_.[0-9]_.json
 src/db/data
 src/db/log
 scratchpad/
+tickets/

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); 
+AuroraClient.login(env.DISCORD_BOT_TOKEN);
+
+// Handle graceful shutdown
+process.on("SIGINT", () => AuroraClient.shutdown());
+process.on("SIGTERM", () => AuroraClient.shutdown());

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] });

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 });
+export const DrizzleClient = drizzle(postgres, { schema });
+
+export const closeDatabase = async () => {
+    await postgres.close();
+};

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);
+    });
+});

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 <T>(
     callback: (tx: Transaction) => Promise<T>,
@@ -8,8 +9,17 @@ export const withTransaction = async <T>(
     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();
+        }
     }
 };

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
+    });
+});

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));
+    }
+};