Support repeatedly calling WhenAny on a progressively smaller list of already scheduled tasks (#446)

Author: davidmrdavid

azure/durable_functions/models/DurableOrchestrationContext.py

@@ -683,7 +683,7 @@ def _get_function_name(self, name: FunctionBuilder,
                 name = name._function._name
                 return name
             else:
-                if(trigger_type == OrchestrationTrigger):
+                if (trigger_type == OrchestrationTrigger):
                     trigger_type = "OrchestrationTrigger"
                 else:
                     trigger_type = "ActivityTrigger"

azure/durable_functions/models/Task.py

@@ -56,6 +56,14 @@ def __init__(self, id_: Union[int, str], actions: Union[List[Action], Action]):
         self.result: Any = None
         self.action_repr: Union[List[Action], Action] = actions
         self.is_played = False
+        self._is_scheduled_flag = False
+
+    @property
+    def _is_scheduled(self) -> bool:
+        return self._is_scheduled_flag
+
+    def _set_is_scheduled(self, is_scheduled: bool):
+        self._is_scheduled_flag = is_scheduled
 
     @property
     def is_completed(self) -> bool:
@@ -158,7 +166,8 @@ def __init__(self, tasks: List[TaskBase], compound_action_constructor=None):
             if isinstance(action_repr, list):
                 child_actions.extend(action_repr)
             else:
-                child_actions.append(action_repr)
+                if not task._is_scheduled:
+                    child_actions.append(action_repr)
         if compound_action_constructor is None:
             self.action_repr = child_actions
         else:  # replay_schema is ReplaySchema.V2
@@ -176,6 +185,10 @@ def __init__(self, tasks: List[TaskBase], compound_action_constructor=None):
             if not (child.state is TaskState.RUNNING):
                 self.handle_completion(child)
 
+    @property
+    def _is_scheduled(self) -> bool:
+        return all([child._is_scheduled for child in self.children])
+
     def handle_completion(self, child: TaskBase):
         """Manage sub-task completion events.
 

azure/durable_functions/models/TaskOrchestrationExecutor.py

@@ -1,4 +1,4 @@
-from azure.durable_functions.models.Task import TaskBase, TaskState, AtomicTask
+from azure.durable_functions.models.Task import TaskBase, TaskState, AtomicTask, CompoundTask
 from azure.durable_functions.models.OrchestratorState import OrchestratorState
 from azure.durable_functions.models.DurableOrchestrationContext import DurableOrchestrationContext
 from typing import Any, List, Optional, Union
@@ -229,7 +229,8 @@ def resume_user_code(self):
             task_succeeded = current_task.state is TaskState.SUCCEEDED
             new_task = self.generator.send(
                 task_value) if task_succeeded else self.generator.throw(task_value)
-            self.context._add_to_open_tasks(new_task)
+            if isinstance(new_task, TaskBase) and not (new_task._is_scheduled):
+                self.context._add_to_open_tasks(new_task)
         except StopIteration as stop_exception:
             # the orchestration returned,
             # flag it as such and capture its output
@@ -245,9 +246,17 @@ def resume_user_code(self):
                 # user yielded the same task multiple times, continue executing code
                 # until a new/not-previously-yielded task is encountered
                 self.resume_user_code()
-            else:
+            elif not (self.current_task._is_scheduled):
                 # new task is received. it needs to be resolved to a value
                 self.context._add_to_actions(self.current_task.action_repr)
+                self._mark_as_scheduled(self.current_task)
+
+    def _mark_as_scheduled(self, task: TaskBase):
+        if isinstance(task, CompoundTask):
+            for task in task.children:
+                self._mark_as_scheduled(task)
+        else:
+            task._set_is_scheduled(True)
 
     def get_orchestrator_state_str(self) -> str:
         """Obtain a JSON-formatted string representing the orchestration's state.

tests/orchestrator/test_sequential_orchestrator.py

@@ -77,6 +77,24 @@ def generator_function_duplicate_yield(context):
 
     return ""
 
+def generator_function_reducing_when_all(context):
+    task1 = context.call_activity("Hello", "Tokyo")
+    task2 = context.call_activity("Hello", "Seattle")
+    pending_tasks = [task1, task2]
+
+    #  Yield until first task is completed
+    finished_task1 = yield context.task_any(pending_tasks)
+
+    #  Remove completed task from pending tasks
+    pending_tasks.remove(finished_task1)
+
+    #  Yield remaining task
+    yield context.task_any(pending_tasks)
+
+    # Ensure we can still schedule new tasks
+    yield context.call_activity("Hello", "London")
+    return ""
+
 def generator_function_compound_tasks(context):
     yield context.call_activity("Hello", "Tokyo")
 
@@ -689,6 +707,30 @@ def test_duplicate_yields_do_not_add_duplicate_actions():
     assert_valid_schema(result)
     assert_orchestration_state_equals(expected, result)
 
+def test_reducing_when_any_pattern():
+    """Tests that a user can call when_any on a progressively smaller list of already scheduled tasks"""
+    context_builder = ContextBuilder('test_reducing_when_any', replay_schema=ReplaySchema.V2)
+    add_hello_completed_events(context_builder, 0, "\"Hello Tokyo!\"")
+    add_hello_completed_events(context_builder, 1, "\"Hello Seattle!\"")
+    add_hello_completed_events(context_builder, 2, "\"Hello London!\"")
+
+    result = get_orchestration_state_result(
+        context_builder, generator_function_reducing_when_all)
+
+    # this scenario is only supported for V2 replay
+    expected_state = base_expected_state("",replay_schema=ReplaySchema.V2)
+    expected_state._actions = [
+        [WhenAnyAction(
+            [CallActivityAction("Hello", "Seattle"), CallActivityAction("Hello", "Tokyo")]),
+            CallActivityAction("Hello", "London")
+        ]
+    ]
+
+    expected_state._is_done = True
+    expected = expected_state.to_json()
+
+    assert_orchestration_state_equals(expected, result)
+
 def test_compound_tasks_return_single_action_in_V2():
     """Tests that compound tasks, in the v2 replay schema, are represented as a single "deep" action"""
     context_builder = ContextBuilder('test_v2_replay_schema', replay_schema=ReplaySchema.V2)