Fix erroneous closing of workflow tasks instead of heartbeating in some edge cases (#1295)

Author: Spikhalskiy

temporal-sdk/src/main/java/io/temporal/internal/Config.java

@@ -0,0 +1,26 @@
+/*
+ * 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;
+
+public class Config {
+  /** Force new workflow task after workflow task timeout multiplied by this coefficient. */
+  public static final double WORKFLOW_TAK_HEARTBEAT_COEFFICIENT = 4d / 5d;
+}

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

@@ -23,6 +23,7 @@
 import static io.temporal.internal.common.ProtobufTimeUtils.toJavaDuration;
 import static io.temporal.serviceclient.CheckedExceptionWrapper.wrap;
 
+import com.google.common.base.Preconditions;
 import com.google.common.base.Throwables;
 import com.google.protobuf.util.Durations;
 import com.google.protobuf.util.Timestamps;
@@ -36,6 +37,7 @@
 import io.temporal.api.query.v1.WorkflowQuery;
 import io.temporal.api.query.v1.WorkflowQueryResult;
 import io.temporal.api.workflowservice.v1.PollWorkflowTaskQueueResponseOrBuilder;
+import io.temporal.internal.Config;
 import io.temporal.internal.statemachines.StatesMachinesCallback;
 import io.temporal.internal.statemachines.WorkflowStateMachines;
 import io.temporal.internal.worker.*;
@@ -61,9 +63,6 @@
  */
 class ReplayWorkflowRunTaskHandler implements WorkflowRunTaskHandler {
 
-  /** Force new decision task after workflow task timeout multiplied by this coefficient. */
-  public static final double FORCED_DECISION_TIME_COEFFICIENT = 4d / 5d;
-
   private final WorkflowServiceStubs service;
 
   private final String namespace;
@@ -86,6 +85,7 @@ class ReplayWorkflowRunTaskHandler implements WorkflowRunTaskHandler {
   private final WorkflowStateMachines workflowStateMachines;
 
   /** Number of non completed local activity tasks */
+  // TODO move and maintain this counter inside workflowStateMachines
   private int localActivityTaskCount;
 
   private final ReplayWorkflowExecutor replayWorkflowExecutor;
@@ -252,54 +252,58 @@ private void processLocalActivityRequests(long startTimeNs) throws InterruptedEx
     long durationUntilWFTHeartbeatNs =
         (long)
             (Durations.toNanos(startedEvent.getWorkflowTaskTimeout())
-                * FORCED_DECISION_TIME_COEFFICIENT);
+                * Config.WORKFLOW_TAK_HEARTBEAT_COEFFICIENT);
 
     long nextWFTHeartbeatTimeNs = startTimeNs + durationUntilWFTHeartbeatNs;
+
     while (true) {
       List<ExecuteLocalActivityParameters> laRequests =
           workflowStateMachines.takeLocalActivityRequests();
+      localActivityTaskCount += laRequests.size();
+
       for (ExecuteLocalActivityParameters laRequest : laRequests) {
         // TODO(maxim): In the presence of workflow task heartbeat this timeout doesn't make
         // much sense. I believe we should add ScheduleToStart timeout for the local activities
         // as well.
-        Duration maxWaitTime = Duration.ofNanos(nextWFTHeartbeatTimeNs - System.nanoTime());
-        if (maxWaitTime.isNegative()) {
-          maxWaitTime = Duration.ZERO;
-        }
+        long maxWaitTimeNs = Math.max(nextWFTHeartbeatTimeNs - System.nanoTime(), 0);
         boolean accepted =
             localActivityTaskPoller.apply(
-                new LocalActivityTask(laRequest, localActivityCompletionSink), maxWaitTime);
-        localActivityTaskCount++;
-        if (!accepted) {
-          throw new IllegalStateException(
-              "Unable to schedule local activity for execution, no more slots available and local activity task queue is full");
-        }
+                new LocalActivityTask(laRequest, localActivityCompletionSink),
+                Duration.ofNanos(maxWaitTimeNs));
+        Preconditions.checkState(
+            accepted,
+            "Unable to schedule local activity for execution, "
+                + "no more slots available and local activity task queue is full");
       }
+
       if (localActivityTaskCount == 0) {
         // No outstanding local activity requests
         break;
       }
-      waitAndProcessLocalActivityCompletion(nextWFTHeartbeatTimeNs);
-      if (nextWFTHeartbeatTimeNs <= System.nanoTime()) {
+
+      long maxWaitTimeTillHeartbeatNs = Math.max(nextWFTHeartbeatTimeNs - System.nanoTime(), 0);
+      ActivityTaskHandler.Result laCompletion =
+          localActivityCompletionQueue.poll(maxWaitTimeTillHeartbeatNs, TimeUnit.NANOSECONDS);
+      if (laCompletion == null) {
+        // Need to force a new task as we are out of time
         break;
       }
-    }
-  }
 
