Fix data races on OutstandingAction result state

pkg/actions/actions.go

@@ -17,6 +17,8 @@ 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/anypb"
 	"google.golang.org/protobuf/types/known/structpb"
 )
 
@@ -31,6 +33,10 @@ type OutstandingAction struct {
 	Err       error
 	StartedAt time.Time
 	sync.Mutex
+
+	// cancelled marks a FAILED status that came from request cancellation
+	// rather than from the handler; the handler's late outcome replaces it.
+	cancelled bool
 }
 
 func NewOutstandingAction(id, name string) *OutstandingAction {
@@ -45,41 +51,167 @@ 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, but a COMPLETE action stays
+// complete. This is a real handler failure, so it clears any provisional
+// cancellation mark.
 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)
+		oa.cancelled = false
+		return
+	}
+	if oa.setStatusLocked(ctx, v2.BatonActionStatus_BATON_ACTION_STATUS_FAILED) {
+		oa.setErrorLocked(err)
+	}
+}
+
+// setCancelled records a cancellation as a provisional failure: unlike other
+// terminal states, the handler's own late outcome replaces it, since the
+// action's side effects can still complete after the request goes away.
+func (oa *OutstandingAction) setCancelled(ctx context.Context, err error) {
+	oa.Lock()
+	defer oa.Unlock()
+	if oa.setStatusLocked(ctx, v2.BatonActionStatus_BATON_ACTION_STATUS_FAILED) {
+		oa.setErrorLocked(err)
+		oa.cancelled = true
+	}
+}
+
+// isProvisional reports whether the action's FAILED status is a provisional
+// cancellation mark that the handler's late outcome may still replace.
+func (oa *OutstandingAction) isProvisional() bool {
+	oa.Lock()
+	defer oa.Unlock()
+	return oa.cancelled
+}
+
+// evictable reports whether the action has reached a state cleanup may
+// remove: terminal, and not a provisional cancellation whose record can
+// still improve.
+func (oa *OutstandingAction) evictable() bool {
+	oa.Lock()
+	defer oa.Unlock()
+	if oa.cancelled {
+		return false
+	}
+	return oa.Status == v2.BatonActionStatus_BATON_ACTION_STATUS_COMPLETE || oa.Status == v2.BatonActionStatus_BATON_ACTION_STATUS_FAILED
+}
+
+// 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, with one
+// exception: a cancellation-FAILED status is provisional, and the handler's
+// own outcome — success or failure — replaces it.
+func (oa *OutstandingAction) setOutcome(ctx context.Context, rv *structpb.Struct, annos annotations.Annotations, err error) {
+	if rv != nil {
+		rv = proto.Clone(rv).(*structpb.Struct)
+	}
+	if annos != nil {
+		annosCopy := make(annotations.Annotations, len(annos))
+		for i, a := range annos {
+			if a != nil {
+				annosCopy[i] = proto.Clone(a).(*anypb.Any)
+			}
+		}
+		annos = annosCopy
+	}
+
+	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)
+			oa.cancelled = false
+		}
+		return
+	}
+
+	if oa.cancelled {
+		// Deliberate cross-terminal replacement: the cancellation was a
+		// transport event, not the action's outcome.
+		ctxzap.Extract(ctx).Debug("replacing provisional cancellation with handler outcome",
+			zap.String("action_id", oa.Id),
+			zap.String("action_name", oa.Name))
+		oa.cancelled = false
+		oa.Status = v2.BatonActionStatus_BATON_ACTION_STATUS_COMPLETE
+		oa.Rv = rv
+		oa.Annos = annos
+		oa.Err = nil
+		return
+	}
+
+	if oa.setStatusLocked(ctx, v2.BatonActionStatus_BATON_ACTION_STATUS_COMPLETE) {
+		oa.Rv = rv
+		oa.Annos = annos
+	}
 }
 
 const maxOldActions = 1000
@@ -163,8 +295,7 @@ 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 {
+		if actionList[i].evictable() {
 			count++
 			delete(a.actions, actionList[i].Id)
 		}
@@ -398,7 +529,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 +589,34 @@ 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()
+		// 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.setCancelled(ctx, ctx.Err())
+		id, st, rv, annos := oa.result()
+		if st == v2.BatonActionStatus_BATON_ACTION_STATUS_COMPLETE {
+			// The handler won the race to the lock; its completed result is
+			// the authoritative pairing, not the cancellation.
+			return id, st, rv, annos, nil
+		}
+		return id, st, rv, annos, ctx.Err()
 	}
 }
 
@@ -547,24 +690,35 @@ 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()
+		// 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.setCancelled(ctx, ctx.Err())
+		id, st, rv, annos := oa.result()
+		if st == v2.BatonActionStatus_BATON_ACTION_STATUS_COMPLETE {
+			// The handler won the race to the lock; its completed result is
+			// the authoritative pairing, not the cancellation.
+			return id, st, rv, annos, nil
+		}
+		return id, st, rv, annos, ctx.Err()
 	}
 }