Reset lastHandledEventId on speculative WFT (#1881)

Author: Quinn-With-Two-Ns

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

@@ -202,7 +202,16 @@ public void setWorklfowStartedEventId(long workflowTaskStartedEventId) {
   }
 
   public void setCurrentStartedEventId(long eventId) {
+    // We have to drop any state machines (which should only be one workflow task machine)
+    // created when handling the speculative workflow task
+    for (long i = this.lastHandledEventId; i > eventId; i--) {
+      stateMachines.remove(i);
+    }
     this.currentStartedEventId = eventId;
+    // When we reset the event ID on a speculative WFT we need to move this counter back
+    // to the last WFT completed to allow new tasks to be processed. Assume the WFT complete
+    // always follows the WFT started.
+    this.lastHandledEventId = eventId + 1;
   }
 
   public long getCurrentStartedEventId() {

temporal-sdk/src/test/java/io/temporal/workflow/shared/TestWorkflows.java

@@ -214,6 +214,9 @@ public interface WorkflowWithUpdateAndSignal {
     @UpdateMethod()
     String update(String value);
 
+    @UpdateValidatorMethod(updateName = "update")
+    void validator(String value);
+
     @UpdateMethod
     void complete();
   }

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

@@ -23,8 +23,10 @@
 import static org.junit.Assert.*;
 
 import io.temporal.api.common.v1.WorkflowExecution;
+import io.temporal.api.enums.v1.EventType;
 import io.temporal.client.WorkflowClient;
 import io.temporal.client.WorkflowOptions;
+import io.temporal.client.WorkflowUpdateException;
 import io.temporal.testing.internal.SDKTestOptions;
 import io.temporal.testing.internal.SDKTestWorkflowRule;
 import io.temporal.worker.WorkerOptions;
@@ -89,6 +91,45 @@ public void testUpdateWithSignal() {
         workflow.execute());
   }
 
+  @Test
+  public void testSpeculativeUpdateWithSignal() {
+    String workflowId = UUID.randomUUID().toString();
+    WorkflowClient workflowClient = testWorkflowRule.getWorkflowClient();
+    WorkflowOptions options =
+        SDKTestOptions.newWorkflowOptionsWithTimeouts(testWorkflowRule.getTaskQueue()).toBuilder()
+            .setWorkflowId(workflowId)
+            .build();
+    TestWorkflows.WorkflowWithUpdateAndSignal workflow =
+        workflowClient.newWorkflowStub(TestWorkflows.WorkflowWithUpdateAndSignal.class, options);
+    // To execute workflow client.execute() would do. But we want to start workflow and immediately
+    // return.
+    WorkflowExecution execution = WorkflowClient.start(workflow::execute);
+
+    assertEquals(workflowId, execution.getWorkflowId());
+
+    SDKTestWorkflowRule.waitForOKQuery(workflow);
+    assertEquals("initial", workflow.getState());
+
+    for (int i = 0; i < 5; i++) {
+      assertThrows(WorkflowUpdateException.class, () -> workflow.update(""));
+    }
+
+    workflow.getState();
+
+    workflow.signal("signal 1");
+
+    assertEquals("update 1", workflow.update("update 1"));
+
+    workflow.signal("signal 2");
+
+    workflow.complete();
+
+    assertEquals(Arrays.asList("signal 1", "update 1", "signal 2"), workflow.execute());
+    SDKTestWorkflowRule.assertNoHistoryEvent(
+        workflowClient.fetchHistory(execution.getWorkflowId()).getHistory(),
+        EventType.EVENT_TYPE_WORKFLOW_TASK_FAILED);
+  }
+
   public static class TestUpdateWithSignalWorkflowImpl
       implements TestWorkflows.WorkflowWithUpdateAndSignal {
     String state = "initial";
@@ -116,10 +157,18 @@ public void signal(String value) {
 
     @Override
     public String update(String value) {
+      Workflow.sleep(100);
       updatesAndSignals.add(value);
       return value;
     }
 
+    @Override
+    public void validator(String value) {
+      if (value.isEmpty()) {
+        throw new RuntimeException("Empty value");
+      }
+    }
+
     @Override
     public void complete() {
       promise.complete(null);