diff --git a/internal/database/asset_test.go b/internal/database/asset_test.go index b9418d9b8..4f40b8980 100644 --- a/internal/database/asset_test.go +++ b/internal/database/asset_test.go @@ -4,6 +4,7 @@ import ( "path/filepath" "strconv" "strings" + "sync" "testing" "time" @@ -447,3 +448,37 @@ func TestAssetListFlexibleFiltersAndOldestScanPagination(t *testing.T) { t.Fatalf("structured filters: total=%d assets=%#v err=%v", total, filtered, err) } } + +func TestUpsertAssetsConcurrentCreatesDoNotLock(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 goroutines = 24 + errs := make([]error, goroutines) + var wg sync.WaitGroup + for i := 0; i < goroutines; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + asset := &Asset{Host: "https://198.51.100." + strconv.Itoa(i+1) + ":443", Source: "concurrent"} + _, errs[i] = db.UpsertAssets([]*Asset{asset}, "") + }(i) + } + wg.Wait() + + for i, err := range errs { + if err != nil { + t.Fatalf("concurrent upsert %d failed: %v", i, err) + } + } + _, total, err := db.ListAssets(100, 0, AssetListFilter{Source: "concurrent"}, RBACListAccess{Scope: RBACScopeAll}) + if err != nil { + t.Fatalf("list assets: %v", err) + } + if total != goroutines { + t.Fatalf("total=%d, want %d", total, goroutines) + } +} diff --git a/internal/database/database.go b/internal/database/database.go index 35884987d..1bc519197 100644 --- a/internal/database/database.go +++ b/internal/database/database.go @@ -1,6 +1,7 @@ package database import ( + "context" "database/sql" "fmt" "os" @@ -47,6 +48,7 @@ func configureSQLitePragmas(db *sql.DB) error { // DB 数据库连接 type DB struct { *sql.DB + writeMu sync.Mutex // 串行化写事务,配合 _txlock=immediate 消除 SQLITE_BUSY logger *zap.Logger conversationArtifactsDir string einoPlantaskBaseDir string // skills_dir + plantask_rel_dir (per-conversation subdirs) @@ -61,6 +63,52 @@ type DB struct { vulnerabilityCreatedHook func(*Vulnerability) } +// serialTx 包装 *sql.Tx:事务结束(Commit/Rollback)时释放 DB.writeMu。 +// 调用点无需感知,仍可把返回值当作 *sql.Tx 使用。 +type serialTx struct { + *sql.Tx + release func() + once sync.Once +} + +func (t *serialTx) Commit() error { + err := t.Tx.Commit() + t.releaseOnce() + return err +} + +func (t *serialTx) Rollback() error { + err := t.Tx.Rollback() + t.releaseOnce() + return err +} + +func (t *serialTx) releaseOnce() { + if t.release != nil { + t.once.Do(t.release) + } +} + +// Begin 启动写事务。写事务在进程内严格串行:获取 writeMu 后才执行 +// BEGIN IMMEDIATE(由 DSN _txlock=immediate 保证),提交/回滚后释放。 +func (db *DB) Begin() (*serialTx, error) { + return db.BeginTx(context.Background(), nil) +} + +// BeginTx 与 Begin 相同,支持显式事务选项。 +func (db *DB) BeginTx(ctx context.Context, opts *sql.TxOptions) (*serialTx, error) { + if db == nil || db.DB == nil { + return nil, fmt.Errorf("数据库未初始化") + } + db.writeMu.Lock() + tx, err := db.DB.BeginTx(ctx, opts) + if err != nil { + db.writeMu.Unlock() + return nil, err + } + return &serialTx{Tx: tx, release: db.writeMu.Unlock}, nil +} + // startPassiveCheckpointLoop 启动后台 PASSIVE checkpoint 循环。 func (db *DB) startPassiveCheckpointLoop(name string) { if sqlitePassiveCheckpointInterval <= 0 || db == nil || db.DB == nil { @@ -122,7 +170,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", dbPath+"?_journal_mode=WAL&_foreign_keys=1&_busy_timeout=5000&_synchronous=NORMAL&_txlock=immediate") if err != nil { return nil, fmt.Errorf("打开数据库失败: %w", err) } @@ -1671,7 +1719,7 @@ func (db *DB) migrateC2ListenersTable() 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", dbPath+"?_journal_mode=WAL&_foreign_keys=1&_busy_timeout=5000&_synchronous=NORMAL&_txlock=immediate") if err != nil { return nil, fmt.Errorf("打开知识库数据库失败: %w", err) } diff --git a/internal/database/rbac.go b/internal/database/rbac.go index 28e865022..3c4cba11b 100644 --- a/internal/database/rbac.go +++ b/internal/database/rbac.go @@ -349,7 +349,7 @@ func (db *DB) BootstrapRBAC(adminPasswordHash string, permissions map[string]str return tx.Commit() } -func grantSystemRolePermissions(tx *sql.Tx, permissions map[string]string) error { +func grantSystemRolePermissions(tx *serialTx, permissions map[string]string) error { now := time.Now() // System roles are immutable and owned by the application. Rebuild their // grants deterministically so policy tightening also removes permissions diff --git a/internal/database/vulnerability.go b/internal/database/vulnerability.go index 2eede0fa4..e2f5072d4 100644 --- a/internal/database/vulnerability.go +++ b/internal/database/vulnerability.go @@ -392,7 +392,7 @@ func (db *DB) DeleteVulnerability(id string) error { return nil } -func collectVulnerabilityConversationIDs(tx *sql.Tx, where string, args []interface{}) ([]string, error) { +func collectVulnerabilityConversationIDs(tx *serialTx, where string, args []interface{}) ([]string, error) { rows, err := tx.Query(`SELECT DISTINCT COALESCE(conversation_id,'') FROM vulnerabilities `+where, args...) if err != nil { return nil, fmt.Errorf("查询受影响漏洞会话失败: %w", err)