kra-oa/internal/modules/system/worker/task_scheduler.go

package worker

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"sync"
	"time"

	"github.com/robfig/cron/v3"
	"kra/internal/modules/system/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
	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 *biz.TaskUsecase, authorities *biz.AuthorityUsecase, executor *TaskExecutor, logger *slog.Logger) *TaskScheduler {
	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) 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.runMu.Lock()
	s.stopping = false
	s.runMu.Unlock()
	s.standard.Start()
	s.seconds.Start()
	if err := s.Reload(ctx); err != nil {
		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.runMu.Lock()
	s.stopping = true
	idle := s.idle
	if idle == nil {
		idle = make(chan struct{})
		close(idle)
	}
	s.runMu.Unlock()
	s.ctxMu.Lock()
	cancel := s.cancel
	s.ctxMu.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 {
	s.mu.Lock()
	defer s.mu.Unlock()

	items, _, err := s.tasks.ListTasks(ctx, 0, 0, nil)
	if err != nil {
		return err
	}
	type preparedTask struct {
		task     *biz.TimedTask
		schedule cron.Schedule
	}
	prepared := make([]preparedTask, 0, len(items))
	for _, task := range items {
		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})
		}
	}

	for id := range s.entries {
		s.removeLocked(id)
	}
	for _, item := range prepared {
		s.scheduleLocked(item.task, item.schedule)
	}
	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) 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 *biz.TimedTask, trigger string, async bool) bool {
	ctx, ok := s.beginRun(task.ID)
	if !ok {
		s.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 *biz.TimedTask, trigger string) {
	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" {
		s.logger.Info("timed task finished", attributes...)
	} else {
		s.logger.Error("timed task finished", attributes...)
	}
	if log.Status != "success" {
		alertCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 3*time.Second)
		defer cancel()
		ids, err := s.authorities.AuthorityUserIDs(alertCtx, 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 {
	if task == nil {
		return errors.New("定时任务不能为空")
	}
	var schedule cron.Schedule
	var err error
	if task.Enabled {
		schedule, err = parseTaskSchedule(task)
		if err != nil {
			return err
		}
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	s.removeLocked(task.ID)
	if !task.Enabled {
		return nil
	}
	s.scheduleLocked(cloneTimedTask(task), schedule)
	return nil
}

func parseTaskSchedule(task *biz.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 *biz.TimedTask) *biz.TimedTask {
	copy := *task
	copy.Params = append([]byte(nil), task.Params...)
	copy.HTTPHeader = append([]byte(nil), task.HTTPHeader...)
	return &copy
}

func (s *TaskScheduler) scheduleLocked(task *biz.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 {
	s.mu.Lock()
	defer s.mu.Unlock()
	task, err := s.tasks.FindTask(ctx, id)
	if err != nil {
		return err
	}
	var schedule cron.Schedule
	if task.Enabled {
		schedule, err = parseTaskSchedule(task)
		if err != nil {
			return err
		}
	}
	s.removeLocked(task.ID)
	if task.Enabled {
		s.scheduleLocked(cloneTimedTask(task), schedule)
	}
	return nil
}

func (s *TaskScheduler) TriggerID(ctx context.Context, id uint) error {
	task, err := s.tasks.FindTask(ctx, id)
	if err != nil {
		return err
	}
	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 *biz.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.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:
			}
		}
	}
}