Fix data races on OutstandingAction result state

pkg/actions/actions.go

@@ -45,6 +45,11 @@ func NewOutstandingAction(id, name string) *OutstandingAction {
 func (oa *OutstandingAction) SetStatus(ctx context.Context, status v2.BatonActionStatus) {
 	oa.Lock()
 	defer oa.Unlock()
+	oa.setStatusLocked(ctx, status)
+}
+
+// setStatusLocked requires oa's mutex to be held.
+func (oa *OutstandingAction) setStatusLocked(ctx context.Context, status v2.BatonActionStatus) {
 	l := ctxzap.Extract(ctx).With(
 		zap.String("action_id", oa.Id),
 		zap.String("action_name", oa.Name),
@@ -60,26 +65,59 @@ func (oa *OutstandingAction) SetStatus(ctx context.Context, status v2.BatonActio
 	oa.Status = status
 }
 
-func (oa *OutstandingAction) setError(_ context.Context, err error) {
-	oa.Lock()
-	defer oa.Unlock()
-	if oa.Rv == nil {
-		oa.Rv = &structpb.Struct{}
-	}
-	if oa.Rv.Fields == nil {
-		oa.Rv.Fields = make(map[string]*structpb.Value)
+// setErrorLocked requires oa's mutex to be held.
+func (oa *OutstandingAction) setErrorLocked(err error) {
+	// Rebuild rather than mutate: a concurrent caller may hold the previously
+	// published struct from a result() snapshot and be marshaling it.
+	fields := make(map[string]*structpb.Value, len(oa.Rv.GetFields())+1)
+	for k, v := range oa.Rv.GetFields() {
+		fields[k] = v
 	}
-	oa.Rv.Fields["error"] = &structpb.Value{
+	fields["error"] = &structpb.Value{
 		Kind: &structpb.Value_StringValue{
 			StringValue: err.Error(),
 		},
 	}
+	oa.Rv = &structpb.Struct{Fields: fields}
 	oa.Err = err
 }
 
+// SetError records the error and marks the action failed in one critical
+// section, so no snapshot can observe the error with a non-terminal status.
 func (oa *OutstandingAction) SetError(ctx context.Context, err error) {
-	oa.setError(ctx, err)
-	oa.SetStatus(ctx, v2.BatonActionStatus_BATON_ACTION_STATUS_FAILED)
+	oa.Lock()
+	defer oa.Unlock()
+	oa.setErrorLocked(err)
+	oa.setStatusLocked(ctx, v2.BatonActionStatus_BATON_ACTION_STATUS_FAILED)
+}
+
+// Result returns the action's identity and current outcome as a consistent
+// snapshot; the handler goroutine may still be mutating the action.
+func (oa *OutstandingAction) Result() (string, v2.BatonActionStatus, *structpb.Struct, annotations.Annotations) {
+	return oa.result()
+}
+
+// result returns the action's identity and current outcome as a consistent
+// snapshot; the handler goroutine may still be mutating the action.
+func (oa *OutstandingAction) result() (string, v2.BatonActionStatus, *structpb.Struct, annotations.Annotations) {
+	oa.Lock()
+	defer oa.Unlock()
+	return oa.Id, oa.Status, oa.Rv, oa.Annos
+}
+
+// setOutcome publishes the handler's result and terminal status in one
+// critical section, so no snapshot observes one without the other.
+func (oa *OutstandingAction) setOutcome(ctx context.Context, rv *structpb.Struct, annos annotations.Annotations, err error) {
+	oa.Lock()
+	defer oa.Unlock()
+	oa.Rv = rv
+	oa.Annos = annos
+	if err != nil {
+		oa.setErrorLocked(err)
+		oa.setStatusLocked(ctx, v2.BatonActionStatus_BATON_ACTION_STATUS_FAILED)
+		return
+	}
+	oa.setStatusLocked(ctx, v2.BatonActionStatus_BATON_ACTION_STATUS_COMPLETE)
 }
 
 const maxOldActions = 1000
@@ -163,8 +201,8 @@ func (a *ActionManager) CleanupOldActions(ctx context.Context) {
 	count := 0
 	// Delete the oldest actions
 	for i := 0; i < len(actionList)-maxOldActions; i++ {
-		action := actionList[i]
-		if action.Status == v2.BatonActionStatus_BATON_ACTION_STATUS_COMPLETE || action.Status == v2.BatonActionStatus_BATON_ACTION_STATUS_FAILED {
+		_, actionStatus, _, _ := actionList[i].result()
+		if actionStatus == v2.BatonActionStatus_BATON_ACTION_STATUS_COMPLETE || actionStatus == v2.BatonActionStatus_BATON_ACTION_STATUS_FAILED {
 			count++
 			delete(a.actions, actionList[i].Id)
 		}
@@ -398,7 +436,8 @@ func (a *ActionManager) GetActionStatus(_ context.Context, actionId string) (v2.
 
 	// Don't return oa.Err here because error is for GetActionStatus, not the action itself.
 	// oa.Rv contains any error.
-	return oa.Status, oa.Name, oa.Rv, oa.Annos, nil
+	_, st, rv, annos := oa.result()
+	return st, oa.Name, rv, annos, nil
 }
 
 // InvokeAction invokes an action. If resourceTypeID is set, it invokes a resource-scoped action.
@@ -457,23 +496,21 @@ func (a *ActionManager) invokeGlobalAction(ctx context.Context, name string, arg
 		bgCtx := trace.ContextWithSpanContext(context.Background(), trace.SpanContextFromContext(ctx))
 		handlerCtx, cancel := context.WithTimeoutCause(bgCtx, 1*time.Hour, errors.New("action handler timed out"))
 		defer cancel()
-		var oaErr error
-		oa.Rv, oa.Annos, oaErr = handler(handlerCtx, args)
-		if oaErr == nil {
-			oa.SetStatus(ctx, v2.BatonActionStatus_BATON_ACTION_STATUS_COMPLETE)
-		} else {
-			oa.SetError(ctx, oaErr)
-		}
+		rv, annos, oaErr := handler(handlerCtx, args)
+		oa.setOutcome(ctx, rv, annos, oaErr)
 	}()
 
 	select {
 	case <-done:
-		return oa.Id, oa.Status, oa.Rv, oa.Annos, nil
+		id, st, rv, annos := oa.result()
+		return id, st, rv, annos, nil
 	case <-time.After(1 * time.Second):
-		return oa.Id, oa.Status, oa.Rv, oa.Annos, nil
+		id, st, rv, annos := oa.result()
+		return id, st, rv, annos, nil
 	case <-ctx.Done():
 		oa.SetError(ctx, ctx.Err())
-		return oa.Id, oa.Status, oa.Rv, oa.Annos, ctx.Err()
+		id, st, rv, annos := oa.result()
+		return id, st, rv, annos, ctx.Err()
 	}
 }
 
@@ -547,24 +584,22 @@ func (a *ActionManager) invokeResourceAction(
 		bgCtx = ctxzap.ToContext(bgCtx, ctxzap.Extract(ctx))
 		handlerCtx, cancel := context.WithTimeoutCause(bgCtx, 1*time.Hour, errors.New("action handler timed out"))
 		defer cancel()
-		var oaErr error
-		oa.Rv, oa.Annos, oaErr = handler(handlerCtx, args)
-		if oaErr == nil {
-			oa.SetStatus(ctx, v2.BatonActionStatus_BATON_ACTION_STATUS_COMPLETE)
-		} else {
-			oa.SetError(ctx, oaErr)
-		}
+		rv, annos, oaErr := handler(handlerCtx, args)
+		oa.setOutcome(ctx, rv, annos, oaErr)
 	}()
 
 	// Wait for completion or timeout
 	select {
 	case <-done:
-		return oa.Id, oa.Status, oa.Rv, oa.Annos, nil
+		id, st, rv, annos := oa.result()
+		return id, st, rv, annos, nil
 	case <-time.After(1 * time.Second):
-		return oa.Id, oa.Status, oa.Rv, oa.Annos, nil
+		id, st, rv, annos := oa.result()
+		return id, st, rv, annos, nil
 	case <-ctx.Done():
 		oa.SetError(ctx, ctx.Err())
-		return oa.Id, oa.Status, oa.Rv, oa.Annos, ctx.Err()
+		id, st, rv, annos := oa.result()
+		return id, st, rv, annos, ctx.Err()
 	}
 }
 

pkg/actions/actions_test.go

@@ -546,3 +546,57 @@ func TestActionHandlerGoroutineLeaks(t *testing.T) {
 		require.LessOrEqual(t, finalCount, initialCount+1, "goroutine leak detected after context cancellation")
 	})
 }
+
+// The data race this test guards against is only detectable under -race,
+// which the plain CI go-test job does not enable for this package; `make
+// race-check` is the out-of-band gate that runs it. The assertions at the end
+// only cover the exported snapshot accessor.
+func TestCleanupOldActionsDuringConcurrentStatusWrites(t *testing.T) {
+	ctx := t.Context()
+	m := NewActionManager(ctx)
+
+	// The cleanup loop only visits the len(actions)-maxOldActions oldest
+	// entries, so the concurrent writer must target the first-created action.
+	// The sort is unstable and StartedAt values can tie, so push the target
+	// strictly earlier to make its position deterministic.
+	oldest := m.GetNewAction("churn")
+	oldest.StartedAt = time.Now().Add(-time.Hour)
+	for i := 0; i < maxOldActions; i++ {
+		m.GetNewAction("churn")
+	}
+
+	stop := make(chan struct{})
+	started := make(chan struct{})
+	writerDone := make(chan struct{})
+	go func() {
+		defer close(writerDone)
+		first := true
+		for {
+			select {
+			case <-stop:
+				return
+			default:
+				oldest.SetStatus(ctx, v2.BatonActionStatus_BATON_ACTION_STATUS_RUNNING)
+				if first {
+					close(started)
+					first = false
+				}
+			}
+		}
+	}()
+
+	// Wait for the writer's first write: otherwise it may only be scheduled
+	// after close(stop), never race cleanup, and leave the action PENDING.
+	<-started
+
+	// Fails under -race if cleanup reads action status without the lock.
+	m.CleanupOldActions(ctx)
+
+	close(stop)
+	<-writerDone
+
+	// The exported snapshot accessor reads the same state race-free.
+	id, actionStatus, _, _ := oldest.Result()
+	require.Equal(t, oldest.Id, id)
+	require.Equal(t, v2.BatonActionStatus_BATON_ACTION_STATUS_RUNNING, actionStatus)
+}