Implement nexus-delivered activity cancels (temporalio#2394)

Author: Sushisource

activity/doc.go

@@ -56,7 +56,8 @@ The first parameter to an activity function can be an optional context.Context.
 * The activity function returns.
 * The context deadline is exceeded. The deadline is calculated based on the minimum of the ScheduleToClose timeout plus
 the activity task scheduled time and the StartToClose timeout plus the activity task start time.
-* The activity calls RecordHeartbeat after being cancelled by the Temporal server.
+* The Temporal server requests cancellation. On supported servers this can be delivered directly to the worker; on older
+servers the activity observes cancellation when it calls RecordHeartbeat after the request.
 
 # Failing the Activity
 
@@ -97,7 +98,8 @@ payload containing progress information.
 
 When an Activity is canceled (or its Workflow execution is completed or failed) the context passed into its function
 is canceled which sets its Done channel’s closed state. So an Activity can use that to perform any necessary cleanup
-and abort its execution. Currently cancellation is delivered only to Activities that call RecordHeartbeat.
+and abort its execution. On supported servers, cancellation can be delivered directly to the worker. Heartbeats are still
+recommended for long-running Activities and may be the cancellation delivery path on older servers.
 
 # Async/Manual Activity Completion
 

internal/activity.go

@@ -289,10 +289,9 @@ func GetWorkerStopChannel(ctx context.Context) <-chan struct{} {
 // If the activity is either canceled or workflow/activity doesn't exist, then we would cancel
 // the context with error context.Canceled.
 //
-//	TODO: Implement automatic heartbeating with cancellation through ctx.
-//
 // details - The details that you provided here can be seen in the workflow when it receives TimeoutError. You
-// can check error TimeoutType()/Details().
+// can check error TimeoutType()/Details(). Heartbeat responses may also deliver server requests such as activity
+// cancellation, pause, and reset to the activity context.
 //
 // Exposed as: [go.temporal.io/sdk/activity.RecordHeartbeat]
 func RecordActivityHeartbeat(ctx context.Context, details ...interface{}) {

internal/internal_task_handlers.go

@@ -145,6 +145,12 @@ type (
 		cache                     *WorkerCache
 		deadlockDetectionTimeout  time.Duration
 		capabilities              *workflowservice.GetSystemInfoResponse_Capabilities
+		workerControlTaskQueue    string
+	}
+
+	activityCancellationCallbacks struct {
+		sync.Mutex
+		cancels map[string]context.CancelCauseFunc
 	}
 
 	activityProvider func(name string) activity
@@ -172,6 +178,7 @@ type (
 		inboundPayloadVisitor            PayloadVisitor
 		outboundPayloadVisitor           PayloadVisitor
 		payloadVisitorConcurrency        int
+		activityCancellationCallbacks    *activityCancellationCallbacks
 	}
 
 	// history wrapper method to help information about events.
@@ -587,6 +594,7 @@ func newWorkflowTaskHandler(params workerExecutionParameters, ppMgr pressurePoin
 		cache:                     params.cache,
 		deadlockDetectionTimeout:  params.DeadlockDetectionTimeout,
 		capabilities:              params.capabilities,
+		workerControlTaskQueue:    params.workerControlTaskQueue,
 	}
 }
 
@@ -2007,6 +2015,7 @@ func (wth *workflowTaskHandlerImpl) completeWorkflow(
 			wth.useBuildIDForVersioning,
 			wth.workerDeploymentVersion,
 		),
+		WorkerControlTaskQueue: wth.workerControlTaskQueue,
 	}
 	if wth.capabilities != nil && wth.capabilities.BuildIdBasedVersioning {
 		//lint:ignore SA1019 ignore deprecated versioning APIs
@@ -2100,12 +2109,40 @@ func newActivityTaskHandlerWithCustomProvider(
 			params.UseBuildIDForVersioning,
 			params.DeploymentOptions.Version,
 		),
-		inboundPayloadVisitor:     params.inboundPayloadVisitor,
-		outboundPayloadVisitor:    params.outboundPayloadVisitor,
-		payloadVisitorConcurrency: params.payloadVisitorConcurrency,
+		inboundPayloadVisitor:         params.inboundPayloadVisitor,
+		outboundPayloadVisitor:        params.outboundPayloadVisitor,
+		payloadVisitorConcurrency:     params.payloadVisitorConcurrency,
+		activityCancellationCallbacks: params.activityCancellationCallbacks,
+	}
+}
+
+func newActivityCancellationCallbacks() *activityCancellationCallbacks {
+	return &activityCancellationCallbacks{cancels: make(map[string]context.CancelCauseFunc)}
+}
+
+func (r *activityCancellationCallbacks) register(taskToken []byte, cancel context.CancelCauseFunc) func() {
+	key := string(taskToken)
+	r.Lock()
+	r.cancels[key] = cancel
+	r.Unlock()
+	return func() {
+		r.Lock()
+		delete(r.cancels, key)
+		r.Unlock()
 	}
 }
 
