Refactored timer handling in ClockDecisionContext (#131)

Author: mfateev

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

@@ -23,17 +23,12 @@
 import com.uber.cadence.TimerFiredEventAttributes;
 import java.util.HashMap;
 import java.util.Map;
-import java.util.SortedMap;
-import java.util.TreeMap;
 import java.util.concurrent.CancellationException;
 import java.util.concurrent.TimeUnit;
 import java.util.function.BiConsumer;
 import java.util.function.Consumer;
 
-/**
- * Clock that must be used inside workflow definition code to ensure replay determinism. TODO:
- * Refactor to become a helper for managing timers instead of the generic clock class.
- */
+/** Clock that must be used inside workflow definition code to ensure replay determinism. */
 final class ClockDecisionContext {
 
   private final class TimerCancellationHandler implements Consumer<Exception> {
@@ -46,24 +41,14 @@ private final class TimerCancellationHandler implements Consumer<Exception> {
 
     @Override
     public void accept(Exception reason) {
-      decisions.cancelTimer(
-          timerId,
-          () -> {
-            OpenRequestInfo<?, ?> scheduled = scheduledTimers.remove(timerId);
-            BiConsumer<?, Exception> context = scheduled.getCompletionCallback();
-            CancellationException exception = new CancellationException("Cancelled by request");
-            exception.initCause(reason);
-            context.accept(null, exception);
-          });
+      decisions.cancelTimer(timerId, () -> timerCancelled(timerId, reason));
     }
   }
 
   private final DecisionsHelper decisions;
 
   private final Map<String, OpenRequestInfo<?, Long>> scheduledTimers = new HashMap<>();
 
-  private final SortedMap<Long, String> timersByFiringTime = new TreeMap<>();
-
   private long replayCurrentTimeMilliseconds;
 
   private boolean replaying = true;
@@ -90,60 +75,18 @@ Consumer<Exception> createTimer(long delaySeconds, Consumer<Exception> callback)
     }
     if (delaySeconds == 0) {
       callback.accept(null);
-      return Exception -> {};
+      return null;
     }
     long firingTime = currentTimeMillis() + TimeUnit.SECONDS.toMillis(delaySeconds);
-    // As the timer resolution is 1 second it doesn't really make sense to update a timer
-    // that is less than one second before the already existing.
-    if (timersByFiringTime.size() > 0) {
-      long nextTimerFiringTime = timersByFiringTime.firstKey();
-      if (firingTime > nextTimerFiringTime
-          || nextTimerFiringTime - firingTime < TimeUnit.SECONDS.toMillis(1)) {
-        return null;
-      }
-    }
-    Consumer<Exception> result = null;
-    if (!timersByFiringTime.containsKey(firingTime)) {
-      final OpenRequestInfo<?, Long> context = new OpenRequestInfo<>(firingTime);
-      final StartTimerDecisionAttributes timer = new StartTimerDecisionAttributes();
-      timer.setStartToFireTimeoutSeconds(delaySeconds);
-      final String timerId = decisions.getNextId();
-      timer.setTimerId(timerId);
-      decisions.startTimer(timer, null);
-      context.setCompletionHandle((ctx, Exception) -> callback.accept(null));
-      scheduledTimers.put(timerId, context);
-      timersByFiringTime.put(firingTime, timerId);
-      result = new ClockDecisionContext.TimerCancellationHandler(timerId);
-    }
-    SortedMap<Long, String> toCancel = timersByFiringTime.subMap(0L, firingTime);
-    for (String timerId : toCancel.values()) {
-      decisions.cancelTimer(
-          timerId,
-          () -> {
-            OpenRequestInfo<?, ?> scheduled = scheduledTimers.remove(timerId);
-            BiConsumer<?, Exception> context = scheduled.getCompletionCallback();
-            CancellationException exception =
-                new CancellationException("Cancelled as next unblock time changed");
-            context.accept(null, exception);
-          });
-    }
-    toCancel.clear();
-    return result;
-  }
-
-  void cancelAllTimers() {
-    for (String timerId : timersByFiringTime.values()) {
-      decisions.cancelTimer(
-          timerId,
-          () -> {
-            OpenRequestInfo<?, ?> scheduled = scheduledTimers.remove(timerId);
-            BiConsumer<?, Exception> context = scheduled.getCompletionCallback();
-            CancellationException exception =
-                new CancellationException("Cancelled as next unblock time changed");
-            context.accept(null, exception);
-          });
-    }
-    timersByFiringTime.clear();
+    final OpenRequestInfo<?, Long> context = new OpenRequestInfo<>(firingTime);
+    final StartTimerDecisionAttributes timer = new StartTimerDecisionAttributes();
+    timer.setStartToFireTimeoutSeconds(delaySeconds);
+    final String timerId = decisions.getNextId();
+    timer.setTimerId(timerId);
+    decisions.startTimer(timer, null);
+    context.setCompletionHandle((ctx, e) -> callback.accept(e));
+    scheduledTimers.put(timerId, context);
+    return new ClockDecisionContext.TimerCancellationHandler(timerId);
   }
 
   void setReplaying(boolean replaying) {
@@ -157,8 +100,6 @@ void handleTimerFired(TimerFiredEventAttributes attributes) {
       if (scheduled != null) {
         BiConsumer<?, Exception> completionCallback = scheduled.getCompletionCallback();
         completionCallback.accept(null, null);
-        long firingTime = scheduled.getUserContext();
-        timersByFiringTime.remove(firingTime);
       }
     }
   }
@@ -167,12 +108,18 @@ void handleTimerCanceled(HistoryEvent event) {
     TimerCanceledEventAttributes attributes = event.getTimerCanceledEventAttributes();
     String timerId = attributes.getTimerId();
     if (decisions.handleTimerCanceled(event)) {
-      OpenRequestInfo<?, ?> scheduled = scheduledTimers.remove(timerId);
-      if (scheduled != null) {
-        BiConsumer<?, Exception> completionCallback = scheduled.getCompletionCallback();
-        CancellationException exception = new CancellationException("Cancelled by request");
-        completionCallback.accept(null, exception);
-      }
+      timerCancelled(timerId, null);
+    }
+  }
+
+  private void timerCancelled(String timerId, Exception reason) {
+    OpenRequestInfo<?, ?> scheduled = scheduledTimers.remove(timerId);
+    if (scheduled == null) {
+      return;
     }
+    BiConsumer<?, Exception> context = scheduled.getCompletionCallback();
+    CancellationException exception = new CancellationException("Cancelled by request");
+    exception.initCause(reason);
+    context.accept(null, exception);
   }
 }

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

