fix: bound SQLite lock waits and strengthen PR 329 coverage

This commit is contained in:
Sycun committed 2026-10-16 12:22:42 +08:00
1 parent 9891b04215
commit 51d78fda91
4 files changed
+197 -58

No files matched your search

-10
View File
@@ -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 {
+93 -17
View File
@@ -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)
}
}
+5 -31
View File
@@ -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)
@@ -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")
}
})
}
}