mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-10-19 22:17:42 +02:00
fix: wait out SQLite busy on concurrent asset upserts
Deep mode fires many create_asset calls at once. UpsertAssets used a DEFERRED transaction (SELECT then INSERT) on a 25-connection pool, so writers hit SQLITE_BUSY immediately and returned "database is locked" in a few milliseconds instead of waiting. Open SQLite with _txlock=immediate and a 30s busy timeout, shrink the pool, and retry locked upserts. Add a 32-way concurrent upsert test.
This commit is contained in:
1 parent
3c0d57fe91
commit
9891b04215
3 files changed
+85
-7
No files matched your search
@@ -304,6 +304,16 @@ 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 {
|
||||
|
||||
@@ -1,9 +1,11 @@
|
||||
package database
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
@@ -86,6 +88,39 @@ func TestFofaAssetIgnoresInvalidOptionalStructuredFields(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestUpsertAssetsConcurrentDoesNotReturnLocked(t *testing.T) {
|
||||
db, err := NewDB(filepath.Join(t.TempDir(), "assets-concurrent.db"), zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
|
||||
const n = 32
|
||||
errCh := make(chan error, n)
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < n; i++ {
|
||||
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
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
close(errCh)
|
||||
for err := range errCh {
|
||||
if err != nil {
|
||||
t.Fatalf("concurrent upsert: %v", 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)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAssetUpsertDeduplicatesAndUpdates(t *testing.T) {
|
||||
db, err := NewDB(filepath.Join(t.TempDir(), "assets.db"), zap.NewNop())
|
||||
if err != nil {
|
||||
|
||||
@@ -14,9 +14,12 @@ import (
|
||||
)
|
||||
|
||||
const (
|
||||
// SQLite 在 WAL 模式下建议使用较保守的连接数,降低长读快照导致 checkpoint 饥饿的概率。
|
||||
sqliteMaxOpenConns = 25
|
||||
sqliteMaxIdleConns = 5
|
||||
// WAL 允许并发读,但同一时间仍只有一个写者。连接池过大时,Deep 模式并行
|
||||
// create_asset 会各自 Begin+SELECT 再 INSERT,触发 SQLITE_BUSY(database is locked)
|
||||
// 且 DEFERRED 事务升级写锁时可能立刻失败、不走 busy_timeout。
|
||||
sqliteMaxOpenConns = 8
|
||||
sqliteMaxIdleConns = 4
|
||||
sqliteBusyTimeoutMS = 30000
|
||||
// 以页为单位的自动 checkpoint 触发阈值(默认 1000 页,约 4MB @ 4KB/page)。
|
||||
sqliteWALAutoCheckpointPages = 1000
|
||||
// 控制 WAL 目标上限,避免异常场景持续膨胀(256MB)。
|
||||
@@ -25,9 +28,39 @@ const (
|
||||
sqlitePassiveCheckpointInterval = 300 * time.Second
|
||||
)
|
||||
|
||||
// configureDBPool 设置 SQLite 连接池参数,提升并发稳定性
|
||||
func sqliteOpenDSN(path string) string {
|
||||
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 路写连接互锁。
|
||||
func configureDBPool(db *sql.DB) {
|
||||
// SQLite 同一时间只允许一个写入者;过高连接数会放大锁竞争和 WAL 回收延迟。
|
||||
db.SetMaxOpenConns(sqliteMaxOpenConns)
|
||||
db.SetMaxIdleConns(sqliteMaxIdleConns)
|
||||
db.SetConnMaxLifetime(30 * time.Minute)
|
||||
@@ -123,7 +156,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 +243,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)
|
||||
}
|
||||
|
||||
Reference in new issue
Block a user