Nexus workflow operation cancellation for CHASM (#9915)

## What changed?

Full implementation of the cancellation path for workflow Nexus
operations on CHASM.
Including:
- State machine logic
- Component logic
- Workflow event handlers
- Integration into workflows
- Executors

I also ended up doing quite a bit of refactoring to reuse as much of the
code as possible across the cancellation and invocation executors.
Also added missing test coverage for the invocation executors.

I tried to structure this into separate commits that make reviews easy
but ended up with a lot of fixes in later commits that touched code in
earlier commits. Best to review in one piece.

Author: bergundy

chasm/lib/nexusoperation/cancellation.go

@@ -1,8 +1,14 @@
 package nexusoperation
 
 import (
+	"time"
+
+	commonpb "go.temporal.io/api/common/v1"
+	failurepb "go.temporal.io/api/failure/v1"
+	"go.temporal.io/api/serviceerror"
 	"go.temporal.io/server/chasm"
-	"go.temporal.io/server/chasm/lib/nexusoperation/gen/nexusoperationpb/v1"
+	nexusoperationpb "go.temporal.io/server/chasm/lib/nexusoperation/gen/nexusoperationpb/v1"
+	"go.temporal.io/server/common/backoff"
 )
 
 var _ chasm.Component = (*Cancellation)(nil)
@@ -14,15 +20,18 @@ type Cancellation struct {
 
 	// Persisted internal state
 	*nexusoperationpb.CancellationState
+
+	// Operation is a pointer to the parent Operation component.
+	Operation chasm.ParentPtr[*Operation]
 }
 
 func newCancellation(state *nexusoperationpb.CancellationState) *Cancellation {
 	return &Cancellation{CancellationState: state}
 }
 
 // LifecycleState maps the cancellation's status to a CHASM lifecycle state.
-func (o *Cancellation) LifecycleState(_ chasm.Context) chasm.LifecycleState {
-	switch o.Status {
+func (c *Cancellation) LifecycleState(_ chasm.Context) chasm.LifecycleState {
+	switch c.Status {
 	case nexusoperationpb.CANCELLATION_STATUS_SUCCEEDED:
 		return chasm.LifecycleStateCompleted
 	case nexusoperationpb.CANCELLATION_STATUS_FAILED,
@@ -34,11 +43,106 @@ func (o *Cancellation) LifecycleState(_ chasm.Context) chasm.LifecycleState {
 }
 
 // StateMachineState returns the current cancellation status.
-func (o *Cancellation) StateMachineState() nexusoperationpb.CancellationStatus {
-	return o.Status
+func (c *Cancellation) StateMachineState() nexusoperationpb.CancellationStatus {
+	return c.Status
 }
 
 // SetStateMachineState sets the cancellation status.
-func (o *Cancellation) SetStateMachineState(status nexusoperationpb.CancellationStatus) {
-	o.Status = status
+func (c *Cancellation) SetStateMachineState(status nexusoperationpb.CancellationStatus) {
+	c.Status = status
+}
+
+// cancelArgs holds the arguments needed to cancel a Nexus operation.
+type cancelArgs struct {
+	service                string
+	operation              string
+	token                  string
+	requestID              string
+	endpointName           string
+	endpointID             string
+	currentTime            time.Time
+	scheduledTime          time.Time
+	startedTime            time.Time
+	scheduleToCloseTimeout time.Duration
+	startToCloseTimeout    time.Duration
+	headers                map[string]string
+	payload                *commonpb.Payload
+}
+
+func (c *Cancellation) onCompleted(ctx chasm.MutableContext) error {
+	op := c.Operation.Get(ctx)
+	if store, ok := op.Store.TryGet(ctx); ok {
+		return store.OnNexusOperationCancellationCompleted(ctx, op)
+	}
+	return TransitionCancellationSucceeded.Apply(c, ctx, EventCancellationSucceeded{})
+}
+
+func (c *Cancellation) onFailed(ctx chasm.MutableContext, failure *failurepb.Failure) error {
+	op := c.Operation.Get(ctx)
+	if store, ok := op.Store.TryGet(ctx); ok {
+		return store.OnNexusOperationCancellationFailed(ctx, op, failure)
+	}
+	return TransitionCancellationFailed.Apply(c, ctx, EventCancellationFailed{
+		Failure: failure,
+	})
+}
+
+// loadArgs loads the cancel arguments from the cancellation and its parent operation.
+func (c *Cancellation) loadArgs(
+	ctx chasm.Context,
+	_ chasm.NoValue,
+) (cancelArgs, error) {
+	op := c.Operation.Get(ctx)
+
+	var invocationData InvocationData
+	if store, ok := op.Store.TryGet(ctx); ok {
+		var err error
+		invocationData, err = store.NexusOperationInvocationData(ctx, op)
+		if err != nil {
+			return cancelArgs{}, err
+		}
+	}
+	// TODO: For standalone operations, load invocation data from the operation state.
+
+	return cancelArgs{
+		service:                op.GetService(),
+		operation:              op.GetOperation(),
+		token:                  op.GetOperationToken(),
+		requestID:              op.GetRequestId(),
+		endpointName:           op.GetEndpoint(),
+		endpointID:             op.GetEndpointId(),
+		currentTime:            ctx.Now(c),
+		scheduledTime:          op.GetScheduledTime().AsTime(),
+		startedTime:            op.GetStartedTime().AsTime(),
+		scheduleToCloseTimeout: op.GetScheduleToCloseTimeout().AsDuration(),
+		startToCloseTimeout:    op.GetStartToCloseTimeout().AsDuration(),
+		headers:                invocationData.Header,
+		payload:                invocationData.Input,
+	}, nil
+}
+
+// saveCancellationResultInput is the input to the Cancellation.saveResult method.
+type saveCancellationResultInput struct {
+	result      cancellationResult
+	retryPolicy func() backoff.RetryPolicy
+}
+
+// saveResult applies the outcome of a cancel operation call to the cancellation state machine.
+func (c *Cancellation) saveResult(
+	ctx chasm.MutableContext,
+	input saveCancellationResultInput,
+) (chasm.NoValue, error) {
+	switch r := input.result.(type) {
+	case cancellationResultOK:
+		return nil, c.onCompleted(ctx)
+	case cancellationResultFail:
+		return nil, c.onFailed(ctx, r.failure)
+	case cancellationResultRetry:
+		return nil, transitionCancellationAttemptFailed.Apply(c, ctx, EventCancellationAttemptFailed{
+			Failure:     r.failure,
+			RetryPolicy: input.retryPolicy(),
+		})
+	default:
+		return nil, serviceerror.NewInternalf("cannot save cancellation result of type %T", r)
+	}
 }

chasm/lib/nexusoperation/cancellation_statemachine.go

@@ -11,13 +11,20 @@ import (
 // EventCancellationScheduled is triggered when cancellation is meant to be scheduled for the first time - immediately
 // after it has been requested.
 type EventCancellationScheduled struct {
+	// Destination is the endpoint name for the cancellation task.
+	// Must be provided by the caller because ParentPtr is not available during inline creation.
+	Destination string
 }
 
 var TransitionCancellationScheduled = chasm.NewTransition(
 	[]nexusoperationpb.CancellationStatus{nexusoperationpb.CANCELLATION_STATUS_UNSPECIFIED},
 	nexusoperationpb.CANCELLATION_STATUS_SCHEDULED,
 	func(c *Cancellation, ctx chasm.MutableContext, event EventCancellationScheduled) error {
-		ctx.AddTask(c, chasm.TaskAttributes{}, &nexusoperationpb.CancellationTask{
+		c.Attempt++
+
+		ctx.AddTask(c, chasm.TaskAttributes{
+			Destination: event.Destination,
+		}, &nexusoperationpb.CancellationTask{
 			Attempt: c.Attempt,
 		})
 
@@ -34,9 +41,12 @@ var transitionCancellationRescheduled = chasm.NewTransition(
 	[]nexusoperationpb.CancellationStatus{nexusoperationpb.CANCELLATION_STATUS_BACKING_OFF},
 	nexusoperationpb.CANCELLATION_STATUS_SCHEDULED,
 	func(c *Cancellation, ctx chasm.MutableContext, event EventCancellationRescheduled) error {
+		c.Attempt++
 		c.NextAttemptScheduleTime = nil
 
-		ctx.AddTask(c, chasm.TaskAttributes{}, &nexusoperationpb.CancellationTask{
+		ctx.AddTask(c, chasm.TaskAttributes{
+			Destination: c.Operation.Get(ctx).GetEndpoint(),
+		}, &nexusoperationpb.CancellationTask{
 			Attempt: c.Attempt,
 		})
 
@@ -56,7 +66,6 @@ var transitionCancellationAttemptFailed = chasm.NewTransition(
 	func(c *Cancellation, ctx chasm.MutableContext, event EventCancellationAttemptFailed) error {
 		currentTime := ctx.Now(c)
 
-		c.Attempt++
 		c.LastAttemptCompleteTime = timestamppb.New(currentTime)
 		c.LastAttemptFailure = event.Failure
 
@@ -76,17 +85,22 @@ var transitionCancellationAttemptFailed = chasm.NewTransition(
 
 // EventCancellationFailed is triggered when a cancellation attempt is failed with a non retryable error.
 type EventCancellationFailed struct {
+	Failure *failurepb.Failure
 }
 
 var TransitionCancellationFailed = chasm.NewTransition(
 	[]nexusoperationpb.CancellationStatus{
 		// We can immediately transition to failed since we don't know how to send a cancellation request for an
 		// unstarted operation.
+		// TODO: This doesn't seem to happen in either the HSM or CHASM implementations.
 		nexusoperationpb.CANCELLATION_STATUS_UNSPECIFIED,
 		nexusoperationpb.CANCELLATION_STATUS_SCHEDULED,
 	},
 	nexusoperationpb.CANCELLATION_STATUS_FAILED,
 	func(c *Cancellation, ctx chasm.MutableContext, event EventCancellationFailed) error {
+		currentTime := ctx.Now(c)
+		c.LastAttemptCompleteTime = timestamppb.New(currentTime)
+		c.LastAttemptFailure = event.Failure
 		// Terminal state - no tasks to emit.
 		return nil
 	},
@@ -100,6 +114,10 @@ var TransitionCancellationSucceeded = chasm.NewTransition(
 	[]nexusoperationpb.CancellationStatus{nexusoperationpb.CANCELLATION_STATUS_SCHEDULED},
 	nexusoperationpb.CANCELLATION_STATUS_SUCCEEDED,
 	func(c *Cancellation, ctx chasm.MutableContext, event EventCancellationSucceeded) error {
+		currentTime := ctx.Now(c)
+		c.LastAttemptCompleteTime = timestamppb.New(currentTime)
+		c.LastAttemptFailure = nil
+
 		// Terminal state - no tasks to emit.
 		return nil
 	},