kra-oa/internal/worker/task_scheduler.go

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(&copy, "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(&copy, "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:
			}
		}
	}
}