优化结构

internal/worker/task_runtime_test.go

@@ -272,3 +272,39 @@ func TestStopCancelsAndWaitsForActiveRuns(t *testing.T) {
 		t.Fatal("scheduler accepted a new run after Stop()")
 	}
 }
+
+func TestStartReturnsAfterStop(t *testing.T) {
+	scheduler := newTestTaskScheduler(&workerTaskRepo{})
+	started := make(chan error, 1)
+	go func() {
+		started <- scheduler.Start(context.Background())
+	}()
+
+	deadline := time.Now().Add(time.Second)
+	for {
+		scheduler.ctxMu.RLock()
+		ready := scheduler.cancel != nil
+		scheduler.ctxMu.RUnlock()
+		if ready {
+			break
+		}
+		if time.Now().After(deadline) {
+			t.Fatal("scheduler did not start")
+		}
+		time.Sleep(time.Millisecond)
+	}
+
+	stopCtx, stopCancel := context.WithTimeout(context.Background(), time.Second)
+	defer stopCancel()
+	if err := scheduler.Stop(stopCtx); err != nil {
+		t.Fatalf("Stop() error = %v", err)
+	}
+	select {
+	case err := <-started:
+		if err != nil {
+			t.Fatalf("Start() error = %v", err)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("Start() did not return after Stop()")
+	}
+}

internal/worker/task_scheduler.go

@@ -40,6 +40,9 @@ type scheduledEntry struct {
 }
 
 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{}{}}
@@ -50,24 +53,42 @@ func NewTaskRuntime(scheduler *TaskScheduler) taskbiz.TaskRuntime { return sched
 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.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)
+	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)
 	}
-	<-ctx.Done()
 	return nil
 }
 
 func (s *TaskScheduler) Stop(ctx context.Context) error {
-	s.closeSubscribers()
+	if ctx == nil {
+		ctx = context.Background()
+	}
 	s.runMu.Lock()
 	s.stopping = true
 	idle := s.idle
@@ -75,6 +96,16 @@ func (s *TaskScheduler) Stop(ctx context.Context) error {
 		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
@@ -329,6 +360,15 @@ func (s *TaskScheduler) Trigger(task *taskbiz.TimedTask) bool {
 
 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{}{}
@@ -342,6 +382,7 @@ func (s *TaskScheduler) Subscribe(userID uint) chan []byte {
 	}
 	s.subscribers[userID][ch] = struct{}{}
 	s.subMu.Unlock()
+	s.runMu.Unlock()
 	return ch
 }
 func (s *TaskScheduler) Unsubscribe(userID uint, ch chan []byte) {