Fix local activity completion synchronization (#474)

Author: meiliang86

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

@@ -39,8 +39,6 @@
 import java.util.concurrent.CancellationException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.locks.Condition;
-import java.util.concurrent.locks.Lock;
-import java.util.concurrent.locks.ReentrantLock;
 import java.util.function.BiConsumer;
 import java.util.function.BiFunction;
 import java.util.function.Consumer;
@@ -88,8 +86,7 @@ public void accept(Exception reason) {
   private final Map<String, ExecuteLocalActivityParameters> unstartedLaTasks = new HashMap<>();
   private final ReplayDecider replayDecider;
   private final DataConverter dataConverter;
-  private final Lock laTaskLock = new ReentrantLock();
-  private final Condition taskCondition = laTaskLock.newCondition();
+  private final Condition taskCondition;
   private boolean taskCompleted = false;
 
   ClockDecisionContext(
@@ -98,6 +95,7 @@ public void accept(Exception reason) {
       ReplayDecider replayDecider,
       DataConverter dataConverter) {
     this.decisions = decisions;
+    this.taskCondition = replayDecider.getLock().newCondition();
     mutableSideEffectHandler =
         new MarkerHandler(decisions, MUTABLE_SIDE_EFFECT_MARKER_NAME, () -> replaying);
     versionHandler = new MarkerHandler(decisions, VERSION_MARKER_NAME, () -> replaying);
@@ -272,13 +270,9 @@ private void handleLocalActivityMarker(MarkerRecordedEventAttributes attributes)
       completionHandle.accept(marker.getResult(), failure);
       setReplayCurrentTimeMilliseconds(marker.getReplayTimeMillis());
 
-      laTaskLock.lock();
-      try {
-        taskCompleted = true;
-        taskCondition.signal();
-      } finally {
-        laTaskLock.unlock();
-      }
+      taskCompleted = true;
+      // This method is already called under the lock.
+      taskCondition.signal();
     }
   }
 
@@ -342,7 +336,7 @@ boolean startUnstartedLaTasks(Duration maxWaitAllowed) {
           laTaskPoller.apply(
               new LocalActivityWorker.Task(
                   params,
-                  replayDecider,
+                  replayDecider.getLocalActivityCompletionSink(),
                   replayDecider.getDecisionTimeoutSeconds(),
                   this::currentTimeMillis,
                   this::replayTimeUpdatedAtMillis),
@@ -360,14 +354,10 @@ int numPendingLaTasks() {
   }
 
   void awaitTaskCompletion(Duration duration) throws InterruptedException {
-    laTaskLock.lock();
-    try {
-      while (!taskCompleted) {
-        taskCondition.awaitNanos(duration.toNanos());
-      }
-      taskCompleted = false;
-    } finally {
-      laTaskLock.unlock();
+    while (!taskCompleted) {
+      // This call is called from already locked object
+      taskCondition.awaitNanos(duration.toNanos());
     }
+    taskCompleted = false;
   }
 }

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

@@ -51,6 +51,8 @@
 import java.util.concurrent.CancellationException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
 import java.util.function.BiFunction;
 import java.util.function.Consumer;
 import org.apache.thrift.TException;
@@ -59,7 +61,7 @@
  * Implements decider that relies on replay of a workflow code. An instance of this class is created
  * per decision.
  */
-class ReplayDecider implements Decider, Consumer<HistoryEvent> {
+class ReplayDecider implements Decider {
 
   private static final int MAXIMUM_PAGE_SIZE = 10000;
 
@@ -75,6 +77,8 @@ class ReplayDecider implements Decider, Consumer<HistoryEvent> {
   private final Scope metricsScope;
   private final long wfStartTimeNanos;
   private final WorkflowExecutionStartedEventAttributes startedEvent;
+  private final Lock lock = new ReentrantLock();
+  private final Consumer<HistoryEvent> localActivityCompletionSink;
 
   ReplayDecider(
       IWorkflowService service,
@@ -100,6 +104,19 @@ class ReplayDecider implements Decider, Consumer<HistoryEvent> {
     context =
         new DecisionContextImpl(
             decisionsHelper, domain, decisionTask, startedEvent, options, laTaskPoller, this);
+    localActivityCompletionSink =
+        historyEvent -> {
+          lock.lock();
+          try {
+            processEvent(historyEvent);
+          } finally {
+            lock.unlock();
+          }
+        };
+  }
+
+  Lock getLock() {
+    return lock;
   }
 
   private void handleWorkflowExecutionStarted(HistoryEvent event) {
@@ -356,8 +373,13 @@ private void handleWorkflowExecutionSignaled(HistoryEvent event) {
 
   @Override
   public DecisionResult decide(PollForDecisionTaskResponse decisionTask) throws Throwable {
-    boolean forceCreateNewDecisionTask = decideImpl(decisionTask, null);
-    return new DecisionResult(decisionsHelper.getDecisions(), forceCreateNewDecisionTask);
+    lock.lock();
+    try {
+      boolean forceCreateNewDecisionTask = decideImpl(decisionTask, null);
+      return new DecisionResult(decisionsHelper.getDecisions(), forceCreateNewDecisionTask);
+    } finally {
+      lock.unlock();
+    }
   }
 
   // Returns boolean to indicate whether we need to force create new decision task for local
@@ -536,19 +558,28 @@ int getDecisionTimeoutSeconds() {
 
   @Override
   public void close() {
-    workflow.close();
+    lock.lock();
+    try {
+      workflow.close();
+    } finally {
+      lock.unlock();
+    }
   }
 
   @Override
   public byte[] query(PollForDecisionTaskResponse response, WorkflowQuery query) throws Throwable {
-    AtomicReference<byte[]> result = new AtomicReference<>();
-    decideImpl(response, () -> result.set(workflow.query(query)));
-    return result.get();
+    lock.lock();
+    try {
+      AtomicReference<byte[]> result = new AtomicReference<>();
+      decideImpl(response, () -> result.set(workflow.query(query)));
+      return result.get();
+    } finally {
+      lock.unlock();
+    }
   }
 
-  @Override
-  public void accept(HistoryEvent event) {
-    processEvent(event);
+  public Consumer<HistoryEvent> getLocalActivityCompletionSink() {
+    return localActivityCompletionSink;
   }
 
   private class DecisionTaskWithHistoryIteratorImpl implements DecisionTaskWithHistoryIterator {
@@ -580,7 +611,12 @@ private final Duration retryServiceOperationExpirationInterval() {
 
     @Override
     public PollForDecisionTaskResponse getDecisionTask() {
-      return task;
+      lock.lock();
+      try {
+        return task;
+      } finally {
+        lock.unlock();
+      }
     }
 
     @Override