Add termination support (#5)

Author: cgillum

CHANGELOG.md

@@ -5,4 +5,4 @@ Initial release, which includes the following features:
 - Orchestrations and activities
 - Durable timers
 - Sub-orchestrations
-
+- Suspend, resume, and terminate client operations

README.md

@@ -117,9 +117,13 @@ Orchestrations can start child orchestrations using the `call_sub_orchestrator`
 
 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
+### Continue-as-new (TODO)
 
-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.
+Orchestrations can be continued as new using the `continue_as_new` API. This API allows an orchestration to restart itself from scratch, optionally with a new input.
+
+### Suspend, resume, and terminate
+
+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. An orchestration can also be terminated using the `terminate_orchestration` client API. Terminated orchestrations will stop processing new events and will discard any buffered events.
 
 ### Retry policies (TODO)
 

durabletask/client.py

@@ -161,8 +161,14 @@ def raise_orchestration_event(self, instance_id: str, event_name: str, *,
         self._logger.info(f"Raising event '{event_name}' for instance '{instance_id}'.")
         self._stub.RaiseEvent(req)
 
-    def terminate_orchestration(self):
-        raise NotImplementedError()
+    def terminate_orchestration(self, instance_id: str, *,
+                                output: Any | None = None):
+        req = pb.TerminateRequest(
+            instanceId=instance_id,
+            output=wrappers_pb2.StringValue(value=shared.to_json(output)) if output else None)
+
+        self._logger.info(f"Terminating instance '{instance_id}'.")
+        self._stub.TerminateInstance(req)
 
     def suspend_orchestration(self, instance_id: str):
         req = pb.SuspendRequest(instanceId=instance_id)

durabletask/internal/helpers.py

@@ -139,6 +139,16 @@ def new_resume_event() -> pb.HistoryEvent:
     )
 
 
+def new_terminated_event(*, encoded_output: str | None = None) -> pb.HistoryEvent:
+    return pb.HistoryEvent(
+        eventId=-1,
+        timestamp=timestamp_pb2.Timestamp(),
+        executionTerminated=pb.ExecutionTerminatedEvent(
+            input=get_string_value(encoded_output)
+        )
+    )
+
+
 def get_string_value(val: str | None) -> wrappers_pb2.StringValue | None:
     if val is None:
         return None

durabletask/worker.py

@@ -277,17 +277,23 @@ def resume(self):
                     if not self._previous_task.is_complete:
                         break
 
-    def set_complete(self, result: Any):
+    def set_complete(self, result: Any, status: pb.OrchestrationStatus, is_result_encoded: bool = False):
+        if self._is_complete:
+            return
+
         self._is_complete = True
         self._result = result
         result_json: str | None = None
         if result is not None:
-            result_json = shared.to_json(result)
+            result_json = result if is_result_encoded else shared.to_json(result)
         action = ph.new_complete_orchestration_action(
-            self.next_sequence_number(), pb.ORCHESTRATION_STATUS_COMPLETED, result_json)
+            self.next_sequence_number(), status, result_json)
         self._pending_actions[action.id] = action
 
     def set_failed(self, ex: Exception):
