Disallow continue as new in update handlers (#2167)

Disallow continue as new in an update handler

Author: Quinn-With-Two-Ns

temporal-sdk/src/main/java/io/temporal/internal/statemachines/UnsupportedContinueAsNewRequest.java

@@ -0,0 +1,34 @@
+/*
+ * Copyright (C) 2022 Temporal Technologies, Inc. All Rights Reserved.
+ *
+ * Copyright (C) 2012-2016 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Modifications copyright (C) 2017 Uber Technologies, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this material except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.temporal.internal.statemachines;
+
+/**
+ * Thrown when {@link io.temporal.workflow.Workflow#continueAsNew} is called from an unsupported
+ * location.
+ *
+ * <p>The reason this class extends Error is for application workflow code to not catch it by
+ * mistake. The default behavior of the SDK is to block workflow execution while Error is thrown.
+ */
+public class UnsupportedContinueAsNewRequest extends Error {
+  public UnsupportedContinueAsNewRequest(String message) {
+    super(message);
+  }
+}

temporal-sdk/src/main/java/io/temporal/internal/sync/WorkflowExecutionHandler.java

@@ -29,6 +29,7 @@
 import io.temporal.failure.CanceledFailure;
 import io.temporal.failure.TemporalFailure;
 import io.temporal.internal.replay.ReplayWorkflowContext;
+import io.temporal.internal.statemachines.UnsupportedContinueAsNewRequest;
 import io.temporal.internal.worker.WorkflowExecutionException;
 import io.temporal.worker.WorkflowImplementationOptions;
 import io.temporal.workflow.Workflow;
@@ -123,6 +124,9 @@ public Optional<Payloads> handleExecuteUpdate(
       io.temporal.api.common.v1.Header header) {
     try {
       return context.handleExecuteUpdate(updateName, input, eventId, new Header(header));
+    } catch (UnsupportedContinueAsNewRequest e) {
+      // Re-throw to fail the workflow task
+      throw e;
     } catch (Throwable e) {
       applyWorkflowFailurePolicyAndRethrow(e);
     }

temporal-sdk/src/main/java/io/temporal/internal/sync/WorkflowInternal.java

@@ -45,6 +45,7 @@
 import io.temporal.internal.common.NonIdempotentHandle;
 import io.temporal.internal.common.SearchAttributesUtil;
 import io.temporal.internal.logging.ReplayAwareLogger;
+import io.temporal.internal.statemachines.UnsupportedContinueAsNewRequest;
 import io.temporal.serviceclient.CheckedExceptionWrapper;
 import io.temporal.workflow.*;
 import io.temporal.workflow.Functions.Func;
@@ -576,6 +577,7 @@ public static <R> R retry(
   public static void continueAsNew(
       @Nullable String workflowType, @Nullable ContinueAsNewOptions options, Object[] args) {
     assertNotReadOnly("continue as new");
+    assertNotInUpdateHandler("ContinueAsNew is not supported in an update handler");
     getWorkflowOutboundInterceptor()
         .continueAsNew(
             new WorkflowOutboundCallsInterceptor.ContinueAsNewInput(
@@ -588,6 +590,7 @@ public static void continueAsNew(
       Object[] args,
       WorkflowOutboundCallsInterceptor outboundCallsInterceptor) {
     assertNotReadOnly("continue as new");
+    assertNotInUpdateHandler("ContinueAsNew is not supported in an update handler");
     outboundCallsInterceptor.continueAsNew(
         new WorkflowOutboundCallsInterceptor.ContinueAsNewInput(
             workflowType, options, args, Header.empty()));
@@ -760,6 +763,12 @@ static void assertNotReadOnly(String action) {
     }
   }
 
+  static void assertNotInUpdateHandler(String message) {
+    if (getCurrentUpdateInfo().isPresent()) {
+      throw new UnsupportedContinueAsNewRequest(message);
+    }
+  }
+
   private static WorkflowThread getWorkflowThread() {
     return DeterministicRunnerImpl.currentThreadInternal();
   }

temporal-sdk/src/test/java/io/temporal/workflow/updateTest/UpdateContinueAsNewInHandlerTest.java

@@ -0,0 +1,114 @@
+/*
+ * Copyright (C) 2022 Temporal Technologies, Inc. All Rights Reserved.
+ *
+ * Copyright (C) 2012-2016 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Modifications copyright (C) 2017 Uber Technologies, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this material except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.temporal.workflow.updateTest;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
+
+import io.temporal.api.common.v1.WorkflowExecution;
+import io.temporal.client.*;
+import io.temporal.failure.ApplicationFailure;
+import io.temporal.internal.statemachines.UnsupportedContinueAsNewRequest;
+import io.temporal.testing.internal.SDKTestOptions;
+import io.temporal.testing.internal.SDKTestWorkflowRule;
+import io.temporal.worker.WorkflowImplementationOptions;
+import io.temporal.workflow.CompletablePromise;
+import io.temporal.workflow.Workflow;
+import io.temporal.workflow.shared.TestWorkflows.WorkflowWithUpdate;
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.UUID;
+import org.junit.Rule;
+import org.junit.Test;
+
+public class UpdateContinueAsNewInHandlerTest {
+  @Rule
+  public SDKTestWorkflowRule testWorkflowRule =
+      SDKTestWorkflowRule.newBuilder()
+          .setWorkflowTypes(
+              WorkflowImplementationOptions.newBuilder()
+                  .setFailWorkflowExceptionTypes(UnsupportedContinueAsNewRequest.class)
+                  .build(),
+              TestUpdateWorkflowImpl.class)
+          .build();
+
+  @Test
+  public void continueAsNewInUpdateHandler() {
+    String workflowId = UUID.randomUUID().toString();
+    WorkflowClient workflowClient = testWorkflowRule.getWorkflowClient();
+    WorkflowOptions options =
+        SDKTestOptions.newWorkflowOptionsWithTimeouts(testWorkflowRule.getTaskQueue()).toBuilder()
+            .setWorkflowId(workflowId)
+            .build();
+    WorkflowWithUpdate workflow = workflowClient.newWorkflowStub(WorkflowWithUpdate.class, options);
+
+    WorkflowExecution execution = WorkflowClient.start(workflow::execute);
+    WorkflowStub stub = WorkflowStub.fromTyped(workflow);
+    stub.startUpdate(
+        UpdateOptions.newBuilder(String.class)
+            .setUpdateName("update")
+            .setWaitForStage(WorkflowUpdateStage.ACCEPTED)
+            .build());
+    WorkflowFailedException e =
+        assertThrows(WorkflowFailedException.class, () -> stub.getResult(String.class));
+    assertEquals(
+        "io.temporal.internal.statemachines.UnsupportedContinueAsNewRequest",
+        ((ApplicationFailure) e.getCause()).getType());
+  }
+
+  public static class TestUpdateWorkflowImpl implements WorkflowWithUpdate {
+    String state = "initial";
+    List<String> updates = new ArrayList<>();
+    CompletablePromise<Void> promise = Workflow.newPromise();
+
+    @Override
+    public String execute() {
+      promise.get();
+      return "";
+    }
+
+    @Override
+    public String getState() {
+      return state;
+    }
+
+    @Override
+    public String update(Integer index, String value) {
+      // Sleep to make sure the update can be accepted before trying to continueAsNew
+      Workflow.sleep(Duration.ofSeconds(1));
+      // This should throw UnsupportedContinueAsNewRequest
+      Workflow.continueAsNew();
+      return "";
+    }
+
+    @Override
+    public void updateValidator(Integer index, String value) {}
+
+    @Override
+    public void complete() {
+      promise.complete(null);
+    }
+
+    @Override
+    public void completeValidator() {}
+  }
+}