feat: refact scheduler
This commit is contained in:
engigu
@@ -0,0 +1,169 @@
|
||||
package executor
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/robfig/cron/v3"
|
||||
)
|
||||
|
||||
// 东八区时区（默认）
|
||||
var defaultLocation = time.FixedZone("CST", 8*3600)
|
||||
|
||||
// CronManager 统一的任务调度管理器
|
||||
type CronManager struct {
|
||||
cron *cron.Cron
|
||||
scheduler *Scheduler
|
||||
entryMap map[string]cron.EntryID // task ID -> cron entry ID
|
||||
mu sync.RWMutex
|
||||
logger SchedulerLogger
|
||||
}
|
||||
|
||||
// NewCronManager 创建一个新的计划任务管理器
|
||||
func NewCronManager(scheduler *Scheduler) *CronManager {
|
||||
// 使用秒级精度的 cron parser
|
||||
c := cron.New(cron.WithSeconds(), cron.WithLocation(defaultLocation))
|
||||
|
||||
m := &CronManager{
|
||||
cron: c,
|
||||
scheduler: scheduler,
|
||||
entryMap: make(map[string]cron.EntryID),
|
||||
logger: &DefaultLogger{},
|
||||
}
|
||||
|
||||
if scheduler != nil && scheduler.logger != nil {
|
||||
m.logger = scheduler.logger
|
||||
}
|
||||
|
||||
return m
|
||||
}
|
||||
|
||||
// SetLogger 设置自定义日志实现
|
||||
func (m *CronManager) SetLogger(logger SchedulerLogger) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
m.logger = logger
|
||||
}
|
||||
|
||||
// Start 启动调度器
|
||||
func (m *CronManager) Start() {
|
||||
m.cron.Start()
|
||||
m.logger.Infof("[CronManager] 调度管理服务已启动")
|
||||
}
|
||||
|
||||
// Stop 停止调度器
|
||||
func (m *CronManager) Stop() {
|
||||
ctx := m.cron.Stop()
|
||||
<-ctx.Done()
|
||||
m.logger.Infof("[CronManager] 调度管理服务已停止")
|
||||
}
|
||||
|
||||
// AddTask 添加或更新计划任务
|
||||
func (m *CronManager) AddTask(task CronTask) error {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
taskID := task.GetID()
|
||||
|
||||
// 如果已存在，先移除旧的
|
||||
if entryID, exists := m.entryMap[taskID]; exists {
|
||||
m.cron.Remove(entryID)
|
||||
delete(m.entryMap, taskID)
|
||||
}
|
||||
|
||||
// 准备任务执行函数
|
||||
cmd := task.GetCommand()
|
||||
name := task.GetName()
|
||||
timeout := task.GetTimeout()
|
||||
|
||||
entryID, err := m.cron.AddFunc(task.GetSchedule(), func() {
|
||||
m.logger.Infof("[CronManager] 触发计划任务 #%s (%s)", taskID, name)
|
||||
|
||||
req := &ExecutionRequest{
|
||||
TaskID: taskID,
|
||||
Name: name,
|
||||
Command: cmd,
|
||||
Type: TaskTypeCron,
|
||||
Timeout: timeout,
|
||||
}
|
||||
|
||||
// 如果有关联的 Scheduler，加入队列执行
|
||||
if m.scheduler != nil {
|
||||
m.scheduler.EnqueueOrExecute(req)
|
||||
}
|
||||
|
||||
// 触发下次运行时间更新事件
|
||||
m.triggerNextRunEvent(taskID, req)
|
||||
})
|
||||
|
||||
if err != nil {
|
||||
m.logger.Errorf("[CronManager] 添加任务失败 #%s: %v", taskID, err)
|
||||
return err
|
||||
}
|
||||
|
||||
m.entryMap[taskID] = entryID
|
||||
m.logger.Infof("[CronManager] 任务已调度 #%s %s (%s)", taskID, name, task.GetSchedule())
|
||||
|
||||
// 初始触发一次下次运行时间通知
|
||||
go func() {
|
||||
req := &ExecutionRequest{TaskID: taskID, Name: name, Type: TaskTypeCron}
|
||||
m.triggerNextRunEvent(taskID, req)
|
||||
}()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// RemoveTask 移除计划任务
|
||||
func (m *CronManager) RemoveTask(taskID string) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
if entryID, exists := m.entryMap[taskID]; exists {
|
||||
m.cron.Remove(entryID)
|
||||
delete(m.entryMap, taskID)
|
||||
m.logger.Infof("[CronManager] 任务已移除 #%s", taskID)
|
||||
}
|
||||
}
|
||||
|
||||
// triggerNextRunEvent 触发下次运行时间更新事件
|
||||
func (m *CronManager) triggerNextRunEvent(taskID string, req *ExecutionRequest) {
|
||||
m.mu.RLock()
|
||||
entryID, exists := m.entryMap[taskID]
|
||||
m.mu.RUnlock()
|
||||
|
||||
if !exists {
|
||||
return
|
||||
}
|
||||
|
||||
entry := m.cron.Entry(entryID)
|
||||
if !entry.Next.IsZero() && m.scheduler != nil && m.scheduler.handler != nil {
|
||||
m.scheduler.handler.OnCronNextRun(req, entry.Next)
|
||||
}
|
||||
}
|
||||
|
||||
// ValidateCron 校验 Cron 表达式
|
||||
func (m *CronManager) ValidateCron(expression string) error {
|
||||
parser := cron.NewParser(cron.Second | cron.Minute | cron.Hour | cron.Dom | cron.Month | cron.Dow | cron.Descriptor)
|
||||
_, err := parser.Parse(expression)
|
||||
return err
|
||||
}
|
||||
|
||||
// GetEntry 获取任务详情
|
||||
func (m *CronManager) GetEntry(taskID string) (cron.Entry, bool) {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
|
||||
entryID, exists := m.entryMap[taskID]
|
||||
if !exists {
|
||||
return cron.Entry{}, false
|
||||
}
|
||||
|
||||
return m.cron.Entry(entryID), true
|
||||
}
|
||||
|
||||
// GetScheduledCount 获取已调度任务总数
|
||||
func (m *CronManager) GetScheduledCount() int {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
return len(m.entryMap)
|
||||
}
|
||||
@@ -0,0 +1,204 @@
|
||||
package executor
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"os"
|
||||
"os/exec"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/engigu/baihu-panel/internal/utils"
|
||||
)
|
||||
|
||||
// Task 任务基础接口
|
||||
type Task interface {
|
||||
GetID() string
|
||||
GetName() string
|
||||
GetCommand() string
|
||||
GetTimeout() int
|
||||
}
|
||||
|
||||
// CronTask 计划任务接口
|
||||
type CronTask interface {
|
||||
Task
|
||||
GetSchedule() string
|
||||
}
|
||||
|
||||
// Request 任务执行请求
|
||||
type Request struct {
|
||||
Command string
|
||||
WorkDir string
|
||||
Envs []string
|
||||
Timeout int // 分钟
|
||||
}
|
||||
|
||||
// Result 任务执行结果
|
||||
type Result struct {
|
||||
Output string
|
||||
Status string // success, failed
|
||||
Duration int64 // 毫秒
|
||||
ExitCode int
|
||||
StartTime time.Time
|
||||
EndTime time.Time
|
||||
}
|
||||
|
||||
// Hooks 执行钩子接口
|
||||
type Hooks interface {
|
||||
// PreExecute 执行前钩子，返回日志ID和错误
|
||||
PreExecute(ctx context.Context, req Request) (logID uint, err error)
|
||||
|
||||
// PostExecute 执行后钩子，处理日志压缩和记录更新
|
||||
PostExecute(ctx context.Context, logID uint, result *Result) error
|
||||
|
||||
// OnHeartbeat 执行中心跳钩子，用于更新实时状态
|
||||
OnHeartbeat(ctx context.Context, logID uint, duration int64) error
|
||||
}
|
||||
|
||||
// Execute 执行命令（基础版本，不带钩子）
|
||||
func Execute(ctx context.Context, req Request, stdout, stderr io.Writer) (*Result, error) {
|
||||
return ExecuteWithHooks(ctx, req, stdout, stderr, nil)
|
||||
}
|
||||
|
||||
// ExecuteWithHooks 执行命令（带钩子支持）
|
||||
func ExecuteWithHooks(ctx context.Context, req Request, stdout, stderr io.Writer, hooks Hooks) (*Result, error) {
|
||||
start := time.Now()
|
||||
|
||||
// 1. 执行前钩子
|
||||
var logID uint
|
||||
if hooks != nil {
|
||||
id, err := hooks.PreExecute(ctx, req)
|
||||
if err != nil {
|
||||
return &Result{
|
||||
Status: "failed",
|
||||
Duration: 0,
|
||||
ExitCode: 1,
|
||||
StartTime: start,
|
||||
EndTime: time.Now(),
|
||||
}, err
|
||||
}
|
||||
logID = id
|
||||
}
|
||||
|
||||
// 2. 执行命令
|
||||
timeout := req.Timeout
|
||||
if timeout <= 0 {
|
||||
timeout = 30
|
||||
}
|
||||
execCtx, cancel := context.WithTimeout(ctx, time.Duration(timeout)*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
finalCommand := req.Command
|
||||
shell, args := utils.GetShellCommand(finalCommand)
|
||||
cmd := exec.CommandContext(execCtx, shell, args...)
|
||||
|
||||
// 设置工作目录
|
||||
// 设置工作目录
|
||||
workDir := strings.TrimSpace(req.WorkDir)
|
||||
if workDir != "" {
|
||||
cmd.Dir = workDir
|
||||
}
|
||||
|
||||
// 设置环境变量（始终继承系统环境变量）
|
||||
cmd.Env = os.Environ()
|
||||
if len(req.Envs) > 0 {
|
||||
cmd.Env = append(cmd.Env, req.Envs...)
|
||||
}
|
||||
|
||||
cmd.Stdout = stdout
|
||||
cmd.Stderr = stderr
|
||||
|
||||
// 使用 cmd.Start() + Wait() 以便在后台处理心跳
|
||||
err := cmd.Start()
|
||||
if err != nil {
|
||||
// Start 失败的处理
|
||||
end := time.Now()
|
||||
result := &Result{
|
||||
Status: "failed",
|
||||
Duration: end.Sub(start).Milliseconds(),
|
||||
ExitCode: 1,
|
||||
StartTime: start, // 修正为 start
|
||||
EndTime: end,
|
||||
}
|
||||
// 执行后钩子
|
||||
if hooks != nil {
|
||||
result.Output += "\n[System Error] " + err.Error()
|
||||
hooks.PostExecute(ctx, logID, result)
|
||||
}
|
||||
return result, err
|
||||
}
|
||||
|
||||
// 启动心跳协程
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
// 每3秒一次心跳
|
||||
ticker := time.NewTicker(3 * time.Second)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-done:
|
||||
return
|
||||
case <-ticker.C:
|
||||
if hooks != nil {
|
||||
hooks.OnHeartbeat(ctx, logID, time.Since(start).Milliseconds())
|
||||
}
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
// 等待命令完成
|
||||
err = cmd.Wait()
|
||||
close(done) // 停止心跳
|
||||
|
||||
end := time.Now()
|
||||
|
||||
result := &Result{
|
||||
StartTime: start,
|
||||
EndTime: end,
|
||||
Duration: end.Sub(start).Milliseconds(),
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
result.Status = "failed"
|
||||
if exitErr, ok := err.(*exec.ExitError); ok {
|
||||
result.ExitCode = exitErr.ExitCode()
|
||||
} else {
|
||||
result.ExitCode = 1
|
||||
}
|
||||
} else {
|
||||
result.Status = "success"
|
||||
result.ExitCode = 0
|
||||
}
|
||||
|
||||
// 3. 执行后钩子
|
||||
if hooks != nil {
|
||||
if hookErr := hooks.PostExecute(ctx, logID, result); hookErr != nil {
|
||||
// 记录钩子错误但不影响执行结果
|
||||
result.Output += "\n[Hook Error] " + hookErr.Error()
|
||||
}
|
||||
}
|
||||
|
||||
return result, err
|
||||
}
|
||||
|
||||
// ParseEnvVars 解析环境变量字符串 "KEY1=VALUE1,KEY2=VALUE2"
|
||||
func ParseEnvVars(envStr string) []string {
|
||||
if envStr == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
pairs := strings.Split(envStr, ",")
|
||||
result := make([]string, 0, len(pairs))
|
||||
|
||||
for _, pair := range pairs {
|
||||
if pair == "" {
|
||||
continue
|
||||
}
|
||||
// 解码特殊字符
|
||||
pair = strings.ReplaceAll(pair, "{{COMMA}}", ",")
|
||||
pair = strings.ReplaceAll(pair, "{{EQUAL}}", "=")
|
||||
result = append(result, pair)
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
@@ -0,0 +1,451 @@
|
||||
package executor
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// SchedulerConfig 调度器配置
|
||||
type SchedulerConfig struct {
|
||||
WorkerCount int // Worker 数量
|
||||
QueueSize int // 队列大小
|
||||
RateInterval time.Duration // 速率限制间隔
|
||||
}
|
||||
|
||||
// TaskType 任务类型
|
||||
type TaskType string
|
||||
|
||||
const (
|
||||
TaskTypeCron TaskType = "cron" // 计划任务
|
||||
TaskTypeManual TaskType = "manual" // 手动任务
|
||||
TaskTypeSystem TaskType = "system" // 系统任务
|
||||
)
|
||||
|
||||
// TaskStatus 任务状态
|
||||
type TaskStatus string
|
||||
|
||||
const (
|
||||
TaskStatusPending TaskStatus = "pending" // 等待中
|
||||
TaskStatusRunning TaskStatus = "running" // 运行中
|
||||
TaskStatusSuccess TaskStatus = "success" // 成功
|
||||
TaskStatusFailed TaskStatus = "failed" // 失败
|
||||
TaskStatusTimeout TaskStatus = "timeout" // 超时
|
||||
TaskStatusCancelled TaskStatus = "cancelled" // 已取消
|
||||
)
|
||||
|
||||
// ExecutionRequest 执行请求（标准接口）
|
||||
type ExecutionRequest struct {
|
||||
TaskID string // 任务 ID
|
||||
LogID uint // 日志 ID
|
||||
Name string // 任务名称
|
||||
Type TaskType // 任务类型
|
||||
Command string // 命令
|
||||
WorkDir string // 工作目录
|
||||
Envs []string // 环境变量
|
||||
Timeout int // 超时时间（分钟）
|
||||
Metadata map[string]interface{} // 额外元数据
|
||||
}
|
||||
|
||||
// ExecutionResult 执行结果（标准接口）
|
||||
type ExecutionResult struct {
|
||||
TaskID string // 任务 ID
|
||||
LogID uint // 日志 ID
|
||||
Success bool // 是否成功
|
||||
Output string // 输出内容
|
||||
Error string // 错误信息
|
||||
Status string // 状态: success, failed, timeout, cancelled
|
||||
Duration int64 // 执行时长（毫秒）
|
||||
ExitCode int // 退出码
|
||||
StartTime time.Time // 开始时间
|
||||
EndTime time.Time // 结束时间
|
||||
}
|
||||
|
||||
// SchedulerEventHandler 调度器事件处理器（标准接口）
|
||||
// 主服务端和 Agent 端通过实现不同的 Handler 来处理事件
|
||||
type SchedulerEventHandler interface {
|
||||
// OnTaskScheduled 任务被调度（加入队列）时触发
|
||||
OnTaskScheduled(req *ExecutionRequest)
|
||||
|
||||
// OnTaskExecuting 任务准备开始执行时触发
|
||||
// 返回 stdout/stderr 写入器用于实时日志推送
|
||||
// 主服务端：返回 TinyLog 写入器（写入本地文件）
|
||||
// Agent 端：返回 WebSocket 写入器（实时推送到主服务）
|
||||
OnTaskExecuting(req *ExecutionRequest) (stdout, stderr io.Writer, err error)
|
||||
|
||||
// OnTaskStarted 任务实际开始运行（已经过了队列等待和速率限制）
|
||||
OnTaskStarted(req *ExecutionRequest)
|
||||
|
||||
// OnTaskCompleted 任务执行完成时触发
|
||||
// 主服务端：压缩日志、更新数据库、清理旧日志
|
||||
// Agent 端：通过 WebSocket 发送执行结果到主服务
|
||||
OnTaskCompleted(req *ExecutionRequest, result *ExecutionResult)
|
||||
|
||||
// OnTaskFailed 任务执行失败时触发
|
||||
OnTaskFailed(req *ExecutionRequest, err error)
|
||||
|
||||
// OnCronNextRun 计划任务下次运行时间更新时触发
|
||||
OnCronNextRun(req *ExecutionRequest, nextRun time.Time)
|
||||
|
||||
// OnTaskHeartbeat 任务执行心跳（用于更新实时耗时等）
|
||||
OnTaskHeartbeat(req *ExecutionRequest, duration int64)
|
||||
}
|
||||
|
||||
// SchedulerLogger 日志接口（允许自定义日志实现）
|
||||
type SchedulerLogger interface {
|
||||
Infof(format string, args ...interface{})
|
||||
Warnf(format string, args ...interface{})
|
||||
Errorf(format string, args ...interface{})
|
||||
}
|
||||
|
||||
// DefaultLogger 默认日志实现（使用 fmt）
|
||||
type DefaultLogger struct{}
|
||||
|
||||
func (l *DefaultLogger) Infof(format string, args ...interface{}) {
|
||||
fmt.Printf("[INFO] "+format+"\n", args...)
|
||||
}
|
||||
func (l *DefaultLogger) Warnf(format string, args ...interface{}) {
|
||||
fmt.Printf("[WARN] "+format+"\n", args...)
|
||||
}
|
||||
func (l *DefaultLogger) Errorf(format string, args ...interface{}) {
|
||||
fmt.Printf("[ERROR] "+format+"\n", args...)
|
||||
}
|
||||
|
||||
// schedulerHooksAdapter 适配器：将 executor.Hooks 映射到 SchedulerEventHandler
|
||||
type schedulerHooksAdapter struct {
|
||||
handler SchedulerEventHandler
|
||||
req *ExecutionRequest
|
||||
}
|
||||
|
||||
func (h *schedulerHooksAdapter) PreExecute(ctx context.Context, req Request) (uint, error) {
|
||||
return h.req.LogID, nil
|
||||
}
|
||||
|
||||
func (h *schedulerHooksAdapter) PostExecute(ctx context.Context, logID uint, result *Result) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (h *schedulerHooksAdapter) OnHeartbeat(ctx context.Context, logID uint, duration int64) error {
|
||||
if h.handler != nil {
|
||||
h.handler.OnTaskHeartbeat(h.req, duration)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// TaskExecutor 定义任务执行函数签名
|
||||
type TaskExecutor func(ctx context.Context, req *ExecutionRequest, stdout, stderr io.Writer) (*Result, error)
|
||||
|
||||
// Scheduler 统一调度器（独立组件，可在主服务和 Agent 中复用）
|
||||
// 调度器本身只负责队列管理和任务调度，具体的执行逻辑和事件处理由 Handler 实现
|
||||
type Scheduler struct {
|
||||
config SchedulerConfig
|
||||
handler SchedulerEventHandler
|
||||
executor TaskExecutor
|
||||
taskQueue chan *ExecutionRequest
|
||||
rateLimiter <-chan time.Time
|
||||
stopCh chan struct{}
|
||||
wg sync.WaitGroup
|
||||
mu sync.RWMutex
|
||||
logger SchedulerLogger
|
||||
runningTasks map[string]context.CancelFunc // 记录运行中的任务，用于停止
|
||||
}
|
||||
|
||||
// NewScheduler 创建调度器
|
||||
func NewScheduler(config SchedulerConfig, handler SchedulerEventHandler) *Scheduler {
|
||||
if config.WorkerCount <= 0 {
|
||||
config.WorkerCount = 4
|
||||
}
|
||||
if config.QueueSize <= 0 {
|
||||
config.QueueSize = 100
|
||||
}
|
||||
if config.RateInterval <= 0 {
|
||||
config.RateInterval = 200 * time.Millisecond
|
||||
}
|
||||
|
||||
s := &Scheduler{
|
||||
config: config,
|
||||
handler: handler,
|
||||
executor: func(ctx context.Context, req *ExecutionRequest, stdout, stderr io.Writer) (*Result, error) {
|
||||
hooks := &schedulerHooksAdapter{handler: handler, req: req}
|
||||
return ExecuteWithHooks(ctx, Request{
|
||||
Command: req.Command,
|
||||
WorkDir: req.WorkDir,
|
||||
Envs: req.Envs,
|
||||
Timeout: req.Timeout,
|
||||
}, stdout, stderr, hooks)
|
||||
},
|
||||
taskQueue: make(chan *ExecutionRequest, config.QueueSize),
|
||||
rateLimiter: time.Tick(config.RateInterval),
|
||||
stopCh: make(chan struct{}),
|
||||
logger: &DefaultLogger{},
|
||||
runningTasks: make(map[string]context.CancelFunc),
|
||||
}
|
||||
|
||||
return s
|
||||
}
|
||||
|
||||
// SetLogger 设置自定义日志实现
|
||||
func (s *Scheduler) SetLogger(logger SchedulerLogger) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.logger = logger
|
||||
}
|
||||
|
||||
// SetExecutor 设置任务执行器
|
||||
func (s *Scheduler) SetExecutor(executor TaskExecutor) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.executor = executor
|
||||
}
|
||||
|
||||
// Start 启动调度器
|
||||
func (s *Scheduler) Start() {
|
||||
for i := 0; i < s.config.WorkerCount; i++ {
|
||||
s.wg.Add(1)
|
||||
go s.worker(i)
|
||||
}
|
||||
s.logger.Infof("[Scheduler] 已启动")
|
||||
}
|
||||
|
||||
// Stop 停止调度器
|
||||
func (s *Scheduler) Stop() {
|
||||
close(s.stopCh)
|
||||
s.wg.Wait()
|
||||
s.logger.Infof("[Scheduler] 已停止")
|
||||
}
|
||||
|
||||
// Enqueue 将任务加入队列
|
||||
func (s *Scheduler) Enqueue(req *ExecutionRequest) error {
|
||||
select {
|
||||
case s.taskQueue <- req:
|
||||
if s.handler != nil {
|
||||
s.handler.OnTaskScheduled(req)
|
||||
}
|
||||
return nil
|
||||
default:
|
||||
// 队列满，返回错误
|
||||
return fmt.Errorf("任务队列已满")
|
||||
}
|
||||
}
|
||||
|
||||
// EnqueueOrExecute 将任务加入队列，如果队列满则直接执行
|
||||
func (s *Scheduler) EnqueueOrExecute(req *ExecutionRequest) {
|
||||
select {
|
||||
case s.taskQueue <- req:
|
||||
// 成功入队
|
||||
if s.handler != nil {
|
||||
s.handler.OnTaskScheduled(req)
|
||||
}
|
||||
default:
|
||||
// 队列满，直接执行（降级处理）
|
||||
s.logger.Warnf("[Scheduler] 任务队列已满，直接执行任务 %s", req.TaskID)
|
||||
go s.executeTask(req)
|
||||
}
|
||||
}
|
||||
|
||||
// ExecuteSync 同步执行任务（不经过队列）
|
||||
func (s *Scheduler) ExecuteSync(req *ExecutionRequest) (*ExecutionResult, error) {
|
||||
return s.executeTask(req)
|
||||
}
|
||||
|
||||
// worker 工作协程
|
||||
func (s *Scheduler) worker(id int) {
|
||||
defer s.wg.Done()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-s.stopCh:
|
||||
return
|
||||
case req := <-s.taskQueue:
|
||||
// 速率限制
|
||||
<-s.rateLimiter
|
||||
s.executeTask(req)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// executeTask 执行任务（本地执行）
|
||||
func (s *Scheduler) executeTask(req *ExecutionRequest) (*ExecutionResult, error) {
|
||||
start := time.Now()
|
||||
|
||||
s.logger.Infof("[Scheduler] 执行任务 %s (名称: %s, 类型: %s)", req.TaskID, req.Name, req.Type)
|
||||
|
||||
// 1. 执行前事件：获取 stdout/stderr 写入器
|
||||
var stdout, stderr io.Writer
|
||||
var err error
|
||||
if s.handler != nil {
|
||||
stdout, stderr, err = s.handler.OnTaskExecuting(req)
|
||||
if err != nil {
|
||||
s.logger.Errorf("[Scheduler] 任务 %s 执行前事件失败: %v", req.TaskID, err)
|
||||
if s.handler != nil {
|
||||
s.handler.OnTaskFailed(req, err)
|
||||
}
|
||||
return &ExecutionResult{
|
||||
TaskID: req.TaskID,
|
||||
Success: false,
|
||||
Status: "failed",
|
||||
Error: err.Error(),
|
||||
Duration: 0,
|
||||
ExitCode: 1,
|
||||
StartTime: start,
|
||||
EndTime: time.Now(),
|
||||
}, err
|
||||
}
|
||||
}
|
||||
|
||||
// 2. 准备输出缓冲区
|
||||
var stdoutBuf, stderrBuf bytes.Buffer
|
||||
var stdoutWriter, stderrWriter io.Writer
|
||||
|
||||
if stdout != nil {
|
||||
stdoutWriter = io.MultiWriter(&stdoutBuf, stdout)
|
||||
} else {
|
||||
stdoutWriter = &stdoutBuf
|
||||
}
|
||||
|
||||
if stderr != nil {
|
||||
stderrWriter = io.MultiWriter(&stderrBuf, stderr)
|
||||
} else {
|
||||
stderrWriter = &stderrBuf
|
||||
}
|
||||
|
||||
// 3. 实际开始执行事件 (经过队列和速率限制之后)
|
||||
if s.handler != nil {
|
||||
s.handler.OnTaskStarted(req)
|
||||
}
|
||||
|
||||
// 4. 执行命令（使用 executor.Execute）
|
||||
// 创建带取消功能的上下文
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
if req.Timeout > 0 {
|
||||
ctx, cancel = context.WithTimeout(ctx, time.Duration(req.Timeout)*time.Minute)
|
||||
}
|
||||
defer cancel()
|
||||
|
||||
// 注册到运行中任务
|
||||
s.mu.Lock()
|
||||
s.runningTasks[req.TaskID] = cancel
|
||||
s.mu.Unlock()
|
||||
|
||||
defer func() {
|
||||
s.mu.Lock()
|
||||
delete(s.runningTasks, req.TaskID)
|
||||
s.mu.Unlock()
|
||||
}()
|
||||
|
||||
execResult, execErr := s.executor(ctx, req, stdoutWriter, stderrWriter)
|
||||
|
||||
// 5. 构建结果
|
||||
result := &ExecutionResult{
|
||||
TaskID: req.TaskID,
|
||||
LogID: req.LogID, // 传递 LogID
|
||||
Success: execResult.Status == "success",
|
||||
Output: stdoutBuf.String(),
|
||||
Status: execResult.Status,
|
||||
Duration: execResult.Duration,
|
||||
ExitCode: execResult.ExitCode,
|
||||
StartTime: execResult.StartTime,
|
||||
EndTime: execResult.EndTime,
|
||||
}
|
||||
|
||||
if execErr != nil {
|
||||
result.Error = execErr.Error()
|
||||
errOutput := stderrBuf.String()
|
||||
if errOutput != "" {
|
||||
result.Output += "\n[ERROR]\n" + errOutput
|
||||
}
|
||||
if ctx.Err() == context.Canceled {
|
||||
result.Status = "cancelled"
|
||||
} else if ctx.Err() == context.DeadlineExceeded {
|
||||
result.Status = "timeout"
|
||||
}
|
||||
}
|
||||
|
||||
// 6. 执行后事件
|
||||
if s.handler != nil {
|
||||
if execErr != nil {
|
||||
s.handler.OnTaskFailed(req, execErr)
|
||||
} else {
|
||||
s.handler.OnTaskCompleted(req, result)
|
||||
}
|
||||
}
|
||||
|
||||
if execErr != nil {
|
||||
s.logger.Errorf("[Scheduler] 任务 %s 执行失败: %v", req.TaskID, execErr)
|
||||
} else {
|
||||
s.logger.Infof("[Scheduler] 任务 %s 执行完成 (状态: %s, 耗时: %dms)",
|
||||
req.TaskID, result.Status, result.Duration)
|
||||
}
|
||||
|
||||
return result, execErr
|
||||
}
|
||||
|
||||
// StopTask 停止正在运行的任务
|
||||
func (s *Scheduler) StopTask(taskID string) bool {
|
||||
s.mu.RLock()
|
||||
cancel, exists := s.runningTasks[taskID]
|
||||
s.mu.RUnlock()
|
||||
|
||||
if exists && cancel != nil {
|
||||
cancel()
|
||||
s.logger.Infof("[Scheduler] 已尝试停止任务 %s", taskID)
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// GetRunningTaskCount 获取正在运行的任务数量
|
||||
func (s *Scheduler) GetRunningTaskCount() int {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
return len(s.runningTasks)
|
||||
}
|
||||
|
||||
// GetRunningTasks 获取所有正在运行的任务 ID
|
||||
func (s *Scheduler) GetRunningTasks() []string {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
ids := make([]string, 0, len(s.runningTasks))
|
||||
for id := range s.runningTasks {
|
||||
ids = append(ids, id)
|
||||
}
|
||||
return ids
|
||||
}
|
||||
|
||||
// Reload 重新加载配置
|
||||
func (s *Scheduler) Reload(config SchedulerConfig) {
|
||||
s.logger.Infof("[Scheduler] 正在重载配置...")
|
||||
|
||||
// 停止现有 workers
|
||||
close(s.stopCh)
|
||||
s.wg.Wait()
|
||||
|
||||
// 更新配置
|
||||
s.mu.Lock()
|
||||
s.config = config
|
||||
s.taskQueue = make(chan *ExecutionRequest, config.QueueSize)
|
||||
s.rateLimiter = time.Tick(config.RateInterval)
|
||||
s.stopCh = make(chan struct{})
|
||||
s.mu.Unlock()
|
||||
|
||||
// 重启 workers
|
||||
s.Start()
|
||||
|
||||
s.logger.Infof("[Scheduler] 配置已重载: workers=%d, queue=%d, rate=%v",
|
||||
config.WorkerCount, config.QueueSize, config.RateInterval)
|
||||
}
|
||||
|
||||
// GetQueueSize 获取当前队列大小
|
||||
func (s *Scheduler) GetQueueSize() int {
|
||||
return len(s.taskQueue)
|
||||
}
|
||||
|
||||
// GetConfig 获取配置
|
||||
func (s *Scheduler) GetConfig() SchedulerConfig {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
return s.config
|
||||
}
|
||||
Reference in New Issue
Block a user