+        if self._is_complete:
+            return
+
         self._is_complete = True
         self._pending_actions.clear()  # Cancel any pending actions
         action = ph.new_complete_orchestration_action(
@@ -409,6 +415,8 @@ def execute(self, instance_id: str, old_events: Sequence[pb.HistoryEvent], new_e
             ctx._is_replaying = False
             for new_event in new_events:
                 self.process_event(ctx, new_event)
+                if ctx._is_complete:
+                    break
         except Exception as ex:
             # Unhandled exceptions fail the orchestration
             ctx.set_failed(ex)
@@ -450,7 +458,7 @@ def process_event(self, ctx: _RuntimeOrchestrationContext, event: pb.HistoryEven
                     ctx.run(result)
                 else:
                     # This is an orchestrator that doesn't schedule any tasks
-                    ctx.set_complete(result)
+                    ctx.set_complete(result, pb.ORCHESTRATION_STATUS_COMPLETED)
             elif event.HasField("timerCreated"):
                 # This history event confirms that the timer was successfully scheduled.
                 # Remove the timerCreated event from the pending action list so we don't schedule it again.
@@ -597,12 +605,17 @@ def process_event(self, ctx: _RuntimeOrchestrationContext, event: pb.HistoryEven
                 for e in self._suspended_events:
                     self.process_event(ctx, e)
                 self._suspended_events = []
+            elif event.HasField("executionTerminated"):
+                if not ctx.is_replaying:
+                    self._logger.info(f"{ctx.instance_id}: Execution terminating.")
+                encoded_output = event.executionTerminated.input.value if not ph.is_empty(event.executionTerminated.input) else None
+                ctx.set_complete(encoded_output, pb.ORCHESTRATION_STATUS_TERMINATED, is_result_encoded=True)
             else:
                 eventType = event.WhichOneof("eventType")
                 raise task.OrchestrationStateError(f"Don't know how to handle event of type '{eventType}'")
         except StopIteration as generatorStopped:
             # The orchestrator generator function completed
-            ctx.set_complete(generatorStopped.value)
+            ctx.set_complete(generatorStopped.value, pb.ORCHESTRATION_STATUS_COMPLETED)
 
 
 class _ActivityExecutor:

tests/test_orchestration_e2e.py

@@ -101,10 +101,10 @@ def parent_orchestrator(ctx: task.OrchestrationContext, count: int):
 
     # Start a worker, which will connect to the sidecar in a background thread
     with worker.TaskHubGrpcWorker() as w:
-        w.start()
         w.add_activity(increment)
         w.add_orchestrator(orchestrator_child)
         w.add_orchestrator(parent_orchestrator)
+        w.start()
 
         task_hub_client = client.TaskHubGrpcClient()
         id = task_hub_client.schedule_new_orchestration(parent_orchestrator, input=10)
@@ -125,8 +125,8 @@ def orchestrator(ctx: task.OrchestrationContext, _):
 
     # Start a worker, which will connect to the sidecar in a background thread
     with worker.TaskHubGrpcWorker() as w:
-        w.start()
         w.add_orchestrator(orchestrator)
+        w.start()
 
         # Start the orchestration and immediately raise events to it.
         task_hub_client = client.TaskHubGrpcClient()
@@ -154,8 +154,8 @@ def orchestrator(ctx: task.OrchestrationContext, _):
 
     # Start a worker, which will connect to the sidecar in a background thread
     with worker.TaskHubGrpcWorker() as w:
-        w.start()
         w.add_orchestrator(orchestrator)
+        w.start()
 
         # Start the orchestration and immediately raise events to it.
         task_hub_client = client.TaskHubGrpcClient()
@@ -209,3 +209,26 @@ def orchestrator(ctx: task.OrchestrationContext, _):
         assert state is not None
         assert state.runtime_status == client.OrchestrationStatus.COMPLETED
         assert state.serialized_output == json.dumps(42)
+
+
+def test_terminate():
+    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
+        assert state.runtime_status == client.OrchestrationStatus.RUNNING
+
+        task_hub_client.terminate_orchestration(id, output="some reason for termination")
+        state = task_hub_client.wait_for_orchestration_completion(id, timeout=30)
+        assert state is not None
+        assert state.runtime_status == client.OrchestrationStatus.TERMINATED
+        assert state.serialized_output == json.dumps("some reason for termination")

tests/test_orchestration_executor.py

@@ -560,6 +560,31 @@ def orchestrator(ctx: task.OrchestrationContext, _):
     assert complete_action.result.value == "42"
 
 
+def test_terminate():
+    """Tests that an orchestration can be terminated before it completes"""
+
+    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_terminated_event(encoded_output=json.dumps("terminated!")),
+        helpers.new_event_raised_event("my_event", encoded_input="42")]
+
+    # Execute the orchestration. It should be in a running state waiting for an external 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_TERMINATED
+    assert complete_action.result.value == json.dumps("terminated!")
+
+
 def test_fan_out():
     """Tests that a fan-out pattern correctly schedules N tasks"""
     def hello(_, name: str):