Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
111 changes: 111 additions & 0 deletions internal/database/asset_test.go
Original file line number Diff line number Diff line change
@@ -1,9 +1,11 @@
package database

import (
"fmt"
"path/filepath"
"strconv"
"strings"
"sync"
"testing"
"time"

Expand Down Expand Up @@ -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 {
Expand Down
21 changes: 14 additions & 7 deletions internal/database/database.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)。
Expand All @@ -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)
Expand Down Expand Up @@ -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)
}
Expand Down Expand Up @@ -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)
}
Expand Down
99 changes: 99 additions & 0 deletions internal/database/transaction_lock_test.go
Original file line number Diff line number Diff line change
@@ -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")
}
})
}
}
Loading