+func (r *activityCancellationCallbacks) cancel(taskToken []byte) bool {
+	r.Lock()
+	cancel, ok := r.cancels[string(taskToken)]
+	r.Unlock()
+	if !ok {
+		return false
+	}
+	cancel(NewCanceledError())
+	return true
+}
+
 // heartbeatVisitorError wraps an outbound payload visitor error from a heartbeat.
 // It is used as the context cancellation cause so Execute() can detect that
 // RespondActivityTaskFailed was already sent proactively and skip sending a second response.
@@ -2120,7 +2157,8 @@ type temporalInvoker struct {
 	service        workflowservice.WorkflowServiceClient
 	metricsHandler metrics.Handler
 	taskToken      []byte
-	// cancelHandler is called when the activity is canceled by a heartbeat request.
+	// cancelHandler is called when the activity is canceled by a heartbeat response
+	// or worker command.
 	cancelHandler context.CancelCauseFunc
 	// Amount of time to wait between each pending heartbeat send
 	heartbeatThrottleInterval time.Duration
@@ -2351,6 +2389,10 @@ func (ath *activityTaskHandlerImpl) Execute(taskQueue string, t *workflowservice
 	}
 	canCtx, cancel := context.WithCancelCause(rootCtx)
 	defer cancel(nil)
+	if ath.activityCancellationCallbacks != nil {
+		unregister := ath.activityCancellationCallbacks.register(t.TaskToken, cancel)
+		defer unregister()
+	}
 
 	if err := visitProtoPayloads(canCtx, ath.inboundPayloadVisitor, t, ath.payloadVisitorConcurrency); err != nil {
 		return ath.visitorErrorToActivityFailure("Activity task preprocess error: ", t, err), nil

internal/internal_task_handlers_test.go

@@ -135,6 +135,19 @@ func TestTaskHandlersTestSuite(t *testing.T) {
 	})
 }
 
+func TestActivityCancellationCallbacksCancel(t *testing.T) {
+	registry := newActivityCancellationCallbacks()
+	taskToken := []byte{1, 2, 3}
+	ctx, cancel := context.WithCancelCause(context.Background())
+
+	unregister := registry.register(taskToken, cancel)
+	require.True(t, registry.cancel([]byte{1, 2, 3}))
+	require.True(t, IsCanceledError(context.Cause(ctx)))
+
+	unregister()
+	require.False(t, registry.cancel(taskToken))
+}
+
 func createTestEventWorkflowExecutionCompleted(eventID int64, attr *historypb.WorkflowExecutionCompletedEventAttributes) *historypb.HistoryEvent {
 	return &historypb.HistoryEvent{
 		EventId:    eventID,

internal/internal_task_pollers.go

@@ -89,6 +89,8 @@ type (
 		pollTimeTracker *pollTimeTracker
 		// Unique identifier for worker
 		workerInstanceKey string
+		// Per-client queue used by the server to send worker commands.
+		workerControlTaskQueue string
 		// Server cancels polls on shutdown
 		workerPollCompleteOnShutdown *atomic.Bool
 	}
@@ -369,6 +371,7 @@ func newWorkflowTaskProcessor(
 			capabilities:                 params.capabilities,
 			pollTimeTracker:              params.pollTimeTracker,
 			workerInstanceKey:            params.workerInstanceKey,
+			workerControlTaskQueue:       params.workerControlTaskQueue,
 			workerPollCompleteOnShutdown: params.workerPollCompleteOnShutdown,
 		},
 		service:                      service,
@@ -1151,7 +1154,8 @@ func (wtp *workflowTaskPoller) getNextPollRequest() (request *workflowservice.Po
 			wtp.useBuildIDVersioning,
 			wtp.workerDeploymentVersion,
 		),
-		WorkerInstanceKey: wtp.workerInstanceKey,
+		WorkerInstanceKey:      wtp.workerInstanceKey,
+		WorkerControlTaskQueue: wtp.workerControlTaskQueue,
 	}
 	if wtp.getCapabilities().BuildIdBasedVersioning {
 		//lint:ignore SA1019 ignore deprecated versioning APIs
@@ -1368,6 +1372,7 @@ func newActivityTaskPoller(taskHandler ActivityTaskHandler, service workflowserv
 			capabilities:                 params.capabilities,
 			pollTimeTracker:              params.pollTimeTracker,
 			workerInstanceKey:            params.workerInstanceKey,
+			workerControlTaskQueue:       params.workerControlTaskQueue,
 			workerPollCompleteOnShutdown: params.workerPollCompleteOnShutdown,
 		},
 		taskHandler:         taskHandler,
@@ -1408,7 +1413,8 @@ func (atp *activityTaskPoller) poll(ctx context.Context) (taskForWorker, error)
 			atp.useBuildIDVersioning,
 			atp.workerDeploymentVersion,
 		),
-		WorkerInstanceKey: atp.workerInstanceKey,
+		WorkerInstanceKey:      atp.workerInstanceKey,
+		WorkerControlTaskQueue: atp.workerControlTaskQueue,
 	}
 
 	response, err := atp.pollActivityTaskQueue(ctx, request)

