Fix activity cancelation check (#2069)

Author: yuandrew

internal/internal_task_handlers.go

@@ -2308,7 +2308,7 @@ func (ath *activityTaskHandlerImpl) Execute(taskQueue string, t *workflowservice
 
 	// 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{})
+	isActivityCanceled := ctx.Err() == context.Canceled && IsCanceledError(context.Cause(ctx))
 
 	dlCancelFunc()
 	if <-ctx.Done(); ctx.Err() == context.DeadlineExceeded {

internal/internal_task_handlers_test.go

@@ -2167,6 +2167,58 @@ func (t *TaskHandlersTestSuite) TestActivityExecutionWorkerStop() {
 	t.NotNil(r)
 }
 
+func (t *TaskHandlersTestSuite) TestActivityCancellationUsesIsCanceledError() {
+	activityName := "activityCancellationIsCanceledError"
+	cancelContextActivity := func(ctx context.Context) error {
+		env := getActivityEnv(ctx)
+		invoker, ok := env.serviceInvoker.(*temporalInvoker)
+		t.Require().True(ok, "expected temporalInvoker")
+		invoker.cancelHandler(NewCanceledError())
+		<-ctx.Done()
+		return ctx.Err()
+	}
+
+	t.registry.RegisterActivityWithOptions(
+		cancelContextActivity,
+		RegisterActivityOptions{Name: activityName, DisableAlreadyRegisteredCheck: true},
+	)
+
+	mockCtrl := gomock.NewController(t.T())
+	mockService := workflowservicemock.NewMockWorkflowServiceClient(mockCtrl)
+	client := WorkflowClient{workflowService: mockService}
+	wep := t.getTestWorkerExecutionParams()
+	activityHandler := newActivityTaskHandler(&client, wep, t.registry)
+	now := time.Now()
+	pats := &workflowservice.PollActivityTaskQueueResponse{
+		Attempt:   1,
+		TaskToken: []byte("token"),
+		WorkflowExecution: &commonpb.WorkflowExecution{
+			WorkflowId: "wID",
+			RunId:      "rID",
+		},
+		ActivityType:           &commonpb.ActivityType{Name: activityName},
+		ActivityId:             uuid.NewString(),
+		ScheduledTime:          timestamppb.New(now),
+		ScheduleToCloseTimeout: durationpb.New(time.Second),
+		StartedTime:            timestamppb.New(now),
+		StartToCloseTimeout:    durationpb.New(time.Second),
+		HeartbeatTimeout:       durationpb.New(time.Second),
+		WorkflowType: &commonpb.WorkflowType{
+			Name: "wType",
+		},
+		WorkflowNamespace: wep.Namespace,
+	}
+
+	result, err := activityHandler.Execute(taskqueue, pats)
+	t.Require().NoError(err)
+
+	canceledReq, ok := result.(*workflowservice.RespondActivityTaskCanceledRequest)
+	t.Require().True(ok, "expected cancel response")
+	t.Equal(pats.TaskToken, canceledReq.TaskToken)
+	t.Equal(wep.Identity, canceledReq.Identity)
+	t.Equal(wep.Namespace, canceledReq.Namespace)
+}
+
 func Test_NonDeterministicCheck(t *testing.T) {
 	unimplementedCommands := []int32{
 		int32(enumspb.COMMAND_TYPE_UNSPECIFIED),

internal/session.go

@@ -419,7 +419,7 @@ func sessionCreationActivity(ctx context.Context, sessionID string) error {
 			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{})) {
+			if !(ctx.Err() == context.Canceled && IsCanceledError(context.Cause(ctx))) {
 				return NewApplicationErrorWithOptions(
 					"session failed due to worker shutdown", "SessionWorkerShutdown",
 					ApplicationErrorOptions{NonRetryable: true, Cause: ctx.Err()})