Verify history is replayed up to StartedEventId (#1916)

Verify history is replayed up to StartedEventId

Author: Quinn-With-Two-Ns

temporal-sdk/src/main/java/io/temporal/internal/replay/ReplayWorkflowRunTaskHandler.java

@@ -228,20 +228,22 @@ private void handleWorkflowTaskImpl(
     workflowStateMachines.setWorklfowStartedEventId(workflowTask.getStartedEventId());
     workflowStateMachines.setReplaying(workflowTask.getPreviousStartedEventId() > 0);
     workflowStateMachines.setMessages(workflowTask.getMessagesList());
-    applyServerHistory(historyIterator);
+    applyServerHistory(workflowTask.getStartedEventId(), historyIterator);
   }
 
-  private void applyServerHistory(WorkflowHistoryIterator historyIterator) {
+  private void applyServerHistory(long lastEventId, WorkflowHistoryIterator historyIterator) {
     Duration expiration = toJavaDuration(startedEvent.getWorkflowTaskTimeout());
     historyIterator.initDeadline(Deadline.after(expiration.toMillis(), TimeUnit.MILLISECONDS));
 
     boolean timerStopped = false;
     Stopwatch sw = metricsScope.timer(MetricsType.WORKFLOW_TASK_REPLAY_LATENCY).start();
+    long currentEventId = 0;
     try {
       while (historyIterator.hasNext()) {
         // iteration itself is intentionally left outside the try-catch below,
         // as gRPC exception happened during history iteration should never ever fail the workflow
         HistoryEvent event = historyIterator.next();
+        currentEventId = event.getEventId();
         boolean hasNext = historyIterator.hasNext();
         try {
           workflowStateMachines.handleEvent(event, hasNext);
@@ -264,13 +266,27 @@ private void applyServerHistory(WorkflowHistoryIterator historyIterator) {
           timerStopped = true;
         }
       }
+      verifyAllEventsProcessed(lastEventId, currentEventId);
     } finally {
       if (!timerStopped) {
         sw.stop();
       }
     }
   }
 
+  // Verify the received and processed all events up to the last one we knew about from the polled
+  // task.
+  // It is possible for the server to send fewer events than required if we are reading history from
+  // a stale node.
+  private void verifyAllEventsProcessed(long lastEventId, long processedEventId) {
+    if (lastEventId != Long.MAX_VALUE && lastEventId > 0 && processedEventId < lastEventId) {
+      throw new IllegalStateException(
+          String.format(
+              "Premature end of stream, expectedLastEventID=%d but no more events after eventID=%d",
+              lastEventId, processedEventId));
+    }
+  }
+
   private Map<String, WorkflowQueryResult> executeQueries(Map<String, WorkflowQuery> queries) {
     Map<String, WorkflowQueryResult> queryResults = new HashMap<>();
     for (Map.Entry<String, WorkflowQuery> entry : queries.entrySet()) {

temporal-sdk/src/test/java/io/temporal/internal/replay/ReplayWorkflowRunTaskHandlerTaskHandlerTests.java

@@ -29,18 +29,27 @@
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
+import com.google.protobuf.ByteString;
 import com.google.protobuf.util.Durations;
 import com.uber.m3.tally.NoopScope;
+import io.temporal.api.enums.v1.EventType;
+import io.temporal.api.history.v1.History;
+import io.temporal.api.history.v1.HistoryEvent;
 import io.temporal.api.taskqueue.v1.StickyExecutionAttributes;
+import io.temporal.api.workflowservice.v1.GetSystemInfoResponse;
+import io.temporal.api.workflowservice.v1.GetWorkflowExecutionHistoryResponse;
 import io.temporal.api.workflowservice.v1.PollWorkflowTaskQueueResponse;
+import io.temporal.api.workflowservice.v1.WorkflowServiceGrpc;
 import io.temporal.internal.common.InternalUtils;
 import io.temporal.internal.worker.SingleWorkerOptions;
 import io.temporal.internal.worker.WorkflowExecutorCache;
 import io.temporal.internal.worker.WorkflowRunLockManager;
 import io.temporal.internal.worker.WorkflowTaskHandler;
+import io.temporal.serviceclient.WorkflowServiceStubs;
 import io.temporal.testUtils.HistoryUtils;
 import io.temporal.testing.internal.SDKTestWorkflowRule;
 import java.time.Duration;
+import java.util.List;
 import java.util.Optional;
 import org.junit.Rule;
 import org.junit.Test;
@@ -77,6 +86,62 @@ public void ifStickyExecutionAttributesAreNotSetThenWorkflowsAreNotCached() thro
     assertFalse(result.getTaskCompleted().hasStickyAttributes());
   }
 
+  @Test
+  public void workflowTaskFailOnIncompleteHistory() throws Throwable {
+    assumeFalse("skipping for docker tests", SDKTestWorkflowRule.useExternalService);
+
+    WorkflowExecutorCache cache =
+        new WorkflowExecutorCache(10, new WorkflowRunLockManager(), new NoopScope());
+    WorkflowServiceStubs client = mock(WorkflowServiceStubs.class);
+    when(client.getServerCapabilities())
+        .thenReturn(() -> GetSystemInfoResponse.Capabilities.newBuilder().build());
+    WorkflowServiceGrpc.WorkflowServiceBlockingStub blockingStub =
+        mock(WorkflowServiceGrpc.WorkflowServiceBlockingStub.class);
+    when(client.blockingStub()).thenReturn(blockingStub);
+    when(blockingStub.withOption(any(), any())).thenReturn(blockingStub);
+
+    // Simulate a stale history node sending a workflow task with an incomplete history
+    List<HistoryEvent> history =
+        HistoryUtils.generateWorkflowTaskWithInitialHistory().getHistory().getEventsList();
+    assertEquals(3, history.size());
+    assertEquals(
+        EventType.EVENT_TYPE_WORKFLOW_TASK_STARTED, history.get(history.size() - 1).getEventType());
+    history = history.subList(0, history.size() - 1);
+    when(blockingStub.getWorkflowExecutionHistory(any()))
+        .thenReturn(
+            GetWorkflowExecutionHistoryResponse.newBuilder()
+                .setHistory(History.newBuilder().addAllEvents(history).build())
+                .build());
+
+    WorkflowTaskHandler taskHandler =
+        new ReplayWorkflowTaskHandler(
+            "namespace",
+            setUpMockWorkflowFactory(),
+            cache,
+            SingleWorkerOptions.newBuilder().build(),
+            null,
+            Duration.ofSeconds(5),
+            client,
+            null);
+
+    // Send a poll with a partial history and no cached execution so the SDK will request a full
+    // history
+    WorkflowTaskHandler.Result result =
+        taskHandler.handleWorkflowTask(
+            HistoryUtils.generateWorkflowTaskWithInitialHistory().toBuilder()
+                .setHistory(History.newBuilder().build())
+                .setNextPageToken(ByteString.EMPTY)
+                .build());
+
+    // Assert
+    assertEquals(0, cache.size());
+    assertNotNull(result.getTaskFailed());
+    assertTrue(result.getTaskFailed().hasFailure());
+    assertEquals(
+        "Premature end of stream, expectedLastEventID=3 but no more events after eventID=2",
+        result.getTaskFailed().getFailure().getMessage());
+  }
+
   @Test
   public void ifStickyExecutionAttributesAreSetThenWorkflowsAreCached() throws Throwable {
     assumeFalse("skipping for docker tests", SDKTestWorkflowRule.useExternalService);

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

@@ -329,18 +329,15 @@ public void sendQueryTask(
       List<HistoryEvent> events = new ArrayList<>(historyStore.getEventsLocked());
       History.Builder history = History.newBuilder();
       PeekingIterator<HistoryEvent> iterator = Iterators.peekingIterator(events.iterator());
-      long startedEventId = 0;
       long previousStaredEventId = 0;
       while (iterator.hasNext()) {
         HistoryEvent event = iterator.next();
         if (event.getEventType() == EventType.EVENT_TYPE_WORKFLOW_TASK_STARTED) {
           if (!iterator.hasNext()
               || iterator.peek().getEventType() == EventType.EVENT_TYPE_WORKFLOW_TASK_COMPLETED) {
             previousStaredEventId = event.getEventId();
-            startedEventId = previousStaredEventId;
           }
         } else if (WorkflowExecutionUtils.isWorkflowExecutionClosedEvent(event)) {
-          startedEventId = 0;
           if (iterator.hasNext()) {
             throw Status.INTERNAL
                 .withDescription("Unexpected event after the completion event: " + iterator.peek())
@@ -351,7 +348,7 @@ public void sendQueryTask(
       task.setPreviousStartedEventId(previousStaredEventId);
       // it's not a real workflow task and the server sends 0 for startedEventId for such a workflow
       // task
-      task.setStartedEventId(startedEventId);
+      task.setStartedEventId(0);
       if (taskQueue.getTaskQueueName().equals(task.getWorkflowExecutionTaskQueue().getName())) {
         history.addAllEvents(events);
       } else {