fix: add graceful database shutdown to prevent WAL file leaks

The SQLite database connection was never explicitly closed when the
process exits. In WAL mode, this can leave behind -wal and -shm files
and risk data loss if the WAL hasn't been checkpointed.

Changes:
- Add closeDb() export for explicit database closure
- Register process shutdown handlers (exit, SIGTERM, SIGINT, SIGHUP)
  that call closeDb() before the process terminates
- Idempotent: multiple calls to closeDb() are safe
- After closeDb(), getDb() transparently re-opens the database
- Add 6 unit tests covering close, idempotency, reopen, data
  persistence, and WAL cleanup (all passing)
This commit is contained in:
liuxiaopai-ai committed 2026-02-08 21:37:35 +08:00
1 parent 4ac1f8cea3
commit 746cea1eb7
2 files changed
+159

No files matched your search

+102
View File
@@ -0,0 +1,102 @@
/**
* Unit tests for database graceful shutdown.
*
* Run with: npx tsx --test src/__tests__/unit/db-shutdown.test.ts
*
* Tests verify that:
* 1. closeDb() closes the database without errors
* 2. closeDb() is idempotent (can be called multiple times)
* 3. getDb() re-opens the database after closeDb()
* 4. Data persists across close/reopen cycles
*/
import { describe, it, before, afterEach } from 'node:test';
import assert from 'node:assert/strict';
import path from 'path';
import os from 'os';
import fs from 'fs';
// Set a temp data dir before importing db module
const tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), 'codepilot-db-test-'));
process.env.CLAUDE_GUI_DATA_DIR = tmpDir;
// Use require to avoid top-level await issues with CJS output
/* eslint-disable @typescript-eslint/no-var-requires */
const { getDb, closeDb, createSession, getSession } = require('../../lib/db') as typeof import('../../lib/db');
describe('closeDb', () => {
afterEach(() => {
// Ensure DB is closed after each test
closeDb();
});
it('should close the database without errors', () => {
// First, ensure DB is open by calling getDb
const db = getDb();
assert.ok(db);
// Close should not throw
assert.doesNotThrow(() => closeDb());
});
it('should be idempotent (safe to call multiple times)', () => {
getDb(); // open
closeDb();
assert.doesNotThrow(() => closeDb()); // second call should be no-op
assert.doesNotThrow(() => closeDb()); // third call too
});
it('should allow re-opening the database after close', () => {
getDb(); // open
closeDb(); // close
// Re-open should work
const db = getDb();
assert.ok(db);
});
it('should persist data across close/reopen cycles', () => {
// Create a session
const session = createSession('Test Session', 'sonnet', '', tmpDir);
assert.ok(session);
assert.equal(session.title, 'Test Session');
const sessionId = session.id;
// Close and reopen
closeDb();
const db = getDb();
assert.ok(db);
// Session should still exist
const retrieved = getSession(sessionId);
assert.ok(retrieved);
assert.equal(retrieved!.title, 'Test Session');
});
it('should clean up WAL files after close', () => {
// Force some writes to create WAL
createSession('WAL Test 1');
createSession('WAL Test 2');
createSession('WAL Test 3');
const dbPath = path.join(tmpDir, 'codepilot.db');
assert.ok(fs.existsSync(dbPath));
// Close the database (should checkpoint WAL)
closeDb();
// After close, WAL file should either not exist or be empty
const walPath = dbPath + '-wal';
const walExists = fs.existsSync(walPath);
if (walExists) {
const walSize = fs.statSync(walPath).size;
assert.equal(walSize, 0, 'WAL file should be empty after graceful close');
}
});
it('cleanup test fixtures', () => {
closeDb();
fs.rmSync(tmpDir, { recursive: true, force: true });
});
});
+57
View File
@@ -483,3 +483,60 @@ export function deactivateAllProviders(): void {
const db = getDb();
db.prepare('UPDATE api_providers SET is_active = 0').run();
}
// ==========================================
// Graceful Shutdown
// ==========================================
/**
* Close the database connection gracefully.
* In WAL mode, this ensures the WAL is checkpointed and the
* -wal/-shm files are cleaned up properly.
*/
export function closeDb(): void {
if (db) {
try {
db.close();
console.log('[db] Database closed gracefully');
} catch (err) {
console.warn('[db] Error closing database:', err);
}
db = null;
}
}
// Register shutdown handlers to close the database when the process exits.
// This prevents WAL file accumulation and potential data loss.
function registerShutdownHandlers(): void {
let shuttingDown = false;
const shutdown = (signal: string) => {
if (shuttingDown) return;
shuttingDown = true;
console.log(`[db] Received ${signal}, closing database...`);
closeDb();
};
// 'exit' fires synchronously when the process is about to exit
process.on('exit', () => shutdown('exit'));
// Handle termination signals (Docker stop, systemd, Ctrl+C, etc.)
process.on('SIGTERM', () => {
shutdown('SIGTERM');
process.exit(0);
});
process.on('SIGINT', () => {
shutdown('SIGINT');
process.exit(0);
});
// Handle Windows-specific close events
if (process.platform === 'win32') {
process.on('SIGHUP', () => {
shutdown('SIGHUP');
process.exit(0);
});
}
}
registerShutdownHandlers();