优化

2026-08-18 15:04:07 +08:00
parent 10ccc005be
commit 1ef8ecb7b2
42 changed files with 3027 additions and 296 deletions
@@ -3,6 +3,7 @@ package worker
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"log/slog"
 	"sync"
@@ -24,6 +25,10 @@ type TaskScheduler struct {
 	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{}
 }
@@ -34,7 +39,9 @@ type scheduledEntry struct {
 }

 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{}{}}
+	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 }
@@ -44,21 +51,12 @@ func (s *TaskScheduler) Start(ctx context.Context) error {
 	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()
-	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 {
+	if err := s.Reload(ctx); err != nil {
 		s.logger.WarnContext(ctx, "timed task table is not ready", "error", err)
 	}
 	<-ctx.Done()
@@ -67,28 +65,38 @@ func (s *TaskScheduler) Start(ctx context.Context) error {

 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()
-	if s.cancel != nil {
-		s.cancel()
-	}
+	cancel := s.cancel
 	s.ctxMu.Unlock()
+	if cancel != nil {
+		cancel()
+	}
 	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()
+	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)
@@ -101,26 +109,33 @@ func (s *TaskScheduler) Remove(id uint) {

 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()
+	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 {
-			if err = s.Schedule(task); err != nil {
-				return err
+			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
 }
@@ -134,8 +149,58 @@ func (s *TaskScheduler) executionContext() context.Context {
 	return context.Background()
 }

-func (s *TaskScheduler) run(task *biz.TimedTask, trigger string) {
-	log := s.executor.Run(s.executionContext(), task, trigger)
+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)
@@ -146,7 +211,9 @@ func (s *TaskScheduler) run(task *biz.TimedTask, trigger string) {
 		s.logger.Error("timed task finished", attributes...)
 	}
 	if log.Status != "success" {
-		ids, err := s.authorities.AuthorityUserIDs(context.Background(), 888)
+		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
@@ -156,36 +223,73 @@ func (s *TaskScheduler) run(task *biz.TimedTask, trigger string) {
 }

 func (s *TaskScheduler) Schedule(task *biz.TimedTask) error {
-	s.Remove(task.ID)
+	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
 	}
-	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()
+	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
 	}
-	return s.Schedule(task)
+	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 {
@@ -193,7 +297,9 @@ func (s *TaskScheduler) TriggerID(ctx context.Context, id uint) error {
 	if err != nil {
 		return err
 	}
-	s.Trigger(task)
+	if !s.Trigger(task) {
+		return errors.New("任务正在执行或调度器正在停止")
+	}
 	return nil
 }

@@ -211,9 +317,11 @@ func (s *TaskScheduler) NextRuns() map[uint]time.Time {
 	return out
 }

-func (s *TaskScheduler) Trigger(task *biz.TimedTask) {
-	copy := *task
-	go s.run(&copy, "manual")
+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 {