Merge pull request #362 from Sycun/codex/fix-328-sqlite-asset-lock

fix: serialize SQLite asset writes without multiplying busy waits

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 {

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

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