Fix data races on OutstandingAction result state

pkg/actions/actions.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"slices"
 	"sort"
 	"sync"
 	"time"
@@ -17,6 +18,7 @@ import (
 	"go.uber.org/zap"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
+	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/structpb"
 )
 
@@ -45,41 +47,111 @@ func NewOutstandingAction(id, name string) *OutstandingAction {
 func (oa *OutstandingAction) SetStatus(ctx context.Context, status v2.BatonActionStatus) {
 	oa.Lock()
 	defer oa.Unlock()
-	l := ctxzap.Extract(ctx).With(
-		zap.String("action_id", oa.Id),
-		zap.String("action_name", oa.Name),
-		zap.String("status", status.String()),
-	)
+	oa.setStatusLocked(ctx, status)
+}
+
+// setStatusLocked applies a lifecycle transition and reports whether it was
+// accepted. Terminal statuses are final and RUNNING is only reachable from
+// PENDING; rejected transitions are dropped. These fire in normal operation
+// (a handler finishing after its request was cancelled), hence debug level.
+// It requires oa's mutex to be held.
+func (oa *OutstandingAction) setStatusLocked(ctx context.Context, status v2.BatonActionStatus) bool {
 	if oa.Status == v2.BatonActionStatus_BATON_ACTION_STATUS_COMPLETE || oa.Status == v2.BatonActionStatus_BATON_ACTION_STATUS_FAILED {
-		l.Error("cannot set status on completed action")
+		ctxzap.Extract(ctx).Debug("dropping status transition on terminal action",
+			zap.String("action_id", oa.Id),
+			zap.String("action_name", oa.Name),
+			zap.String("status", oa.Status.String()),
+			zap.String("requested_status", status.String()))
+		return false
 	}
 	if status == v2.BatonActionStatus_BATON_ACTION_STATUS_RUNNING && oa.Status != v2.BatonActionStatus_BATON_ACTION_STATUS_PENDING {
-		l.Error("cannot set status to running unless action is pending")
+		ctxzap.Extract(ctx).Debug("dropping running transition on non-pending action",
+			zap.String("action_id", oa.Id),
+			zap.String("action_name", oa.Name),
+			zap.String("status", oa.Status.String()))
+		return false
 	}
 
 	oa.Status = status
+	return true
 }
 
-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.
+// A FAILED action may refresh its error payload — a late handler error is a
+// richer account than an earlier cancellation — but a COMPLETE action stays
+// complete.
 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()
+	if oa.Status == v2.BatonActionStatus_BATON_ACTION_STATUS_FAILED {
+		oa.setErrorLocked(err)
+		return
+	}
+	if oa.setStatusLocked(ctx, v2.BatonActionStatus_BATON_ACTION_STATUS_FAILED) {
+		oa.setErrorLocked(err)
+	}
+}
+
+// Result returns the action's identity and current outcome. The snapshot is
+// internally consistent; the returned message and annotations are owned by
+// the action and must not be modified.
+func (oa *OutstandingAction) Result() (string, v2.BatonActionStatus, *structpb.Struct, annotations.Annotations) {
+	return oa.result()
+}
+
+// result is the unexported form of Result.
+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. The
+// values are cloned at this publication seam: the handler owns what it
+// returned and may keep mutating it. Terminal statuses are final — a late
+// completion after a cancellation is dropped, except that a FAILED action
+// takes a late handler failure as a replacement account.
+func (oa *OutstandingAction) setOutcome(ctx context.Context, rv *structpb.Struct, annos annotations.Annotations, err error) {
+	if rv != nil {
+		rv = proto.Clone(rv).(*structpb.Struct)
+	}
+	annos = slices.Clone(annos)
+
+	oa.Lock()
+	defer oa.Unlock()
+
+	if err != nil {
+		if oa.Status == v2.BatonActionStatus_BATON_ACTION_STATUS_FAILED || oa.setStatusLocked(ctx, v2.BatonActionStatus_BATON_ACTION_STATUS_FAILED) {
+			oa.Rv = rv
+			oa.Annos = annos
+			oa.setErrorLocked(err)
+		}
+		return
+	}
+
+	if oa.setStatusLocked(ctx, v2.BatonActionStatus_BATON_ACTION_STATUS_COMPLETE) {
+		oa.Rv = rv
+		oa.Annos = annos
+	}
 }
 
 const maxOldActions = 1000
@@ -163,8 +235,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 +470,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 +530,29 @@ 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():
+		// The handler may have finished in the same instant; prefer its
+		// completed result over a spurious cancellation return.
+		select {
+		case <-done:
+			id, st, rv, annos := oa.result()
+			return id, st, rv, annos, nil
+		default:
+		}
 		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 +626,30 @@ 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():
+		// The handler may have finished in the same instant; prefer its
+		// completed result over a spurious cancellation return.
+		select {
+		case <-done:
+			id, st, rv, annos := oa.result()
+			return id, st, rv, annos, nil
+		default:
+		}
 		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()
 	}
 }