git.iko.soy/crush
on main help sign in
fix(db): use connection pool to avoid corrupted writes
61ee2d2Kieran Klukas
internal/app/app.go+4−2
compare with main
unifiedsplit
@@ -113,10 +113,12 @@ func New(ctx context.Context, conn *sql.DB, store *config.ConfigStore) (*App, er
113
113
114 go mcp.Initialize(ctx, app.Permissions, store)
114 go mcp.Initialize(ctx, app.Permissions, store)
115
115
116− // cleanup database upon app shutdown
116+ // Release the shared database connection on shutdown. The pool
117+ // closes the underlying *sql.DB when the last reference is released.
118+ dataDir := cfg.Options.DataDirectory
117 app.cleanupFuncs = append(
119 app.cleanupFuncs = append(
118 app.cleanupFuncs,
120 app.cleanupFuncs,
119− func(context.Context) error { return conn.Close() },
121+ func(context.Context) error { return db.Release(dataDir) },
120 func(ctx context.Context) error { return mcp.Close(ctx) },
122 func(ctx context.Context) error { return mcp.Close(ctx) },
121 )
123 )
122
124
internal/db/connect.go+84−12
@@ -38,41 +38,113 @@ func init() {
38 }
38 }
39}
39}
40
40
41+// connEntry holds a shared database connection and its reference count.
42+type connEntry struct {
43+ db *sql.DB
44+ refCount int
45+}
46+
47+var (
48+ pool = make(map[string]*connEntry)
49+ poolMu sync.Mutex
50+)
51+
41−// Connect opens a SQLite database connection and runs migrations.
52+// Connect opens a SQLite database connection for the given data
53+// directory and runs migrations. If a connection to the same database
54+// file already exists, the existing connection is returned with its
55+// reference count incremented. Callers must pair each Connect with a
56+// [Release] when they no longer need the connection.
42func Connect(ctx context.Context, dataDir string) (*sql.DB, error) {
57func Connect(ctx context.Context, dataDir string) (*sql.DB, error) {
43 if dataDir == "" {
58 if dataDir == "" {
44 return nil, fmt.Errorf("data.dir is not set")
59 return nil, fmt.Errorf("data.dir is not set")
45 }
60 }
61+
46 dbPath := filepath.Join(dataDir, "crush.db")
62 dbPath := filepath.Join(dataDir, "crush.db")
47
63
64+ // Resolve to an absolute path so that different relative paths to
65+ // the same file share a single connection.
66+ absPath, err := filepath.Abs(dbPath)
67+ if err != nil {
68+ absPath = dbPath
69+ }
70+
71+ poolMu.Lock()
72+ defer poolMu.Unlock()
73+
74+ if entry, ok := pool[absPath]; ok {
75+ entry.refCount++
76+ return entry.db, nil
77+ }
78+
48− db, err := openDB(dbPath)
79+ conn, err := openDB(dbPath)
49 if err != nil {
80 if err != nil {
50 return nil, err
81 return nil, err
51 }
82 }
52
83
53− // Serialize all access through a single connection. SQLite serializes
84+ // Serialize all access through a single connection. SQLite
54− // writes at the file level anyway, and allowing multiple pool
85+ // serializes writes at the file level anyway, and allowing multiple
55− // connections to interleave writes/checkpoints (especially under
86+ // pool connections to interleave writes/checkpoints (especially
56− // concurrent sub-agents) has caused WAL/header desync resulting in
87+ // under concurrent sub-agents) has caused WAL/header desync
57− // SQLITE_NOTADB (26) on the next open.
88+ // resulting in SQLITE_NOTADB (26) on the next open.
58− db.SetMaxOpenConns(1)
89+ conn.SetMaxOpenConns(1)
59
90
60− if err = db.PingContext(ctx); err != nil {
91+ if err = conn.PingContext(ctx); err != nil {
61− db.Close()
92+ conn.Close()
62 return nil, fmt.Errorf("failed to connect to database: %w", err)
93 return nil, fmt.Errorf("failed to connect to database: %w", err)
63 }
94 }
64
95
65 if err := initGoose(); err != nil {
96 if err := initGoose(); err != nil {
97+ conn.Close()
66 slog.Error("Failed to initialize goose", "error", err)
98 slog.Error("Failed to initialize goose", "error", err)
67 return nil, fmt.Errorf("failed to initialize goose: %w", err)
99 return nil, fmt.Errorf("failed to initialize goose: %w", err)
68 }
100 }
69
101
70− if err := goose.Up(db, "migrations"); err != nil {
102+ if err := goose.Up(conn, "migrations"); err != nil {
103+ conn.Close()
71 slog.Error("Failed to apply migrations", "error", err)
104 slog.Error("Failed to apply migrations", "error", err)
72 return nil, fmt.Errorf("failed to apply migrations: %w", err)
105 return nil, fmt.Errorf("failed to apply migrations: %w", err)
73 }
106 }
74
107
108+ pool[absPath] = &connEntry{db: conn, refCount: 1}
75− return db, nil
109+ return conn, nil
110+}
111+
112+// Release decrements the reference count for the database at the given
113+// data directory. When the count reaches zero the underlying connection
114+// is closed and removed from the pool.
115+func Release(dataDir string) error {
116+ dbPath := filepath.Join(dataDir, "crush.db")
117+ absPath, err := filepath.Abs(dbPath)
118+ if err != nil {
119+ absPath = dbPath
120+ }
121+
122+ poolMu.Lock()
123+ defer poolMu.Unlock()
124+
125+ entry, ok := pool[absPath]
126+ if !ok {
127+ return nil
128+ }
129+
130+ entry.refCount--
131+ if entry.refCount > 0 {
132+ return nil
133+ }
134+
135+ delete(pool, absPath)
136+ return entry.db.Close()
137+}
138+
139+// ResetPool closes all pooled connections and clears the pool. This is
140+// intended for use in tests to ensure a clean state between test cases.
141+func ResetPool() {
142+ poolMu.Lock()
143+ defer poolMu.Unlock()
144+ for path, entry := range pool {
145+ entry.db.Close()
146+ delete(pool, path)
147+ }
76}
148}
77
149
78func initGoose() error {
150func initGoose() error {
internal/db/connect_test.gonew file+54
@@ -0,0 +1,54 @@
1+package db
2+
3+import (
4+ "context"
5+ "testing"
6+
7+ "github.com/stretchr/testify/require"
8+)
9+
10+func TestConnect_SharesConnectionForSameDataDir(t *testing.T) {
11+ t.Cleanup(ResetPool)
12+
13+ dataDir := t.TempDir()
14+
15+ conn1, err := Connect(context.Background(), dataDir)
16+ require.NoError(t, err)
17+
18+ conn2, err := Connect(context.Background(), dataDir)
19+ require.NoError(t, err)
20+
21+ require.Same(t, conn1, conn2, "should return the same *sql.DB for the same data dir")
22+
23+ // Releasing once should not close the connection.
24+ require.NoError(t, Release(dataDir))
25+ require.NoError(t, conn1.PingContext(context.Background()), "connection should still be usable after partial release")
26+
27+ // Releasing again should close it.
28+ require.NoError(t, Release(dataDir))
29+ require.Error(t, conn1.PingContext(context.Background()), "connection should be closed after final release")
30+}
31+
32+func TestConnect_SeparateConnectionsForDifferentDataDirs(t *testing.T) {
33+ t.Cleanup(ResetPool)
34+
35+ dir1 := t.TempDir()
36+ dir2 := t.TempDir()
37+
38+ conn1, err := Connect(context.Background(), dir1)
39+ require.NoError(t, err)
40+
41+ conn2, err := Connect(context.Background(), dir2)
42+ require.NoError(t, err)
43+
44+ require.NotSame(t, conn1, conn2, "different data dirs should get different connections")
45+
46+ require.NoError(t, Release(dir1))
47+ require.NoError(t, Release(dir2))
48+}
49+
50+func TestRelease_NoopForUnknownDataDir(t *testing.T) {
51+ t.Cleanup(ResetPool)
52+
53+ require.NoError(t, Release("/nonexistent/path"), "releasing unknown data dir should not error")
54+}