package worker import ( "context" "encoding/json" "fmt" "log/slog" "sync" "time" "github.com/robfig/cron/v3" "kra/internal/biz" ) type TaskScheduler struct { tasks *biz.TaskUsecase authorities *biz.AuthorityUsecase executor *TaskExecutor logger *slog.Logger standard *cron.Cron seconds *cron.Cron mu sync.Mutex entries map[uint]scheduledEntry ctxMu sync.RWMutex runContext context.Context cancel context.CancelFunc subMu sync.RWMutex subscribers map[uint]map[chan []byte]struct{} } type scheduledEntry struct { seconds bool entry cron.EntryID } func NewTaskScheduler(tasks *biz.TaskUsecase, authorities *biz.AuthorityUsecase, executor *TaskExecutor, logger *slog.Logger) *TaskScheduler { 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{}, subscribers: map[uint]map[chan []byte]struct{}{}} } func NewTaskRuntime(scheduler *TaskScheduler) biz.TaskRuntime { return scheduler } func (s *TaskScheduler) Start(ctx context.Context) error { runContext, cancel := context.WithCancel(ctx) s.ctxMu.Lock() s.runContext, s.cancel = runContext, cancel s.ctxMu.Unlock() s.standard.Start() s.seconds.Start() items, _, err := s.tasks.ListTasks(ctx, 0, 0, nil) if err == nil { for _, task := range items { if task.Enabled { if scheduleErr := s.Schedule(task); scheduleErr != nil { s.logger.ErrorContext(ctx, "restore timed task failed", "id", task.ID, "error", scheduleErr) } } } // Report the number of rows read from the database, not only // the enabled rows that were scheduled successfully. s.logger.InfoContext(ctx, "定时任务加载完成", "task_count", len(items)) } else { s.logger.WarnContext(ctx, "timed task table is not ready", "error", err) } <-ctx.Done() return nil } func (s *TaskScheduler) Stop(ctx context.Context) error { s.closeSubscribers() s.ctxMu.Lock() if s.cancel != nil { s.cancel() } s.ctxMu.Unlock() standardDone, secondsDone := s.standard.Stop().Done(), s.seconds.Stop().Done() select { case <-standardDone: case <-ctx.Done(): return ctx.Err() } select { case <-secondsDone: return nil case <-ctx.Done(): return ctx.Err() } } func (s *TaskScheduler) Remove(id uint) { s.mu.Lock() defer s.mu.Unlock() 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 { s.mu.Lock() for id, old := range s.entries { if old.seconds { s.seconds.Remove(old.entry) } else { s.standard.Remove(old.entry) } delete(s.entries, id) } s.mu.Unlock() items, _, err := s.tasks.ListTasks(ctx, 0, 0, nil) if err != nil { return err } for _, task := range items { if task.Enabled { if err = s.Schedule(task); err != nil { return err } } } s.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) run(task *biz.TimedTask, trigger string) { log := s.executor.Run(s.executionContext(), 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" { s.logger.Info("timed task finished", attributes...) } else { s.logger.Error("timed task finished", attributes...) } if log.Status != "success" { ids, err := s.authorities.AuthorityUserIDs(context.Background(), 888) if err != nil { s.logger.Error("query timed task alert recipients failed", "error", err) return } s.PublishToUsers(ids, map[string]any{"taskId": task.ID, "name": task.Name, "error": log.ErrorMsg, "time": time.Now().Format(time.RFC3339)}) } } func (s *TaskScheduler) Schedule(task *biz.TimedTask) error { s.Remove(task.ID) if !task.Enabled { return nil } copy := *task run := func() { s.run(©, "auto") } var id cron.EntryID var err error if task.WithSeconds { id, err = s.seconds.AddFunc(task.Spec, run) } else { id, err = s.standard.AddFunc(task.Spec, run) } if err != nil { return err } s.mu.Lock() s.entries[task.ID] = scheduledEntry{seconds: task.WithSeconds, entry: id} s.mu.Unlock() return nil } func (s *TaskScheduler) ScheduleID(ctx context.Context, id uint) error { task, err := s.tasks.FindTask(ctx, id) if err != nil { return err } return s.Schedule(task) } func (s *TaskScheduler) TriggerID(ctx context.Context, id uint) error { task, err := s.tasks.FindTask(ctx, id) if err != nil { return err } s.Trigger(task) 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 *biz.TimedTask) { copy := *task go s.run(©, "manual") } func (s *TaskScheduler) Subscribe(userID uint) chan []byte { ch := make(chan []byte, 16) s.subMu.Lock() 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() 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: } } } }