Guard scheduler cancel context fields and honor early aborts

pkg/scheduler/scheduler.go

@@ -23,11 +23,19 @@ type TaskScheduler struct {
 	rootTasks        []types.TaskIndex
 	allCleanupTasks  []types.TaskIndex
 	rootCleanupTasks []types.TaskIndex
-	testRunCtx       context.Context
 	taskStateMutex   sync.RWMutex
 	taskStateMap     map[types.TaskIndex]*taskState
-	cancelTaskCtx    context.CancelFunc
-	cancelCleanupCtx context.CancelFunc
+
+	// runCtxMutex guards testRunCtx, cancelTaskCtx, cancelCleanupCtx, and the
+	// pending-cancel flags below. RunTasks (the test worker goroutine) writes
+	// them; CancelTasks (an HTTP abort handler) and GetTestRunCtx (task
+	// goroutines) read them from other goroutines.
+	runCtxMutex            sync.Mutex
+	testRunCtx             context.Context
+	cancelTaskCtx          context.CancelFunc
+	cancelCleanupCtx       context.CancelFunc
+	cancelRequested        bool
+	cancelRequestedCleanup bool
 
 	testResultMutex sync.Mutex
 	testResultDir   string
@@ -56,6 +64,9 @@ func (ts *TaskScheduler) GetTestRunID() uint64 {
 }
 
 func (ts *TaskScheduler) GetTestRunCtx() context.Context {
+	ts.runCtxMutex.Lock()
+	defer ts.runCtxMutex.Unlock()
+
 	return ts.testRunCtx
 }
 
@@ -167,11 +178,9 @@ func (ts *TaskScheduler) AddCleanupTask(options *types.TaskOptions) (types.TaskI
 func (ts *TaskScheduler) RunTasks(testRunCtx context.Context, timeout time.Duration) error {
 	var cleanupCtx, tasksCtx context.Context
 
-	cleanupCtx, ts.cancelCleanupCtx = context.WithCancel(testRunCtx)
-
-	defer ts.cleanupTestResult()
-	defer ts.runCleanupTasks(cleanupCtx)
+	ts.runCtxMutex.Lock()
 
+	cleanupCtx, ts.cancelCleanupCtx = context.WithCancel(testRunCtx)
 	ts.testRunCtx = testRunCtx
 
 	if timeout > 0 {
@@ -180,7 +189,27 @@ func (ts *TaskScheduler) RunTasks(testRunCtx context.Context, timeout time.Durat
 		tasksCtx, ts.cancelTaskCtx = context.WithCancel(testRunCtx)
 	}
 
-	defer ts.cancelTaskCtx()
+	cancelTaskCtx, cancelCleanupCtx := ts.cancelTaskCtx, ts.cancelCleanupCtx
+	cancelRequested, cancelRequestedCleanup := ts.cancelRequested, ts.cancelRequestedCleanup
+	ts.cancelRequested, ts.cancelRequestedCleanup = false, false
+
+	ts.runCtxMutex.Unlock()
+
+	defer ts.cleanupTestResult()
+	defer ts.runCleanupTasks(cleanupCtx)
+	defer cancelTaskCtx()
+
+	// Honor an abort that was requested before these contexts existed.
+	// Without this, CancelTasks had nothing to cancel yet and silently
+	// dropped the request, letting every root task run to completion
+	// despite the operator having already asked to stop.
+	if cancelRequested {
+		cancelTaskCtx()
+
+		if cancelRequestedCleanup {
+			cancelCleanupCtx()
+		}
+	}
 
 	for _, task := range ts.rootTasks {
 		err := ts.ExecuteTask(tasksCtx, task, nil)
@@ -211,12 +240,26 @@ func (ts *TaskScheduler) runCleanupTasks(ctx context.Context) {
 }
 
 func (ts *TaskScheduler) CancelTasks(cancelCleanup bool) {
-	if ts.cancelTaskCtx != nil {
-		ts.cancelTaskCtx()
+	ts.runCtxMutex.Lock()
+	defer ts.runCtxMutex.Unlock()
+
+	if ts.cancelTaskCtx == nil {
+		// RunTasks hasn't reached its context setup yet. Record the
+		// request so RunTasks honors it as soon as the contexts exist,
+		// instead of it being silently dropped.
+		ts.cancelRequested = true
 
 		if cancelCleanup {
-			ts.cancelCleanupCtx()
+			ts.cancelRequestedCleanup = true
 		}
+
+		return
+	}
+
+	ts.cancelTaskCtx()
+
+	if cancelCleanup {
+		ts.cancelCleanupCtx()
 	}
 }
 

pkg/scheduler/scheduler_test.go

@@ -0,0 +1,191 @@
+package scheduler
+
+import (
+	"context"
+	"io"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/ethpandaops/assertoor/pkg/clients"
+	"github.com/ethpandaops/assertoor/pkg/db"
+	"github.com/ethpandaops/assertoor/pkg/events"
+	"github.com/ethpandaops/assertoor/pkg/logger"
+	"github.com/ethpandaops/assertoor/pkg/names"
+	"github.com/ethpandaops/assertoor/pkg/txmgr"
+	"github.com/ethpandaops/assertoor/pkg/types"
+	"github.com/ethpandaops/assertoor/pkg/vars"
+	"github.com/sirupsen/logrus"
+)
+
+type fakeServices struct{}
+
+func (f *fakeServices) Database() *db.Database                { return nil }
+func (f *fakeServices) ClientPool() *clients.ClientPool       { return nil }
+func (f *fakeServices) WalletManager() *txmgr.Spamoor         { return nil }
+func (f *fakeServices) ValidatorNames() *names.ValidatorNames { return nil }
+func (f *fakeServices) EventBus() *events.EventBus            { return nil }
+
+// blockingTask waits for its context to be cancelled, or gives up after a
+// short deadline and reports that it ran to completion uninterrupted.
+type blockingTask struct {
+	wasCancelled    atomic.Bool
+	ranToCompletion atomic.Bool
+}
+
+func (b *blockingTask) Config() interface{}    { return nil }
+func (b *blockingTask) Timeout() time.Duration { return 0 }
+func (b *blockingTask) LoadConfig() error      { return nil }
+
+func (b *blockingTask) Execute(ctx context.Context) error {
+	select {
+	case <-ctx.Done():
+		b.wasCancelled.Store(true)
+		return ctx.Err()
+	case <-time.After(150 * time.Millisecond):
+		b.ranToCompletion.Store(true)
+		return nil
+	}
+}
+
+func addBlockingRootTask(ts *TaskScheduler) *blockingTask {
+	task := &blockingTask{}
+	descriptor := &types.TaskDescriptor{
+		Name: "test-blocking-task",
+		NewTask: func(_ *types.TaskContext, _ *types.TaskOptions) (types.Task, error) {
+			return task, nil
+		},
+	}
+
+	ts.taskStateMutex.Lock()
+	ts.taskCount++
+	idx := ts.taskCount
+
+	state := &taskState{
+		ts:             ts,
+		index:          idx,
+		options:        &types.TaskOptions{},
+		descriptor:     descriptor,
+		taskVars:       vars.NewVariables(nil),
+		logger:         logger.NewLogger(&logger.ScopeOptions{Parent: ts.logger}),
+		taskOutputs:    vars.NewVariables(nil),
+		taskStatusVars: vars.NewVariables(nil),
+	}
+
+	ts.taskStateMap[idx] = state
+	ts.allTasks = append(ts.allTasks, idx)
+	ts.rootTasks = append(ts.rootTasks, idx)
+	ts.taskStateMutex.Unlock()
+
+	return task
+}
+
+func newTestScheduler() *TaskScheduler {
+	log := logrus.New()
+	log.SetOutput(io.Discard)
+
+	return NewTaskScheduler(log, &fakeServices{}, vars.NewVariables(nil), 1)
+}
+
+// TestCancelTasksBeforeRunTasksIsHonored covers the abort-endpoint race: an
+// operator can request cancellation between a test being registered and its
+// worker goroutine reaching RunTasks. Before the fix, CancelTasks found
+// cancelTaskCtx still nil and silently did nothing, so every root task ran
+// to completion despite the abort.
+func TestCancelTasksBeforeRunTasksIsHonored(t *testing.T) {
+	ts := newTestScheduler()
+	task := addBlockingRootTask(ts)
+
+	ts.CancelTasks(true)
+
+	if err := ts.RunTasks(context.Background(), 0); err == nil {
+		t.Fatalf("expected RunTasks to return the cancellation error")
+	}
+
+	if !task.wasCancelled.Load() {
+		t.Fatalf("expected the root task to observe cancellation requested before RunTasks started")
+	}
+
+	if task.ranToCompletion.Load() {
+		t.Fatalf("expected the root task to be interrupted, not run to completion")
+	}
+}
+
+// TestCancelTasksAfterRunTasksStartedIsHonored is the pre-existing case:
+// once RunTasks has assigned cancelTaskCtx, CancelTasks correctly
+// interrupts the running root task.
+func TestCancelTasksAfterRunTasksStartedIsHonored(t *testing.T) {
+	ts := newTestScheduler()
+	task := addBlockingRootTask(ts)
+
+	var wg sync.WaitGroup
+
+	wg.Add(1)
+
+	go func() {
+		defer wg.Done()
+
+		_ = ts.RunTasks(context.Background(), 0)
+	}()
+
+	time.Sleep(30 * time.Millisecond)
+
+	ts.CancelTasks(true)
+	wg.Wait()
+
+	if !task.wasCancelled.Load() {
+		t.Fatalf("expected the root task to observe cancellation")
+	}
+
+	if task.ranToCompletion.Load() {
+		t.Fatalf("expected the root task to be interrupted, not run to completion")
+	}
+}
+
+// TestRunTasksNoAbortRunsNormally guards against the pending-cancel flag
+// leaking into a run that was never aborted.
+func TestRunTasksNoAbortRunsNormally(t *testing.T) {
+	ts := newTestScheduler()
+	task := addBlockingRootTask(ts)
+
+	if err := ts.RunTasks(context.Background(), 0); err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	if task.wasCancelled.Load() {
+		t.Fatalf("task should not have been cancelled")
+	}
+
+	if !task.ranToCompletion.Load() {
+		t.Fatalf("task should have run to completion")
+	}
+}
+
+// TestCancelCtxFieldsNoRace exercises RunTasks, CancelTasks, and
+// GetTestRunCtx concurrently under -race to guard the field-level locking.
+func TestCancelCtxFieldsNoRace(t *testing.T) {
+	ts := newTestScheduler()
+	addBlockingRootTask(ts)
+
+	var wg sync.WaitGroup
+
+	wg.Add(2)
+
+	go func() {
+		defer wg.Done()
+
+		_ = ts.RunTasks(context.Background(), 0)
+	}()
+
+	go func() {
+		defer wg.Done()
+
+		for i := 0; i < 200; i++ {
+			ts.CancelTasks(true)
+			_ = ts.GetTestRunCtx()
+		}
+	}()
+
+	wg.Wait()
+}