admin/TaskPool
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add startup task and fix backup import

Browse Source
This commit is contained in:
engigu
2026-02-26 10:43:49 +08:00
parent afb412e71c
commit acc65dce64
12 changed files with 131 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/services/backup_service.go
+49 -53
View File
@@ -227,6 +227,32 @@ func (s *BackupService) Restore(zipPath string) error {
})
}
func restoreStreamBatch[T any](tx *gorm.DB, decoder *json.Decoder) error {
batchSize := 1000
var batch []*T
for decoder.More() {
var m T
if err := decoder.Decode(&m); err != nil {
return err
}
batch = append(batch, &m)
if len(batch) >= batchSize {
if err := tx.CreateInBatches(batch, batchSize).Error; err != nil {
return err
}
batch = nil // reset batch
}
}
if len(batch) > 0 {
return tx.CreateInBatches(batch, len(batch)).Error
}
return nil
}
func (s *BackupService) restoreFromZipFile(tx *gorm.DB, f *zip.File, filename string) error {
rc, err := f.Open()
if err != nil {
@@ -254,60 +280,30 @@ func (s *BackupService) restoreFromZipFile(tx *gorm.DB, f *zip.File, filename st
return fmt.Errorf("invalid json format: expected %s", filename)
}
batchSize := 1000
var batch []any
// 根据文件名确定模型类型
getModel := func() any {
switch filename {
case "tasks.json":
return &models.Task{}
case "task_logs.json":
return &models.TaskLog{}
case "envs.json":
return &models.EnvironmentVariable{}
case "scripts.json":
return &models.Script{}
case "send_stats.json":
return &models.SendStats{}
case "login_logs.json":
return &models.LoginLog{}
case "agents.json":
return &models.Agent{}
case "tokens.json":
return &models.AgentToken{}
case "languages.json":
return &models.Language{}
case "deps.json":
return &models.Dependency{}
default:
return nil
}
switch filename {
case "tasks.json":
return restoreStreamBatch[models.Task](tx, decoder)
case "task_logs.json":
return restoreStreamBatch[models.TaskLog](tx, decoder)
case "envs.json":
return restoreStreamBatch[models.EnvironmentVariable](tx, decoder)
case "scripts.json":
return restoreStreamBatch[models.Script](tx, decoder)
case "send_stats.json":
return restoreStreamBatch[models.SendStats](tx, decoder)
case "login_logs.json":
return restoreStreamBatch[models.LoginLog](tx, decoder)
case "agents.json":
return restoreStreamBatch[models.Agent](tx, decoder)
case "tokens.json":
return restoreStreamBatch[models.AgentToken](tx, decoder)
case "languages.json":
return restoreStreamBatch[models.Language](tx, decoder)
case "deps.json":
return restoreStreamBatch[models.Dependency](tx, decoder)
default:
return nil
}
for decoder.More() {
m := getModel()
if m == nil {
break
}
if err := decoder.Decode(m); err != nil {
return err
}
batch = append(batch, m)
if len(batch) >= batchSize {
if err := tx.CreateInBatches(batch, batchSize).Error; err != nil {
return err
}
batch = nil
}
}
if len(batch) > 0 {
return tx.CreateInBatches(batch, batchSize).Error
}
return nil
}
// insertRecords, restoreFromData 方法已合并入 restoreFromZipFile，此处删除冗余方法
internal/services/tasks/executor_service.go
+17 -2
View File
@@ -416,6 +416,10 @@ func (es *ExecutorService) StopCron() {
// AddCronTask 添加计划任务
func (es *ExecutorService) AddCronTask(task *models.Task) error {
if task.TriggerType != constant.TriggerTypeCron {
es.RemoveCronTask(task.ID) // 如果不是cron类型，确保从调度器移除
return nil
}
return es.cronManager.AddTask(task)
}
@@ -439,8 +443,19 @@ func (es *ExecutorService) loadCronTasks() {
tasks := es.taskService.GetTasks()
count := 0
for _, task := range tasks {
// 只调度本地任务（agent_id 为空或 0）
if task.Enabled && (task.AgentID == nil || *task.AgentID == 0) {
if !task.Enabled {
continue
}
if task.TriggerType == constant.TriggerTypeBaihuStartup {
go func(t models.Task) {
// 延迟一点时间再触发，确保系统完全启动
time.Sleep(3 * time.Second)
logger.Infof("[Executor] 触发开机服务启动任务 #%d: %s", t.ID, t.Name)
es.ExecuteTask(int(t.ID), nil)
}(task)
} else if task.TriggerType == constant.TriggerTypeCron && task.Schedule != "" && (task.AgentID == nil || *task.AgentID == 0) {
// 只调度本地任务（agent_id 为空或 0）的定时任务
err := es.cronManager.AddTask(&task)
if err != nil {
continue
internal/services/tasks/task_service.go
+16 -2
View File
@@ -1,6 +1,7 @@
package tasks
import (
"github.com/engigu/baihu-panel/internal/constant"
"github.com/engigu/baihu-panel/internal/database"
"github.com/engigu/baihu-panel/internal/models"
)
@@ -11,14 +12,18 @@ func NewTaskService() *TaskService {
return &TaskService{}
}
func (ts *TaskService) CreateTask(name, command, schedule string, timeout int, workDir, cleanConfig, envs, taskType, config string, agentID *uint, languages []map[string]string) *models.Task {
func (ts *TaskService) CreateTask(name, command, schedule string, timeout int, workDir, cleanConfig, envs, taskType, config string, agentID *uint, languages []map[string]string, triggerType string) *models.Task {
if taskType == "" {
taskType = "task"
}
if triggerType == "" {
triggerType = constant.TriggerTypeCron
}
task := &models.Task{
Name: name,
Command: command,
Type: taskType,
TriggerType: triggerType,
Config: config,
Schedule: schedule,
Timeout: timeout,
@@ -29,6 +34,9 @@ func (ts *TaskService) CreateTask(name, command, schedule string, timeout int, w
AgentID: agentID,
Enabled: true,
}
if triggerType != constant.TriggerTypeCron {
task.NextRun = nil
}
database.DB.Create(task)
return task
}
@@ -66,7 +74,7 @@ func (ts *TaskService) GetTaskByID(id int) *models.Task {
return &task
}
func (ts *TaskService) UpdateTask(id int, name, command, schedule string, timeout int, workDir, cleanConfig, envs string, enabled bool, taskType, config string, agentID *uint, languages []map[string]string) *models.Task {
func (ts *TaskService) UpdateTask(id int, name, command, schedule string, timeout int, workDir, cleanConfig, envs string, enabled bool, taskType, config string, agentID *uint, languages []map[string]string, triggerType string) *models.Task {
var task models.Task
if err := database.DB.First(&task, id).Error; err != nil {
return nil
@@ -84,6 +92,12 @@ func (ts *TaskService) UpdateTask(id int, name, command, schedule string, timeou
if taskType != "" {
task.Type = taskType
}
if triggerType != "" {
task.TriggerType = triggerType
}
if task.TriggerType != constant.TriggerTypeCron {
task.NextRun = nil
}
task.Config = config
database.DB.Save(&task)
return &task