package worker import ( "context" "encoding/json" "errors" "fmt" "kra/internal/biz/system" taskbiz "kra/internal/biz/task" "log/slog" "sync" "time" "github.com/robfig/cron/v3" ) type TaskScheduler struct { tasks *taskbiz.TaskUsecase authorities *system.AuthorityUsecase executor *TaskExecutor logger *slog.Logger standard *cron.Cron seconds *cron.Cron mu sync.Mutex entries map[uint]scheduledEntry lifecycleMu sync.Mutex started bool ctxMu sync.RWMutex runContext context.Context cancel context.CancelFunc runMu sync.Mutex running map[uint]struct{} idle chan struct{} stopping bool subMu sync.RWMutex subscribers map[uint]map[chan []byte]struct{} } type scheduledEntry struct { seconds bool entry cron.EntryID } func NewTaskScheduler(tasks *taskbiz.TaskUsecase, authorities *system.AuthorityUsecase, executor *TaskExecutor, logger *slog.Logger) *TaskScheduler { if logger == nil { logger = slog.Default() } idle := make(chan struct{}) close(idle) return &TaskScheduler{tasks: tasks, authorities: authorities, executor: executor, logger: logger.With("mod", "timedTask"), standard: cron.New(), seconds: cron.New(cron.WithSeconds()), entries: map[uint]scheduledEntry{}, running: map[uint]struct{}{}, idle: idle, subscribers: map[uint]map[chan []byte]struct{}{}} } func NewTaskRuntime(scheduler *TaskScheduler) taskbiz.TaskRuntime { if scheduler == nil { return nil } return scheduler } func NewTaskReloader(scheduler *TaskScheduler) system.TaskReloader { if scheduler == nil { return nil } return scheduler } func (s *TaskScheduler) Start(ctx context.Context) error { if ctx == nil { ctx = context.Background() } runContext, cancel := context.WithCancel(ctx) logger := s.logger if logger == nil { logger = slog.Default() } s.lifecycleMu.Lock() s.runMu.Lock() if s.started || s.stopping { s.runMu.Unlock() s.lifecycleMu.Unlock() cancel() return nil } s.started = true s.stopping = false s.runMu.Unlock() s.ctxMu.Lock() s.runContext, s.cancel = runContext, cancel s.ctxMu.Unlock() s.standard.Start() s.seconds.Start() s.lifecycleMu.Unlock() if err := s.Reload(runContext); err != nil { logger.WarnContext(runContext, "timed task table is not ready", "error", err) } // Stop cancels runContext. Waiting on the child context is important: the // application lifecycle passes a long-lived parent context to Start and // invokes Stop separately during graceful shutdown. <-runContext.Done() if ctx.Err() != nil { cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 5*time.Second) defer cleanupCancel() _ = s.Stop(cleanupCtx) } return nil } func (s *TaskScheduler) Stop(ctx context.Context) error { if ctx == nil { ctx = context.Background() } s.lifecycleMu.Lock() s.runMu.Lock() s.stopping = true s.started = false idle := s.idle if idle == nil { idle = make(chan struct{}) close(idle) } // Keep the lifecycle lock while closing subscriber channels so a new SSE // subscription cannot arrive between the shutdown flag and the close. s.closeSubscribers() s.runMu.Unlock() s.ctxMu.Lock() cancel := s.cancel s.ctxMu.Unlock() s.lifecycleMu.Unlock() if cancel != nil { cancel() } standardDone, secondsDone := s.standard.Stop().Done(), s.seconds.Stop().Done() for _, done := range []<-chan struct{}{standardDone, secondsDone, idle} { select { case <-done: case <-ctx.Done(): return ctx.Err() } } return nil } func (s *TaskScheduler) Remove(id uint) { s.mu.Lock() defer s.mu.Unlock() s.removeLocked(id) } func (s *TaskScheduler) removeLocked(id uint) { if old, ok := s.entries[id]; ok { if old.seconds { s.seconds.Remove(old.entry) } else { s.standard.Remove(old.entry) } delete(s.entries, id) } } func (s *TaskScheduler) Reload(ctx context.Context) error { if ctx == nil { ctx = context.Background() } if s.tasks == nil { return errors.New("任务用例未初始化") } s.mu.Lock() defer s.mu.Unlock() items, _, err := s.tasks.ListTasks(ctx, 0, 0, nil) if err != nil { return err } if err := ctx.Err(); err != nil { return err } type preparedTask struct { task *taskbiz.TimedTask schedule cron.Schedule } prepared := make([]preparedTask, 0, len(items)) for _, task := range items { if task == nil { return errors.New("任务列表包含空任务") } if task.Enabled { schedule, parseErr := parseTaskSchedule(task) if parseErr != nil { return fmt.Errorf("定时任务 %d 的 cron 表达式非法: %w", task.ID, parseErr) } prepared = append(prepared, preparedTask{task: cloneTimedTask(task), schedule: schedule}) } } s.runMu.Lock() if err := ctx.Err(); err != nil { s.runMu.Unlock() return err } if s.stopping { s.runMu.Unlock() return errors.New("任务调度器正在停止") } for id := range s.entries { s.removeLocked(id) } for _, item := range prepared { s.scheduleLocked(item.task, item.schedule) } s.runMu.Unlock() logger := s.logger if logger == nil { logger = slog.Default() } logger.InfoContext(ctx, "定时任务重载完成", "task_count", len(items)) return nil } func (s *TaskScheduler) executionContext() context.Context { s.ctxMu.RLock() defer s.ctxMu.RUnlock() if s.runContext != nil { return s.runContext } return context.Background() } func (s *TaskScheduler) beginRun(id uint) (context.Context, bool) { ctx := s.executionContext() s.runMu.Lock() defer s.runMu.Unlock() if s.stopping || ctx.Err() != nil { return nil, false } if s.running == nil { s.running = map[uint]struct{}{} } if _, exists := s.running[id]; exists { return nil, false } if len(s.running) == 0 { s.idle = make(chan struct{}) } s.running[id] = struct{}{} return ctx, true } func (s *TaskScheduler) finishRun(id uint) { s.runMu.Lock() defer s.runMu.Unlock() if _, exists := s.running[id]; !exists { return } delete(s.running, id) if len(s.running) == 0 && s.idle != nil { close(s.idle) } } func (s *TaskScheduler) dispatch(task *taskbiz.TimedTask, trigger string, async bool) bool { if task == nil { return false } ctx, ok := s.beginRun(task.ID) if !ok { logger := s.logger if logger == nil { logger = slog.Default() } logger.Warn("timed task skipped because it is already running or the scheduler is stopping", "task_id", task.ID, "task_name", task.Name, "trigger_type", trigger) return false } run := func() { defer s.finishRun(task.ID) s.run(ctx, task, trigger) } if async { go run() return true } run() return true } func (s *TaskScheduler) run(ctx context.Context, task *taskbiz.TimedTask, trigger string) { if ctx == nil { ctx = context.Background() } logger := s.logger if logger == nil { logger = slog.Default() } var taskID uint var taskName string if task != nil { taskID, taskName = task.ID, task.Name } if s.executor == nil { logger.Error("timed task execution skipped because executor is unavailable", "task_id", taskID, "task_name", taskName, "trigger_type", trigger) return } log := s.executor.Run(ctx, task, trigger) attributes := []any{"task_id", log.TaskID, "task_name", log.TaskName, "trigger_type", log.TriggerType, "status", log.Status, "duration_ms", log.DurationMS, "started_at", log.StartedAt, "finished_at", log.FinishedAt} if log.ErrorMsg != "" { attributes = append(attributes, "error", log.ErrorMsg) } if log.Status == "success" { logger.Info("timed task finished", attributes...) } else { logger.Error("timed task finished", attributes...) } if log.Status != "success" { if s.authorities == nil { logger.Warn("timed task alert skipped because authority service is unavailable", attributes...) return } alertCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 3*time.Second) defer cancel() ids, err := s.authorities.AuthorityUserIDs(alertCtx, system.SuperAdminAuthorityID) if err != nil { logger.Error("query timed task alert recipients failed", "error", err) return } s.PublishToUsers(ids, map[string]any{"taskId": log.TaskID, "name": log.TaskName, "error": log.ErrorMsg, "time": time.Now().Format(time.RFC3339)}) } } func parseTaskSchedule(task *taskbiz.TimedTask) (cron.Schedule, error) { if task.WithSeconds { return cron.NewParser(cron.Second | cron.Minute | cron.Hour | cron.Dom | cron.Month | cron.Dow | cron.Descriptor).Parse(task.Spec) } return cron.ParseStandard(task.Spec) } func cloneTimedTask(task *taskbiz.TimedTask) *taskbiz.TimedTask { clone := *task clone.Params = append([]byte(nil), task.Params...) clone.HTTPHeader = append([]byte(nil), task.HTTPHeader...) return &clone } func (s *TaskScheduler) scheduleLocked(task *taskbiz.TimedTask, schedule cron.Schedule) { run := func() { s.dispatch(task, "auto", false) } if task.WithSeconds { id := s.seconds.Schedule(schedule, cron.FuncJob(run)) s.entries[task.ID] = scheduledEntry{seconds: true, entry: id} } else { id := s.standard.Schedule(schedule, cron.FuncJob(run)) s.entries[task.ID] = scheduledEntry{entry: id} } } func (s *TaskScheduler) ScheduleID(ctx context.Context, id uint) error { if ctx == nil { ctx = context.Background() } if s.tasks == nil { return errors.New("任务用例未初始化") } s.mu.Lock() defer s.mu.Unlock() task, err := s.tasks.FindTask(ctx, id) if err != nil { return err } if task == nil { return errors.New("任务不存在") } var schedule cron.Schedule if task.Enabled { schedule, err = parseTaskSchedule(task) if err != nil { return err } } s.runMu.Lock() stopping := s.stopping s.runMu.Unlock() if stopping { return errors.New("任务调度器正在停止") } s.removeLocked(task.ID) if task.Enabled { s.scheduleLocked(cloneTimedTask(task), schedule) } return nil } func (s *TaskScheduler) TriggerID(ctx context.Context, id uint) error { if ctx == nil { ctx = context.Background() } if s.tasks == nil { return errors.New("任务用例未初始化") } task, err := s.tasks.FindTask(ctx, id) if err != nil { return err } if task == nil { return errors.New("任务不存在") } if !s.Trigger(task) { return errors.New("任务正在执行或调度器正在停止") } return nil } func (s *TaskScheduler) NextRuns() map[uint]time.Time { s.mu.Lock() defer s.mu.Unlock() out := map[uint]time.Time{} for id, item := range s.entries { if item.seconds { out[id] = s.seconds.Entry(item.entry).Next } else { out[id] = s.standard.Entry(item.entry).Next } } return out } func (s *TaskScheduler) Trigger(task *taskbiz.TimedTask) bool { if task == nil { return false } return s.dispatch(cloneTimedTask(task), "manual", true) } func (s *TaskScheduler) Subscribe(userID uint) chan []byte { ch := make(chan []byte, 16) s.runMu.Lock() if s.stopping { close(ch) s.runMu.Unlock() return ch } s.subMu.Lock() if s.subscribers == nil { s.subscribers = make(map[uint]map[chan []byte]struct{}) } if s.subscribers[userID] == nil { s.subscribers[userID] = map[chan []byte]struct{}{} } if len(s.subscribers[userID]) >= 10 { for old := range s.subscribers[userID] { delete(s.subscribers[userID], old) close(old) break } } s.subscribers[userID][ch] = struct{}{} s.subMu.Unlock() s.runMu.Unlock() return ch } func (s *TaskScheduler) Unsubscribe(userID uint, ch chan []byte) { s.subMu.Lock() if subscribers := s.subscribers[userID]; subscribers != nil { if _, ok := subscribers[ch]; !ok { s.subMu.Unlock() return } delete(subscribers, ch) if len(subscribers) == 0 { delete(s.subscribers, userID) } close(ch) } s.subMu.Unlock() } func (s *TaskScheduler) closeSubscribers() { s.subMu.Lock() defer s.subMu.Unlock() for userID, subscribers := range s.subscribers { for ch := range subscribers { close(ch) } delete(s.subscribers, userID) } } func (s *TaskScheduler) PublishToUsers(userIDs []uint, value any) { raw, err := json.Marshal(value) if err != nil { raw = []byte(fmt.Sprint(value)) } s.subMu.RLock() defer s.subMu.RUnlock() for _, userID := range userIDs { for ch := range s.subscribers[userID] { select { case ch <- raw: default: } } } }