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