Suspend and resume support (#4)

Author: cgillum

README.md

@@ -9,15 +9,6 @@ This repo contains a Python client SDK for use with the [Durable Task Framework
 
 > Note that this project is **not** currently affiliated with the [Durable Functions](https://docs.microsoft.com/azure/azure-functions/durable/durable-functions-overview) project for Azure Functions. If you are looking for a Python SDK for Durable Functions, please see [this repo](https://github.com/Azure/azure-functions-durable-python).
 
-## Features
-
-- [x] Orchestrations
-- [x] Activities
-- [x] Durable timers
-- [x] Sub-orchestrations
-- [ ] External events
-- [ ] Suspend and resume
-- [ ] Retry policies
 
 ## Supported patterns
 
@@ -102,6 +93,38 @@ As an aside, you'll also notice that the example orchestration above works with
 
 You can find the full sample [here](./examples/human_interaction.py).
 
+## Feature overview
+
+The following features are currently supported:
+
+### Orchestrations
+
+Orchestrations are implemented using ordinary Python functions that take an `OrchestrationContext` as their first parameter. The `OrchestrationContext` provides APIs for starting child orchestrations, scheduling activities, and waiting for external events, among other things. Orchestrations are fault-tolerant and durable, meaning that they can automatically recover from failures and rebuild their local execution state. Orchestrator functions must be deterministic, meaning that they must always produce the same output given the same input.
+
+### Activities
+
+Activities are implemented using ordinary Python functions that take an `ActivityContext` as their first parameter. Activity functions are scheduled by orchestrations and have at-least-once execution guarantees, meaning that they will be executed at least once but may be executed multiple times in the event of a transient failure. Activity functions are where the real "work" of any orchestration is done.
+
+### Durable timers
+
+Orchestrations can schedule durable timers using the `create_timer` API. These timers are durable, meaning that they will survive orchestrator restarts and will fire even if the orchestrator is not actively in memory. Durable timers can be of any duration, from milliseconds to months.
+
+### Sub-orchestrations
+
+Orchestrations can start child orchestrations using the `call_sub_orchestrator` API. Child orchestrations are useful for encapsulating complex logic and for breaking up large orchestrations into smaller, more manageable pieces.
+
+### External events
+
+Orchestrations can wait for external events using the `wait_for_external_event` API. External events are useful for implementing human interaction patterns, such as waiting for a user to approve an order before continuing.
+
+### Suspend and resume
+
+Orchestrations can be suspended using the `suspend_orchestration` client API and will remain suspended until resumed using the `resume_orchestration` client API. A suspended orchestration will stop processing new events, but will continue to buffer any that happen to arrive until resumed, ensuring that no data is lost.
+
+### Retry policies (TODO)
+
+Orchestrations can specify retry policies for activities and sub-orchestrations. These policies control how many times and how frequently an activity or sub-orchestration will be retried in the event of a transient error.
+
 ## Getting Started
 
 ### Prerequisites

durabletask/client.py

@@ -162,10 +162,14 @@ def raise_orchestration_event(self, instance_id: str, event_name: str, *,
         self._stub.RaiseEvent(req)
 
     def terminate_orchestration(self):
-        pass
+        raise NotImplementedError()
 
-    def suspend_orchestration(self):
-        pass
+    def suspend_orchestration(self, instance_id: str):
+        req = pb.SuspendRequest(instanceId=instance_id)
+        self._logger.info(f"Suspending instance '{instance_id}'.")
+        self._stub.SuspendInstance(req)
 
-    def resume_orchestration(self):
-        pass
+    def resume_orchestration(self, instance_id: str):
+        req = pb.ResumeRequest(instanceId=instance_id)
+        self._logger.info(f"Resuming instance '{instance_id}'.")
+        self._stub.ResumeInstance(req)

durabletask/internal/helpers.py

@@ -123,6 +123,22 @@ def new_event_raised_event(name: str, encoded_input: str | None = None) -> pb.Hi
     )
 
 
+def new_suspend_event() -> pb.HistoryEvent:
+    return pb.HistoryEvent(
+        eventId=-1,
+        timestamp=timestamp_pb2.Timestamp(),
+        executionSuspended=pb.ExecutionSuspendedEvent()
+    )
+
+
+def new_resume_event() -> pb.HistoryEvent:
+    return pb.HistoryEvent(
+        eventId=-1,
+        timestamp=timestamp_pb2.Timestamp(),
+        executionResumed=pb.ExecutionResumedEvent()
+    )
+
+
 def get_string_value(val: str | None) -> wrappers_pb2.StringValue | None:
     if val is None:
         return None

durabletask/worker.py

@@ -93,6 +93,7 @@ def __init__(self, *,
         self._logger = shared.get_logger(log_handler, log_formatter)
         self._shutdown = Event()
         self._response_stream = None
+        self._is_running = False
 
     def __enter__(self):
         return self
@@ -102,17 +103,24 @@ def __exit__(self, type, value, traceback):
 
     def add_orchestrator(self, fn: task.Orchestrator) -> str:
         """Registers an orchestrator function with the worker."""
+        if self._is_running:
+            raise RuntimeError('Orchestrators cannot be added while the worker is running.')
         return self._registry.add_orchestrator(fn)
 
     def add_activity(self, fn: task.Activity) -> str:
         """Registers an activity function with the worker."""
+        if self._is_running:
+            raise RuntimeError('Activities cannot be added while the worker is running.')
         return self._registry.add_activity(fn)
 
     def start(self):
         """Starts the worker on a background thread and begins listening for work items."""
         channel = shared.get_grpc_channel(self._host_address)
         stub = stubs.TaskHubSidecarServiceStub(channel)
 
+        if self._is_running:
+            raise RuntimeError('The worker is already running.')
+
         def run_loop():
             # TODO: Investigate whether asyncio could be used to enable greater concurrency for async activity
             #       functions. We'd need to know ahead of time whether a function is async or not.
@@ -158,16 +166,21 @@ def run_loop():
         self._logger.info(f"starting gRPC worker that connects to {self._host_address}")
         self._runLoop = Thread(target=run_loop)
         self._runLoop.start()
+        self._is_running = True
 
     def stop(self):
         """Stops the worker and waits for any pending work items to complete."""
+        if not self._is_running:
+            return
+
         self._logger.info("Stopping gRPC worker...")
         self._shutdown.set()
         if self._response_stream is not None:
             self._response_stream.cancel()
         if self._runLoop is not None:
             self._runLoop.join(timeout=30)
         self._logger.info("Worker shutdown completed")
+        self._is_running = False
 
     def _execute_orchestrator(self, req: pb.OrchestratorRequest, stub: stubs.TaskHubSidecarServiceStub):
         try:
@@ -374,6 +387,8 @@ def __init__(self, registry: _Registry, logger: logging.Logger):
         self._registry = registry
         self._logger = logger
         self._generator = None
+        self._is_suspended = False
+        self._suspended_events: List[pb.HistoryEvent] = []
 
     def execute(self, instance_id: str, old_events: Sequence[pb.HistoryEvent], new_events: Sequence[pb.HistoryEvent]) -> List[pb.OrchestratorAction]:
         if not new_events:
@@ -408,6 +423,12 @@ def execute(self, instance_id: str, old_events: Sequence[pb.HistoryEvent], new_e
         return actions
 
     def process_event(self, ctx: _RuntimeOrchestrationContext, event: pb.HistoryEvent) -> None:
+        if self._is_suspended and _is_suspendable(event):
+            # We are suspended, so we need to buffer this event until we are resumed
+            self._suspended_events.append(event)
+            return
+
+        # CONSIDER: change to a switch statement with event.WhichOneof("eventType")
         try:
             if event.HasField("orchestratorStarted"):
                 ctx.current_utc_datetime = event.timestamp.ToDatetime()
@@ -445,8 +466,9 @@ def process_event(self, ctx: _RuntimeOrchestrationContext, event: pb.HistoryEven
                 timer_task = ctx._pending_tasks.pop(timer_id, None)
                 if not timer_task:
                     # TODO: Should this be an error? When would it ever happen?
-                    self._logger.warning(
-                        f"Ignoring unexpected timerFired event for '{ctx.instance_id}' with ID = {timer_id}.")
+                    if not ctx._is_replaying:
+                        self._logger.warning(
+                            f"{ctx.instance_id}: Ignoring unexpected timerFired event with ID = {timer_id}.")
                     return
                 timer_task.complete(None)
                 ctx.resume()
@@ -472,8 +494,9 @@ def process_event(self, ctx: _RuntimeOrchestrationContext, event: pb.HistoryEven
                 activity_task = ctx._pending_tasks.pop(task_id, None)
                 if not activity_task:
                     # TODO: Should this be an error? When would it ever happen?
-                    self._logger.warning(
-                        f"Ignoring unexpected taskCompleted event for '{ctx.instance_id}' with ID = {task_id}.")
+                    if not ctx.is_replaying:
+                        self._logger.warning(
+                            f"{ctx.instance_id}: Ignoring unexpected taskCompleted event with ID = {task_id}.")
                     return
                 result = None
                 if not ph.is_empty(event.taskCompleted.result):
@@ -485,11 +508,12 @@ def process_event(self, ctx: _RuntimeOrchestrationContext, event: pb.HistoryEven
                 activity_task = ctx._pending_tasks.pop(task_id, None)
                 if not activity_task:
                     # TODO: Should this be an error? When would it ever happen?
-                    self._logger.warning(
-                        f"Ignoring unexpected taskFailed event for '{ctx.instance_id}' with ID = {task_id}.")
+                    if not ctx.is_replaying:
+                        self._logger.warning(
+                            f"{ctx.instance_id}: Ignoring unexpected taskFailed event with ID = {task_id}.")
                     return
                 activity_task.fail(
-                    f"Activity task #{task_id} failed: {event.taskFailed.failureDetails.errorMessage}",
+                    f"{ctx.instance_id}: Activity task #{task_id} failed: {event.taskFailed.failureDetails.errorMessage}",
                     event.taskFailed.failureDetails)
                 ctx.resume()
             elif event.HasField("subOrchestrationInstanceCreated"):
@@ -513,8 +537,9 @@ def process_event(self, ctx: _RuntimeOrchestrationContext, event: pb.HistoryEven
                 sub_orch_task = ctx._pending_tasks.pop(task_id, None)
                 if not sub_orch_task:
                     # TODO: Should this be an error? When would it ever happen?
-                    self._logger.warning(
-                        f"Ignoring unexpected subOrchestrationInstanceCompleted event for '{ctx.instance_id}' with ID = {task_id}.")
+                    if not ctx.is_replaying:
+                        self._logger.warning(
+                            f"{ctx.instance_id}: Ignoring unexpected subOrchestrationInstanceCompleted event with ID = {task_id}.")
                     return
                 result = None
                 if not ph.is_empty(event.subOrchestrationInstanceCompleted.result):
@@ -527,8 +552,9 @@ def process_event(self, ctx: _RuntimeOrchestrationContext, event: pb.HistoryEven
                 sub_orch_task = ctx._pending_tasks.pop(task_id, None)
                 if not sub_orch_task:
                     # TODO: Should this be an error? When would it ever happen?
-                    self._logger.warning(
-                        f"Ignoring unexpected subOrchestrationInstanceFailed event for '{ctx.instance_id}' with ID = {task_id}.")
+                    if not ctx.is_replaying:
+                        self._logger.warning(
+                            f"{ctx.instance_id}: Ignoring unexpected subOrchestrationInstanceFailed event with ID = {task_id}.")
                     return
                 sub_orch_task.fail(
                     f"Sub-orchestration task #{task_id} failed: {failedEvent.failureDetails.errorMessage}",
@@ -559,7 +585,18 @@ def process_event(self, ctx: _RuntimeOrchestrationContext, event: pb.HistoryEven
                         decoded_result = shared.from_json(event.eventRaised.input.value)
                     event_list.append(_ExternalEvent(event.eventRaised.name, decoded_result))
                     if not ctx.is_replaying:
-                        self._logger.info(f"Event '{event_name}' has been buffered as there are no tasks waiting for it.")
+                        self._logger.info(f"{ctx.instance_id}: Event '{event_name}' has been buffered as there are no tasks waiting for it.")
+            elif event.HasField("executionSuspended"):
+                if not self._is_suspended and not ctx.is_replaying:
+                    self._logger.info(f"{ctx.instance_id}: Execution suspended.")
+                self._is_suspended = True
+            elif event.HasField("executionResumed") and self._is_suspended:
+                if not ctx.is_replaying:
+                    self._logger.info(f"{ctx.instance_id}: Resuming execution.")
+                self._is_suspended = False
+                for e in self._suspended_events:
+                    self.process_event(ctx, e)
+                self._suspended_events = []
             else:
                 eventType = event.WhichOneof("eventType")
                 raise task.OrchestrationStateError(f"Don't know how to handle event of type '{eventType}'")
@@ -667,3 +704,8 @@ def _get_action_summary(new_actions: Sequence[pb.OrchestratorAction]) -> str:
             action_type = action.WhichOneof('orchestratorActionType')
             counts[action_type] = counts.get(action_type, 0) + 1
         return f"[{', '.join(f'{name}={count}' for name, count in counts.items())}]"
+
+
+def _is_suspendable(event: pb.HistoryEvent) -> bool:
+    """Returns true if the event is one that can be suspended and resumed."""
+    return event.WhichOneof("eventType") not in ["executionResumed", "executionTerminated"]

tests/test_orchestration_e2e.py

@@ -3,6 +3,7 @@
 
 import json
 import threading
+import time
 from datetime import timedelta
 
 import pytest
@@ -169,3 +170,42 @@ def orchestrator(ctx: task.OrchestrationContext, _):
         assert state.serialized_output == json.dumps("approved")
     else:
         assert state.serialized_output == json.dumps("timed out")
+
+
+def test_suspend_and_resume():
+    def orchestrator(ctx: task.OrchestrationContext, _):
+        result = yield ctx.wait_for_external_event("my_event")
+        return result
+
+    # Start a worker, which will connect to the sidecar in a background thread
+    with worker.TaskHubGrpcWorker() as w:
+        w.add_orchestrator(orchestrator)
+        w.start()
+
+        task_hub_client = client.TaskHubGrpcClient()
+        id = task_hub_client.schedule_new_orchestration(orchestrator)
+        state = task_hub_client.wait_for_orchestration_start(id, timeout=30)
+        assert state is not None
+
+        # Suspend the orchestration and wait for it to go into the SUSPENDED state
+        task_hub_client.suspend_orchestration(id)
+        while state.runtime_status == client.OrchestrationStatus.RUNNING:
+            time.sleep(0.1)
+            state = task_hub_client.get_orchestration_state(id)
+            assert state is not None
+        assert state.runtime_status == client.OrchestrationStatus.SUSPENDED
+
+        # Raise an event to the orchestration and confirm that it does NOT complete
+        task_hub_client.raise_orchestration_event(id, "my_event", data=42)
+        try:
+            state = task_hub_client.wait_for_orchestration_completion(id, timeout=3)
+            assert False, "Orchestration should not have completed"
+        except TimeoutError:
+            pass
+
+        # Resume the orchestration and wait for it to complete
+        task_hub_client.resume_orchestration(id)
+        state = task_hub_client.wait_for_orchestration_completion(id, timeout=30)
+        assert state is not None
+        assert state.runtime_status == client.OrchestrationStatus.COMPLETED
+        assert state.serialized_output == json.dumps(42)

tests/test_orchestration_executor.py

@@ -486,6 +486,7 @@ def orchestrator(ctx: task.OrchestrationContext, _):
     # the orchestration should complete.
     old_events = new_events
     new_events = [helpers.new_event_raised_event("my_event", encoded_input="42")]
+    executor = worker._OrchestrationExecutor(registry, TEST_LOGGER)
     actions = executor.execute(TEST_INSTANCE_ID, old_events, new_events)
     complete_action = get_and_validate_single_complete_orchestration_action(actions)
     assert complete_action.orchestrationStatus == pb.ORCHESTRATION_STATUS_COMPLETED
@@ -519,6 +520,40 @@ def orchestrator(ctx: task.OrchestrationContext, _):
     timer_due_time = datetime.utcnow() + timedelta(days=1)
     old_events = new_events + [helpers.new_timer_created_event(1, timer_due_time)]
     new_events = [helpers.new_timer_fired_event(1, timer_due_time)]
+    executor = worker._OrchestrationExecutor(registry, TEST_LOGGER)
+    actions = executor.execute(TEST_INSTANCE_ID, old_events, new_events)
+    complete_action = get_and_validate_single_complete_orchestration_action(actions)
+    assert complete_action.orchestrationStatus == pb.ORCHESTRATION_STATUS_COMPLETED
+    assert complete_action.result.value == "42"
+
+
+def test_suspend_resume():
+    """Tests that an orchestration can be suspended and resumed"""
+
+    def orchestrator(ctx: task.OrchestrationContext, _):
+        result = yield ctx.wait_for_external_event("my_event")
+        return result
+
+    registry = worker._Registry()
+    orchestrator_name = registry.add_orchestrator(orchestrator)
+
+    old_events = [
+        helpers.new_orchestrator_started_event(),
+        helpers.new_execution_started_event(orchestrator_name, TEST_INSTANCE_ID)]
+    new_events = [
+        helpers.new_suspend_event(),
+        helpers.new_event_raised_event("my_event", encoded_input="42")]
+
+    # Execute the orchestration. It should remain in a running state because it was suspended prior
+    # to processing the event raised event.
+    executor = worker._OrchestrationExecutor(registry, TEST_LOGGER)
+    actions = executor.execute(TEST_INSTANCE_ID, old_events, new_events)
+    assert len(actions) == 0
+
+    # Resume the orchestration. It should complete successfully.
+    old_events = old_events + new_events
+    new_events = [helpers.new_resume_event()]
+    executor = worker._OrchestrationExecutor(registry, TEST_LOGGER)
     actions = executor.execute(TEST_INSTANCE_ID, old_events, new_events)
     complete_action = get_and_validate_single_complete_orchestration_action(actions)
     assert complete_action.orchestrationStatus == pb.ORCHESTRATION_STATUS_COMPLETED