From 9891b042155a0bb4ad94b4db7959fe054402e5ef Mon Sep 17 00:00:00 2001 From: Ordinary Date: Fri, 18 Sep 2026 16:02:12 +0800 Subject: [PATCH 1/2] fix: wait out SQLite busy on concurrent asset upserts Deep mode fires many create_asset calls at once. UpsertAssets used a DEFERRED transaction (SELECT then INSERT) on a 25-connection pool, so writers hit SQLITE_BUSY immediately and returned "database is locked" in a few milliseconds instead of waiting. Open SQLite with _txlock=immediate and a 30s busy timeout, shrink the pool, and retry locked upserts. Add a 32-way concurrent upsert test. --- internal/database/asset.go | 10 +++++++ internal/database/asset_test.go | 35 ++++++++++++++++++++++++ internal/database/database.go | 47 ++++++++++++++++++++++++++++----- 3 files changed, 85 insertions(+), 7 deletions(-) diff --git a/internal/database/asset.go b/internal/database/asset.go index 183bdd9f..b142a101 100644 --- a/internal/database/asset.go +++ b/internal/database/asset.go @@ -304,6 +304,16 @@ func appendAssetAccess(query string, args []interface{}, access RBACListAccess, } func (db *DB) UpsertAssets(assets []*Asset, ownerUserID string, allowGlobal ...bool) (AssetImportResult, error) { + var result AssetImportResult + err := retrySQLiteBusy(8, func() error { + var err error + result, err = db.upsertAssetsOnce(assets, ownerUserID, allowGlobal...) + return err + }) + return result, err +} + +func (db *DB) upsertAssetsOnce(assets []*Asset, ownerUserID string, allowGlobal ...bool) (AssetImportResult, error) { result := AssetImportResult{} tx, err := db.Begin() if err != nil { diff --git a/internal/database/asset_test.go b/internal/database/asset_test.go index 18cc637b..b75c9531 100644 --- a/internal/database/asset_test.go +++ b/internal/database/asset_test.go @@ -1,9 +1,11 @@ package database import ( + "fmt" "path/filepath" "strconv" "strings" + "sync" "testing" "time" @@ -86,6 +88,39 @@ func TestFofaAssetIgnoresInvalidOptionalStructuredFields(t *testing.T) { } } +func TestUpsertAssetsConcurrentDoesNotReturnLocked(t *testing.T) { + db, err := NewDB(filepath.Join(t.TempDir(), "assets-concurrent.db"), zap.NewNop()) + if err != nil { + t.Fatal(err) + } + defer db.Close() + + const n = 32 + errCh := make(chan error, n) + var wg sync.WaitGroup + for i := 0; i < n; i++ { + wg.Add(1) + i := i + go func() { + defer wg.Done() + asset := &Asset{Host: fmt.Sprintf("h%d.example.com", i), Port: 443, Protocol: "https", Title: fmt.Sprintf("t%d", i)} + _, err := db.UpsertAssets([]*Asset{asset}, "user-a") + errCh <- err + }() + } + wg.Wait() + close(errCh) + for err := range errCh { + if err != nil { + t.Fatalf("concurrent upsert: %v", err) + } + } + _, total, err := db.ListAssets(100, 0, AssetListFilter{}, RBACListAccess{Scope: RBACScopeAll}) + if err != nil || total != n { + t.Fatalf("want %d assets, got total=%d err=%v", n, total, err) + } +} + func TestAssetUpsertDeduplicatesAndUpdates(t *testing.T) { db, err := NewDB(filepath.Join(t.TempDir(), "assets.db"), zap.NewNop()) if err != nil { diff --git a/internal/database/database.go b/internal/database/database.go index 97aefafc..b0d3bcaa 100644 --- a/internal/database/database.go +++ b/internal/database/database.go @@ -14,9 +14,12 @@ import ( ) const ( - // SQLite 在 WAL 模式下建议使用较保守的连接数，降低长读快照导致 checkpoint 饥饿的概率。 - sqliteMaxOpenConns = 25 - sqliteMaxIdleConns = 5 + // WAL 允许并发读，但同一时间仍只有一个写者。连接池过大时，Deep 模式并行 + // create_asset 会各自 Begin+SELECT 再 INSERT，触发 SQLITE_BUSY（database is locked） + // 且 DEFERRED 事务升级写锁时可能立刻失败、不走 busy_timeout。 + sqliteMaxOpenConns = 8 + sqliteMaxIdleConns = 4 + sqliteBusyTimeoutMS = 30000 // 以页为单位的自动 checkpoint 触发阈值（默认 1000 页，约 4MB @ 4KB/page）。 sqliteWALAutoCheckpointPages = 1000 // 控制 WAL 目标上限，避免异常场景持续膨胀（256MB）。 @@ -25,9 +28,39 @@ const ( sqlitePassiveCheckpointInterval = 300 * time.Second ) -// configureDBPool 设置 SQLite 连接池参数，提升并发稳定性 +func sqliteOpenDSN(path string) string { + return fmt.Sprintf("%s?_journal_mode=WAL&_foreign_keys=1&_busy_timeout=%d&_synchronous=NORMAL&_txlock=immediate", path, sqliteBusyTimeoutMS) +} + +func isSQLiteBusy(err error) bool { + if err == nil { + return false + } + msg := strings.ToLower(err.Error()) + return strings.Contains(msg, "database is locked") || strings.Contains(msg, "sqlite_busy") +} + +func retrySQLiteBusy(attempts int, op func() error) error { + if attempts < 1 { + attempts = 1 + } + var err error + backoff := 5 * time.Millisecond + for i := 0; i < attempts; i++ { + err = op() + if err == nil || !isSQLiteBusy(err) { + return err + } + time.Sleep(backoff) + if backoff < 100*time.Millisecond { + backoff *= 2 + } + } + return err +} + +// configureDBPool 设置 SQLite 连接池：保留少量连接给 WAL 并发读，避免 25 路写连接互锁。 func configureDBPool(db *sql.DB) { - // SQLite 同一时间只允许一个写入者；过高连接数会放大锁竞争和 WAL 回收延迟。 db.SetMaxOpenConns(sqliteMaxOpenConns) db.SetMaxIdleConns(sqliteMaxIdleConns) db.SetConnMaxLifetime(30 * time.Minute) @@ -123,7 +156,7 @@ func (db *DB) runPassiveCheckpoint(trigger string) { // NewDB 创建数据库连接 func NewDB(dbPath string, logger *zap.Logger) (*DB, error) { - db, err := sql.Open("sqlite3", dbPath+"?_journal_mode=WAL&_foreign_keys=1&_busy_timeout=5000&_synchronous=NORMAL") + db, err := sql.Open("sqlite3", sqliteOpenDSN(dbPath)) if err != nil { return nil, fmt.Errorf("打开数据库失败: %w", err) } @@ -210,7 +243,7 @@ func (db *DB) initTables() error { // NewKnowledgeDB 创建知识库数据库连接（只包含知识库相关的表） func NewKnowledgeDB(dbPath string, logger *zap.Logger) (*DB, error) { - sqlDB, err := sql.Open("sqlite3", dbPath+"?_journal_mode=WAL&_foreign_keys=1&_busy_timeout=5000&_synchronous=NORMAL") + sqlDB, err := sql.Open("sqlite3", sqliteOpenDSN(dbPath)) if err != nil { return nil, fmt.Errorf("打开知识库数据库失败: %w", err) } From 51d78fda91a5690e4cdec43e464bc090ab128e79 Mon Sep 17 00:00:00 2001 From: Sycun <165354365+Sycun@users.noreply.github.com> Date: Fri, 16 Oct 2026 12:22:42 +0800 Subject: [PATCH 2/2] fix: bound SQLite lock waits and strengthen PR 329 coverage --- internal/database/asset.go | 10 -- internal/database/asset_test.go | 110 +++++++++++++++++---- internal/database/database.go | 36 +------ internal/database/transaction_lock_test.go | 99 +++++++++++++++++++ 4 files changed, 197 insertions(+), 58 deletions(-) create mode 100644 internal/database/transaction_lock_test.go diff --git a/internal/database/asset.go b/internal/database/asset.go index b142a101..183bdd9f 100644 --- a/internal/database/asset.go +++ b/internal/database/asset.go @@ -304,16 +304,6 @@ func appendAssetAccess(query string, args []interface{}, access RBACListAccess, } func (db *DB) UpsertAssets(assets []*Asset, ownerUserID string, allowGlobal ...bool) (AssetImportResult, error) { - var result AssetImportResult - err := retrySQLiteBusy(8, func() error { - var err error - result, err = db.upsertAssetsOnce(assets, ownerUserID, allowGlobal...) - return err - }) - return result, err -} - -func (db *DB) upsertAssetsOnce(assets []*Asset, ownerUserID string, allowGlobal ...bool) (AssetImportResult, error) { result := AssetImportResult{} tx, err := db.Begin() if err != nil { diff --git a/internal/database/asset_test.go b/internal/database/asset_test.go index b75c9531..9039483c 100644 --- a/internal/database/asset_test.go +++ b/internal/database/asset_test.go @@ -88,36 +88,112 @@ func TestFofaAssetIgnoresInvalidOptionalStructuredFields(t *testing.T) { } } -func TestUpsertAssetsConcurrentDoesNotReturnLocked(t *testing.T) { - db, err := NewDB(filepath.Join(t.TempDir(), "assets-concurrent.db"), zap.NewNop()) +// Start more workers than the pool can hold, with separate Asset inputs just as +// parallel create_asset calls use. Check both returned counts and persisted IDs. +func TestUpsertAssetsConcurrent(t *testing.T) { + for _, sameTarget := range []bool{false, true} { + t.Run(fmt.Sprintf("same_target=%t", sameTarget), func(t *testing.T) { + db, err := NewDB(filepath.Join(t.TempDir(), "assets.db"), zap.NewNop()) + if err != nil { + t.Fatal(err) + } + defer db.Close() + + const workers = 32 + start := make(chan struct{}) + var wg sync.WaitGroup + assets := make([]*Asset, workers) + results := make([]AssetImportResult, workers) + errs := make([]error, workers) + for i := range workers { + host := fmt.Sprintf("host-%d.example.com", i) + if sameTarget { + host = "shared.example.com" + } + assets[i] = &Asset{Host: host, Port: 443, Protocol: "https", Title: fmt.Sprintf("title-%d", i)} + wg.Add(1) + go func() { + defer wg.Done() + <-start + results[i], errs[i] = db.UpsertAssets([]*Asset{assets[i]}, "user-a") + }() + } + close(start) + wg.Wait() + + created, updated := 0, 0 + for i, err := range errs { + if err != nil { + t.Errorf("worker %d: %v", i, err) + continue + } + created += results[i].Created + updated += results[i].Updated + saved, err := db.GetAsset(assets[i].ID, RBACListAccess{Scope: RBACScopeAll}) + if err != nil || saved.Host != assets[i].Host { + t.Errorf("worker %d asset was not persisted: %v", i, err) + } + if sameTarget && assets[i].ID != assets[0].ID { + t.Errorf("worker %d returned a different ID for the same target", i) + } + } + wantCreated, wantUpdated := workers, 0 + if sameTarget { + wantCreated, wantUpdated = 1, workers-1 + } + if created != wantCreated || updated != wantUpdated { + t.Errorf("created/updated = %d/%d, want %d/%d", created, updated, wantCreated, wantUpdated) + } + var count int + if err := db.QueryRow("SELECT COUNT(*) FROM assets").Scan(&count); err != nil || count != wantCreated { + t.Fatalf("persisted count = %d, want %d: %v", count, wantCreated, err) + } + }) + } +} + +func TestAssetConcurrentCreateAndUpdate(t *testing.T) { + db, err := NewDB(filepath.Join(t.TempDir(), "assets.db"), zap.NewNop()) if err != nil { t.Fatal(err) } defer db.Close() - const n = 32 - errCh := make(chan error, n) + const workers = 32 + assets := make([]*Asset, workers) + for i := range workers { + assets[i] = &Asset{Host: fmt.Sprintf("existing-%d.example.com", i)} + } + if _, err := db.UpsertAssets(assets, "user-a"); err != nil { + t.Fatal(err) + } + start := make(chan struct{}) var wg sync.WaitGroup - for i := 0; i < n; i++ { + for i := range workers { wg.Add(1) - i := i go func() { defer wg.Done() - asset := &Asset{Host: fmt.Sprintf("h%d.example.com", i), Port: 443, Protocol: "https", Title: fmt.Sprintf("t%d", i)} - _, err := db.UpsertAssets([]*Asset{asset}, "user-a") - errCh <- err + <-start + if i%2 == 0 { + assets[i].Title = "updated" + if err := db.UpdateAsset(assets[i].ID, assets[i], RBACListAccess{Scope: RBACScopeAll}); err != nil { + t.Errorf("update %d: %v", i, err) + } + } else { + if _, err := db.UpsertAssets([]*Asset{{Host: fmt.Sprintf("new-%d.example.com", i)}}, "user-a"); err != nil { + t.Errorf("create %d: %v", i, err) + } + } }() } + close(start) wg.Wait() - close(errCh) - for err := range errCh { - if err != nil { - t.Fatalf("concurrent upsert: %v", err) - } + var total, updated int + if err := db.QueryRow("SELECT COUNT(*), COUNT(CASE WHEN title='updated' THEN 1 END) FROM assets").Scan(&total, &updated); err != nil { + t.Fatal(err) } - _, total, err := db.ListAssets(100, 0, AssetListFilter{}, RBACListAccess{Scope: RBACScopeAll}) - if err != nil || total != n { - t.Fatalf("want %d assets, got total=%d err=%v", n, total, err) + if total != workers+workers/2 || updated != workers/2 { + t.Fatalf("persisted total/updated = %d/%d, want %d/%d", total, updated, workers+workers/2, workers/2) } } diff --git a/internal/database/database.go b/internal/database/database.go index b0d3bcaa..02fb807a 100644 --- a/internal/database/database.go +++ b/internal/database/database.go @@ -14,9 +14,7 @@ import ( ) const ( - // WAL 允许并发读，但同一时间仍只有一个写者。连接池过大时，Deep 模式并行 - // create_asset 会各自 Begin+SELECT 再 INSERT，触发 SQLITE_BUSY（database is locked） - // 且 DEFERRED 事务升级写锁时可能立刻失败、不走 busy_timeout。 + // WAL permits concurrent readers but only one writer; keep the pool modest. sqliteMaxOpenConns = 8 sqliteMaxIdleConns = 4 sqliteBusyTimeoutMS = 30000 @@ -29,37 +27,13 @@ const ( ) func sqliteOpenDSN(path string) string { + // Reserve the writer at Begin, before read-modify-write transactions take a + // snapshot. DEFERRED upgrades may fail immediately without busy_timeout. + // Let SQLite wait once; retrying after each timeout multiplies lock latency. return fmt.Sprintf("%s?_journal_mode=WAL&_foreign_keys=1&_busy_timeout=%d&_synchronous=NORMAL&_txlock=immediate", path, sqliteBusyTimeoutMS) } -func isSQLiteBusy(err error) bool { - if err == nil { - return false - } - msg := strings.ToLower(err.Error()) - return strings.Contains(msg, "database is locked") || strings.Contains(msg, "sqlite_busy") -} - -func retrySQLiteBusy(attempts int, op func() error) error { - if attempts < 1 { - attempts = 1 - } - var err error - backoff := 5 * time.Millisecond - for i := 0; i < attempts; i++ { - err = op() - if err == nil || !isSQLiteBusy(err) { - return err - } - time.Sleep(backoff) - if backoff < 100*time.Millisecond { - backoff *= 2 - } - } - return err -} - -// configureDBPool 设置 SQLite 连接池：保留少量连接给 WAL 并发读，避免 25 路写连接互锁。 +// configureDBPool 设置 SQLite 连接池，保留少量连接给 WAL 并发读。 func configureDBPool(db *sql.DB) { db.SetMaxOpenConns(sqliteMaxOpenConns) db.SetMaxIdleConns(sqliteMaxIdleConns) diff --git a/internal/database/transaction_lock_test.go b/internal/database/transaction_lock_test.go new file mode 100644 index 00000000..c60c2033 --- /dev/null +++ b/internal/database/transaction_lock_test.go @@ -0,0 +1,99 @@ +package database + +import ( + "database/sql" + "errors" + "path/filepath" + "testing" + "time" + + "github.com/mattn/go-sqlite3" + "go.uber.org/zap" +) + +func TestSQLiteBeginReservesWriterWithoutBlockingReaders(t *testing.T) { + for name, open := range map[string]func(string, *zap.Logger) (*DB, error){ + "conversations": NewDB, + "knowledge": NewKnowledgeDB, + } { + t.Run(name, func(t *testing.T) { + path := filepath.Join(t.TempDir(), "transactions.db") + db, err := open(path, zap.NewNop()) + if err != nil { + t.Fatal(err) + } + defer db.Close() + // A separate connection with no busy wait makes the lock assertion + // deterministic, without relying on goroutine scheduling or sleeps. + other, err := sql.Open("sqlite3", path+"?_busy_timeout=0") + if err != nil { + t.Fatal(err) + } + defer other.Close() + other.SetMaxOpenConns(1) + if err := other.Ping(); err != nil { + t.Fatal(err) + } + + tx, err := db.Begin() + if err != nil { + t.Fatal(err) + } + defer tx.Rollback() + // Begin must already own the writer reservation, even before SELECT. + _, err = other.Exec("BEGIN IMMEDIATE") + if err == nil { + _, _ = other.Exec("ROLLBACK") + t.Fatal("Begin did not reserve the writer lock") + } + var sqliteErr sqlite3.Error + if !errors.As(err, &sqliteErr) || sqliteErr.Code != sqlite3.ErrBusy { + t.Fatalf("competing writer: want SQLITE_BUSY, got %v", err) + } + var tables int + if err := other.QueryRow("SELECT COUNT(*) FROM sqlite_master").Scan(&tables); err != nil { + t.Fatalf("WAL reader blocked by writer reservation: %v", err) + } + if err := tx.Rollback(); err != nil { + t.Fatal(err) + } + if _, err := other.Exec("BEGIN IMMEDIATE"); err != nil { + t.Fatalf("writer reservation was not released: %v", err) + } + if _, err := other.Exec("ROLLBACK"); err != nil { + t.Fatal(err) + } + + // A temporarily busy writer should make Begin wait, then succeed + // when that writer releases its reservation. + if _, err := other.Exec("BEGIN IMMEDIATE"); err != nil { + t.Fatal(err) + } + defer other.Exec("ROLLBACK") + done := make(chan error, 1) + go func() { + pending, err := db.Begin() + if err == nil { + err = pending.Rollback() + } + done <- err + }() + select { + case err := <-done: + t.Fatalf("Begin returned before the competing writer released its lock: %v", err) + case <-time.After(50 * time.Millisecond): + } + if _, err := other.Exec("ROLLBACK"); err != nil { + t.Fatal(err) + } + select { + case err := <-done: + if err != nil { + t.Fatalf("Begin did not recover after writer released its lock: %v", err) + } + case <-time.After(5 * time.Second): + t.Fatal("Begin did not resume after writer released its lock") + } + }) + } +}