Support sticky queries (#205)

Author: halakaraki

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

@@ -39,16 +39,17 @@
 import com.uber.cadence.workflow.Functions;
 import com.uber.m3.tally.Scope;
 import com.uber.m3.tally.Stopwatch;
+import org.apache.thrift.TException;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
 import java.time.Duration;
 import java.util.Iterator;
 import java.util.Objects;
 import java.util.concurrent.CancellationException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
-import org.apache.thrift.TException;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 
 /**
  * Implements decider that relies on replay of a worklfow code. An instance of this class is created
@@ -381,7 +382,8 @@ private void decideImpl(PollForDecisionTaskResponse decisionTask, Functions.Proc
               != historyHelper.getPreviousStartedEventId()
                   + 2) // getNextDecisionEventId() skips over completed.
           && (decisionsHelper.getNextDecisionEventId() != 0
-              && historyHelper.getPreviousStartedEventId() != 0)) {
+              && historyHelper.getPreviousStartedEventId() != 0)
+          && (decisionTask.getHistory().getEventsSize() > 0)) {
         throw new IllegalStateException(
             String.format(
                 "ReplayDecider expects next event id at %d. History's previous started event id is %d",

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

@@ -106,16 +106,49 @@ public DecisionTaskHandler.Result handleDecisionTask(PollForDecisionTaskResponse
 
   private Result handleDecisionTaskImpl(PollForDecisionTaskResponse decisionTask) throws Throwable {
 
+    if (decisionTask.isSetQuery()) {
+      return processQuery(decisionTask);
+    } else {
+
+      return processDecision(decisionTask);
+    }
+  }
+
+  private Result processDecision(PollForDecisionTaskResponse decisionTask) throws Throwable {
     Decider decider =
         stickyTaskListName == null
             ? createDecider(decisionTask)
             : cache.getOrCreate(decisionTask, this::createDecider);
     try {
-      if (decisionTask.isSetQuery()) {
-        return processQuery(decisionTask, (ReplayDecider) decider);
-      } else {
-        return processDecision(decisionTask, (ReplayDecider) decider);
+      decider.decide(decisionTask);
+      DecisionsHelper decisionsHelper = ((ReplayDecider) decider).getDecisionsHelper();
+      List<Decision> decisions = decisionsHelper.getDecisions();
+
+      if (log.isTraceEnabled()) {
+        WorkflowExecution execution = decisionTask.getWorkflowExecution();
+        log.trace(
+            "WorkflowTask startedEventId="
+                + decisionTask.getStartedEventId()
+                + ", WorkflowID="
+                + execution.getWorkflowId()
+                + ", RunID="
+                + execution.getRunId()
+                + " completed with "
+                + WorkflowExecutionUtils.prettyPrintDecisions(decisions));
+      } else if (log.isDebugEnabled()) {
+        WorkflowExecution execution = decisionTask.getWorkflowExecution();
+        log.debug(
+            "WorkflowTask startedEventId="
+                + decisionTask.getStartedEventId()
+                + ", WorkflowID="
+                + execution.getWorkflowId()
+                + ", RunID="
+                + execution.getRunId()
+                + " completed with "
+                + decisions.size()
+                + " new decisions");
       }
+      return createCompletedRequest(decisionTask, decisionsHelper, decisions);
     } catch (Exception e) {
       if (stickyTaskListName != null) {
         cache.invalidate(decisionTask);
@@ -128,44 +161,14 @@ private Result handleDecisionTaskImpl(PollForDecisionTaskResponse decisionTask)
     }
   }
 
-  private Result processDecision(PollForDecisionTaskResponse decisionTask, ReplayDecider decider)
-      throws Throwable {
-    decider.decide(decisionTask);
-    DecisionsHelper decisionsHelper = decider.getDecisionsHelper();
-    List<Decision> decisions = decisionsHelper.getDecisions();
-
-    if (log.isTraceEnabled()) {
-      WorkflowExecution execution = decisionTask.getWorkflowExecution();
-      log.trace(
-          "WorkflowTask startedEventId="
-              + decisionTask.getStartedEventId()
-              + ", WorkflowID="
-              + execution.getWorkflowId()
-              + ", RunID="
-              + execution.getRunId()
-              + " completed with "
-              + WorkflowExecutionUtils.prettyPrintDecisions(decisions));
-    } else if (log.isDebugEnabled()) {
-      WorkflowExecution execution = decisionTask.getWorkflowExecution();
-      log.debug(
-          "WorkflowTask startedEventId="
-              + decisionTask.getStartedEventId()
-              + ", WorkflowID="
-              + execution.getWorkflowId()
-              + ", RunID="
-              + execution.getRunId()
-              + " completed with "
-              + decisions.size()
-              + " new decisions");
-    }
-
-    return createCompletedRequest(decisionTask, decisionsHelper, decisions);
-  }
-
-  private Result processQuery(PollForDecisionTaskResponse decisionTask, ReplayDecider decider) {
+  private Result processQuery(PollForDecisionTaskResponse decisionTask) {
     RespondQueryTaskCompletedRequest queryCompletedRequest = new RespondQueryTaskCompletedRequest();
     queryCompletedRequest.setTaskToken(decisionTask.getTaskToken());
     try {
+      Decider decider =
+          stickyTaskListName == null
+              ? createDecider(decisionTask)
+              : cache.getOrCreate(decisionTask, this::createDecider);
       byte[] queryResult = decider.query(decisionTask, decisionTask.getQuery());
       queryCompletedRequest.setQueryResult(queryResult);
       queryCompletedRequest.setCompletedType(QueryTaskCompletedType.COMPLETED);

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

@@ -82,6 +82,10 @@
 import com.uber.cadence.internal.testservice.StateMachines.TimerData;
 import com.uber.cadence.internal.testservice.StateMachines.WorkflowData;
 import com.uber.cadence.internal.testservice.TestWorkflowStore.TaskListId;
+import org.apache.thrift.TException;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
 import java.io.DataInputStream;
@@ -102,9 +106,6 @@
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantLock;
 import java.util.function.LongSupplier;
-import org.apache.thrift.TException;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 
 class TestWorkflowMutableStateImpl implements TestWorkflowMutableState {
 
@@ -138,6 +139,7 @@ void apply(RequestContext ctx)
   private long lastNonFailedDecisionStartEventId;
   private final Map<String, CompletableFuture<QueryWorkflowResponse>> queries =
       new ConcurrentHashMap<>();
+  private final Map<String, PollForDecisionTaskResponse> queryRequests = new ConcurrentHashMap<>();
   public StickyExecutionAttributes stickyExecutionAttributes;
 
   /** @param parentChildInitiatedEventId id of the child initiated event in the parent history */
@@ -1183,15 +1185,22 @@ public void requestCancelWorkflowExecution(RequestCancelWorkflowExecutionRequest
   @Override
   public QueryWorkflowResponse query(QueryWorkflowRequest queryRequest) throws TException {
     QueryId queryId = new QueryId(executionId);
+
     PollForDecisionTaskResponse task =
         new PollForDecisionTaskResponse()
             .setTaskToken(queryId.toBytes())
             .setWorkflowExecution(executionId.getExecution())
             .setWorkflowType(startRequest.getWorkflowType())
-            .setQuery(queryRequest.getQuery());
+            .setQuery(queryRequest.getQuery())
+            .setWorkflowExecutionTaskList(startRequest.getTaskList());
     TaskListId taskListId =
-        new TaskListId(queryRequest.getDomain(), startRequest.getTaskList().getName());
+        new TaskListId(
+            queryRequest.getDomain(),
+            stickyExecutionAttributes == null
+                ? startRequest.getTaskList().getName()
+                : stickyExecutionAttributes.getWorkerTaskList().getName());
     CompletableFuture<QueryWorkflowResponse> result = new CompletableFuture<>();
+    queryRequests.put(queryId.getQueryId(), task);
     queries.put(queryId.getQueryId(), result);
     store.sendQueryTask(executionId, taskListId, task);
     try {
@@ -1218,6 +1227,13 @@ public void completeQuery(QueryId queryId, RespondQueryTaskCompletedRequest comp
       QueryWorkflowResponse response =
           new QueryWorkflowResponse().setQueryResult(completeRequest.getQueryResult());
       result.complete(response);
+    } else if (stickyExecutionAttributes != null) {
+      stickyExecutionAttributes = null;
+      PollForDecisionTaskResponse task = queryRequests.remove(queryId.getQueryId());
+
+      TaskListId taskListId =
+          new TaskListId(startRequest.getDomain(), startRequest.getTaskList().getName());
+      store.sendQueryTask(executionId, taskListId, task);
     } else {
       QueryFailedError error = new QueryFailedError().setMessage(completeRequest.getErrorMessage());
       result.completeExceptionally(error);

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

@@ -34,6 +34,9 @@
 import com.uber.cadence.WorkflowExecutionInfo;
 import com.uber.cadence.internal.common.WorkflowExecutionUtils;
 import com.uber.cadence.internal.testservice.RequestContext.Timer;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.HashMap;
@@ -46,8 +49,6 @@
 import java.util.concurrent.locks.Condition;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantLock;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 
 class TestWorkflowStoreImpl implements TestWorkflowStore {
 
@@ -316,9 +317,15 @@ public void sendQueryTask(
       throws EntityNotExistsError {
     lock.lock();
     try {
-      HistoryStore history = getHistoryStore(executionId);
-      List<HistoryEvent> events = new ArrayList<>(history.getEventsLocked());
-      task.setHistory(new History().setEvents(events));
+      HistoryStore historyStore = getHistoryStore(executionId);
+      List<HistoryEvent> events = new ArrayList<>(historyStore.getEventsLocked());
+      History history = new History();
+      if (taskList.getTaskListName().equals(task.getWorkflowExecutionTaskList().getName())) {
+        history.setEvents(events);
+      } else {
+        history.setEvents(new ArrayList<>());
+      }
+      task.setHistory(history);
     } finally {
       lock.unlock();
     }