Lock decider while processing (#325)

* Lock decider while processing
* Start decision task should not look at sticky schedule to start timeout
* Move workflow run lock into workflow worker
* Add unit test for local activity heartbeating and query

Author: meiliang86

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

@@ -405,7 +405,11 @@ private boolean decideImpl(PollForDecisionTaskResponse decisionTask, Functions.P
         }
 
         forceCreateNewDecisionTask =
-            processEventLoop(startTime, startedEvent.getTaskStartToCloseTimeoutSeconds(), decision);
+            processEventLoop(
+                startTime,
+                startedEvent.getTaskStartToCloseTimeoutSeconds(),
+                decision,
+                decisionTask.getQuery() != null);
 
         mayBeCompleteWorkflow();
         if (decision.isReplay()) {
@@ -443,11 +447,12 @@ private boolean decideImpl(PollForDecisionTaskResponse decisionTask, Functions.P
     }
   }
 
-  private boolean processEventLoop(long startTime, int decisionTimeoutSecs, DecisionEvents decision)
+  private boolean processEventLoop(
+      long startTime, int decisionTimeoutSecs, DecisionEvents decision, boolean isQuery)
       throws Throwable {
     eventLoop();
 
-    if (decision.isReplay()) {
+    if (decision.isReplay() || isQuery) {
       return replayLocalActivities(decision);
     } else {
       return executeLocalActivities(startTime, decisionTimeoutSecs);

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

@@ -88,7 +88,9 @@ public SyncWorkflowWorker(
             service,
             laWorker.getLocalActivityTaskPoller());
 
-    workflowWorker = new WorkflowWorker(service, domain, taskList, workflowOptions, taskHandler);
+    workflowWorker =
+        new WorkflowWorker(
+            service, domain, taskList, workflowOptions, taskHandler, stickyTaskListName);
   }
 
   public void setWorkflowImplementationTypes(

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

@@ -216,9 +216,7 @@ public void startDecisionTask(
             long scheduledEventId = decision.getData().scheduledEventId;
             decision.action(StateMachines.Action.START, ctx, pollRequest, 0);
             ctx.addTimer(
-                stickyExecutionAttributes != null
-                    ? stickyExecutionAttributes.getScheduleToStartTimeoutSeconds()
-                    : startRequest.getTaskStartToCloseTimeoutSeconds(),
+                startRequest.getTaskStartToCloseTimeoutSeconds(),
                 () -> timeoutDecisionTask(scheduledEventId),
                 "DecisionTask StartToCloseTimeout");
           });

src/main/java/com/uber/cadence/internal/worker/WorkflowRunLockManager.java

@@ -0,0 +1,105 @@
+/*
+ *  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.worker;
+
+import java.util.HashMap;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
+
+final class WorkflowRunLockManager {
+
+  private static class CountableLock {
+    private final Lock lock = new ReentrantLock();
+    private int count = 1;
+
+    void incrementCount() {
+      count++;
+    }
+
+    void decrementCount() {
+      count--;
+    }
+
+    int getCount() {
+      return count;
+    }
+
+    Lock getLock() {
+      return lock;
+    }
+  }
+
+  private final Lock mapLock = new ReentrantLock();
+  private final HashMap<String, CountableLock> perRunLock = new HashMap<>();
+
+  /**
+   * This method returns a lock that can be used to serialize decision task processing for a
+   * particular workflow run. This is used to make sure that query tasks and real decision tasks are
+   * serialized when sticky is on.
+   *
+   * @param runId
+   * @return a lock to be used during decision task processing
+   */
+  Lock getLockForLocking(String runId) {
+    mapLock.lock();
+
+    try {
+      CountableLock cl = perRunLock.get(runId);
+      if (cl == null) {
+        cl = new CountableLock();
+        perRunLock.put(runId, cl);
+      } else {
+        cl.incrementCount();
+      }
+
+      return cl.getLock();
+    } finally {
+      mapLock.unlock();
+    }
+  }
+
+  void unlock(String runId) {
+    mapLock.lock();
+
+    try {
+      CountableLock cl = perRunLock.get(runId);
+      if (cl == null) {
+        throw new RuntimeException("lock for run " + runId + " does not exist.");
+      }
+
+      cl.decrementCount();
+      if (cl.getCount() == 0) {
+        perRunLock.remove(runId);
+      }
+
+      cl.getLock().unlock();
+    } finally {
+      mapLock.unlock();
+    }
+  }
+
+  int totalLocks() {
+    mapLock.lock();
+
+    try {
+      return perRunLock.size();
+    } finally {
+      mapLock.unlock();
+    }
+  }
+}

src/main/java/com/uber/cadence/internal/worker/WorkflowWorker.java

@@ -17,6 +17,7 @@
 
 package com.uber.cadence.internal.worker;
 
+import com.google.common.base.Strings;
 import com.uber.cadence.BadRequestError;
 import com.uber.cadence.DomainNotActiveError;
 import com.uber.cadence.EntityNotExistsError;
@@ -46,6 +47,7 @@
 import java.util.List;
 import java.util.Objects;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.locks.Lock;
 import java.util.function.Consumer;
 import org.apache.thrift.TException;
 import org.slf4j.MDC;
@@ -62,17 +64,21 @@ public final class WorkflowWorker
   private final String domain;
   private final String taskList;
   private final SingleWorkerOptions options;
+  private final String stickyTaskListName;
+  private final WorkflowRunLockManager runLocks = new WorkflowRunLockManager();
 
   public WorkflowWorker(
       IWorkflowService service,
       String domain,
       String taskList,
       SingleWorkerOptions options,
-      DecisionTaskHandler handler) {
+      DecisionTaskHandler handler,
+      String stickyTaskListName) {
     this.service = Objects.requireNonNull(service);
     this.domain = Objects.requireNonNull(domain);
     this.taskList = Objects.requireNonNull(taskList);
     this.handler = handler;
+    this.stickyTaskListName = stickyTaskListName;
 
     PollerOptions pollerOptions = options.getPollerOptions();
     if (pollerOptions.getPollThreadNamePrefix() == null) {
@@ -240,6 +246,13 @@ public void handle(PollForDecisionTaskResponse task) throws Exception {
       MDC.put(LoggerTag.WORKFLOW_ID, task.getWorkflowExecution().getWorkflowId());
       MDC.put(LoggerTag.WORKFLOW_TYPE, task.getWorkflowType().getName());
       MDC.put(LoggerTag.RUN_ID, task.getWorkflowExecution().getRunId());
+
+      Lock runLock = null;
+      if (!Strings.isNullOrEmpty(stickyTaskListName)) {
+        runLock = runLocks.getLockForLocking(task.getWorkflowExecution().getRunId());
+        runLock.lock();
+      }
+
       try {
         Stopwatch sw = metricsScope.timer(MetricsType.DECISION_EXECUTION_LATENCY).start();
         DecisionTaskHandler.Result response = handler.handleDecisionTask(task);
@@ -254,6 +267,10 @@ public void handle(PollForDecisionTaskResponse task) throws Exception {
         MDC.remove(LoggerTag.WORKFLOW_ID);
         MDC.remove(LoggerTag.WORKFLOW_TYPE);
         MDC.remove(LoggerTag.RUN_ID);
+
+        if (runLock != null) {
+          runLocks.unlock(task.getWorkflowExecution().getRunId());
+        }
       }
     }
 

src/test/java/com/uber/cadence/internal/worker/WorkflowRunLockManagerTest.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.worker;
+
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
+
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.locks.Lock;
+import org.junit.Test;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class WorkflowRunLockManagerTest {
+  private static final Logger log = LoggerFactory.getLogger(WorkflowRunLockManagerTest.class);
+  private WorkflowRunLockManager runLockManager = new WorkflowRunLockManager();
+
+  @Test
+  public void lockAndUnlockTest() throws ExecutionException, InterruptedException {
+    ExecutorService executor = Executors.newFixedThreadPool(4);
+    ConcurrentLinkedQueue<String> finishedTasks = new ConcurrentLinkedQueue<>();
+    Future<?> f1 = executor.submit(() -> finishedTasks.add(processTask("run1", 1)));
+    Thread.sleep(100);
+    Future<?> f3 = executor.submit(() -> finishedTasks.add(processTask("run1", 2)));
+    Future<?> f2 = executor.submit(() -> finishedTasks.add(processTask("run2", 1)));
+    Thread.sleep(100);
+    Future<?> f4 = executor.submit(() -> finishedTasks.add(processTask("run1", 3)));
+
+    f1.get();
+    f2.get();
+    f3.get();
+    f4.get();
+
+    log.info("All done.");
+    assertEquals(0, runLockManager.totalLocks());
+    String[] expectedTasks = {"run1.1", "run2.1", "run1.2", "run1.3"};
+    String[] processedTasks = new String[4];
+    assertArrayEquals(expectedTasks, finishedTasks.toArray(processedTasks));
+  }
+
+  private String processTask(String runId, int taskId) {
+    Lock runLock = runLockManager.getLockForLocking(runId);
+    runLock.lock();
+
+    log.info("Got lock runId " + runId + " taskId " + taskId);
+    try {
+      Thread.sleep(1000);
+    } catch (InterruptedException e) {
+      throw new RuntimeException("interrupted");
+    } finally {
+      runLockManager.unlock(runId);
+    }
+    log.info("Finished processing runId " + runId + " taskId " + taskId);
+    return runId + "." + taskId;
+  }
+}