Nexus Standalone: Poll (#9780)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: stephanos

chasm/lib/nexusoperation/config.go

@@ -12,6 +12,21 @@ import (
 	"go.temporal.io/server/common/rpc/interceptor"
 )
 
+var LongPollTimeout = dynamicconfig.NewNamespaceDurationSetting(
+	"nexusoperation.longPollTimeout",
+	20*time.Second,
+	`Maximum timeout for nexus operation long-poll requests. Actual wait may be shorter to leave
+longPollBuffer before the caller deadline.`,
+)
+
+var LongPollBuffer = dynamicconfig.NewNamespaceDurationSetting(
+	"nexusoperation.longPollBuffer",
+	time.Second,
+	`A buffer used to adjust the nexus operation long-poll timeouts.
+ Specifically, nexus operation long-poll requests are timed out at a time which leaves at least the buffer's duration
+ remaining before the caller's deadline, if permitted by the caller's deadline.`,
+)
+
 var Enabled = dynamicconfig.NewNamespaceBoolSetting(
 	"nexusoperation.enableStandalone",
 	false,
@@ -185,6 +200,8 @@ type Config struct {
 	ChasmEnabled                        dynamicconfig.BoolPropertyFnWithNamespaceFilter
 	ChasmNexusEnabled                   dynamicconfig.BoolPropertyFnWithNamespaceFilter
 	NumHistoryShards                    int32
+	LongPollBuffer                      dynamicconfig.DurationPropertyFnWithNamespaceFilter
+	LongPollTimeout                     dynamicconfig.DurationPropertyFnWithNamespaceFilter
 	RequestTimeout                      dynamicconfig.DurationPropertyFnWithDestinationFilter
 	MinRequestTimeout                   dynamicconfig.DurationPropertyFnWithNamespaceFilter
 	MaxConcurrentOperations             dynamicconfig.IntPropertyFnWithNamespaceFilter
@@ -212,6 +229,8 @@ func configProvider(dc *dynamicconfig.Collection, cfg *config.Persistence) *Conf
 		ChasmEnabled:                       dynamicconfig.EnableChasm.Get(dc),
 		ChasmNexusEnabled:                  ChasmNexusEnabled.Get(dc),
 		NumHistoryShards:                   cfg.NumHistoryShards,
+		LongPollBuffer:                     LongPollBuffer.Get(dc),
+		LongPollTimeout:                    LongPollTimeout.Get(dc),
 		RequestTimeout:                     RequestTimeout.Get(dc),
 		MinRequestTimeout:                  MinRequestTimeout.Get(dc),
 		MaxConcurrentOperations:            MaxConcurrentOperations.Get(dc),

chasm/lib/nexusoperation/frontend.go

@@ -110,7 +110,7 @@ func (h *frontendHandler) DescribeNexusOperationExecution(
 		return nil, err
 	}
 
-	if err := validateAndNormalizeDescribeRequest(req, h.config); err != nil {
+	if err := validateAndNormalizeDescribeRequest(req, namespaceID.String(), h.config); err != nil {
 		return nil, err
 	}
 
@@ -121,11 +121,29 @@ func (h *frontendHandler) DescribeNexusOperationExecution(
 	return resp.GetFrontendResponse(), err
 }
 
-func (h *frontendHandler) PollNexusOperationExecution(_ context.Context, req *workflowservice.PollNexusOperationExecutionRequest) (*workflowservice.PollNexusOperationExecutionResponse, error) {
+// PollNexusOperationExecution long-polls for a Nexus operation to reach a specific stage.
+func (h *frontendHandler) PollNexusOperationExecution(
+	ctx context.Context,
+	req *workflowservice.PollNexusOperationExecutionRequest,
+) (*workflowservice.PollNexusOperationExecutionResponse, error) {
 	if !h.isStandaloneNexusOperationEnabled(req.GetNamespace()) {
 		return nil, ErrStandaloneNexusOperationDisabled
 	}
-	return nil, serviceerror.NewUnimplemented("PollNexusOperationExecution not implemented")
+
+	if err := validateAndNormalizePollRequest(req, h.config); err != nil {
+		return nil, err
+	}
+
+	namespaceID, err := h.namespaceRegistry.GetNamespaceID(namespace.Name(req.GetNamespace()))
+	if err != nil {
+		return nil, err
+	}
+
+	resp, err := h.client.PollNexusOperation(ctx, &nexusoperationpb.PollNexusOperationRequest{
+		NamespaceId:     namespaceID.String(),
+		FrontendRequest: req,
+	})
+	return resp.GetFrontendResponse(), err
 }
 
 func (h *frontendHandler) ListNexusOperationExecutions(

chasm/lib/nexusoperation/handler.go

@@ -2,11 +2,14 @@ package nexusoperation
 
 import (
 	"context"
+	"errors"
 
 	enumspb "go.temporal.io/api/enums/v1"
+	"go.temporal.io/api/serviceerror"
 	"go.temporal.io/api/workflowservice/v1"
 	"go.temporal.io/server/chasm"
 	nexusoperationpb "go.temporal.io/server/chasm/lib/nexusoperation/gen/nexusoperationpb/v1"
+	"go.temporal.io/server/common/contextutil"
 	"go.temporal.io/server/common/log"
 )
 
@@ -59,7 +62,14 @@ func (h *handler) StartNexusOperation(
 	}, nil
 }
 
-// TODO: Add long-poll support.
+// DescribeNexusOperation queries current operation state, optionally as a long-poll that waits
+// for any state change.
+//
+// When used to long-poll, it returns an empty non-error response on context
+// deadline expiry, to indicate that the state being waited for was not reached. Callers should
+// interpret this as an invitation to resubmit their long-poll request. This response is sent before
+// the caller's deadline (see nexusoperation.longPollBuffer) so that it is likely that the caller
+// does indeed receive the non-error response.
 func (h *handler) DescribeNexusOperation(
 	ctx context.Context,
 	req *nexusoperationpb.DescribeNexusOperationRequest,
@@ -72,7 +82,102 @@ func (h *handler) DescribeNexusOperation(
 		RunID:       req.GetFrontendRequest().GetRunId(),
 	})
 
-	return chasm.ReadComponent(ctx, ref, (*Operation).buildDescribeResponse, req, nil)
+	token := req.GetFrontendRequest().GetLongPollToken()
+	if len(token) == 0 {
+		// No long poll.
+		return chasm.ReadComponent(ctx, ref, (*Operation).buildDescribeResponse, req)
+	}
+
+	// Determine the long poll timeout and buffer.
+	ns := req.GetFrontendRequest().GetNamespace()
+	ctx, cancel := contextutil.WithDeadlineBuffer(
+		ctx,
+		h.config.LongPollTimeout(ns),
+		h.config.LongPollBuffer(ns),
+	)
+	defer cancel()
+
+	// Poll for the operation state to change.
+	response, _, err = chasm.PollComponent(ctx, ref, func(
+		o *Operation,
+		ctx chasm.Context,
+		req *nexusoperationpb.DescribeNexusOperationRequest,
+	) (*nexusoperationpb.DescribeNexusOperationResponse, bool, error) {
+		changed, err := chasm.ExecutionStateChanged(o, ctx, token)
+		if err != nil {
+			if errors.Is(err, chasm.ErrMalformedComponentRef) {
+				return nil, false, serviceerror.NewInvalidArgument("invalid long poll token")
+			}
+			if errors.Is(err, chasm.ErrInvalidComponentRef) {
+				return nil, false, serviceerror.NewInvalidArgument("long poll token does not match execution")
+			}
+			return nil, false, err
+		}
+		if changed {
+			response, err := o.buildDescribeResponse(ctx, req)
+			return response, true, err
+		}
+		return nil, false, nil
+	}, req)
+
+	if err != nil && ctx.Err() != nil {
+		// Send empty non-error response on deadline expiry: caller should continue long-polling.
+		return &nexusoperationpb.DescribeNexusOperationResponse{
+			FrontendResponse: &workflowservice.DescribeNexusOperationExecutionResponse{},
+		}, nil
+	}
+	return response, err
+}
+
+// PollNexusOperation long-polls for a Nexus operation to reach a specific stage.
+//
+// It returns an empty non-error response on context deadline expiry, to indicate that the state
+// being waited for was not reached. Callers should interpret this as an invitation to resubmit
+// their long-poll request. This response is sent before the caller's
+// deadline (see nexusoperation.longPollBuffer) so that it is likely that the caller
+// does indeed receive the non-error response.
+func (h *handler) PollNexusOperation(
+	ctx context.Context,
+	req *nexusoperationpb.PollNexusOperationRequest,
+) (response *nexusoperationpb.PollNexusOperationResponse, err error) {
+	defer log.CapturePanic(h.logger, &err)
+
+	ref := chasm.NewComponentRef[*Operation](chasm.ExecutionKey{
+		NamespaceID: req.GetNamespaceId(),
+		BusinessID:  req.GetFrontendRequest().GetOperationId(),
+		RunID:       req.GetFrontendRequest().GetRunId(),
+	})
+
+	// Determine the long poll timeout and buffer.
+	ns := req.GetFrontendRequest().GetNamespace()
+	ctx, cancel := contextutil.WithDeadlineBuffer(
+		ctx,
+		h.config.LongPollTimeout(ns),
+		h.config.LongPollBuffer(ns),
+	)
+	defer cancel()
+
+	// Poll for the wait stage to be reached.
+	waitStage := req.GetFrontendRequest().GetWaitStage()
+	response, _, err = chasm.PollComponent(ctx, ref, func(
+		o *Operation,
+		ctx chasm.Context,
+		req *nexusoperationpb.PollNexusOperationRequest,
+	) (*nexusoperationpb.PollNexusOperationResponse, bool, error) {
+		if o.isWaitStageReached(ctx, waitStage) {
+			response := o.buildPollResponse(ctx)
+			return response, true, nil
+		}
+		return nil, false, nil
+	}, req)
+
+	if err != nil && ctx.Err() != nil {
+		// Send an empty non-error response as an invitation to resubmit the long-poll.
+		return &nexusoperationpb.PollNexusOperationResponse{
+			FrontendResponse: &workflowservice.PollNexusOperationExecutionResponse{},
+		}, nil
+	}
+	return response, err
 }
 
 // RequestCancelNexusOperation requests cancellation of a standalone Nexus operation via CHASM.

chasm/lib/nexusoperation/operation.go

@@ -336,10 +336,15 @@ func (o *Operation) buildDescribeResponse(
 	ctx chasm.Context,
 	req *nexusoperationpb.DescribeNexusOperationRequest,
 ) (*nexusoperationpb.DescribeNexusOperationResponse, error) {
+	token, err := ctx.Ref(o)
+	if err != nil {
+		return nil, err
+	}
+
 	resp := &workflowservice.DescribeNexusOperationExecutionResponse{
-		RunId: ctx.ExecutionKey().RunID,
-		Info:  o.buildExecutionInfo(ctx),
-		// TODO: Add LongPollToken support.
+		RunId:         ctx.ExecutionKey().RunID,
+		Info:          o.buildExecutionInfo(ctx),
+		LongPollToken: token,
 	}
 
 	// Include input, if requested
@@ -348,14 +353,12 @@ func (o *Operation) buildDescribeResponse(
 	}
 
 	// Include output, if available and requested
-	// TODO: get failure from last attempt for running operation, if available
-	if req.GetFrontendRequest().GetIncludeOutcome() && o.LifecycleState(ctx).IsClosed() {
-		outcome := o.Outcome.Get(ctx)
-		if successful := outcome.GetSuccessful(); successful != nil {
+	if req.GetFrontendRequest().GetIncludeOutcome() && o.isClosed() {
+		if successful, failure := o.describeOutcome(ctx); successful != nil {
 			resp.Outcome = &workflowservice.DescribeNexusOperationExecutionResponse_Result{
-				Result: successful.GetResult(),
+				Result: successful,
 			}
-		} else if failure := outcome.GetFailed().GetFailure(); failure != nil {
+		} else if failure != nil {
 			resp.Outcome = &workflowservice.DescribeNexusOperationExecutionResponse_Failure{
 				Failure: failure,
 			}
@@ -367,6 +370,63 @@ func (o *Operation) buildDescribeResponse(
 	}, nil
 }
 
+func (o *Operation) buildPollResponse(
+	ctx chasm.Context,
+) *nexusoperationpb.PollNexusOperationResponse {
+	resp := &workflowservice.PollNexusOperationExecutionResponse{
+		RunId:          ctx.ExecutionKey().RunID,
+		OperationToken: o.OperationToken,
+	}
+
+	if o.isClosed() {
+		resp.WaitStage = enumspb.NEXUS_OPERATION_WAIT_STAGE_CLOSED
+		if successful, failure := o.describeOutcome(ctx); successful != nil {
+			resp.Outcome = &workflowservice.PollNexusOperationExecutionResponse_Result{
+				Result: successful,
+			}
+		} else if failure != nil {
+			resp.Outcome = &workflowservice.PollNexusOperationExecutionResponse_Failure{
+				Failure: failure,
+			}
+		}
+	} else {
+		resp.WaitStage = enumspb.NEXUS_OPERATION_WAIT_STAGE_STARTED
+	}
+
+	return &nexusoperationpb.PollNexusOperationResponse{
+		FrontendResponse: resp,
+	}
+}
+
+func (o *Operation) describeOutcome(ctx chasm.Context) (*commonpb.Payload, *failurepb.Failure) {
+	outcome := o.Outcome.Get(ctx)
+	if successful := outcome.GetSuccessful(); successful != nil {
+		return successful.GetResult(), nil
+	}
+	if failure := outcome.GetFailed().GetFailure(); failure != nil {
+		return nil, failure
+	}
+	return nil, o.LastAttemptFailure
+}
+
+// isWaitStageReached checks if the operation has reached the requested wait stage.
+// NOTE: WaitStage.UNSPECIFIED is normalized to CLOSED by the frontend validator.
+func (o *Operation) isWaitStageReached(_ chasm.Context, waitStage enumspb.NexusOperationWaitStage) bool {
+	switch waitStage {
+	case enumspb.NEXUS_OPERATION_WAIT_STAGE_STARTED:
+		return o.Status == nexusoperationpb.OPERATION_STATUS_STARTED || o.isClosed()
+	case enumspb.NEXUS_OPERATION_WAIT_STAGE_CLOSED:
+		return o.isClosed()
+	default:
+		return false
+	}
+}
+
+// isClosed returns true if the operation is in a terminal state.
+func (o *Operation) isClosed() bool {
+	return o.LifecycleState(nil).IsClosed()
+}
+
 func (o *Operation) buildExecutionInfo(ctx chasm.Context) *nexuspb.NexusOperationExecutionInfo {
 	requestData := o.RequestData.Get(ctx)