Fix scheduler cleanup edge cases

agent/internal/scheduler/scheduler.go

@@ -202,13 +202,24 @@ func (w *checkWorker) Wait() {
 }
 
 func (w *checkWorker) loop(initial config.CheckConfig) {
+	cfg := initial
+	stoppedNotified := false
+	notifyStopped := func() {
+		if stoppedNotified {
+			return
+		}
+		stoppedNotified = true
+		if w.hooks.OnStopped != nil {
+			w.hooks.OnStopped(cfg.ID)
+		}
+	}
 	defer func() {
+		notifyStopped()
 		if w.onDone != nil {
 			w.onDone(w)
 		}
 		w.wg.Done()
 	}()
-	cfg := initial
 	t := time.NewTimer(time.Hour)
 	stopTimer(t)
 	defer t.Stop()
@@ -249,9 +260,6 @@ func (w *checkWorker) loop(initial config.CheckConfig) {
 	}
 
 	if !arm(cfg, true) {
-		if w.hooks.OnStopped != nil {
-			w.hooks.OnStopped(cfg.ID)
-		}
 		return
 	}
 
@@ -261,18 +269,12 @@ func (w *checkWorker) loop(initial config.CheckConfig) {
 			stopping = true
 			ctxDone = nil
 			if !running {
-				if w.hooks.OnStopped != nil {
-					w.hooks.OnStopped(cfg.ID)
-				}
 				return
 			}
 		case <-stopCh:
 			stopping = true
 			stopCh = nil
 			if !running {
-				if w.hooks.OnStopped != nil {
-					w.hooks.OnStopped(cfg.ID)
-				}
 				return
 			}
 		case next := <-w.updateCh:
@@ -307,9 +309,6 @@ func (w *checkWorker) loop(initial config.CheckConfig) {
 				}
 			}
 			if stopping {
-				if w.hooks.OnStopped != nil {
-					w.hooks.OnStopped(cfg.ID)
-				}
 				return
 			}
 			if runAfterCurrent {

agent/internal/scheduler/scheduler_test.go

@@ -2,6 +2,7 @@ package scheduler
 
 import (
 	"context"
+	"errors"
 	"strings"
 	"sync/atomic"
 	"testing"
@@ -107,6 +108,38 @@ func TestCronEmitsOnDueSlot(t *testing.T) {
 	m.Wait()
 }
 
+func TestCronArmErrorCallsStoppedHook(t *testing.T) {
+	old := parseCronSchedule
+	var calls int32
+	parseCronSchedule = func(string) (cron.Schedule, error) {
+		if atomic.AddInt32(&calls, 1) == 1 {
+			return fakeCronSchedule{delay: 10 * time.Millisecond}, nil
+		}
+		return nil, errors.New("boom")
+	}
+	t.Cleanup(func() { parseCronSchedule = old })
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+	out := make(chan Event, 10)
+	stopped := make(chan string, 1)
+	m := NewManager(ctx, "a1", out, Hooks{
+		OnStopped: func(id string) { stopped <- id },
+	})
+	m.Update([]config.CheckConfig{mkCronCheck("cron", "true", 100*time.Millisecond)})
+
+	select {
+	case id := <-stopped:
+		if id != "cron" {
+			t.Fatalf("unexpected stopped id: %q", id)
+		}
+	case <-ctx.Done():
+		t.Fatal("stopped hook was not called")
+	}
+	m.Stop()
+	m.Wait()
+}
+
 func TestNoOverlapSkips(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), 600*time.Millisecond)
 	defer cancel()