Guard task state lifecycle and result fields with the result mutex

pkg/scheduler/scheduler.go

@@ -247,6 +247,9 @@ func (ts *TaskScheduler) GetTaskCount() uint64 {
 		return 0
 	}
 
+	ts.taskStateMutex.RLock()
+	defer ts.taskStateMutex.RUnlock()
+
 	return uint64(len(ts.allTasks))
 }
 

pkg/scheduler/task_execution.go

@@ -19,8 +19,10 @@ func (ts *TaskScheduler) ExecuteTask(ctx context.Context, taskIndex types.TaskIn
 
 	taskLogger := taskState.logger.GetLogger()
 
-	// check if task has already been started/executed
+	// check if the task has already been started/executed, and mark it started
+	taskState.resultMutex.Lock()
 	if taskState.isStarted {
+		taskState.resultMutex.Unlock()
 		return fmt.Errorf("task has already been executed")
 	}
 
@@ -30,20 +32,26 @@ func (ts *TaskScheduler) ExecuteTask(ctx context.Context, taskIndex types.TaskIn
 	taskState.taskStatusVars.SetVar("started", true)
 	taskState.taskStatusVars.SetVar("running", true)
 
-	if err := taskState.updateTaskState(); err != nil {
-		taskLogger.Errorf("task state update on db failed: %v", err)
+	updateErr := taskState.updateTaskStateLocked()
+	taskState.resultMutex.Unlock()
+
+	if updateErr != nil {
+		taskLogger.Errorf("task state update on db failed: %v", updateErr)
 	}
 
 	// emit task started event
 	ts.emitTaskStarted(taskState)
 
 	defer func() {
+		taskState.resultMutex.Lock()
 		taskState.isRunning = false
 		taskState.stopTime = time.Now()
 		taskState.taskStatusVars.SetVar("running", false)
+		updateErr := taskState.updateTaskStateLocked()
+		taskState.resultMutex.Unlock()
 
-		if err := taskState.updateTaskState(); err != nil {
-			taskLogger.Errorf("task state update on db failed: %v", err)
+		if updateErr != nil {
+			taskLogger.Errorf("task state update on db failed: %v", updateErr)
 		}
 
 		taskState.logger.Flush()
@@ -67,7 +75,9 @@ func (ts *TaskScheduler) ExecuteTask(ctx context.Context, taskIndex types.TaskIn
 		if !isValid {
 			taskLogger.Infof("task condition not met, skipping task")
 
+			taskState.resultMutex.Lock()
 			taskState.isSkipped = true
+			taskState.resultMutex.Unlock()
 			taskState.setTaskResult(types.TaskResultNone, false)
 
 			return nil
@@ -136,7 +146,9 @@ func (ts *TaskScheduler) ExecuteTask(ctx context.Context, taskIndex types.TaskIn
 		go func() {
 			select {
 			case <-time.After(taskTimeout):
+				taskState.resultMutex.Lock()
 				taskState.isTimeout = true
+				taskState.resultMutex.Unlock()
 				taskState.taskStatusVars.SetVar("timeout", true)
 
 				taskLogger.Warnf("task timed out")
@@ -152,7 +164,9 @@ func (ts *TaskScheduler) ExecuteTask(ctx context.Context, taskIndex types.TaskIn
 		if r := recover(); r != nil {
 			pErr, ok := r.(error)
 			if ok {
+				taskState.resultMutex.Lock()
 				taskState.taskError = pErr
+				taskState.resultMutex.Unlock()
 			}
 
 			taskLogger.Errorf("task execution panic: %v, stack: %v", r, string(debug.Stack()))
@@ -172,9 +186,11 @@ func (ts *TaskScheduler) ExecuteTask(ctx context.Context, taskIndex types.TaskIn
 	if err != nil {
 		taskLogger.Errorf("task execution returned error: %v", err)
 
+		taskState.resultMutex.Lock()
 		if taskState.taskError == nil {
 			taskState.taskError = err
 		}
+		taskState.resultMutex.Unlock()
 	}
 
 	// Persist the shared test-result file (if any task in this run wrote
@@ -187,20 +203,30 @@ func (ts *TaskScheduler) ExecuteTask(ctx context.Context, taskIndex types.TaskIn
 	}
 
 	// set task result
-	if !taskState.updatedResult || taskState.taskResult == types.TaskResultNone {
+	taskState.resultMutex.RLock()
+	resultUnset := !taskState.updatedResult || taskState.taskResult == types.TaskResultNone
+	timedOut := taskState.isTimeout
+	taskState.resultMutex.RUnlock()
+
+	if resultUnset {
 		// set task result if not already done by task
-		if taskState.isTimeout || err != nil {
+		if timedOut || err != nil {
 			taskState.setTaskResult(types.TaskResultFailure, false)
 		} else {
 			taskState.setTaskResult(types.TaskResultSuccess, false)
 		}
 	}
 
-	if taskState.taskResult == types.TaskResultFailure {
-		taskLogger.Warnf("task failed with failure result: %v", taskState.taskError)
+	taskState.resultMutex.RLock()
+	finalResult := taskState.taskResult
+	finalError := taskState.taskError
+	taskState.resultMutex.RUnlock()
+
+	if finalResult == types.TaskResultFailure {
+		taskLogger.Warnf("task failed with failure result: %v", finalError)
 		ts.emitTaskFailed(taskState)
 
-		return fmt.Errorf("task failed: %w", taskState.taskError)
+		return fmt.Errorf("task failed: %w", finalError)
 	}
 
 	taskLogger.Infof("task completed")

pkg/scheduler/task_state.go

@@ -172,7 +172,10 @@ func (ts *TaskScheduler) newTaskState(options *types.TaskOptions, parentState *t
 	return taskState, nil
 }
 
-func (ts *taskState) updateTaskState() error {
+// updateTaskStateLocked persists the current task state to the database.
+// The caller must hold resultMutex: this method reads the volatile lifecycle
+// fields and mutates dbTaskState.
+func (ts *taskState) updateTaskStateLocked() error {
 	if ts.dbTaskState == nil {
 		return nil
 	}
@@ -289,7 +292,7 @@ func (ts *taskState) setTaskResult(result types.TaskResult, setUpdated bool) {
 	ts.taskResult = result
 	ts.taskStatusVars.SetVar("result", uint8(result))
 
-	if err := ts.updateTaskState(); err != nil {
+	if err := ts.updateTaskStateLocked(); err != nil {
 		ts.logger.GetLogger().Errorf("failed to update task state in db: %v", err)
 	}
 
@@ -300,6 +303,9 @@ func (ts *taskState) setTaskResult(result types.TaskResult, setUpdated bool) {
 }
 
 func (ts *taskState) GetTaskStatus() *types.TaskStatus {
+	ts.resultMutex.RLock()
+	defer ts.resultMutex.RUnlock()
+
 	taskStatus := &types.TaskStatus{
 		Index:           ts.index,
 		ParentIndex:     0,
@@ -350,8 +356,8 @@ func (ts *taskState) GetScopeOwner() types.TaskIndex {
 }
 
 func (ts *taskState) GetTaskResultUpdateChan(oldResult types.TaskResult) <-chan bool {
-	ts.resultMutex.RLock()
-	defer ts.resultMutex.RUnlock()
+	ts.resultMutex.Lock()
+	defer ts.resultMutex.Unlock()
 
 	if ts.taskResult != oldResult {
 		return nil

pkg/scheduler/task_state_test.go

@@ -0,0 +1,69 @@
+package scheduler
+
+import (
+	"sync"
+	"testing"
+
+	"github.com/ethpandaops/assertoor/pkg/types"
+	"github.com/ethpandaops/assertoor/pkg/vars"
+)
+
+// TestTaskStateStatusNoRace exercises the concurrent access pattern between the
+// web/watcher readers (GetTaskStatus) and the writers (setTaskResult, SetProgress)
+// on a single task state. All of them must go through resultMutex, otherwise the
+// shared result/progress fields race (run with -race).
+func TestTaskStateStatusNoRace(t *testing.T) {
+	state := &taskState{
+		index:          1,
+		taskStatusVars: vars.NewVariables(nil),
+	}
+
+	var wg sync.WaitGroup
+
+	stop := make(chan struct{})
+
+	// writer: task result transitions
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+
+		for {
+			select {
+			case <-stop:
+				return
+			default:
+				state.setTaskResult(types.TaskResultSuccess, false)
+				state.setTaskResult(types.TaskResultFailure, false)
+			}
+		}
+	}()
+
+	// writer: progress updates
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+
+		for {
+			select {
+			case <-stop:
+				return
+			default:
+				state.SetProgress(50, "half way")
+			}
+		}
+	}()
+
+	// reader: the status snapshot served to watchers and the web api
+	for i := 0; i < 200000; i++ {
+		status := state.GetTaskStatus()
+		_ = status.Result
+		_ = status.Progress
+
+		if status.Error != nil {
+			_ = status.Error.Error()
+		}
+	}
+
+	close(stop)
+	wg.Wait()
+}