From f4f955d721cd5a45e543b88e2fae672fddcc02c5 Mon Sep 17 00:00:00 2001
From: engigu <sayheya@qq.com>
Date: Fri, 6 Mar 2026 15:59:03 +0800
Subject: [PATCH] fix: sqlite atabase is locked #20

---
 internal/database/database.go               | 13 ++-
 internal/services/tasks/executor_service.go | 97 ++++++++++++---------
 2 files changed, 69 insertions(+), 41 deletions(-)

diff --git a/internal/database/database.go b/internal/database/database.go
index 54bafc2..3490467 100644
--- a/internal/database/database.go
+++ b/internal/database/database.go
@@ -36,7 +36,7 @@ func Init(cfg *Config) error {
 
 	switch cfg.Type {
 	case "sqlite":
-		dialector = sqlite.Open(cfg.Path)
+		dialector = sqlite.Open(cfg.Path + "?_busy_timeout=5000")
 	case "mysql":
 		dsn := fmt.Sprintf("%s:%s@tcp(%s:%d)/%s?charset=utf8mb4&parseTime=True&loc=Asia%%2FShanghai",
 			cfg.User, cfg.Password, cfg.Host, cfg.Port, cfg.DBName)
@@ -60,6 +60,17 @@ func Init(cfg *Config) error {
 	}
 
 	logger.Infof("[Database] 已连接 %s 数据库 (时区: Asia/Shanghai)", cfg.Type)
+
+	// SQLite 特殊优化：开启 WAL 模式，提升并发性能
+	if cfg.Type == "sqlite" {
+		sqlDB, _ := DB.DB()
+		if sqlDB != nil {
+			sqlDB.SetMaxOpenConns(1) // SQLite 只允许单写连接
+			sqlDB.Exec("PRAGMA journal_mode=WAL")
+			sqlDB.Exec("PRAGMA synchronous=NORMAL")
+		}
+	}
+
 	return nil
 }
 
diff --git a/internal/services/tasks/executor_service.go b/internal/services/tasks/executor_service.go
index d39d4cb..cabd05e 100644
--- a/internal/services/tasks/executor_service.go
+++ b/internal/services/tasks/executor_service.go
@@ -19,7 +19,6 @@ import (
 	"github.com/engigu/baihu-panel/internal/utils"
 
 	"gorm.io/gorm"
-	"gorm.io/gorm/clause"
 )
 
 // AgentWSManager 接口定义（避免循环依赖）
@@ -768,54 +767,72 @@ func (es *ExecutorService) CheckConcurrency(taskID string) error {
 // AddRunningGo 添加当前 goroutine ID 到任务的 running_go 字段
 func (es *ExecutorService) AddRunningGo(taskID string) (int64, error) {
 	goid := utils.GetGoroutineID()
-	err := database.DB.Transaction(func(tx *gorm.DB) error {
-		var task models.Task
-		if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id = ?", taskID).First(&task).Error; err != nil {
-			return err
-		}
-		var goids []int64
-		if task.RunningGo != "" {
-			_ = json.Unmarshal([]byte(task.RunningGo), &goids)
-		}
+	var lastErr error
+	for attempt := 0; attempt < 3; attempt++ {
+		lastErr = database.DB.Transaction(func(tx *gorm.DB) error {
+			var task models.Task
+			if err := tx.Where("id = ?", taskID).First(&task).Error; err != nil {
+				return err
+			}
+			var goids []int64
+			if task.RunningGo != "" {
+				_ = json.Unmarshal([]byte(task.RunningGo), &goids)
+			}
 
-		// 解析配置以获取并发设置
-		var config models.TaskConfig
-		if task.Config != "" {
-			_ = json.Unmarshal([]byte(task.Config), &config)
-		}
+			// 解析配置以获取并发设置
+			var config models.TaskConfig
+			if task.Config != "" {
+				_ = json.Unmarshal([]byte(task.Config), &config)
+			}
 
-		// 如果并发为0(禁用)且已有执行中的任务，返回错误
-		if config.Concurrency == 0 && len(goids) > 0 {
-			return fmt.Errorf("task is running")
-		}
+			// 如果并发为0(禁用)且已有执行中的任务，返回错误
+			if config.Concurrency == 0 && len(goids) > 0 {
+				return fmt.Errorf("task is running")
+			}
 
-		goids = append(goids, goid)
-		data, _ := json.Marshal(goids)
-		return tx.Model(&task).Update("running_go", string(data)).Error
-	})
-	return goid, err
+			goids = append(goids, goid)
+			data, _ := json.Marshal(goids)
+			return tx.Model(&task).Update("running_go", string(data)).Error
+		})
+		if lastErr == nil {
+			return goid, nil
+		}
+		// 如果是业务错误（任务正在运行），不重试
+		if lastErr.Error() == "task is running" {
+			return goid, lastErr
+		}
+		// 数据库锁错误，等待后重试
+		time.Sleep(100 * time.Millisecond)
+	}
+	return goid, fmt.Errorf("任务并发限制: %v", lastErr)
 }
 
 // RemoveRunningGo 从任务的 running_go 字段移除指定 goroutine ID
 func (es *ExecutorService) RemoveRunningGo(taskID string, goid int64) {
-	database.DB.Transaction(func(tx *gorm.DB) error {
-		var task models.Task
-		if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id = ?", taskID).First(&task).Error; err != nil {
-			return err
-		}
-		var goids []int64
-		if task.RunningGo != "" {
-			_ = json.Unmarshal([]byte(task.RunningGo), &goids)
-		}
-		newGoids := make([]int64, 0)
-		for _, id := range goids {
-			if id != goid {
-				newGoids = append(newGoids, id)
+	for attempt := 0; attempt < 3; attempt++ {
+		err := database.DB.Transaction(func(tx *gorm.DB) error {
+			var task models.Task
+			if err := tx.Where("id = ?", taskID).First(&task).Error; err != nil {
+				return err
 			}
+			var goids []int64
+			if task.RunningGo != "" {
+				_ = json.Unmarshal([]byte(task.RunningGo), &goids)
+			}
+			newGoids := make([]int64, 0)
+			for _, id := range goids {
+				if id != goid {
+					newGoids = append(newGoids, id)
+				}
+			}
+			data, _ := json.Marshal(newGoids)
+			return tx.Model(&task).Update("running_go", string(data)).Error
+		})
+		if err == nil {
+			return
 		}
-		data, _ := json.Marshal(newGoids)
-		return tx.Model(&task).Update("running_go", string(data)).Error
-	})
+		time.Sleep(100 * time.Millisecond)
+	}
 }
 
 // ExecuteRemoteForScheduler 供 Scheduler 调用，执行远程任务并等待结果