diff --git a/internal/database/asset_test.go b/internal/database/asset_test.go index 18cc637b..9039483c 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,115 @@ func TestFofaAssetIgnoresInvalidOptionalStructuredFields(t *testing.T) { } } +// 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 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 := range workers { + wg.Add(1) + go func() { + defer wg.Done() + <-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() + 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) + } + 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) + } +} + 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..02fb807a 100644 --- a/internal/database/database.go +++ b/internal/database/database.go @@ -14,9 +14,10 @@ import ( ) const ( - // SQLite 在 WAL 模式下建议使用较保守的连接数,降低长读快照导致 checkpoint 饥饿的概率。 - sqliteMaxOpenConns = 25 - sqliteMaxIdleConns = 5 + // WAL permits concurrent readers but only one writer; keep the pool modest. + sqliteMaxOpenConns = 8 + sqliteMaxIdleConns = 4 + sqliteBusyTimeoutMS = 30000 // 以页为单位的自动 checkpoint 触发阈值(默认 1000 页,约 4MB @ 4KB/page)。 sqliteWALAutoCheckpointPages = 1000 // 控制 WAL 目标上限,避免异常场景持续膨胀(256MB)。 @@ -25,9 +26,15 @@ const ( sqlitePassiveCheckpointInterval = 300 * time.Second ) -// configureDBPool 设置 SQLite 连接池参数,提升并发稳定性 +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) +} + +// configureDBPool 设置 SQLite 连接池,保留少量连接给 WAL 并发读。 func configureDBPool(db *sql.DB) { - // SQLite 同一时间只允许一个写入者;过高连接数会放大锁竞争和 WAL 回收延迟。 db.SetMaxOpenConns(sqliteMaxOpenConns) db.SetMaxIdleConns(sqliteMaxIdleConns) db.SetConnMaxLifetime(30 * time.Minute) @@ -123,7 +130,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 +217,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) } 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") + } + }) + } +}