ContinueAsNew works, needed to mark the orchestrator as completed (#158)

Author: davidmrdavid

azure/durable_functions/models/DurableOrchestrationContext.py

@@ -32,6 +32,7 @@ def __init__(self,
         self._custom_status: Any = None
         self._new_uuid_counter: int = 0
         self._sub_orchestrator_counter: int = 0
+        self._continue_as_new_flag: bool = False
         self.call_activity = lambda n, i=None: call_activity_task(
             state=self.histories,
             name=n,
@@ -65,7 +66,7 @@ def __init__(self,
             state=self.histories,
             name=n)
         self.new_uuid = lambda: new_uuid(context=self)
-        self.continue_as_new = lambda i: continue_as_new(input_=i)
+        self.continue_as_new = lambda i: continue_as_new(context=self, input_=i)
         self.task_any = lambda t: task_any(tasks=t)
         self.task_all = lambda t: task_all(tasks=t)
         self.create_timer = lambda d: create_timer_task(state=self.histories, fire_at=d)
@@ -344,3 +345,8 @@ def function_context(self) -> FunctionContext:
             Object containing function level attributes not used by durable orchestrator.
         """
         return self._function_context
+
+    @property
+    def will_continue_as_new(self) -> bool:
+        """Return true if continue_as_new was called."""
+        return self._continue_as_new_flag

azure/durable_functions/orchestrator.py

@@ -49,57 +49,66 @@ def handle(self, context: DurableOrchestrationContext):
         suspended = False
 
         fn_output = self.fn(self.durable_context)
+
         # If `fn_output` is not an Iterator, then the orchestrator
         # function does not make use of its context parameter. If so,
         # `fn_output` is the return value instead of a generator
-        if isinstance(fn_output, Iterator):
-            self.generator = fn_output
-
-        else:
+        if not isinstance(fn_output, Iterator):
             orchestration_state = OrchestratorState(
                 is_done=True,
                 output=fn_output,
                 actions=self.durable_context.actions,
                 custom_status=self.durable_context.custom_status)
-            return orchestration_state.to_json_string()
-        try:
-            generation_state = self._generate_next(None)
-
-            while not suspended:
-                self._add_to_actions(generation_state)
-
-                if should_suspend(generation_state):
-                    orchestration_state = OrchestratorState(
-                        is_done=False,
-                        output=None,
-                        actions=self.durable_context.actions,
-                        custom_status=self.durable_context.custom_status)
-                    suspended = True
-                    continue
-
-                if (isinstance(generation_state, Task)
-                    or isinstance(generation_state, TaskSet)) and (
-                        generation_state.is_faulted):
-                    generation_state = self.generator.throw(
-                        generation_state.exception)
-                    continue
-
-                self._reset_timestamp()
-                generation_state = self._generate_next(generation_state)
-
-        except StopIteration as sie:
-            orchestration_state = OrchestratorState(
-                is_done=True,
-                output=sie.value,
-                actions=self.durable_context.actions,
-                custom_status=self.durable_context.custom_status)
-        except Exception as e:
-            orchestration_state = OrchestratorState(
-                is_done=False,
-                output=None,  # Should have no output, after generation range
-                actions=self.durable_context.actions,
-                error=str(e),
-                custom_status=self.durable_context.custom_status)
+
+        else:
+            self.generator = fn_output
+            try:
+                generation_state = self._generate_next(None)
+
+                while not suspended:
+                    self._add_to_actions(generation_state)
+
+                    if should_suspend(generation_state):
+
+                        # The `is_done` field should be False here unless
+                        # `continue_as_new` was called. Therefore,
+                        # `will_continue_as_new` essentially "tracks"
+                        # whether or not the orchestration is done.
+                        orchestration_state = OrchestratorState(
+                            is_done=self.durable_context.will_continue_as_new,
+                            output=None,
+                            actions=self.durable_context.actions,
+                            custom_status=self.durable_context.custom_status)
+                        suspended = True
+                        continue
+
+                    if (isinstance(generation_state, Task)
+                        or isinstance(generation_state, TaskSet)) and (
+                            generation_state.is_faulted):
+                        generation_state = self.generator.throw(
+                            generation_state.exception)
+                        continue
+
+                    self._reset_timestamp()
+                    generation_state = self._generate_next(generation_state)
+
+            except StopIteration as sie:
+                orchestration_state = OrchestratorState(
+                    is_done=True,
+                    output=sie.value,
+                    actions=self.durable_context.actions,
+                    custom_status=self.durable_context.custom_status)
+            except Exception as e:
+                orchestration_state = OrchestratorState(
+                    is_done=False,
+                    output=None,  # Should have no output, after generation range
+                    actions=self.durable_context.actions,
+                    error=str(e),
+                    custom_status=self.durable_context.custom_status)
+
+        # No output if continue_as_new was called
+        if self.durable_context.will_continue_as_new:
+            orchestration_state._output = None
 
         return orchestration_state.to_json_string()
 
@@ -108,9 +117,13 @@ def _generate_next(self, partial_result):
             gen_result = self.generator.send(partial_result.result)
         else:
             gen_result = self.generator.send(None)
+
         return gen_result
 
     def _add_to_actions(self, generation_state):
+        # Do not add new tasks to action if continue_as_new was called
+        if self.durable_context.will_continue_as_new:
+            return
         if (isinstance(generation_state, Task)
                 and hasattr(generation_state, "action")):
             self.durable_context.actions.append([generation_state.action])

azure/durable_functions/tasks/continue_as_new.py

@@ -1,24 +1,19 @@
 from typing import Any
 
-from ..models.Task import (
-    Task)
 from ..models.actions.ContinueAsNewAction import ContinueAsNewAction
 
 
 def continue_as_new(
-        input_: Any = None) -> Task:
+        context,
+        input_: Any = None):
     """Create a new continue as new action.
 
     Parameters
     ----------
     input_: Any
         The JSON-serializable input to pass to the activity function.
-
-    Returns
-    -------
-    Task
-        A Durable Task that causes the orchestrator reset and start as a new orchestration.
     """
     new_action = ContinueAsNewAction(input_)
 
-    return Task(is_completed=False, is_faulted=False, action=new_action)
+    context.actions.append([new_action])
+    context._continue_as_new_flag = True

tests/orchestrator/test_continue_as_new.py

@@ -10,7 +10,7 @@
 
 def generator_function(context):
     yield context.call_activity("Hello", "Tokyo")
-    yield context.continue_as_new("Cause I can")
+    context.continue_as_new("Cause I can")
 
 
 def base_expected_state(output=None) -> OrchestratorState:
@@ -25,6 +25,7 @@ def add_hello_action(state: OrchestratorState, input_: str):
 def add_continue_as_new_action(state: OrchestratorState, input_: str):
     action = ContinueAsNewAction(input_=input_)
     state.actions.append([action])
+    state._is_done = True
 
 
 def add_hello_completed_events(