431 lines
10 KiB
Go
431 lines
10 KiB
Go
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
|
|
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 { return scheduler }
|
|
|
|
func NewTaskReloader(scheduler *TaskScheduler) system.TaskReloader { return scheduler }
|
|
|
|
func (s *TaskScheduler) Start(ctx context.Context) error {
|
|
if ctx == nil {
|
|
ctx = context.Background()
|
|
}
|
|
runContext, cancel := context.WithCancel(ctx)
|
|
s.ctxMu.Lock()
|
|
s.runContext, s.cancel = runContext, cancel
|
|
s.ctxMu.Unlock()
|
|
s.runMu.Lock()
|
|
if s.stopping {
|
|
s.runMu.Unlock()
|
|
cancel()
|
|
return nil
|
|
}
|
|
s.stopping = false
|
|
s.standard.Start()
|
|
s.seconds.Start()
|
|
if err := s.Reload(runContext); err != nil {
|
|
s.logger.WarnContext(runContext, "timed task table is not ready", "error", err)
|
|
}
|
|
s.runMu.Unlock()
|
|
// 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.runMu.Lock()
|
|
s.stopping = true
|
|
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.subMu.Lock()
|
|
for userID, subscribers := range s.subscribers {
|
|
for ch := range subscribers {
|
|
close(ch)
|
|
}
|
|
delete(s.subscribers, userID)
|
|
}
|
|
s.subMu.Unlock()
|
|
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 *taskbiz.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 *taskbiz.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 *taskbiz.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 *taskbiz.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 *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 {
|
|
copy := *task
|
|
copy.Params = append([]byte(nil), task.Params...)
|
|
copy.HTTPHeader = append([]byte(nil), task.HTTPHeader...)
|
|
return ©
|
|
}
|
|
|
|
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 {
|
|
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 *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
|
|
}
|
|
if s.subscribers == nil {
|
|
s.subscribers = make(map[uint]map[chan []byte]struct{})
|
|
}
|
|
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()
|
|
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:
|
|
}
|
|
}
|
|
}
|
|
}
|