Added WorkflowTestEnvironment.sleep (#130)

Author: mfateev

src/main/java/com/uber/cadence/client/WorkflowClientInterceptor.java

@@ -26,7 +26,7 @@ UntypedWorkflowStub newUntypedWorkflowStub(
       String workflowType, WorkflowOptions options, UntypedWorkflowStub next);
 
   UntypedWorkflowStub newUntypedWorkflowStub(
-      WorkflowExecution execution, Optional<String> workflowType, UntypedWorkflowStub result);
+      WorkflowExecution execution, Optional<String> workflowType, UntypedWorkflowStub next);
 
   ActivityCompletionClient newActivityCompletionClient(ActivityCompletionClient next);
 }

src/main/java/com/uber/cadence/internal/replay/ActivityDecisionContext.java

@@ -49,6 +49,10 @@ private ActivityCancellationHandler(
 
     @Override
     public void accept(Exception cause) {
+      if (!scheduledActivities.containsKey(activityId)) {
+        // Cancellation handlers are not deregistered. So they fire after an activity completion.
+        return;
+      }
       decisions.requestCancelActivityTask(
           activityId,
           () -> {

src/main/java/com/uber/cadence/internal/replay/WorkflowDecisionContext.java

@@ -62,6 +62,10 @@ private ChildWorkflowCancellationHandler(
 
     @Override
     public void accept(Exception cause) {
+      if (!scheduledExternalWorkflows.containsKey(workflowId)) {
+        // Cancellation handlers are not deregistered. So they fire after a child completion.
+        return;
+      }
       RequestCancelExternalWorkflowExecutionDecisionAttributes cancelAttributes =
           new RequestCancelExternalWorkflowExecutionDecisionAttributes();
       cancelAttributes.setWorkflowId(workflowId);
@@ -175,6 +179,10 @@ Consumer<Exception> signalWorkflowExecution(
     final String finalSignalId = new String(attributes.getControl(), StandardCharsets.UTF_8);
     scheduledSignals.put(finalSignalId, context);
     return (e) -> {
+      if (!scheduledSignals.containsKey(finalSignalId)) {
+        // Cancellation handlers are not deregistered. So they fire after a signal completion.
+        return;
+      }
       decisions.cancelSignalExternalWorkflowExecution(finalSignalId, null);
       OpenRequestInfo<Void, Void> scheduled = scheduledSignals.remove(finalSignalId);
       if (scheduled == null) {

src/main/java/com/uber/cadence/internal/sync/SyncDecisionContext.java

@@ -121,7 +121,7 @@ private Promise<byte[]> executeActivityOnce(String name, ActivityOptions options
                     () -> result.completeExceptionally(mapActivityException(failure)));
               } else {
                 runner.executeInWorkflowThread(
-                    "activity failure callback", () -> result.complete(output));
+                    "activity completion callback", () -> result.complete(output));
               }
             });
     CancellationScope.current()

src/main/java/com/uber/cadence/internal/sync/TestWorkflowEnvironmentInternal.java

@@ -62,10 +62,14 @@
 import com.uber.cadence.TerminateWorkflowExecutionRequest;
 import com.uber.cadence.UpdateDomainRequest;
 import com.uber.cadence.UpdateDomainResponse;
+import com.uber.cadence.WorkflowExecution;
 import com.uber.cadence.WorkflowExecutionAlreadyStartedError;
+import com.uber.cadence.client.ActivityCompletionClient;
+import com.uber.cadence.client.UntypedWorkflowStub;
 import com.uber.cadence.client.WorkflowClient;
+import com.uber.cadence.client.WorkflowClientInterceptor;
 import com.uber.cadence.client.WorkflowClientOptions;
-import com.uber.cadence.client.WorkflowClientOptions.Builder;
+import com.uber.cadence.client.WorkflowOptions;
 import com.uber.cadence.internal.testservice.TestWorkflowService;
 import com.uber.cadence.serviceclient.IWorkflowService;
 import com.uber.cadence.testing.TestEnvironmentOptions;
@@ -76,6 +80,10 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.Optional;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import org.apache.thrift.TException;
 import org.apache.thrift.async.AsyncMethodCallback;
 import org.slf4j.Logger;
@@ -96,6 +104,7 @@ class TestWorkflowEnvironmentInternal implements TestWorkflowEnvironment {
       this.testEnvironmentOptions = options;
     }
     service = new WorkflowServiceWrapper();
+    service.lockTimeSkipping();
   }
 
   @Override
@@ -113,7 +122,10 @@ public Worker newWorker(String taskList) {
   @Override
   public WorkflowClient newWorkflowClient() {
     WorkflowClientOptions options =
-        new Builder().setDataConverter(testEnvironmentOptions.getDataConverter()).build();
+        new WorkflowClientOptions.Builder()
+            .setDataConverter(testEnvironmentOptions.getDataConverter())
+            .setInterceptors(new TimeLockingInterceptor(service))
+            .build();
     return WorkflowClientInternal.newInstance(service, testEnvironmentOptions.getDomain(), options);
   }
 
@@ -127,6 +139,11 @@ public long currentTimeMillis() {
     return service.currentTimeMillis();
   }
 
+  @Override
+  public void sleep(Duration duration) {
+    service.sleep(duration);
+  }
+
   @Override
   public void registerDelayedCallback(Duration delay, Runnable r) {
     service.registerDelayedCallback(delay, r);
@@ -523,5 +540,124 @@ public void close() {
     public void registerDelayedCallback(Duration delay, Runnable r) {
       impl.registerDelayedCallback(delay, r);
     }
+
+    public void lockTimeSkipping() {
+      impl.lockTimeSkipping();
+    }
+
+    public void unlockTimeSkipping() {
+      impl.unlockTimeSkipping();
+    }
+
+    public void sleep(Duration duration) {
+      impl.sleep(duration);
+    }
+  }
+
+  private static class TimeLockingInterceptor implements WorkflowClientInterceptor {
+
+    private final WorkflowServiceWrapper service;
+
+    TimeLockingInterceptor(WorkflowServiceWrapper service) {
+      this.service = service;
+    }
+
+    @Override
+    public UntypedWorkflowStub newUntypedWorkflowStub(
+        String workflowType, WorkflowOptions options, UntypedWorkflowStub next) {
+      return new TimeLockingWorkflowStub(service, next);
+    }
+
+    @Override
+    public UntypedWorkflowStub newUntypedWorkflowStub(
+        WorkflowExecution execution, Optional<String> workflowType, UntypedWorkflowStub next) {
+      return new TimeLockingWorkflowStub(service, next);
+    }
+
+    @Override
+    public ActivityCompletionClient newActivityCompletionClient(ActivityCompletionClient next) {
+      return next;
+    }
+
+    private class TimeLockingWorkflowStub implements UntypedWorkflowStub {
+
+      private final WorkflowServiceWrapper service;
+      private final UntypedWorkflowStub next;
+
+      TimeLockingWorkflowStub(WorkflowServiceWrapper service, UntypedWorkflowStub next) {
+        this.service = service;
+        this.next = next;
+      }
+
+      @Override
+      public void signal(String signalName, Object... args) {
+        next.signal(signalName, args);
+      }
+
+      @Override
+      public WorkflowExecution start(Object... args) {
+        return next.start(args);
+      }
+
+      @Override
+      public Optional<String> getWorkflowType() {
+        return next.getWorkflowType();
+      }
+
+      @Override
+      public WorkflowExecution getExecution() {
+        return next.getExecution();
+      }
+
+      @Override
+      public <R> R getResult(Class<R> returnType) {
+        service.unlockTimeSkipping();
+        try {
+          return next.getResult(returnType);
+        } finally {
+          service.lockTimeSkipping();
+        }
+      }
+
+      @Override
+      public <R> CompletableFuture<R> getResultAsync(Class<R> returnType) {
+        service.unlockTimeSkipping();
+        return next.getResultAsync(returnType).whenComplete((r, e) -> service.lockTimeSkipping());
+      }
+
+      @Override
+      public <R> R getResult(long timeout, TimeUnit unit, Class<R> returnType)
+          throws TimeoutException {
+        service.unlockTimeSkipping();
+        try {
+          return next.getResult(timeout, unit, returnType);
+        } finally {
+          service.lockTimeSkipping();
+        }
+      }
+
+      @Override
+      public <R> CompletableFuture<R> getResultAsync(
+          long timeout, TimeUnit unit, Class<R> returnType) {
+        service.unlockTimeSkipping();
+        return next.getResultAsync(timeout, unit, returnType)
+            .whenComplete((r, e) -> service.lockTimeSkipping());
+      }
+
+      @Override
+      public <R> R query(String queryType, Class<R> returnType, Object... args) {
+        return next.query(queryType, returnType, args);
+      }
+
+      @Override
+      public void cancel() {
+        next.cancel();
+      }
+
+      @Override
+      public Optional<WorkflowOptions> getOptions() {
+        return next.getOptions();
+      }
+    }
   }
 }

src/main/java/com/uber/cadence/internal/testservice/RequestContext.java

@@ -28,6 +28,8 @@
 import java.util.List;
 import java.util.Objects;
 import java.util.function.LongSupplier;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 final class RequestContext {
 

src/main/java/com/uber/cadence/internal/testservice/TestWorkflowMutableStateImpl.java

@@ -974,6 +974,7 @@ public RecordActivityTaskHeartbeatResponse heartbeatActivityTask(
   }
 
   private void timeoutActivity(String activityId, TimeoutType timeoutType) {
+    boolean unlockTimer = true;
     try {
       update(
           ctx -> {
@@ -996,9 +997,14 @@ private void timeoutActivity(String activityId, TimeoutType timeoutType) {
           });
     } catch (EntityNotExistsError e) {
       // Expected as timers are not removed
+      unlockTimer = false;
     } catch (Exception e) {
       // Cannot fail to timer threads
       log.error("Failure trying to timeout an activity", e);
+    } finally {
+      if (unlockTimer) {
+        selfAdvancingTimer.unlockTimeSkipping();
+      }
     }
   }
 

src/main/java/com/uber/cadence/internal/testservice/TestWorkflowService.java

@@ -73,6 +73,8 @@
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Optional;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ForkJoinPool;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantLock;
@@ -681,4 +683,33 @@ public long currentTimeMillis() {
   public void registerDelayedCallback(Duration delay, Runnable r) {
     store.registerDelayedCallback(delay, r);
   }
+
+  public void lockTimeSkipping() {
+    store.getTimer().lockTimeSkipping();
+  }
+
+  public void unlockTimeSkipping() {
+    store.getTimer().unlockTimeSkipping();
+  }
+
+  public void sleep(Duration duration) {
+    long start = store.currentTimeMillis();
+    CompletableFuture<Void> result = new CompletableFuture<>();
+    store
+        .getTimer()
+        .schedule(
+            duration,
+            () -> {
+              store.getTimer().lockTimeSkipping();
+              result.complete(null);
+            });
+    store.getTimer().unlockTimeSkipping();
+    try {
+      result.get();
+    } catch (InterruptedException e) {
+      throw new RuntimeException(e);
+    } catch (ExecutionException e) {
+      throw new RuntimeException(e);
+    }
+  }
 }

src/main/java/com/uber/cadence/internal/testservice/TestWorkflowStore.java

@@ -125,6 +125,8 @@ public PollForActivityTaskResponse getTask() {
 
   SelfAdvancingTimer getTimer();
 
+  long currentTimeMillis();
+
   long save(RequestContext requestContext) throws InternalServiceError, EntityNotExistsError;
 
   void registerDelayedCallback(Duration delay, Runnable r);

src/main/java/com/uber/cadence/internal/testservice/TestWorkflowStoreImpl.java

@@ -150,6 +150,11 @@ public SelfAdvancingTimer getTimer() {
     return timerService;
   }
 
+  @Override
+  public long currentTimeMillis() {
+    return timerService.getClock().getAsLong();
+  }
+
   @Override
   public long save(RequestContext ctx) throws InternalServiceError, EntityNotExistsError {
     long result;