diff --git a/pkg/scheduler/scheduler.go b/pkg/scheduler/scheduler.go index 5846e345..435802e7 100644 --- a/pkg/scheduler/scheduler.go +++ b/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() } } diff --git a/pkg/scheduler/scheduler_test.go b/pkg/scheduler/scheduler_test.go new file mode 100644 index 00000000..ecc584b2 --- /dev/null +++ b/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() +}