Improve continueAsNew (#150)

kaibocai · web-flow · commit b919929a0503 · 2023-07-12T06:52:12.000-05:00
* improve continueAsNew

* update CHANGELOG.md

* fix events got dropped

* rethrow ContinueAsNewInterruption in await method

* update integration tests

* fix dropping event

* exclude potential duplicated event
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,7 +1,10 @@
 ## v1.2.0
+
+### Updates
 * Add `thenAccept` and `thenApply` to `Task` interface ([#148](https://github.com/microsoft/durabletask-java/pull/148))
 * Support Suspend and Resume Client APIs ([#151](https://github.com/microsoft/durabletask-java/pull/151))
 * Support restartInstance and pass restartPostUri in HttpManagementPayload ([#108](https://github.com/microsoft/durabletask-java/issues/108))
+* Improve `TaskOrchestrationContext#continueAsNew` method so it doesn't require `return` statement right after it anymore ([#149](https://github.com/microsoft/durabletask-java/pull/149))
 
 ## v1.1.1
 
diff --git a/azurefunctions/src/main/java/com/microsoft/durabletask/azurefunctions/internal/middleware/OrchestrationMiddleware.java b/azurefunctions/src/main/java/com/microsoft/durabletask/azurefunctions/internal/middleware/OrchestrationMiddleware.java
@@ -10,7 +10,8 @@
 import com.microsoft.azure.functions.internal.spi.middleware.MiddlewareChain;
 import com.microsoft.azure.functions.internal.spi.middleware.MiddlewareContext;
 import com.microsoft.durabletask.OrchestrationRunner;
-import com.microsoft.durabletask.OrchestratorBlockedException;
+import com.microsoft.durabletask.interruption.ContinueAsNewInterruption;
+import com.microsoft.durabletask.interruption.OrchestratorBlockedException;
 
 /**
  * Durable Function Orchestration Middleware
@@ -39,9 +40,14 @@ public void invoke(MiddlewareContext context, MiddlewareChain chain) throws Exce
                 // The OrchestratorBlockedEvent will be wrapped into InvocationTargetException by using reflection to
                 // invoke method. Thus get the cause to check if it's OrchestratorBlockedEvent.
                 Throwable cause = e.getCause();
-                if (cause instanceof OrchestratorBlockedException){
+                if (cause instanceof OrchestratorBlockedException) {
                     throw (OrchestratorBlockedException) cause;
                 }
+                // The ContinueAsNewInterruption will be wrapped into InvocationTargetException by using reflection to
+                // invoke method. Thus get the cause to check if it's ContinueAsNewInterruption.
+                if (cause instanceof ContinueAsNewInterruption) {
+                    throw (ContinueAsNewInterruption) cause;
+                }
                 throw new RuntimeException("Unexpected failure in the task execution", e);
             }
         });
diff --git a/client/src/main/java/com/microsoft/durabletask/Task.java b/client/src/main/java/com/microsoft/durabletask/Task.java
@@ -2,6 +2,8 @@
 // Licensed under the MIT License.
 package com.microsoft.durabletask;
 
+import com.microsoft.durabletask.interruption.OrchestratorBlockedException;
+
 import java.util.concurrent.CompletableFuture;
 import java.util.function.Consumer;
 import java.util.function.Function;
diff --git a/client/src/main/java/com/microsoft/durabletask/TaskOrchestrationExecutor.java b/client/src/main/java/com/microsoft/durabletask/TaskOrchestrationExecutor.java
@@ -4,6 +4,8 @@
 
 import com.google.protobuf.StringValue;
 import com.google.protobuf.Timestamp;
+import com.microsoft.durabletask.interruption.ContinueAsNewInterruption;
+import com.microsoft.durabletask.interruption.OrchestratorBlockedException;
 import com.microsoft.durabletask.implementation.protobuf.OrchestratorService.*;
 import com.microsoft.durabletask.implementation.protobuf.OrchestratorService.ScheduleTaskAction.Builder;
 
@@ -18,8 +20,8 @@
 import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.IntFunction;
-import java.util.function.Supplier;
 import java.util.logging.Logger;
+import java.util.stream.Collectors;
 
 final class TaskOrchestrationExecutor {
 
@@ -51,6 +53,9 @@ public TaskOrchestratorResult execute(List<HistoryEvent> pastEvents, List<Histor
             completed = true;
         } catch (OrchestratorBlockedException orchestratorBlockedException) {
             logger.fine("The orchestrator has yielded and will await for new events.");
+        } catch (ContinueAsNewInterruption continueAsNewInterruption) {
+            logger.fine("The orchestrator has continued as new.");
+            context.complete(null);
         } catch (Exception e) {
             // The orchestrator threw an unhandled exception - fail it
             // TODO: What's the right way to log this?
@@ -289,6 +294,13 @@ public void continueAsNew(Object input, boolean preserveUnprocessedEvents) {
             this.continuedAsNew = true;
             this.continuedAsNewInput = input;
             this.preserveUnprocessedEvents = preserveUnprocessedEvents;
+
+            // The ContinueAsNewInterruption exception allows the orchestration to complete immediately and return back
+            // to the sidecar.
+            // We can send the current set of actions back to the worker and wait for new events to come in.
+            // This is *not* an exception - it's a normal part of orchestrator control flow.
+            throw new ContinueAsNewInterruption(
+                    "The orchestrator invoked continueAsNew. This Throwable should never be caught by user code.");
         }
 
         @Override
@@ -739,9 +751,7 @@ private void completeInternal(
             }
 
             if (this.continuedAsNew && this.preserveUnprocessedEvents) {
-                for (HistoryEvent e : this.unprocessedEvents) {
-                    builder.addCarryoverEvents(e);
-                }
+                addCarryoverEvents(builder);
             }
 
             if (!this.isReplaying) {
@@ -755,6 +765,21 @@ private void completeInternal(
             this.pendingActions.put(id, action);
             this.isComplete = true;
         }
+
+        private void addCarryoverEvents(CompleteOrchestrationAction.Builder builder) {
+            // Add historyEvent in the unprocessedEvents buffer
+            // Add historyEvent in the new event list that haven't been added to the buffer.
+            // We don't check the event in the pass event list to avoid duplicated events.
+            Set<HistoryEvent> externalEvents = new HashSet<>(this.unprocessedEvents);
+            List<HistoryEvent> newEvents = this.historyEventPlayer.getNewEvents();
+
+            Set<HistoryEvent> filteredEvents = newEvents.stream()
+                    .filter(e -> e.getEventTypeCase() == HistoryEvent.EventTypeCase.EVENTRAISED)
+                    .collect(Collectors.toSet());
+
+            externalEvents.addAll(filteredEvents);
+            externalEvents.forEach(builder::addCarryoverEvents);
+        }
         
         private boolean waitingForEvents() {
             return this.outstandingEvents.size() > 0;
@@ -904,6 +929,10 @@ public boolean moveNext() {
                 ContextImplTask.this.processEvent(next);
                 return true;
             }
+
+            List<HistoryEvent> getNewEvents() {
+                return newEvents;
+            }
         }
 
         private class ExternalEventTask<V> extends CompletableTask<V> {
@@ -1090,6 +1119,10 @@ public V await() {
                         try {
                             return this.future.get();
                         } catch (ExecutionException e) {
+                            // rethrow if it's ContinueAsNewInterruption
+                            if (e.getCause() instanceof ContinueAsNewInterruption) {
+                                throw (ContinueAsNewInterruption) e.getCause();
+                            }
                             this.handleException(e.getCause());
                         } catch (Exception e) {
                             this.handleException(e);
diff --git a/client/src/main/java/com/microsoft/durabletask/interruption/ContinueAsNewInterruption.java b/client/src/main/java/com/microsoft/durabletask/interruption/ContinueAsNewInterruption.java
@@ -0,0 +1,18 @@
+package com.microsoft.durabletask.interruption;
+
+import com.microsoft.durabletask.TaskOrchestrationContext;
+
+/**
+ * Control flow {@code Throwable} class for orchestrator when invoke {@link TaskOrchestrationContext#continueAsNew}.
+ * This {@code Throwable} must never be caught by user
+ * code.
+ * <p>
+ * {@code ContinueAsNewInterruption} is thrown when an orchestrator calls {@link TaskOrchestrationContext#continueAsNew}.
+ * Catching {@code ContinueAsNewInterruption} in user code could prevent the orchestration from saving
+ * state and scheduling new tasks, resulting in the orchestration getting stuck.
+ */
+public class ContinueAsNewInterruption extends RuntimeException {
+    public ContinueAsNewInterruption(String message) {
+        super(message);
+    }
+}
diff --git a/client/src/main/java/com/microsoft/durabletask/interruption/OrchestratorBlockedException.java b/client/src/main/java/com/microsoft/durabletask/interruption/OrchestratorBlockedException.java
@@ -1,6 +1,8 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
-package com.microsoft.durabletask;
+package com.microsoft.durabletask.interruption;
+
+import com.microsoft.durabletask.Task;
 
 /**
  * Control flow {@code Throwable} class for orchestrator functions. This {@code Throwable} must never be caught by user
@@ -12,7 +14,7 @@
  * state and scheduling new tasks, resulting in the orchestration getting stuck.
  */
 public final class OrchestratorBlockedException extends RuntimeException {
-    OrchestratorBlockedException(String message) {
+    public OrchestratorBlockedException(String message) {
         super(message);
     }
 }
diff --git a/client/src/test/java/com/microsoft/durabletask/IntegrationTests.java b/client/src/test/java/com/microsoft/durabletask/IntegrationTests.java
@@ -356,7 +356,7 @@ void continueAsNew() throws TimeoutException {
     }
 
     @Test
-    void continueAsNewWithExternalEvents() throws TimeoutException {
+    void continueAsNewWithExternalEvents() throws TimeoutException, InterruptedException{
         final String orchestratorName = "continueAsNewWithExternalEvents";
         final String eventName = "MyEvent";
         final int expectedEventCount = 10;
@@ -1194,7 +1194,7 @@ void thenApply() throws IOException, InterruptedException, TimeoutException {
     }
 
     @Test
-    void externalEventThenAccept() throws IOException, InterruptedException, TimeoutException {
+    void externalEventThenAccept() throws InterruptedException, TimeoutException {
         final String orchestratorName = "continueAsNewWithExternalEvents";
         final String eventName = "MyEvent";
         final int expectedEventCount = 10;
