feat: add scheduler config

This commit is contained in:
engigu
2025-12-22 16:05:07 +08:00
parent 03f4fa7b5c
commit be3527cb61
12 changed files with 387 additions and 52 deletions
+143 -26
View File
@@ -1,19 +1,19 @@
package services
import (
"bytes"
"context"
"encoding/json"
"os"
"os/exec"
"sync"
"time"
"baihu/internal/constant"
"baihu/internal/database"
"baihu/internal/logger"
"baihu/internal/models"
"baihu/internal/utils"
"bytes"
"context"
"encoding/json"
"fmt"
"os"
"os/exec"
"sync"
"time"
)
// ExecutionResult represents the result of a task execution
@@ -29,6 +29,11 @@ type ExecutionResult struct {
// ExecutionCallback 任务执行完成后的回调函数类型
type ExecutionCallback func(taskID uint, command string, result *ExecutionResult)
// taskJob 任务队列项
type taskJob struct {
taskID int
}
// ExecutorService handles task execution
type ExecutorService struct {
taskService *TaskService
@@ -36,23 +41,119 @@ type ExecutorService struct {
runningTasks map[int]bool
callbacks []ExecutionCallback
mu sync.RWMutex
resultsMu sync.RWMutex
// 任务队列和 worker pool
taskQueue chan taskJob
workerCount int
rateLimiter <-chan time.Time
stopCh chan struct{}
wg sync.WaitGroup
}
// NewExecutorService creates a new executor service
func NewExecutorService(taskService *TaskService) *ExecutorService {
// 从设置中读取调度配置
settingsService := NewSettingsService()
workerCount := getIntSetting(settingsService, constant.SectionScheduler, constant.KeyWorkerCount, 4)
queueSize := getIntSetting(settingsService, constant.SectionScheduler, constant.KeyQueueSize, 100)
rateInterval := getIntSetting(settingsService, constant.SectionScheduler, constant.KeyRateInterval, 200)
logger.Infof("Executor service config: workers=%d, queue=%d, rate=%dms", workerCount, queueSize, rateInterval)
es := &ExecutorService{
taskService: taskService,
results: make([]ExecutionResult, 0),
results: make([]ExecutionResult, 0, 100),
runningTasks: make(map[int]bool),
callbacks: make([]ExecutionCallback, 0),
taskQueue: make(chan taskJob, queueSize),
workerCount: workerCount,
rateLimiter: time.Tick(time.Duration(rateInterval) * time.Millisecond),
stopCh: make(chan struct{}),
}
// 注册默认回调
es.RegisterCallback(es.saveTaskLogCallback)
es.RegisterCallback(es.updateStatsCallback)
es.RegisterCallback(es.cleanLogsCallback)
// 启动 worker pool
es.startWorkers()
return es
}
// getIntSetting 从设置中获取整数值
func getIntSetting(s *SettingsService, section, key string, defaultVal int) int {
val := s.Get(section, key)
if val == "" {
return defaultVal
}
var result int
if _, err := fmt.Sscanf(val, "%d", &result); err != nil {
return defaultVal
}
return result
}
// startWorkers 启动 worker pool
func (es *ExecutorService) startWorkers() {
for i := 0; i < es.workerCount; i++ {
es.wg.Add(1)
go es.worker(i)
}
}
// worker 从队列中取任务执行
func (es *ExecutorService) worker(id int) {
defer es.wg.Done()
for {
select {
case <-es.stopCh:
return
case job := <-es.taskQueue:
// 速率限制
<-es.rateLimiter
es.executeTaskInternal(job.taskID)
}
}
}
// Stop 停止 executor service
func (es *ExecutorService) Stop() {
close(es.stopCh)
es.wg.Wait()
}
// Reload 重新加载配置并重建 worker pool
func (es *ExecutorService) Reload() {
logger.Info("Reloading executor service...")
// 停止现有 workers
close(es.stopCh)
es.wg.Wait()
logger.Info("Stopped executor service...")
// 从设置中读取新配置
settingsService := NewSettingsService()
workerCount := getIntSetting(settingsService, constant.SectionScheduler, constant.KeyWorkerCount, 4)
queueSize := getIntSetting(settingsService, constant.SectionScheduler, constant.KeyQueueSize, 100)
rateInterval := getIntSetting(settingsService, constant.SectionScheduler, constant.KeyRateInterval, 200)
// 重建 channel 和配置
es.mu.Lock()
es.taskQueue = make(chan taskJob, queueSize)
es.workerCount = workerCount
es.rateLimiter = time.Tick(time.Duration(rateInterval) * time.Millisecond)
es.stopCh = make(chan struct{})
es.mu.Unlock()
// 启动新的 workers
es.startWorkers()
logger.Infof("Executor service reloaded: workers=%d, queue=%d, rate=%dms", workerCount, queueSize, rateInterval)
}
// RegisterCallback 注册执行完成回调
func (es *ExecutorService) RegisterCallback(cb ExecutionCallback) {
es.mu.Lock()
@@ -60,19 +161,21 @@ func (es *ExecutorService) RegisterCallback(cb ExecutionCallback) {
es.mu.Unlock()
}
// executeCallbacks 执行所有回调
func (es *ExecutorService) executeCallbacks(taskID uint, command string, result *ExecutionResult) {
// executeCallbacksAsync 异步执行所有回调
func (es *ExecutorService) executeCallbacksAsync(taskID uint, command string, result *ExecutionResult) {
es.mu.RLock()
callbacks := make([]ExecutionCallback, len(es.callbacks))
copy(callbacks, es.callbacks)
es.mu.RUnlock()
for _, cb := range callbacks {
cb(taskID, command, result)
}
go func() {
for _, cb := range callbacks {
cb(taskID, command, result)
}
}()
}
// saveTaskLogCallback 保存任务日志的回调
// saveTaskLogCallback 保存任务日志的回调(异步执行,压缩在此处进行)
func (es *ExecutorService) saveTaskLogCallback(taskID uint, command string, result *ExecutionResult) {
output := result.Output
if result.Error != "" {
@@ -141,17 +244,13 @@ func (es *ExecutorService) cleanLogsCallback(taskID uint, _ string, _ *Execution
var deleted int64
switch config.Type {
case "day":
// 按天清理:删除 N 天前的日志
cutoff := time.Now().AddDate(0, 0, -config.Keep)
result := database.DB.Where("task_id = ? AND created_at < ?", taskID, cutoff).Delete(&models.TaskLog{})
deleted = result.RowsAffected
case "count":
// 按条数清理:使用子查询删除超出保留数量的旧日志
// 先获取第 N 条的 ID 作为边界
var boundaryLog models.TaskLog
err := database.DB.Where("task_id = ?", taskID).Order("id DESC").Offset(config.Keep - 1).Limit(1).First(&boundaryLog).Error
if err == nil {
// 删除 ID 小于边界的所有日志
result := database.DB.Where("task_id = ? AND id < ?", taskID, boundaryLog.ID).Delete(&models.TaskLog{})
deleted = result.RowsAffected
}
@@ -162,8 +261,25 @@ func (es *ExecutorService) cleanLogsCallback(taskID uint, _ string, _ *Execution
}
}
// ExecuteTask executes a task by ID
// EnqueueTask 将任务加入队列(供 cron 调度器调用)
func (es *ExecutorService) EnqueueTask(taskID int) {
select {
case es.taskQueue <- taskJob{taskID: taskID}:
// 成功入队
default:
// 队列满,直接执行(降级处理)
logger.Warnf("Task queue full, executing task %d directly", taskID)
go es.executeTaskInternal(taskID)
}
}
// ExecuteTask executes a task by ID(同步执行,供 API 调用)
func (es *ExecutorService) ExecuteTask(taskID int) *ExecutionResult {
return es.executeTaskInternal(taskID)
}
// executeTaskInternal 内部执行任务逻辑
func (es *ExecutorService) executeTaskInternal(taskID int) *ExecutionResult {
task := es.taskService.GetTaskByID(taskID)
if task == nil {
return &ExecutionResult{
@@ -197,8 +313,8 @@ func (es *ExecutorService) ExecuteTask(taskID int) *ExecutionResult {
delete(es.runningTasks, taskID)
es.mu.Unlock()
// 执行回调
es.executeCallbacks(uint(taskID), task.Command, result)
// 异步执行回调(日志压缩、统计更新、日志清理)
es.executeCallbacksAsync(uint(taskID), task.Command, result)
return result
}
@@ -255,20 +371,21 @@ func (es *ExecutorService) ExecuteCommandWithEnv(command string, timeout time.Du
result.Success = true
}
es.mu.Lock()
// 使用独立锁保存结果
es.resultsMu.Lock()
es.results = append(es.results, *result)
if len(es.results) > 100 {
es.results = es.results[1:]
}
es.mu.Unlock()
es.resultsMu.Unlock()
return result
}
// GetLastResults returns the last execution results
func (es *ExecutorService) GetLastResults(count int) []ExecutionResult {
es.mu.RLock()
defer es.mu.RUnlock()
es.resultsMu.RLock()
defer es.resultsMu.RUnlock()
start := 0
if len(es.results) > count {