-  private void waitAndProcessLocalActivityCompletion(long nextForcedDecisionTimeNs)
-      throws InterruptedException {
-    long maxWaitTimeNs = nextForcedDecisionTimeNs - System.nanoTime();
-    if (maxWaitTimeNs <= 0) {
-      return;
-    }
-    ActivityTaskHandler.Result laCompletion;
-    laCompletion = localActivityCompletionQueue.poll(maxWaitTimeNs, TimeUnit.NANOSECONDS);
-    if (laCompletion == null) {
-      // Need to force a new task as nextForcedDecisionTime has passed.
-      return;
+      localActivityTaskCount--;
+      workflowStateMachines.handleLocalActivityCompletion(laCompletion);
+      // handleLocalActivityCompletion triggers eventLoop.
+      // After this call, there may be new local activity requests available in
+      // workflowStateMachines.takeLocalActivityRequests()
+      // These requests need to be processed and accounted for, otherwise we may end up not
+      // heartbeating and completing workflow task instead. So we have to make another iteration.
     }
-    localActivityTaskCount--;
-    workflowStateMachines.handleLocalActivityCompletion(laCompletion);
+
+    // it's safe to call and discard the result of takeLocalActivityRequests() here, because if it's
+    // not empty - we are in trouble anyway
+    Preconditions.checkState(
+        workflowStateMachines.takeLocalActivityRequests().isEmpty(),
+        "[BUG] Local activities requests from the last event loop were not drained "
+            + "and accounted in the outstanding local activities counter");
   }
 
   private class StatesMachinesCallbackImpl implements StatesMachinesCallback {

temporal-sdk/src/test/java/io/temporal/workflow/activityTests/LocalActivityGettingScheduledRightBeforeWorkflowTaskHeartbeatTest.java

@@ -0,0 +1,112 @@
+/*
+ * 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.activityTests;
+
+import static org.junit.Assert.assertEquals;
+
+import io.temporal.client.WorkflowOptions;
+import io.temporal.client.WorkflowStub;
+import io.temporal.internal.Config;
+import io.temporal.internal.Issue;
+import io.temporal.testing.internal.SDKTestOptions;
+import io.temporal.testing.internal.SDKTestWorkflowRule;
+import io.temporal.workflow.Workflow;
+import io.temporal.workflow.shared.TestActivities.TestActivitiesImpl;
+import io.temporal.workflow.shared.TestActivities.VariousTestActivities;
+import io.temporal.workflow.shared.TestWorkflows.TestWorkflow1;
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
+import org.junit.Rule;
+import org.junit.Test;
+
+@Issue("https://github.com/temporalio/sdk-java/issues/1262")
+public class LocalActivityGettingScheduledRightBeforeWorkflowTaskHeartbeatTest {
+  private static final Duration WORKFLOW_TASK_TIMEOUT = Duration.ofSeconds(5);
+  private static final Duration SLEEP_DURATION =
+      Duration.ofMillis(800); // << 1000 to avoid deadlock detection
+
+  private final TestActivitiesImpl activitiesImpl = new TestActivitiesImpl();
+
+  @Rule
+  public SDKTestWorkflowRule testWorkflowRule =
+      SDKTestWorkflowRule.newBuilder()
+          .setWorkflowTypes(HeartbeatingWorkflowImpl.class)
+          .setActivityImplementations(activitiesImpl)
+          .build();
+
+  @Test(timeout = 15_000)
+  public void testLocalActivitiesWorkflowTaskHeartbeat() {
+    WorkflowOptions options =
+        WorkflowOptions.newBuilder()
+            .setWorkflowRunTimeout(WORKFLOW_TASK_TIMEOUT.multipliedBy(2))
+            .setWorkflowTaskTimeout(WORKFLOW_TASK_TIMEOUT)
+            .setTaskQueue(testWorkflowRule.getTaskQueue())
+            .build();
+
+    List<WorkflowStub> stubs = new ArrayList<>();
+
+    // this test is actually pretty stable,
+    // but run several instances to increase the chances of the right timing being hit
+    for (int i = 0; i < 5; i++) {
+      TestWorkflow1 workflow =
+          testWorkflowRule.getWorkflowClient().newWorkflowStub(TestWorkflow1.class, options);
+      WorkflowStub stub = WorkflowStub.fromTyped(workflow);
+      stub.start(testWorkflowRule.getTaskQueue());
+      stubs.add(stub);
+    }
+
+    for (WorkflowStub stub : stubs) {
+      assertEquals("done", stub.getResult(String.class));
+    }
+  }
+
+  public static class HeartbeatingWorkflowImpl implements TestWorkflow1 {
+    @Override
+    public String execute(String taskQueue) {
+      VariousTestActivities localActivities =
+          Workflow.newLocalActivityStub(
+              VariousTestActivities.class, SDKTestOptions.newLocalActivityOptions());
+
+      long firstLocalActivityDurationMs =
+          (long) (WORKFLOW_TASK_TIMEOUT.toMillis() * Config.WORKFLOW_TAK_HEARTBEAT_COEFFICIENT)
+              - SLEEP_DURATION.toMillis() / 2;
+      localActivities.sleepActivity(firstLocalActivityDurationMs, 0);
+
+      // It is very important for reproduction that the workflow heartbeat timeout is reached DURING
+      // this sleep / workflow code execution.
+      // So the first local activity is done, eventLoop is triggered and heartbeat timeout is
+      // reached at the end of
+      // this workflow code event loop call with the next activity scheduled.
+      try {
+        Thread.sleep(SLEEP_DURATION.toMillis());
+      } catch (InterruptedException e) {
+        Thread.currentThread().interrupt();
+        throw new RuntimeException(e);
+      }
+
+      localActivities.sleepActivity(TimeUnit.SECONDS.toMillis(1), 0);
+
+      return "done";
+    }
+  }
+}

temporal-testing/src/main/java/io/temporal/internal/Issue.java

@@ -0,0 +1,33 @@
+/*
+ * 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;
+
+/**
+ * Test annotated with {@link Issue} covers a specific problem or edge case and is usually crafted
+ * carefully to hit the right conditions to reproduce the problem. Please specify a link to a ticket
+ * that describes what this test intends to cover.
+ */
+public @interface Issue {
+  /**
+   * @return Link to a ticket that describes what this test is intended to cover.
+   */
+  String value();
+}