internal/internal_task_pollers_test.go

@@ -4,12 +4,12 @@ import (
 	"context"
 	"encoding/binary"
 	"errors"
-	"github.com/google/uuid"
 	"sync/atomic"
 	"testing"
 	"time"
 
 	"github.com/golang/mock/gomock"
+	"github.com/google/uuid"
 	"github.com/stretchr/testify/require"
 	commonpb "go.temporal.io/api/common/v1"
 	enumspb "go.temporal.io/api/enums/v1"
@@ -40,6 +40,63 @@ func (wth *countingTaskHandler) ProcessWorkflowTask(
 	return wth.WorkflowTaskHandler.ProcessWorkflowTask(task, wfctx, hb)
 }
 
+func TestPollRequestsIncludeWorkerControlTaskQueue(t *testing.T) {
+	t.Parallel()
+
+	ctrl := gomock.NewController(t)
+	service := workflowservicemock.NewMockWorkflowServiceClient(ctrl)
+	const (
+		namespace    = "test-ns"
+		taskQueue    = "test-task-queue"
+		identity     = "test-worker"
+		controlQueue = "temporal-sys/worker-commands/test-ns/grouping-key"
+	)
+
+	service.EXPECT().PollWorkflowTaskQueue(gomock.Any(), gomock.Any(), gomock.Any()).
+		DoAndReturn(func(_ context.Context, req *workflowservice.PollWorkflowTaskQueueRequest, _ ...grpc.CallOption) (*workflowservice.PollWorkflowTaskQueueResponse, error) {
+			require.Equal(t, controlQueue, req.WorkerControlTaskQueue)
+			require.Equal(t, taskQueue, req.TaskQueue.GetName())
+			return &workflowservice.PollWorkflowTaskQueueResponse{}, nil
+		})
+
+	base := basePoller{
+		metricsHandler:         metrics.NopHandler,
+		workerBuildID:          "test-build-id",
+		workerControlTaskQueue: controlQueue,
+	}
+	wtp := &workflowTaskPoller{
+		basePoller:            base,
+		mode:                  NonSticky,
+		namespace:             namespace,
+		taskQueueName:         taskQueue,
+		identity:              identity,
+		service:               service,
+		logger:                ilog.NewDefaultLogger(),
+		numNormalPollerMetric: newNumPollerMetric(metrics.NopHandler, metrics.PollerTypeWorkflowTask),
+	}
+	_, err := wtp.poll(context.Background())
+	require.NoError(t, err)
+
+	service.EXPECT().PollActivityTaskQueue(gomock.Any(), gomock.Any(), gomock.Any()).
+		DoAndReturn(func(_ context.Context, req *workflowservice.PollActivityTaskQueueRequest, _ ...grpc.CallOption) (*workflowservice.PollActivityTaskQueueResponse, error) {
+			require.Equal(t, controlQueue, req.WorkerControlTaskQueue)
+			require.Equal(t, taskQueue, req.TaskQueue.GetName())
+			return &workflowservice.PollActivityTaskQueueResponse{}, nil
+		})
+
+	atp := &activityTaskPoller{
+		basePoller:      base,
+		namespace:       namespace,
+		taskQueueName:   taskQueue,
+		identity:        identity,
+		service:         service,
+		logger:          ilog.NewDefaultLogger(),
+		numPollerMetric: newNumPollerMetric(metrics.NopHandler, metrics.PollerTypeActivityTask),
+	}
+	_, err = atp.poll(context.Background())
+	require.NoError(t, err)
+}
+
 func TestWFTRacePrevention(t *testing.T) {
 	params := workerExecutionParameters{cache: NewWorkerCache()}
 	ensureRequiredParams(&params)

internal/internal_worker.go

@@ -228,6 +228,10 @@ type (
 
 		workerInstanceKey string
 
+		workerControlTaskQueue string
+
+		activityCancellationCallbacks *activityCancellationCallbacks
+
 		workerPollCompleteOnShutdown *atomic.Bool
 
 		// Set to true during start() when the namespace has the poller_autoscaling capability.
@@ -314,6 +318,10 @@ func (params *workerExecutionParameters) isInternalWorker() bool {
 	return params.Namespace == "temporal-system" || params.TaskQueue == "temporal-sys-per-ns-tq"
 }
 
+func workerControlTaskQueue(namespace, groupingKey string) string {
+	return fmt.Sprintf("temporal-sys/worker-commands/%s/%s", namespace, groupingKey)
+}
+
 func newWorkflowWorkerInternal(client *WorkflowClient, params workerExecutionParameters, ppMgr pressurePointMgr, overrides *workerOverrides, registry *registry) *workflowWorker {
 	workerStopChannel := make(chan struct{})
 	params.WorkerStopChannel = getReadOnlyChannel(workerStopChannel)
@@ -1224,7 +1232,7 @@ type AggregatedWorker struct {
 	// stopC is created in NewAggregatedWorker and closed by AggregatedWorker.Stop()
 	// to mark the aggregated worker stopped, unblock Run(), and prevent restart.
 	// Child worker stop channels are closed later by their own Stop methods.
-	stopC chan struct{}
+	stopC        chan struct{}
 	fatalErr     error
 	fatalErrLock sync.Mutex
 	capabilities *workflowservice.GetSystemInfoResponse_Capabilities
@@ -2269,6 +2277,12 @@ func NewAggregatedWorker(client *WorkflowClient, taskQueue string, options Worke
 	// All worker systems that depend on the capabilities to process workflow/activity tasks
 	// should take a pointer to this struct and wait for it to be populated when the worker is run.
 	var capabilities workflowservice.GetSystemInfoResponse_Capabilities
+	var activityCancellationCallbacks *activityCancellationCallbacks
+	if client.heartbeatManager != nil {
+		activityCancellationCallbacks = client.heartbeatManager.
+			sharedNamespaceWorkerFor(client.namespace).
+			activityCancellationCallbacks
+	}
 
 	baseMetricsHandler := client.metricsHandler.WithTags(metrics.TaskQueueTags(taskQueue))
 	var metricsHandler metrics.Handler
@@ -2338,12 +2352,14 @@ func NewAggregatedWorker(client *WorkflowClient, taskQueue string, options Worke
 			taskQueue:     taskQueue,
 			maxConcurrent: options.MaxConcurrentEagerActivityExecutionSize,
 		}),
-		capabilities:                 &capabilities,
-		pollTimeTracker:              &pollTimeTracker{},
-		workerInstanceKey:            workerInstanceKey,
-		workerPollCompleteOnShutdown: workerPollCompleteOnShutdown,
-		serverSupportsAutoscaling:    &atomic.Bool{},
-		inboundPayloadVisitor:        extstore.NewExternalRetrievalVisitor(client.storageParams),
+		capabilities:                  &capabilities,
+		pollTimeTracker:               &pollTimeTracker{},
+		workerInstanceKey:             workerInstanceKey,
+		workerControlTaskQueue:        workerControlTaskQueue(client.namespace, client.workerGroupingKey),
+		activityCancellationCallbacks: activityCancellationCallbacks,
+		workerPollCompleteOnShutdown:  workerPollCompleteOnShutdown,
+		serverSupportsAutoscaling:     &atomic.Bool{},
+		inboundPayloadVisitor:         extstore.NewExternalRetrievalVisitor(client.storageParams),
 		outboundPayloadVisitor: newCompositePayloadVisitor(
 			extstore.NewExternalStorageVisitor(client.storageParams),
 			payloadLimitVisitor,