Accept CancelTimer for a timer that fired during a workflow task

temporal-test-server/src/main/java/io/temporal/internal/testservice/StateMachines.java

@@ -2440,6 +2440,27 @@ private static void cancelTimer(
     ctx.addEvent(event);
   }
 
+  /**
+   * Records a TIMER_CANCELED event for a timer that has already fired while a workflow task was in
+   * progress. Such a timer has no state machine associated with it anymore, and its buffered
+   * TIMER_FIRED event is expected to be discarded by the caller instead of being committed to the
+   * history.
+   */
+  static void cancelFiredTimer(
+      RequestContext ctx, String timerId, long startedEventId, long workflowTaskCompletedEventId) {
+    TimerCanceledEventAttributes.Builder a =
+        TimerCanceledEventAttributes.newBuilder()
+            .setWorkflowTaskCompletedEventId(workflowTaskCompletedEventId)
+            .setTimerId(timerId)
+            .setStartedEventId(startedEventId);
+    HistoryEvent event =
+        HistoryEvent.newBuilder()
+            .setEventType(EventType.EVENT_TYPE_TIMER_CANCELED)
+            .setTimerCanceledEventAttributes(a)
+            .build();
+    ctx.addEvent(event);
+  }
+
   private static void initiateExternalSignal(
       RequestContext ctx,
       SignalExternalData data,

temporal-test-server/src/main/java/io/temporal/internal/testservice/TestWorkflowMutableStateImpl.java

@@ -661,15 +661,35 @@ private void failWorkflowTaskWithAReason(
 
   private boolean unhandledCommand(RespondWorkflowTaskCompletedRequest request) {
     boolean newEvents = false;
+    outer:
     for (RequestContext ctx2 : workflowTaskStateMachine.getData().bufferedEvents) {
-      if (!ctx2.getEvents().isEmpty()) {
+      for (HistoryEvent event : ctx2.getEvents()) {
+        if (event.getEventType() == EventType.EVENT_TYPE_TIMER_FIRED
+            && isCancelledByCommands(event.getTimerFiredEventAttributes().getTimerId(), request)) {
+          // A CancelTimer command from this request replaces this buffered TIMER_FIRED event with a
+          // TIMER_CANCELED event, so it does not make the workflow completion command unhandled.
+          // A workflow completion command is always the last one, so it is processed after the
+          // cancel, like on the real server.
+          continue;
+        }
         newEvents = true;
-        break;
+        break outer;
       }
     }
     return (newEvents && hasCompletionCommand(request.getCommandsList()));
   }
 
+  private boolean isCancelledByCommands(
+      String timerId, RespondWorkflowTaskCompletedRequest request) {
+    for (Command command : request.getCommandsList()) {
+      if (command.getCommandType() == CommandType.COMMAND_TYPE_CANCEL_TIMER
+          && command.getCancelTimerCommandAttributes().getTimerId().equals(timerId)) {
+        return true;
+      }
+    }
+    return false;
+  }
+
   private boolean unhandledMessages(RespondWorkflowTaskCompletedRequest request) {
     return (!workflowTaskStateMachine.getData().updateRequestBuffer.isEmpty()
         && hasCompletionCommand(request.getCommandsList()));
@@ -1004,14 +1024,50 @@ private void processCancelTimer(
     String timerId = d.getTimerId();
     StateMachine<TimerData> timer = timers.get(timerId);
     if (timer == null) {
-      throw Status.INVALID_ARGUMENT
-          .withDescription("invalid history builder state for action")
-          .asRuntimeException();
+      // The timer may have fired while the current workflow task was still in progress. In that
+      // case its TIMER_FIRED event is buffered until the workflow task completion and this
+      // command should replace it with a TIMER_CANCELED event instead of failing the workflow
+      // task. This mirrors the real server, which removes the buffered TimerFired event.
+      HistoryEvent timerFiredEvent = removeBufferedTimerFiredEvent(timerId);
+      if (timerFiredEvent == null) {
+        throw Status.INVALID_ARGUMENT
+            .withDescription("invalid history builder state for action")
+            .asRuntimeException();
+      }
+      long startedEventId = timerFiredEvent.getTimerFiredEventAttributes().getStartedEventId();
+      StateMachines.cancelFiredTimer(ctx, timerId, startedEventId, workflowTaskCompletedId);
+      // The removal of the buffered TIMER_FIRED event may leave the buffered context empty, so
+      // request the next workflow task explicitly instead of relying on the leftover buffered
+      // context to trigger the scheduling.
+      ctx.setNeedWorkflowTask(true);
+      return;
     }
     timer.action(StateMachines.Action.CANCEL, ctx, d, workflowTaskCompletedId);
     timers.remove(timerId);
   }
 
+  /**
+   * Removes a buffered TIMER_FIRED event of the specified timer, if any. Events are buffered only
+   * if the timer fired while a workflow task was in progress and the workflow task completion has
+   * not flushed them yet.
+   *
+   * @return the removed TIMER_FIRED event or {@code null} if there is no such buffered event.
+   */
+  private @Nullable HistoryEvent removeBufferedTimerFiredEvent(String timerId) {
+    List<RequestContext> bufferedEvents = workflowTaskStateMachine.getData().bufferedEvents;
+    for (RequestContext bufferedCtx : bufferedEvents) {
+      List<HistoryEvent> events = bufferedCtx.getEvents();
+      for (int i = 0; i < events.size(); i++) {
+        HistoryEvent event = events.get(i);
+        if (event.getEventType() == EventType.EVENT_TYPE_TIMER_FIRED
+            && event.getTimerFiredEventAttributes().getTimerId().equals(timerId)) {
+          return events.remove(i);
+        }
+      }
+    }
+    return null;
+  }
+
   private void processRequestCancelActivityTask(
       RequestContext ctx,
       RequestCancelActivityTaskCommandAttributes a,

temporal-test-server/src/test/java/io/temporal/testserver/functional/CancelFiredTimerTest.java

@@ -0,0 +1,254 @@
+package io.temporal.testserver.functional;
+
+import static io.temporal.internal.common.InternalUtils.createNormalTaskQueue;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+import com.google.protobuf.ByteString;
+import io.temporal.api.command.v1.CancelTimerCommandAttributes;
+import io.temporal.api.command.v1.Command;
+import io.temporal.api.command.v1.CompleteWorkflowExecutionCommandAttributes;
+import io.temporal.api.command.v1.ScheduleActivityTaskCommandAttributes;
+import io.temporal.api.command.v1.StartTimerCommandAttributes;
+import io.temporal.api.common.v1.ActivityType;
+import io.temporal.api.common.v1.Payloads;
+import io.temporal.api.common.v1.WorkflowExecution;
+import io.temporal.api.common.v1.WorkflowType;
+import io.temporal.api.enums.v1.CommandType;
+import io.temporal.api.enums.v1.EventType;
+import io.temporal.api.history.v1.HistoryEvent;
+import io.temporal.api.taskqueue.v1.TaskQueue;
+import io.temporal.api.workflowservice.v1.GetWorkflowExecutionHistoryRequest;
+import io.temporal.api.workflowservice.v1.PollWorkflowTaskQueueResponse;
+import io.temporal.api.workflowservice.v1.RespondWorkflowTaskCompletedRequest;
+import io.temporal.api.workflowservice.v1.StartWorkflowExecutionRequest;
+import io.temporal.internal.common.ProtobufTimeUtils;
+import io.temporal.serviceclient.WorkflowServiceStubs;
+import io.temporal.serviceclient.WorkflowServiceStubsOptions;
+import io.temporal.testing.internal.TestServiceUtils;
+import io.temporal.testserver.TestServer;
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import java.util.UUID;
+import java.util.concurrent.TimeUnit;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+/**
+ * Verifies that a timer that fired while a workflow task was in progress can be cancelled by that
+ * same workflow task. The buffered TIMER_FIRED event is replaced with a TIMER_CANCELED event
+ * instead of failing the workflow task completion with INVALID_ARGUMENT.
+ *
+ * @see <a href="https://github.com/temporalio/sdk-java/issues/2606">Issue 2606</a>
+ */
+public class CancelFiredTimerTest {
+
+  private static final Duration TIMER_DURATION = Duration.ofSeconds(2);
+  private static final long TIMER_FIRING_WAIT_MILLIS = 4000;
+
+  private static final String NAMESPACE = "namespace";
+  private static final String TASK_QUEUE = "taskQueue";
+  private static final String WORKFLOW_TYPE = "wfType";
+  private static final String TIMER_ID = "timer";
+  private static final String ACTIVITY_ID = "activity";
+
+  private TestServer.InProcessTestServer testServer;
+  private WorkflowServiceStubs workflowServiceStubs;
+
+  @Before
+  public void setUp() {
+    this.testServer = TestServer.createServer(true);
+    this.workflowServiceStubs =
+        WorkflowServiceStubs.newServiceStubs(
+            WorkflowServiceStubsOptions.newBuilder()
+                .setChannel(testServer.getChannel())
+                .validateAndBuildWithDefaults());
+  }
+
+  @After
+  public void tearDown() {
+    this.workflowServiceStubs.shutdownNow();
+    this.workflowServiceStubs.awaitTermination(1, TimeUnit.SECONDS);
+    this.testServer.close();
+  }
+
+  /**
+   * Verifies that the workflow task that cancels the fired timer can also complete the workflow in
+   * the same response. Previously the completion failed with INVALID_ARGUMENT "UnhandledCommand".
+   */
+  @Test
+  public void cancelFiredTimerWithWorkflowCompletion() throws Exception {
+    String workflowId = UUID.randomUUID().toString();
+    PollWorkflowTaskQueueResponse task = startWorkflowAndPollTaskWithFiredTimer(workflowId);
+
+    respondWorkflowTaskCompleted(
+        task.getTaskToken(), cancelTimerCommand(), completeWorkflowCommand());
+
+    List<HistoryEvent> history = getHistory(workflowId);
+    assertCompleted(history);
+    assertTimerCancelledAndNotFired(history);
+  }
+
+  /**
+   * Verifies that a workflow task that cancels the fired timer without completing the workflow is
+   * accepted, and the workflow is not wedged and keeps making progress. Previously the completion
+   * failed with INVALID_ARGUMENT "invalid history builder state for action".
+   */
+  @Test
+  public void cancelFiredTimerWithoutWorkflowCompletion() throws Exception {
+    String workflowId = UUID.randomUUID().toString();
+    PollWorkflowTaskQueueResponse task = startWorkflowAndPollTaskWithFiredTimer(workflowId);
+
+    respondWorkflowTaskCompleted(task.getTaskToken(), cancelTimerCommand());
+    List<HistoryEvent> history = getHistory(workflowId);
+    assertTimerCancelledAndNotFired(history);
+
+    // The workflow must still be able to progress and complete after the race.
+    PollWorkflowTaskQueueResponse completionTask =
+        TestServiceUtils.pollWorkflowTaskQueue(
+            NAMESPACE, createNormalTaskQueue(TASK_QUEUE), workflowServiceStubs);
+    respondWorkflowTaskCompleted(completionTask.getTaskToken(), completeWorkflowCommand());
+    assertCompleted(getHistory(workflowId));
+  }
+
+  /**
+   * Starts a workflow and brings it to a state where a timer fired while a workflow task is in
+   * progress.
+   *
+   * <p>The first workflow task schedules an activity and starts a short timer. The outstanding
+   * activity holds the time skipping lock, so the timer can only fire in real time. A signal then
+   * schedules the second workflow task, and the method polls and starts it. While the second
+   * workflow task is in progress, the timer fires in real time and its TIMER_FIRED event gets
+   * buffered until the workflow task completion.
+   *
+   * @return the response of the started workflow task to complete with commands.
+   */
+  private PollWorkflowTaskQueueResponse startWorkflowAndPollTaskWithFiredTimer(String workflowId)
+      throws Exception {
+    StartWorkflowExecutionRequest startRequest =
+        StartWorkflowExecutionRequest.newBuilder()
+            .setRequestId(UUID.randomUUID().toString())
+            .setNamespace(NAMESPACE)
+            .setWorkflowId(workflowId)
+            .setTaskQueue(createNormalTaskQueue(TASK_QUEUE))
+            .setWorkflowRunTimeout(ProtobufTimeUtils.toProtoDuration(Duration.ofSeconds(100)))
+            .setWorkflowTaskTimeout(ProtobufTimeUtils.toProtoDuration(Duration.ofSeconds(100)))
+            .setWorkflowType(WorkflowType.newBuilder().setName(WORKFLOW_TYPE))
+            .build();
+    workflowServiceStubs.blockingStub().startWorkflowExecution(startRequest);
+
+    PollWorkflowTaskQueueResponse firstTask =
+        TestServiceUtils.pollWorkflowTaskQueue(
+            NAMESPACE, createNormalTaskQueue(TASK_QUEUE), workflowServiceStubs);
+    respondWorkflowTaskCompleted(
+        firstTask.getTaskToken(), startTimerCommand(), scheduleActivityTaskCommand());
+
+    // The signal schedules the second workflow task while the activity keeps the time skipping
+    // locked, so the clock tracks real time from here on.
+    TestServiceUtils.signalWorkflow(
+        WorkflowExecution.newBuilder().setWorkflowId(workflowId).build(),
+        NAMESPACE,
+        workflowServiceStubs);
+    PollWorkflowTaskQueueResponse secondTask =
+        TestServiceUtils.pollWorkflowTaskQueue(
+            NAMESPACE, createNormalTaskQueue(TASK_QUEUE), workflowServiceStubs);
+
+    // Wait for the timer to fire in real time while the second workflow task is in progress. The
+    // timer has a two second timeout and the sleep is four seconds long to leave a margin both
+    // before the timer fires and for the test server to buffer the fired event.
+    Thread.sleep(TIMER_FIRING_WAIT_MILLIS);
+    return secondTask;
+  }
+
+  private void assertCompleted(List<HistoryEvent> history) {
+    assertTrue(
+        "Expected the workflow to complete",
+        history.stream()
+            .anyMatch(
+                event ->
+                    event.getEventType() == EventType.EVENT_TYPE_WORKFLOW_EXECUTION_COMPLETED));
+  }
+
+  private void assertTimerCancelledAndNotFired(List<HistoryEvent> history) {
+    assertTrue(
+        "Expected a TIMER_CANCELED event for the cancelled timer",
+        history.stream()
+            .anyMatch(
+                event ->
+                    event.getEventType() == EventType.EVENT_TYPE_TIMER_CANCELED
+                        && event.getTimerCanceledEventAttributes().getTimerId().equals(TIMER_ID)));
+    assertFalse(
+        "The buffered TIMER_FIRED event of the cancelled timer should be replaced with the TIMER_CANCELED event",
+        history.stream()
+            .anyMatch(
+                event ->
+                    event.getEventType() == EventType.EVENT_TYPE_TIMER_FIRED
+                        && event.getTimerFiredEventAttributes().getTimerId().equals(TIMER_ID)));
+  }
+
+  private List<HistoryEvent> getHistory(String workflowId) {
+    GetWorkflowExecutionHistoryRequest request =
+        GetWorkflowExecutionHistoryRequest.newBuilder()
+            .setNamespace(NAMESPACE)
+            .setExecution(WorkflowExecution.newBuilder().setWorkflowId(workflowId))
+            .build();
+    return new ArrayList<>(
+        workflowServiceStubs
+            .blockingStub()
+            .getWorkflowExecutionHistory(request)
+            .getHistory()
+            .getEventsList());
+  }
+
+  private void respondWorkflowTaskCompleted(ByteString taskToken, Command... commands) {
+    RespondWorkflowTaskCompletedRequest request =
+        RespondWorkflowTaskCompletedRequest.newBuilder()
+            .setTaskToken(taskToken)
+            .addAllCommands(Arrays.asList(commands))
+            .build();
+    workflowServiceStubs.blockingStub().respondWorkflowTaskCompleted(request);
+  }
+
+  private Command startTimerCommand() {
+    return Command.newBuilder()
+        .setCommandType(CommandType.COMMAND_TYPE_START_TIMER)
+        .setStartTimerCommandAttributes(
+            StartTimerCommandAttributes.newBuilder()
+                .setTimerId(TIMER_ID)
+                .setStartToFireTimeout(ProtobufTimeUtils.toProtoDuration(TIMER_DURATION)))
+        .build();
+  }
+
+  private Command cancelTimerCommand() {
+    return Command.newBuilder()
+        .setCommandType(CommandType.COMMAND_TYPE_CANCEL_TIMER)
+        .setCancelTimerCommandAttributes(
+            CancelTimerCommandAttributes.newBuilder().setTimerId(TIMER_ID))
+        .build();
+  }
+
+  private Command completeWorkflowCommand() {
+    return Command.newBuilder()
+        .setCommandType(CommandType.COMMAND_TYPE_COMPLETE_WORKFLOW_EXECUTION)
+        .setCompleteWorkflowExecutionCommandAttributes(
+            CompleteWorkflowExecutionCommandAttributes.newBuilder()
+                .setResult(Payloads.getDefaultInstance()))
+        .build();
+  }
+
+  private Command scheduleActivityTaskCommand() {
+    return Command.newBuilder()
+        .setCommandType(CommandType.COMMAND_TYPE_SCHEDULE_ACTIVITY_TASK)
+        .setScheduleActivityTaskCommandAttributes(
+            ScheduleActivityTaskCommandAttributes.newBuilder()
+                .setActivityId(ACTIVITY_ID)
+                .setActivityType(ActivityType.newBuilder().setName("activity"))
+                .setTaskQueue(TaskQueue.newBuilder().setName(TASK_QUEUE))
+                .setScheduleToCloseTimeout(
+                    ProtobufTimeUtils.toProtoDuration(Duration.ofSeconds(60))))
+        .build();
+  }
+}