Fix test server time-skipping locks after workflow closure

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

@@ -344,6 +344,10 @@ private void update(boolean completeWorkflowTaskUpdate, UpdateProcedure updater,
         // we don't buffer the events and don't wait till the finish of the workflow task
         // in-progress even if there is one,
         // but instead we apply them to the history immediately.
+        State taskState = workflowTaskStateMachine.getState();
+        if (taskState == State.INITIATED || taskState == State.STARTED) {
+          ctx.unlockTimer("forcefullyCompleteWorkflowTask");
+        }
         nextEventId = ctx.commitChanges(store);
       } else if (concurrentWorkflowTask) {
         // if there is a concurrent workflow task in progress and the workflow wasn't terminated and
@@ -370,6 +374,15 @@ private void update(boolean completeWorkflowTaskUpdate, UpdateProcedure updater,
     }
   }
 
+  private void updateWithTimerUnlock(String caller, UpdateProcedure updater) {
+    try {
+      update(updater);
+    } finally {
+      // A child notification owns a lock even if the parent has already closed.
+      timerService.unlockTimeSkipping(caller);
+    }
+  }
+
   @Override
   public ExecutionId getExecutionId() {
     return executionId;
@@ -1365,79 +1378,77 @@ private void timeoutWorkflowTask(long scheduledEventId) {
 
   @Override
   public void childWorkflowStarted(ChildWorkflowExecutionStartedEventAttributes a) {
-    update(
+    updateWithTimerUnlock(
+        "childWorkflowStarted",
         ctx -> {
           StateMachine<ChildWorkflowData> child = getChildWorkflow(a.getInitiatedEventId());
           child.action(StateMachines.Action.START, ctx, a, 0);
           scheduleWorkflowTask(ctx);
-          // No need to lock until completion as child workflow might skip
-          // time as well
-          ctx.unlockTimer("childWorkflowStarted");
         });
   }
 
   @Override
   public void childWorkflowFailed(
       String activityId, ChildWorkflowExecutionFailedEventAttributes a) {
-    update(
+    updateWithTimerUnlock(
+        "childWorkflowFailed",
         ctx -> {
           StateMachine<ChildWorkflowData> child = getChildWorkflow(a.getInitiatedEventId());
           child.action(StateMachines.Action.FAIL, ctx, a, 0);
           childWorkflows.remove(a.getInitiatedEventId());
           scheduleWorkflowTask(ctx);
-          ctx.unlockTimer("childWorkflowFailed");
         });
   }
 
   @Override
   public void childWorkflowTimedOut(
       String activityId, ChildWorkflowExecutionTimedOutEventAttributes a) {
-    update(
+    updateWithTimerUnlock(
+        "childWorkflowTimedOut",
         ctx -> {
           StateMachine<ChildWorkflowData> child = getChildWorkflow(a.getInitiatedEventId());
           child.action(Action.TIME_OUT, ctx, a.getRetryState(), 0);
           childWorkflows.remove(a.getInitiatedEventId());
           scheduleWorkflowTask(ctx);
-          ctx.unlockTimer("childWorkflowTimedOut");
         });
   }
 
   @Override
   public void failStartChildWorkflow(
       String childId, StartChildWorkflowExecutionFailedEventAttributes a) {
-    update(
+    updateWithTimerUnlock(
+        "failStartChildWorkflow",
         ctx -> {
           StateMachine<ChildWorkflowData> child = getChildWorkflow(a.getInitiatedEventId());
           child.action(StateMachines.Action.FAIL, ctx, a, 0);
           childWorkflows.remove(a.getInitiatedEventId());
           scheduleWorkflowTask(ctx);
-          ctx.unlockTimer("failStartChildWorkflow");
         });
   }
 
   @Override
   public void childWorkflowCompleted(
       String activityId, ChildWorkflowExecutionCompletedEventAttributes a) {
-    update(
+    updateWithTimerUnlock(
+        "childWorkflowCompleted",
         ctx -> {
           StateMachine<ChildWorkflowData> child = getChildWorkflow(a.getInitiatedEventId());
           child.action(StateMachines.Action.COMPLETE, ctx, a, 0);
           childWorkflows.remove(a.getInitiatedEventId());
           scheduleWorkflowTask(ctx);
-          ctx.unlockTimer("childWorkflowCompleted");
         });
   }
 
   @Override
   public void childWorkflowCanceled(
       String activityId, ChildWorkflowExecutionCanceledEventAttributes a) {
-    update(
+    updateWithTimerUnlock(
+        "childWorkflowCanceled",
         ctx -> {
           StateMachine<ChildWorkflowData> child = getChildWorkflow(a.getInitiatedEventId());
           child.action(StateMachines.Action.CANCEL, ctx, a, 0);
           childWorkflows.remove(a.getInitiatedEventId());
           scheduleWorkflowTask(ctx);
-          ctx.unlockTimer("childWorkflowCanceled");
         });
   }
 

temporal-test-server/src/test/java/io/temporal/internal/testservice/TestWorkflowMutableStateTimeSkippingTest.java

@@ -0,0 +1,193 @@
+package io.temporal.internal.testservice;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
+
+import io.grpc.Status;
+import io.grpc.StatusRuntimeException;
+import io.temporal.api.common.v1.Payloads;
+import io.temporal.api.common.v1.WorkflowType;
+import io.temporal.api.enums.v1.WorkflowExecutionStatus;
+import io.temporal.api.history.v1.ChildWorkflowExecutionCanceledEventAttributes;
+import io.temporal.api.history.v1.ChildWorkflowExecutionCompletedEventAttributes;
+import io.temporal.api.history.v1.ChildWorkflowExecutionFailedEventAttributes;
+import io.temporal.api.history.v1.ChildWorkflowExecutionStartedEventAttributes;
+import io.temporal.api.history.v1.ChildWorkflowExecutionTimedOutEventAttributes;
+import io.temporal.api.history.v1.StartChildWorkflowExecutionFailedEventAttributes;
+import io.temporal.api.taskqueue.v1.TaskQueue;
+import io.temporal.api.workflowservice.v1.PollWorkflowTaskQueueRequest;
+import io.temporal.api.workflowservice.v1.PollWorkflowTaskQueueResponse;
+import io.temporal.api.workflowservice.v1.RespondWorkflowTaskCompletedRequest;
+import io.temporal.api.workflowservice.v1.StartWorkflowExecutionRequest;
+import io.temporal.api.workflowservice.v1.TerminateWorkflowExecutionRequest;
+import java.time.Clock;
+import java.time.Duration;
+import java.util.Optional;
+import java.util.OptionalLong;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+public class TestWorkflowMutableStateTimeSkippingTest {
+  private SelfAdvancingTimerImpl timer;
+  private LockHandle testLock;
+  private TestWorkflowMutableStateImpl state;
+  private TestWorkflowService service;
+
+  @Before
+  public void setUp() {
+    timer = new SelfAdvancingTimerImpl(1, Clock.systemUTC());
+    testLock = timer.lockTimeSkipping("test setup");
+    TestWorkflowStore store = new TestWorkflowStoreImpl(timer);
+    TestVisibilityStore visibilityStore = new TestVisibilityStoreImpl();
+    TestNexusEndpointStore nexusStore = new TestNexusEndpointStoreImpl();
+    service = new TestWorkflowService(store, visibilityStore, nexusStore, timer);
+    state =
+        new TestWorkflowMutableStateImpl(
+            StartWorkflowExecutionRequest.newBuilder()
+                .setNamespace("test")
+                .setWorkflowId("workflow")
+                .setRequestId("request")
+                .setWorkflowType(WorkflowType.newBuilder().setName("workflow"))
+                .setTaskQueue(TaskQueue.newBuilder().setName("queue"))
+                .build(),
+            "run",
+            "run",
+            Optional.empty(),
+            Duration.ZERO,
+            Payloads.getDefaultInstance(),
+            Optional.empty(),
+            Optional.empty(),
+            OptionalLong.empty(),
+            Optional.empty(),
+            service,
+            store,
+            visibilityStore,
+            nexusStore,
+            timer);
+    state.startWorkflow(false, null, null);
+  }
+
+  @After
+  public void tearDown() {
+    service.close();
+  }
+
+  @Test
+  public void terminateScheduledWorkflowTask() throws InterruptedException {
+    state.terminateWorkflowExecution(TerminateWorkflowExecutionRequest.getDefaultInstance());
+    assertTimeSkips();
+  }
+
+  @Test
+  public void terminateStartedWorkflowTask() throws InterruptedException {
+    startTask();
+    state.terminateWorkflowExecution(TerminateWorkflowExecutionRequest.getDefaultInstance());
+    StatusRuntimeException failure =
+        assertThrows(
+            StatusRuntimeException.class,
+            () ->
+                state.completeWorkflowTask(
+                    3, RespondWorkflowTaskCompletedRequest.getDefaultInstance()));
+    assertEquals(Status.Code.NOT_FOUND, failure.getStatus().getCode());
+    assertTimeSkips();
+  }
+
+  @Test
+  public void terminateWithoutWorkflowTask() throws InterruptedException {
+    finishTask();
+    state.terminateWorkflowExecution(TerminateWorkflowExecutionRequest.getDefaultInstance());
+    assertTimeSkips();
+  }
+
+  @Test
+  public void timeoutStartedWorkflowTask() throws InterruptedException {
+    startTask();
+    timer.skip(Duration.ofDays(3651));
+    assertTimeSkips();
+    assertEquals(
+        WorkflowExecutionStatus.WORKFLOW_EXECUTION_STATUS_TIMED_OUT,
+        state.getWorkflowExecutionStatus());
+  }
+
+  @Test
+  public void childStartedAfterParentClosed() throws InterruptedException {
+    assertClosedParentReleasesLock(
+        parent ->
+            parent.childWorkflowStarted(
+                ChildWorkflowExecutionStartedEventAttributes.getDefaultInstance()));
+  }
+
+  @Test
+  public void childStartFailedAfterParentClosed() throws InterruptedException {
+    assertClosedParentReleasesLock(
+        parent ->
+            parent.failStartChildWorkflow(
+                "child", StartChildWorkflowExecutionFailedEventAttributes.getDefaultInstance()));
+  }
+
+  @Test
+  public void childCompletedAfterParentClosed() throws InterruptedException {
+    assertClosedParentReleasesLock(
+        parent ->
+            parent.childWorkflowCompleted(
+                "child", ChildWorkflowExecutionCompletedEventAttributes.getDefaultInstance()));
+  }
+
+  @Test
+  public void childFailedAfterParentClosed() throws InterruptedException {
+    assertClosedParentReleasesLock(
+        parent ->
+            parent.childWorkflowFailed(
+                "child", ChildWorkflowExecutionFailedEventAttributes.getDefaultInstance()));
+  }
+
+  @Test
+  public void childCanceledAfterParentClosed() throws InterruptedException {
+    assertClosedParentReleasesLock(
+        parent ->
+            parent.childWorkflowCanceled(
+                "child", ChildWorkflowExecutionCanceledEventAttributes.getDefaultInstance()));
+  }
+
+  @Test
+  public void childTimedOutAfterParentClosed() throws InterruptedException {
+    assertClosedParentReleasesLock(
+        parent ->
+            parent.childWorkflowTimedOut(
+                "child", ChildWorkflowExecutionTimedOutEventAttributes.getDefaultInstance()));
+  }
+
+  private void startTask() {
+    state.startWorkflowTask(
+        PollWorkflowTaskQueueResponse.newBuilder(),
+        PollWorkflowTaskQueueRequest.getDefaultInstance());
+  }
+
+  private void finishTask() {
+    startTask();
+    state.completeWorkflowTask(3, RespondWorkflowTaskCompletedRequest.getDefaultInstance());
+  }
+
+  private void assertClosedParentReleasesLock(Consumer<TestWorkflowMutableState> notification)
+      throws InterruptedException {
+    finishTask();
+    state.terminateWorkflowExecution(TerminateWorkflowExecutionRequest.getDefaultInstance());
+    timer.lockTimeSkipping("child notification");
+    StatusRuntimeException failure =
+        assertThrows(StatusRuntimeException.class, () -> notification.accept(state));
+    assertEquals(Status.Code.NOT_FOUND, failure.getStatus().getCode());
+    assertTimeSkips();
+  }
+
+  private void assertTimeSkips() throws InterruptedException {
+    CountDownLatch fired = new CountDownLatch(1);
+    timer.schedule(Duration.ofDays(1), fired::countDown, "probe");
+    testLock.unlock();
+    assertTrue("Clock must skip after workflow closure", fired.await(5, TimeUnit.SECONDS));
+  }
+}

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

@@ -32,7 +32,7 @@ public class ChildLivesLongerThanParentTest {
 
   private static final Signal timerFired = new Signal();
 
-  @Test
+  @Test(timeout = 10000)
   public void testAbandonChild() throws InterruptedException {
     testWorkflowRule
         .getWorkflowClient()
@@ -44,6 +44,10 @@ public void testAbandonChild() throws InterruptedException {
                 .build())
         .execute();
     assertTrue(timerFired.waitForSignal(SLEEPING_BY.plus(Duration.ofSeconds(2))));
+    if (!testWorkflowRule.isUseExternalService()) {
+      // A closed parent's notification must not prevent later time skips.
+      testWorkflowRule.getTestEnvironment().sleep(Duration.ofDays(1));
+    }
   }
 
   public static class TestWorkflowImpl implements TestWorkflows.PrimitiveWorkflow {