TestWorkflowEnvironment implemented (#132)

* Fixed TestWorkflowEnvironment time skipping

* Added Javadoc to TestWorkflowEnvironment

Author: mfateev

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

@@ -87,14 +87,14 @@
 import org.apache.thrift.TException;
 import org.apache.thrift.async.AsyncMethodCallback;
 
-class TestActivityEnvironmentInternal implements TestActivityEnvironment {
+public final class TestActivityEnvironmentInternal implements TestActivityEnvironment {
 
   private final POJOActivityTaskHandler activityTaskHandler;
   private final TestEnvironmentOptions testEnvironmentOptions;
   private final AtomicInteger idSequencer = new AtomicInteger();
   private ClassConsumerPair<Object> activityHeartbetListener;
 
-  TestActivityEnvironmentInternal(TestEnvironmentOptions options) {
+  public TestActivityEnvironmentInternal(TestEnvironmentOptions options) {
     if (options == null) {
       this.testEnvironmentOptions = new TestEnvironmentOptions.Builder().build();
     } else {

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

@@ -82,22 +82,23 @@
 import java.util.List;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
 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;
 import org.slf4j.LoggerFactory;
 
-class TestWorkflowEnvironmentInternal implements TestWorkflowEnvironment {
+public final class TestWorkflowEnvironmentInternal implements TestWorkflowEnvironment {
 
   private static final Logger log = LoggerFactory.getLogger(TestWorkflowEnvironmentInternal.class);
 
   private final TestEnvironmentOptions testEnvironmentOptions;
   private final WorkflowServiceWrapper service;
   private final List<Worker> workers = Collections.synchronizedList(new ArrayList<>());
 
-  TestWorkflowEnvironmentInternal(TestEnvironmentOptions options) {
+  public TestWorkflowEnvironmentInternal(TestEnvironmentOptions options) {
     if (options == null) {
       this.testEnvironmentOptions = new TestEnvironmentOptions.Builder().build();
     } else {
@@ -149,19 +150,22 @@ public void registerDelayedCallback(Duration delay, Runnable r) {
     service.registerDelayedCallback(delay, r);
   }
 
+  @Override
   public IWorkflowService getWorkflowService() {
     return service;
   }
 
+  @Override
   public String getDiagnostics() {
     StringBuilder result = new StringBuilder();
     service.getDiagnostics(result);
     return result.toString();
   }
 
-  public void shutdown() {
+  @Override
+  public void close() {
     for (Worker w : workers) {
-      w.shutdown(Duration.ofMillis(1));
+      w.shutdown(Duration.ofMillis(10));
     }
     service.close();
   }
@@ -621,8 +625,7 @@ public <R> R getResult(Class<R> returnType) {
 
       @Override
       public <R> CompletableFuture<R> getResultAsync(Class<R> returnType) {
-        service.unlockTimeSkipping();
-        return next.getResultAsync(returnType).whenComplete((r, e) -> service.lockTimeSkipping());
+        return new TimeLockingFuture<>(next.getResultAsync(returnType));
       }
 
       @Override
@@ -639,9 +642,7 @@ public <R> R getResult(long timeout, TimeUnit unit, Class<R> returnType)
       @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());
+        return new TimeLockingFuture<>(next.getResultAsync(timeout, unit, returnType));
       }
 
       @Override
@@ -658,6 +659,47 @@ public void cancel() {
       public Optional<WorkflowOptions> getOptions() {
         return next.getOptions();
       }
+
+      /** Unlocks time skipping before blocking calls and locks back after completion. */
+      private class TimeLockingFuture<R> extends CompletableFuture<R> {
+
+        public TimeLockingFuture(CompletableFuture<R> resultAsync) {
+          CompletableFuture<R> ignored =
+              resultAsync.whenComplete(
+                  (r, e) -> {
+                    service.lockTimeSkipping();
+                    if (e == null) {
+                      this.complete(r);
+                    } else {
+                      this.completeExceptionally(e);
+                    }
+                  });
+        }
+
+        @Override
+        public R get() throws InterruptedException, ExecutionException {
+          service.unlockTimeSkipping();
+          return super.get();
+        }
+
+        @Override
+        public R get(long timeout, TimeUnit unit)
+            throws InterruptedException, ExecutionException, TimeoutException {
+          service.unlockTimeSkipping();
+          try {
+            return super.get(timeout, unit);
+          } catch (TimeoutException | InterruptedException e) {
+            service.lockTimeSkipping();
+            throw e;
+          }
+        }
+
+        @Override
+        public R join() {
+          service.unlockTimeSkipping();
+          return super.join();
+        }
+      }
     }
   }
 }

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

@@ -0,0 +1,75 @@
+/*
+ *  Copyright 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 file except in compliance with the License. A copy of the License is
+ *  located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ *  or in the "license" file accompanying this file. This file 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 com.uber.cadence.internal.testservice;
+
+import com.google.common.base.Throwables;
+import com.uber.cadence.InternalServiceError;
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.DataInputStream;
+import java.io.DataOutputStream;
+import java.io.IOException;
+import java.util.Objects;
+
+final class DecisionTaskToken {
+
+  private final ExecutionId executionId;
+  private final int historySize;
+
+  DecisionTaskToken(ExecutionId executionId, int historySize) {
+    this.executionId = Objects.requireNonNull(executionId);
+    this.historySize = Objects.requireNonNull(historySize);
+  }
+
+  ExecutionId getExecutionId() {
+    return executionId;
+  }
+
+  int getHistorySize() {
+    return historySize;
+  }
+
+  /** Used for task tokens. */
+  byte[] toBytes() throws InternalServiceError {
+    ByteArrayOutputStream bout = new ByteArrayOutputStream();
+    DataOutputStream out = new DataOutputStream(bout);
+    try {
+      addBytes(out);
+    } catch (IOException e) {
+      throw new InternalServiceError(Throwables.getStackTraceAsString(e));
+    }
+    return bout.toByteArray();
+  }
+
+  private void addBytes(DataOutputStream out) throws IOException {
+    executionId.addBytes(out);
+    out.writeInt(historySize);
+  }
+
+  static DecisionTaskToken fromBytes(byte[] serialized) throws InternalServiceError {
+    ByteArrayInputStream bin = new ByteArrayInputStream(serialized);
+    DataInputStream in = new DataInputStream(bin);
+    try {
+      ExecutionId executionId = ExecutionId.readFromBytes(in);
+      int historySize = in.readInt();
+      return new DecisionTaskToken(executionId, historySize);
+    } catch (IOException e) {
+      throw new InternalServiceError(Throwables.getStackTraceAsString(e));
+    }
+  }
+}

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

@@ -34,7 +34,7 @@ final class RequestContext {
   @FunctionalInterface
   interface CommitCallback {
 
-    void apply() throws InternalServiceError;
+    void apply(int historySize) throws InternalServiceError;
   }
 
   static final class Timer {
@@ -199,9 +199,9 @@ long commitChanges(TestWorkflowStore store) throws InternalServiceError, EntityN
   }
 
   /** Called by {@link TestWorkflowStore#save(RequestContext)} */
-  void fireCallbacks() throws InternalServiceError {
+  void fireCallbacks(int historySize) throws InternalServiceError {
     for (CommitCallback callback : commitCallbacks) {
-      callback.apply();
+      callback.apply(historySize);
     }
   }
 

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

@@ -24,7 +24,7 @@
  * Timer service that automatically forwards current time to the next task time when is not locked
  * through {@link #lockTimeSkipping()}.
  */
-public interface SelfAdvancingTimer {
+interface SelfAdvancingTimer {
 
   /**
    * Schedule a task with a specified delay. The actual wait time is defined by the internal clock