@@ -116,6 +116,4 @@ Consumer<Exception> signalWorkflowExecution(
    * @return cancellation handle. Invoke {@link Consumer#accept(Object)} to cancel timer.
    */
   Consumer<Exception> createTimer(long delaySeconds, Consumer<Exception> callback);
-
-  void cancelAllTimers();
 }

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

@@ -163,11 +163,6 @@ public Consumer<Exception> createTimer(long delaySeconds, Consumer<Exception> ca
     return workflowClock.createTimer(delaySeconds, callback);
   }
 
-  @Override
-  public void cancelAllTimers() {
-    workflowClock.cancelAllTimers();
-  }
-
   @Override
   public long currentTimeMillis() {
     return workflowClock.currentTimeMillis();

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

@@ -256,6 +256,10 @@ void startTimer(StartTimerDecisionAttributes request, Object createTimerUserCont
 
   boolean cancelTimer(String timerId, Runnable immediateCancellationCallback) {
     DecisionStateMachine decision = getDecision(new DecisionId(DecisionTarget.TIMER, timerId));
+    if (decision.isDone()) {
+      // Cancellation callbacks are not deregistered and might be invoked after timer firing
+      return true;
+    }
     decision.cancel(immediateCancellationCallback);
     return decision.isDone();
   }

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

@@ -32,6 +32,7 @@
 import java.util.List;
 import java.util.concurrent.CancellationException;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Consumer;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -59,6 +60,9 @@ class ReplayDecider {
 
   private WorkflowExecutionException failure;
 
+  private long wakeUpTime;
+  private Consumer<Exception> timerCancellationHandler;
+
   ReplayDecider(
       String domain,
       ReplayWorkflow workflow,
@@ -229,33 +233,58 @@ private void eventLoop() {
     }
   }
 
-  private void completeWorkflow() {
+  private void mayBeCompleteWorkflow() {
     if (completed) {
-      if (failure != null) {
-        decisionsHelper.failWorkflowExecution(failure);
-      } else if (cancelRequested) {
-        decisionsHelper.cancelWorkflowExecution();
+      completeWorkflow();
+    } else {
+      updateTimers();
+    }
+  }
+
+  private void completeWorkflow() {
+    if (failure != null) {
+      decisionsHelper.failWorkflowExecution(failure);
+    } else if (cancelRequested) {
+      decisionsHelper.cancelWorkflowExecution();
+    } else {
+      ContinueAsNewWorkflowExecutionParameters continueAsNewOnCompletion =
+          context.getContinueAsNewOnCompletion();
+      if (continueAsNewOnCompletion != null) {
+        decisionsHelper.continueAsNewWorkflowExecution(continueAsNewOnCompletion);
       } else {
-        ContinueAsNewWorkflowExecutionParameters continueAsNewOnCompletion =
-            context.getContinueAsNewOnCompletion();
-        if (continueAsNewOnCompletion != null) {
-          decisionsHelper.continueAsNewWorkflowExecution(continueAsNewOnCompletion);
-        } else {
-          byte[] workflowOutput = workflow.getOutput();
-          decisionsHelper.completeWorkflowExecution(workflowOutput);
-        }
+        byte[] workflowOutput = workflow.getOutput();
+        decisionsHelper.completeWorkflowExecution(workflowOutput);
       }
-    } else {
-      long nextWakeUpTime = workflow.getNextWakeUpTime();
-      if (nextWakeUpTime == 0) { // No time based waiting
-        context.cancelAllTimers();
+    }
+  }
+
+  private void updateTimers() {
+    long nextWakeUpTime = workflow.getNextWakeUpTime();
+    if (nextWakeUpTime == 0) {
+      if (timerCancellationHandler != null) {
+        timerCancellationHandler.accept(null);
+        timerCancellationHandler = null;
       }
-      long delayMilliseconds = nextWakeUpTime - context.currentTimeMillis();
-      if (nextWakeUpTime > context.currentTimeMillis()) {
-        // Round up to the nearest second as we don't want to deliver a timer
-        // earlier than requested.
-        long delaySeconds =
-            InternalUtils.roundUpToSeconds(Duration.ofMillis(delayMilliseconds)).getSeconds();
+      wakeUpTime = nextWakeUpTime;
+      return;
+    }
+    if (wakeUpTime == nextWakeUpTime && timerCancellationHandler != null) {
+      return; // existing timer
+    }
+    long delayMilliseconds = nextWakeUpTime - context.currentTimeMillis();
+    if (delayMilliseconds < 0) {
+      throw new IllegalStateException("Negative delayMilliseconds=" + delayMilliseconds);
+    }
+    // Round up to the nearest second as we don't want to deliver a timer
+    // earlier than requested.
+    long delaySeconds =
+        InternalUtils.roundUpToSeconds(Duration.ofMillis(delayMilliseconds)).getSeconds();
+    if (timerCancellationHandler != null) {
+      timerCancellationHandler.accept(null);
+      timerCancellationHandler = null;
+    }
+    wakeUpTime = nextWakeUpTime;
+    timerCancellationHandler =
         context.createTimer(
             delaySeconds,
             (t) -> {
@@ -264,8 +293,6 @@ private void completeWorkflow() {
               // But no specific timer related action is necessary as Workflow.sleep is just a
               // Workflow.await with a time based condition.
             });
-      }
-    }
   }
 
   private void handleDecisionTaskStarted(HistoryEvent event) {}
@@ -359,7 +386,7 @@ private void decideImpl(Functions.Proc query) throws Throwable {
           processEvent(event, eventType);
         }
         eventLoop();
-        completeWorkflow();
+        mayBeCompleteWorkflow();
       } while (eventsIterator.hasNext());
     } finally {
       if (query != null) {

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

@@ -28,8 +28,6 @@
 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/test/java/com/uber/cadence/internal/testing/WorkflowTestingTest.java

@@ -485,6 +485,8 @@ public String workflow1(String input) {
       try {
         activity.activity1("input");
         Workflow.sleep(Duration.ofDays(3));
+      } catch (CancellationException e) {
+        return "cancelled";
       } finally {
         s.get();
       }