Keep time skipping after activity timeouts

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

@@ -477,9 +477,7 @@ private void unlockTimeSkippingLocked(String caller) {
   }
 
   private void unlockTimeSkippingLockedInternal(String caller) {
-    if (lockCount == 0) {
-      throw new IllegalStateException("Unbalanced lock and unlock calls: \n" + getDiagnostics());
-    }
+    // An activity can time out while time-skipping sleep has borrowed its lock.
     lockCount--;
     if (caller == null) {
       log.trace("---");
@@ -490,12 +488,6 @@ private void unlockTimeSkippingLockedInternal(String caller) {
     }
   }
 
-  private String getDiagnostics() {
-    StringBuilder result = new StringBuilder();
-    getDiagnostics(result);
-    return result.toString();
-  }
-
   /**
    * TimerPump is very conservative with resources usage and if time is locked, it calculates the
    * time needed for the first scheduled task and waits for this calculated period to do the next

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

@@ -1,6 +1,7 @@
 package io.temporal.internal.testservice;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
@@ -40,6 +41,13 @@ public void tearDown() throws Exception {
     fixedTimer.shutdown();
   }
 
+  @Test
+  public void testLockHandleCannotBeUnlockedTwice() {
+    LockHandle handle = fixedTimer.lockTimeSkipping("unit test");
+    handle.unlock();
+    assertThrows(IllegalStateException.class, handle::unlock);
+  }
+
   @Test
   public void testSchedule() throws InterruptedException {
     AtomicLong captured = new AtomicLong();

temporal-test-server/src/test/java/io/temporal/testserver/functional/timeskipping/ActivityTimeoutTimeSkippingTest.java

@@ -0,0 +1,86 @@
+package io.temporal.testserver.functional.timeskipping;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+import io.temporal.activity.ActivityOptions;
+import io.temporal.client.WorkflowClient;
+import io.temporal.client.WorkflowClientOptions;
+import io.temporal.client.WorkflowStub;
+import io.temporal.testing.internal.SDKTestWorkflowRule;
+import io.temporal.testserver.functional.common.TestWorkflows;
+import io.temporal.workflow.Workflow;
+import java.time.Duration;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+
+@RunWith(Parameterized.class)
+public class ActivityTimeoutTimeSkippingTest {
+  @Parameterized.Parameters(name = "sleep={0}s")
+  public static Object[][] sleepDurations() {
+    return new Object[][] {{4}, {6}};
+  }
+
+  @Parameterized.Parameter public int sleepSeconds;
+
+  private final AtomicBoolean waited = new AtomicBoolean();
+  private final AtomicInteger attempts = new AtomicInteger();
+
+  @Rule
+  public SDKTestWorkflowRule testWorkflowRule =
+      SDKTestWorkflowRule.newBuilder()
+          .setWorkflowClientOptions(
+              WorkflowClientOptions.newBuilder()
+                  .setNamespace(SDKTestWorkflowRule.NAMESPACE)
+                  .setIdentity("activity-timeout-repro")
+                  .build())
+          .setWorkflowTypes(TestWorkflowImpl.class)
+          .setActivityImplementations(new SleepingActivityImpl())
+          .build();
+
+  @Test
+  public void activityCompletesAfterTimeSkipping() throws TimeoutException {
+    TestWorkflows.PrimitiveWorkflow workflow =
+        testWorkflowRule.newWorkflowStub(TestWorkflows.PrimitiveWorkflow.class);
+    WorkflowClient.start(workflow::execute);
+    try {
+      WorkflowStub.fromTyped(workflow).getResult(3, TimeUnit.SECONDS, Void.class);
+    } catch (TimeoutException e) {
+      System.err.println(testWorkflowRule.getTestEnvironment().getDiagnostics());
+      throw e;
+    }
+    if (sleepSeconds < 5) {
+      assertEquals(1, attempts.get());
+    } else {
+      assertTrue(attempts.get() >= 2);
+    }
+  }
+
+  public class SleepingActivityImpl implements SleepingActivity {
+    @Override
+    public void sleep() {
+      attempts.incrementAndGet();
+      if (!waited.get()) {
+        testWorkflowRule.getTestEnvironment().sleep(Duration.ofSeconds(sleepSeconds));
+        waited.set(true);
+      }
+    }
+  }
+
+  public static class TestWorkflowImpl implements TestWorkflows.PrimitiveWorkflow {
+    @Override
+    public void execute() {
+      Workflow.newActivityStub(
+              SleepingActivity.class,
+              ActivityOptions.newBuilder().setStartToCloseTimeout(Duration.ofSeconds(5)).build())
+          .sleep();
+      Workflow.sleep(Duration.ofHours(1));
+    }
+  }
+}