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
	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[chan []byte]struct{}
}

type scheduledEntry struct {
	seconds bool
	entry   cron.EntryID
}

func NewTaskScheduler(tasks *biz.TaskUsecase, executor *TaskExecutor, logger *slog.Logger) *TaskScheduler {
	return &TaskScheduler{tasks: tasks, executor: executor, logger: logger, standard: cron.New(), seconds: cron.New(cron.WithSeconds()), entries: map[uint]scheduledEntry{}, subscribers: map[chan []byte]struct{}{}}
}

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)
				}
			}
		}
	} 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.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
			}
		}
	}
	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)
	if log.Status != "success" {
		s.Broadcast(map[string]any{"taskId": task.ID, "taskName": task.Name, "status": log.Status, "errorMsg": log.ErrorMsg, "time": time.Now()})
	}
}

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() chan []byte {
	ch := make(chan []byte, 16)
	s.subMu.Lock()
	s.subscribers[ch] = struct{}{}
	s.subMu.Unlock()
	return ch
}
func (s *TaskScheduler) Unsubscribe(ch chan []byte) {
	s.subMu.Lock()
	if _, ok := s.subscribers[ch]; ok {
		delete(s.subscribers, ch)
		close(ch)
	}
	s.subMu.Unlock()
}
func (s *TaskScheduler) Broadcast(value any) {
	raw, err := json.Marshal(value)
	if err != nil {
		raw = []byte(fmt.Sprint(value))
	}
	s.subMu.RLock()
	defer s.subMu.RUnlock()
	for ch := range s.subscribers {
		select {
		case ch <- raw:
		default:
		}
	}
}