chasm.PollComponent and PollActivityExecution

chasm/engine.go

@@ -42,9 +42,7 @@ type Engine interface {
 	PollComponent(
 		context.Context,
 		ComponentRef,
-		func(Context, Component) (any, bool, error),
-		func(MutableContext, Component, any) error,
-		...TransitionOption,
+		func(Context, Component) (bool, error),
 	) ([]byte, error)
 }
 
@@ -176,6 +174,9 @@ func UpdateWithNewEntity[C Component, I any, O1 any, O2 any](
 //   - consider remove ComponentRef from the return value and allow components to get
 //     the ref in the transition function. There are some caveats there, check the
 //     comment of the NewRef method in MutableContext.
+//
+// UpdateComponent applies updateFn to the component identified by the supplied component reference.
+// It returns the result, along with the new component reference.
 func UpdateComponent[C Component, R []byte | ComponentRef, I any, O any](
 	ctx context.Context,
 	r R,
@@ -207,6 +208,8 @@ func UpdateComponent[C Component, R []byte | ComponentRef, I any, O any](
 	return output, newSerializedRef, err
 }
 
+// ReadComponent returns the result of evaluating readFn against the current state of the component
+// identified by the supplied component reference.
 func ReadComponent[C Component, R []byte | ComponentRef, I any, O any](
 	ctx context.Context,
 	r R,
@@ -234,20 +237,16 @@ func ReadComponent[C Component, R []byte | ComponentRef, I any, O any](
 	return output, err
 }
 
-type PollComponentRequest[C Component, I any, O any] struct {
-	Ref         ComponentRef
-	PredicateFn func(C, Context, I) bool
-	OperationFn func(C, MutableContext, I) (O, error)
-	Input       I
-}
-
-func PollComponent[C Component, R []byte | ComponentRef, I any, O any, T any](
+// PollComponent waits until the predicate is true when evaluated against the component identified
+// by the supplied component reference. If it times out due to a server-imposed long-poll timeout
+// then it returns (nil, nil, nil). Otherwise it returns (output, ref, err), where output is the
+// output of the predicate function, and ref is a component reference identifying the state at which
+// the predicate was satisfied. If the predicate is true at the outset then it returns immediately.
+func PollComponent[C Component, R []byte | ComponentRef, I any, O any](
 	ctx context.Context,
 	r R,
-	predicateFn func(C, Context, I) (T, bool, error),
-	operationFn func(C, MutableContext, I, T) (O, error),
+	predicateFn func(C, Context, I) (O, bool, error),
 	input I,
-	opts ...TransitionOption,
 ) (O, []byte, error) {
 	var output O
 
@@ -259,15 +258,13 @@ func PollComponent[C Component, R []byte | ComponentRef, I any, O any, T any](
 	newSerializedRef, err := engineFromContext(ctx).PollComponent(
 		ctx,
 		ref,
-		func(ctx Context, c Component) (any, bool, error) {
-			return predicateFn(c.(C), ctx, input)
+		func(ctx Context, c Component) (bool, error) {
+			out, satisfied, err := predicateFn(c.(C), ctx, input)
+			if satisfied {
+				output = out
+			}
+			return satisfied, err
 		},
-		func(ctx MutableContext, c Component, t any) error {
-			var err error
-			output, err = operationFn(c.(C), ctx, input, t.(T))
-			return err
-		},
-		opts...,
 	)
 	if err != nil {
 		return output, nil, err

chasm/lib/activity/activity.go

@@ -1,14 +1,17 @@
 package activity
 
 import (
+	"errors"
 	"fmt"
 	"time"
 
+	"go.temporal.io/api/activity/v1"
 	commonpb "go.temporal.io/api/common/v1"
 	deploymentpb "go.temporal.io/api/deployment/v1"
 	enumspb "go.temporal.io/api/enums/v1"
 	failurepb "go.temporal.io/api/failure/v1"
 	historypb "go.temporal.io/api/history/v1"
+	"go.temporal.io/api/serviceerror"
 	"go.temporal.io/api/workflowservice/v1"
 	"go.temporal.io/server/api/historyservice/v1"
 	"go.temporal.io/server/api/matchingservice/v1"
@@ -37,14 +40,12 @@ type Activity struct {
 
 	*activitypb.ActivityState
 
-	// Standalone only
 	Visibility    chasm.Field[*chasm.Visibility]
 	Attempt       chasm.Field[*activitypb.ActivityAttemptState]
 	LastHeartbeat chasm.Field[*activitypb.ActivityHeartbeatState]
+	Outcome       chasm.Field[*activitypb.ActivityOutcome]
 	// Standalone only
 	RequestData chasm.Field[*activitypb.ActivityRequestData]
-	Outcome     chasm.Field[*activitypb.ActivityOutcome]
-
 	// Pointer to an implementation of the "store". for a workflow activity this would be a parent pointer back to
 	// the workflow. For a standalone activity this would be nil.
 	// TODO: revisit a standalone activity pointing to itself once we handle storing it more efficiently.
@@ -314,3 +315,123 @@ func (a *Activity) RecordHeartbeat(ctx chasm.MutableContext, details *commonpb.P
 	})
 	return nil, nil
 }
+
+func (a *Activity) buildActivityExecutionInfo(ctx chasm.Context) (*activity.ActivityExecutionInfo, error) {
+	if a.ActivityState == nil {
+		return nil, errors.New("activity state is nil")
+	}
+
+	// TODO(dan): support pause states
+	var status enumspb.ActivityExecutionStatus
+	var runState enumspb.PendingActivityState
+	switch a.GetStatus() {
+	case activitypb.ACTIVITY_EXECUTION_STATUS_UNSPECIFIED:
+		status = enumspb.ACTIVITY_EXECUTION_STATUS_UNSPECIFIED
+		runState = enumspb.PENDING_ACTIVITY_STATE_UNSPECIFIED
+	case activitypb.ACTIVITY_EXECUTION_STATUS_SCHEDULED:
+		status = enumspb.ACTIVITY_EXECUTION_STATUS_RUNNING
+		runState = enumspb.PENDING_ACTIVITY_STATE_SCHEDULED
+	case activitypb.ACTIVITY_EXECUTION_STATUS_STARTED:
+		status = enumspb.ACTIVITY_EXECUTION_STATUS_RUNNING
+		runState = enumspb.PENDING_ACTIVITY_STATE_STARTED
+	case activitypb.ACTIVITY_EXECUTION_STATUS_CANCEL_REQUESTED:
+		status = enumspb.ACTIVITY_EXECUTION_STATUS_RUNNING
+		runState = enumspb.PENDING_ACTIVITY_STATE_CANCEL_REQUESTED
+	case activitypb.ACTIVITY_EXECUTION_STATUS_COMPLETED:
+		status = enumspb.ACTIVITY_EXECUTION_STATUS_COMPLETED
+		runState = enumspb.PENDING_ACTIVITY_STATE_UNSPECIFIED
+	case activitypb.ACTIVITY_EXECUTION_STATUS_FAILED:
+		status = enumspb.ACTIVITY_EXECUTION_STATUS_FAILED
+		runState = enumspb.PENDING_ACTIVITY_STATE_UNSPECIFIED
+	case activitypb.ACTIVITY_EXECUTION_STATUS_CANCELED:
+		status = enumspb.ACTIVITY_EXECUTION_STATUS_CANCELED
+		runState = enumspb.PENDING_ACTIVITY_STATE_UNSPECIFIED
+	case activitypb.ACTIVITY_EXECUTION_STATUS_TERMINATED:
+		status = enumspb.ACTIVITY_EXECUTION_STATUS_TERMINATED
+		runState = enumspb.PENDING_ACTIVITY_STATE_UNSPECIFIED
+	case activitypb.ACTIVITY_EXECUTION_STATUS_TIMED_OUT:
+		status = enumspb.ACTIVITY_EXECUTION_STATUS_TIMED_OUT
+		runState = enumspb.PENDING_ACTIVITY_STATE_UNSPECIFIED
+	default:
+		return nil, serviceerror.NewInternalf("unknown activity execution status: %s", a.GetStatus())
+	}
+
+	requestData, err := a.RequestData.Get(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	key := ctx.ExecutionKey()
+
+	info := &activity.ActivityExecutionInfo{
+		ActivityId:    key.BusinessID,
+		RunId:         key.EntityID,
+		ActivityType:  a.GetActivityType(),
+		Status:        status,
+		RunState:      runState,
+		ScheduledTime: a.GetScheduledTime(),
+		Priority:      a.GetPriority(),
+		Header:        requestData.GetHeader(),
+		// TODO(dan): populate remaining fields
+	}
+
+	return info, nil
+}
+
+func (a *Activity) buildPollActivityExecutionResponse(
+	ctx chasm.Context,
+	req *activitypb.PollActivityExecutionRequest,
+) (*activitypb.PollActivityExecutionResponse, error) {
+	request := req.GetFrontendRequest()
+
+	// TODO(dan): pass ref into this function?
+	token, err := ctx.Ref(a)
+	if err != nil {
+		return nil, err
+	}
+
+	var info *activity.ActivityExecutionInfo
+	if request.GetIncludeInfo() {
+		info, err = a.buildActivityExecutionInfo(ctx)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	var input *commonpb.Payloads
+	if request.GetIncludeInput() {
+		activityRequest, err := a.RequestData.Get(ctx)
+		if err != nil {
+			return nil, err
+		}
+		input = activityRequest.GetInput()
+	}
+
+	response := &workflowservice.PollActivityExecutionResponse{
+		Info:                     info,
+		RunId:                    ctx.ExecutionKey().EntityID,
+		Input:                    input,
+		StateChangeLongPollToken: token,
+	}
+
+	if request.GetIncludeOutcome() {
+		activityOutcome, err := a.Outcome.Get(ctx)
+		if err != nil {
+			return nil, err
+		}
+		switch v := activityOutcome.GetVariant().(type) {
+		case *activitypb.ActivityOutcome_Failed_:
+			response.Outcome = &workflowservice.PollActivityExecutionResponse_Failure{
+				Failure: v.Failed.GetFailure(),
+			}
+		case *activitypb.ActivityOutcome_Successful_:
+			response.Outcome = &workflowservice.PollActivityExecutionResponse_Result{
+				Result: v.Successful.GetOutput(),
+			}
+		}
+	}
+
+	return &activitypb.PollActivityExecutionResponse{
+		FrontendResponse: response,
+	}, nil
+}

chasm/lib/activity/frontend.go

@@ -85,6 +85,24 @@ func (h *frontendHandler) StartActivityExecution(ctx context.Context, req *workf
 	return resp.GetFrontendResponse(), err
 }
 
+// PollActivityExecution handles PollActivityExecutionRequest. This method supports querying current
+// activity state, optionally as a long-poll that waits for certain state changes. It is used by
+// clients to poll for activity state and/or result.
+func (h *frontendHandler) PollActivityExecution(
+	ctx context.Context,
+	req *workflowservice.PollActivityExecutionRequest,
+) (*workflowservice.PollActivityExecutionResponse, error) {
+	namespaceID, err := h.namespaceRegistry.GetNamespaceID(namespace.Name(req.GetNamespace()))
+	if err != nil {
+		return nil, err
+	}
+	resp, err := h.client.PollActivityExecution(ctx, &activitypb.PollActivityExecutionRequest{
+		NamespaceId:     namespaceID.String(),
+		FrontendRequest: req,
+	})
+	return resp.GetFrontendResponse(), err
+}
+
 func (h *frontendHandler) validateAndPopulateStartRequest(
 	req *workflowservice.StartActivityExecutionRequest,
 	namespaceID namespace.ID,