Handle activity cancel due to worker shutdown properly (#1910)

* Differentiate context cancel due to worker shutdown

* change cancelAllowed to cancelReason to handle more nuanced scenario of cancel source

* go back to using isActivityCanceled

* Fix sessions for TestSessionStateFailedWorkerFailed

* PR feedback

* Add local activity test

* Add local activity test for worker shutdown, as well as for no heartbeating due to shutdown

* Only avoid heartbeating early due to non-server cancel

* Remove redundant tests, remove old LA fix now that fix is in main,

* Add new test, activity actually handles ctx cancel

* increase time to reduce race condition change

Author: yuandrew

internal/internal_task_handlers.go

@@ -2283,8 +2283,10 @@ func (ath *activityTaskHandlerImpl) Execute(taskQueue string, t *workflowservice
 
 	output, err := activityImplementation.Execute(ctx, t.Input)
 	// Check if context canceled at a higher level before we cancel it ourselves
-	// TODO : check if the cause of the context cancellation is from the server
-	isActivityCanceled := ctx.Err() == context.Canceled
+
+	// Cancels that don't originate from the server will have separate cancel reasons, like
+	// ErrWorkerShutdown or ErrActivityPaused
+	isActivityCanceled := ctx.Err() == context.Canceled && errors.Is(context.Cause(ctx), &CanceledError{})
 
 	dlCancelFunc()
 	if <-ctx.Done(); ctx.Err() == context.DeadlineExceeded {

internal/session.go

@@ -286,7 +286,7 @@ func createSession(ctx Context, creationTaskqueue string, options *SessionOption
 
 	taskqueueChan := GetSignalChannel(ctx, sessionID) // use sessionID as channel name
 	// Retry is only needed when creating new session and the error returned is
-	// NewApplicationError(errTooManySessionsMsg). Therefore we make sure to
+	// NewApplicationError(errTooManySessionsMsg). Therefore, we make sure to
 	// disable retrying for start-to-close and heartbeat timeouts which can occur
 	// when attempting to retry a create-session on a different worker.
 	retryPolicy := &RetryPolicy{
@@ -417,6 +417,14 @@ func sessionCreationActivity(ctx context.Context, sessionID string) error {
 		select {
 		case <-ctx.Done():
 			sessionEnv.CompleteSession(sessionID)
+			// Because of how session creation configures retryPolicy, we need to wrap context cancels that don't
+			// originate from the server as non-retryable errors. See retrypolicy in createSession() above.
+			if !(ctx.Err() == context.Canceled && errors.Is(context.Cause(ctx), &CanceledError{})) {
+				return NewApplicationErrorWithOptions(
+					"session failed due to worker shutdown", "SessionWorkerShutdown",
+					ApplicationErrorOptions{NonRetryable: true, Cause: ctx.Err()})
+
+			}
 			return ctx.Err()
 		case <-ticker.C:
 			heartbeatOp := func() error {

test/activity_test.go

@@ -461,3 +461,13 @@ func (a *Activities) RawValueActivity(ctx context.Context, value converter.RawVa
 	activity.GetLogger(ctx).Info("RawValue value", value.Payload())
 	return value, nil
 }
+
+func (a *Activities) CancelActivity(ctx context.Context) error {
+	t := time.NewTicker(200 * time.Millisecond)
+	defer t.Stop()
+	select {
+	case <-ctx.Done():
+	case <-t.C:
+	}
+	return ctx.Err()
+}

test/integration_test.go

@@ -7709,6 +7709,66 @@ func (ts *IntegrationTestSuite) TestLocalActivityFailureMetric_BenignHandling()
 	ts.assertMetricCount(metrics.LocalActivityExecutionFailedCounter, currCount)
 }
 
+func (ts *IntegrationTestSuite) TestActivityCancelFromWorkerShutdown() {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	run, err := ts.client.ExecuteWorkflow(ctx, ts.startWorkflowOptions("test-activity-cancel"), ts.workflows.WorkflowReactToCancel, false)
+	ts.NoError(err)
+
+	// Give the workflow time to run and run activity
+	time.Sleep(100 * time.Millisecond)
+	ts.worker.Stop()
+	ts.workerStopped = true
+	// Now create a new worker on that same task queue to resume the work of the
+	// activity retry
+	nextWorker := worker.New(ts.client, ts.taskQueueName, worker.Options{})
+	ts.registerWorkflowsAndActivities(nextWorker)
+	ts.NoError(nextWorker.Start())
+	defer nextWorker.Stop()
+
+	err = run.Get(ctx, nil)
+	ts.NoError(err)
+}
+
+func (ts *IntegrationTestSuite) TestLocalActivityCancelFromWorkerShutdown() {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	run, err := ts.client.ExecuteWorkflow(ctx, ts.startWorkflowOptions("test-local-activity-cancel"), ts.workflows.WorkflowReactToCancel, true)
+	ts.NoError(err)
+
+	// Give the workflow time to run and run activity
+	time.Sleep(100 * time.Millisecond)
+	ts.worker.Stop()
+	ts.workerStopped = true
+	// Now create a new worker on that same task queue to resume the work of the
+	// activity retry
+	nextWorker := worker.New(ts.client, ts.taskQueueName, worker.Options{})
+	ts.registerWorkflowsAndActivities(nextWorker)
+	ts.NoError(nextWorker.Start())
+	defer nextWorker.Stop()
+
+	err = run.Get(ctx, nil)
+	ts.NoError(err)
+
+	timeout_count := 0
+	iter := ts.client.GetWorkflowHistory(ctx, run.GetID(), run.GetRunID(), true, enumspb.HISTORY_EVENT_FILTER_TYPE_ALL_EVENT)
+	for iter.HasNext() {
+		event, err := iter.Next()
+		if err != nil {
+			break
+		}
+
+		// WFT timeout should come from first worker stopping and LA being canceled
+		if event.EventType == enumspb.EVENT_TYPE_WORKFLOW_TASK_TIMED_OUT {
+			timeout_count++
+		}
+	}
+
+	ts.Equal(1, timeout_count)
+}
+
 func (ts *IntegrationTestSuite) TestLocalActivityWorkerShutdownNoHeartbeat() {
 	// FYI, setup of this test allows the worker to wait to stop for 10 seconds
 	ctx, cancel := context.WithCancel(context.Background())

test/workflow_test.go

@@ -3500,6 +3500,35 @@ func (w *Workflows) WorkflowRawValue(ctx workflow.Context, value converter.RawVa
 	return returnVal, err
 }
 
+func (w *Workflows) WorkflowReactToCancel(ctx workflow.Context, localActivity bool) error {
+	var activities *Activities
+	var err error
+	// Allow for 2 attempts so when a worker shuts down and a 2nd one is created,
+	// it can use the 2nd attempt to complete the activity.
+	retryPolicy := temporal.RetryPolicy{
+		MaximumAttempts: 2,
+	}
+
+	if localActivity {
+		ctx = workflow.WithLocalActivityOptions(ctx, workflow.LocalActivityOptions{
+			ScheduleToCloseTimeout: 2 * time.Second,
+			RetryPolicy:            &retryPolicy,
+		})
+		err = workflow.ExecuteLocalActivity(ctx, activities.CancelActivity).Get(ctx, nil)
+	} else {
+		ctx = workflow.WithActivityOptions(ctx, workflow.ActivityOptions{
+			ScheduleToCloseTimeout: 2 * time.Second,
+			RetryPolicy:            &retryPolicy,
+		})
+		err = workflow.ExecuteActivity(ctx, activities.CancelActivity).Get(ctx, nil)
+	}
+
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
 func (w *Workflows) register(worker worker.Worker) {
 	worker.RegisterWorkflow(w.ActivityCancelRepro)
 	worker.RegisterWorkflow(w.ActivityCompletionUsingID)
@@ -3648,6 +3677,7 @@ func (w *Workflows) register(worker worker.Worker) {
 	worker.RegisterWorkflow(w.WorkflowClientFromActivity)
 	worker.RegisterWorkflow(w.WorkflowTemporalPrefixSignal)
 	worker.RegisterWorkflow(w.WorkflowRawValue)
+	worker.RegisterWorkflow(w.WorkflowReactToCancel)
 }
 
 func (w *Workflows) defaultActivityOptions() workflow.ActivityOptions {