fix: cron scheduler exec context to avoid canceled watcher ctx (#6977)

* Fix cron scheduler exec context to avoid canceled watcher ctx

* add context cancel tests

Author: povilasv

cmd/api-server/main.go

@@ -784,11 +784,8 @@ func main() {
 		leaderTasks = append(leaderTasks, leader.Task{
 			Name: "cron-scheduler",
 			Start: func(taskCtx context.Context) error {
+				scheduleManager.Start()
 				go func() {
-					// Start the schedule manager.
-					scheduleManager.Start()
-					// If we're no longer the leader then stop the manager.
-					// This probably won't happen as losing leadership likely means we died.
 					<-taskCtx.Done()
 					scheduleManager.Stop()
 				}()

internal/cronjob/robfig/manager.go

@@ -28,28 +28,36 @@ type Manager struct {
 	cron           *cron.Cron
 	cronEntries    map[string]map[string]cron.EntryID
 	executor       Executor
+	execCtx        context.Context
+	execCancel     context.CancelFunc
 }
 
-func New(logger *zap.SugaredLogger, executor Executor, proModeEnabled bool) Manager {
-	return Manager{
+func New(logger *zap.SugaredLogger, executor Executor, proModeEnabled bool) *Manager {
+	ctx, cancel := context.WithCancel(context.Background())
+	return &Manager{
 		proModeEnabled: proModeEnabled,
 		logger:         logger,
 		cron:           cron.New(),
 		cronEntries:    make(map[string]map[string]cron.EntryID),
 		executor:       executor,
+		execCtx:        ctx,
+		execCancel:     cancel,
 	}
 }
 
 // Start the cron manager in its own goroutine, or no-op if already started.
-func (m Manager) Start() {
+func (m *Manager) Start() {
 	m.logger.Infow("cron manager starting")
+	m.execCancel()
+	m.execCtx, m.execCancel = context.WithCancel(context.Background())
 	m.cron.Start()
 	m.logger.Infow("cron manager started")
 }
 
 // Stop stops the cron manager if it is running; otherwise it does nothing.
-func (m Manager) Stop() {
+func (m *Manager) Stop() {
 	m.logger.Infow("cron manager stopping")
+	m.execCancel()
 	m.cron.Stop()
 	m.logger.Infow("cron manager stopped")
 }
@@ -62,7 +70,7 @@ func cronSpec(config testkube.TestWorkflowCronJobConfig) string {
 	return spec
 }
 
-func (m Manager) ReplaceWorkflowSchedules(ctx context.Context, workflow cronjob.Workflow, configs []testkube.TestWorkflowCronJobConfig) error {
+func (m *Manager) ReplaceWorkflowSchedules(ctx context.Context, workflow cronjob.Workflow, configs []testkube.TestWorkflowCronJobConfig) error {
 	log := m.logger.With("workflow", workflow.Name)
 	// Delete all existing schedules for this workflow.
 	// This is because we may not know when a schedule is removed from
@@ -93,7 +101,7 @@ func (m Manager) ReplaceWorkflowSchedules(ctx context.Context, workflow cronjob.
 				"cron", config.Cron,
 			)
 		}
-		entryId, err := m.cron.AddJob(spec, m.testWorkflowExecuteJob(ctx, workflow.Name, spec, config))
+		entryId, err := m.cron.AddJob(spec, m.testWorkflowExecuteJob(workflow.Name, spec, config))
 		if err != nil {
 			m.logger.Errorw("Error adding cron for workflow, continuing processing",
 				"cron", spec,
@@ -115,8 +123,9 @@ func (m Manager) ReplaceWorkflowSchedules(ctx context.Context, workflow cronjob.
 	return nil
 }
 
-func (m Manager) testWorkflowExecuteJob(ctx context.Context, workflow, cronSpec string, config testkube.TestWorkflowCronJobConfig) cron.FuncJob {
+func (m *Manager) testWorkflowExecuteJob(workflow, cronSpec string, config testkube.TestWorkflowCronJobConfig) cron.FuncJob {
 	return cron.FuncJob(func() {
+		execCtx := m.execCtx
 		var targets []*cloud.ExecutionTarget
 		if config.Target != nil {
 			targets = commonmapper.MapAllTargetsApiToGrpc([]testkube.ExecutionTarget{*config.Target})
@@ -141,7 +150,7 @@ func (m Manager) testWorkflowExecuteJob(ctx context.Context, workflow, cronSpec
 		)
 		log.Info("executing scheduled workflow")
 
-		results, err := m.executor.Execute(ctx, request)
+		results, err := m.executor.Execute(execCtx, request)
 		if err != nil {
 			log.Errorw("unable to execute scheduled workflow",
 				"error", err)

internal/cronjob/robfig/manager_test.go

@@ -0,0 +1,90 @@
+package robfig
+
+import (
+	"context"
+	"testing"
+
+	"go.uber.org/zap"
+
+	"github.com/kubeshop/testkube/pkg/api/v1/testkube"
+	"github.com/kubeshop/testkube/pkg/cloud"
+	"github.com/kubeshop/testkube/pkg/cronjob"
+)
+
+type recordingExecutor struct {
+	calls int
+	ctx   context.Context
+	req   *cloud.ScheduleRequest
+}
+
+func (e *recordingExecutor) Execute(ctx context.Context, req *cloud.ScheduleRequest) ([]testkube.TestWorkflowExecution, error) {
+	e.calls++
+	e.ctx = ctx
+	e.req = req
+	return nil, nil
+}
+
+func TestManagerStartStopResetsExecContext(t *testing.T) {
+	logger := zap.NewNop().Sugar()
+	executor := &recordingExecutor{}
+	manager := New(logger, executor, false)
+
+	oldCtx := manager.execCtx
+	manager.Start()
+
+	select {
+	case <-oldCtx.Done():
+	default:
+		t.Fatal("expected previous exec context to be canceled on start")
+	}
+
+	if err := manager.execCtx.Err(); err != nil {
+		t.Fatalf("expected new exec context to be active, got %v", err)
+	}
+
+	currentCtx := manager.execCtx
+	manager.Stop()
+
+	select {
+	case <-currentCtx.Done():
+	default:
+		t.Fatal("expected exec context to be canceled on stop")
+	}
+}
+
+func TestJobUsesManagerContext(t *testing.T) {
+	logger := zap.NewNop().Sugar()
+	executor := &recordingExecutor{}
+	manager := New(logger, executor, false)
+
+	canceledCtx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	workflow := cronjob.Workflow{Name: "workflow-a", EnvId: "env-1"}
+	config := testkube.TestWorkflowCronJobConfig{Cron: "* * * * *"}
+
+	if err := manager.ReplaceWorkflowSchedules(canceledCtx, workflow, []testkube.TestWorkflowCronJobConfig{config}); err != nil {
+		t.Fatalf("unexpected error setting schedules: %v", err)
+	}
+
+	entryID := manager.cronEntries[workflow.Name][cronSpec(config)]
+	entry := manager.cron.Entry(entryID)
+	if entry.Job == nil {
+		t.Fatal("expected scheduled job to be registered")
+	}
+
+	entry.Job.Run()
+
+	if executor.calls != 1 {
+		t.Fatalf("expected executor to be called once, got %d", executor.calls)
+	}
+	if executor.ctx == nil {
+		t.Fatal("expected executor context to be set")
+	}
+	if err := executor.ctx.Err(); err != nil {
+		t.Fatalf("expected executor context to be active, got %v", err)
+	}
+	if executor.ctx == canceledCtx {
+		t.Fatal("expected executor to use manager context, got canceled caller context")
+	}
+}