Add continue_as_new orchestration API (#9)

Compatible with durabletask-go v0.2.1 or newer.

Author: cgillum

durabletask/client.py

@@ -97,7 +97,7 @@ def __init__(self, *,
                  log_formatter: Union[logging.Formatter, None] = None):
         channel = shared.get_grpc_channel(host_address)
         self._stub = stubs.TaskHubSidecarServiceStub(channel)
-        self._logger = shared.get_logger(log_handler, log_formatter)
+        self._logger = shared.get_logger("client", log_handler, log_formatter)
 
     def schedule_new_orchestration(self, orchestrator: Union[task.Orchestrator[TInput, TOutput], str], *,
                                    input: Union[TInput, None] = None,
@@ -109,7 +109,7 @@ def schedule_new_orchestration(self, orchestrator: Union[task.Orchestrator[TInpu
         req = pb.CreateInstanceRequest(
             name=name,
             instanceId=instance_id if instance_id else uuid.uuid4().hex,
-            input=wrappers_pb2.StringValue(value=shared.to_json(input)) if input else None,
+            input=wrappers_pb2.StringValue(value=shared.to_json(input)) if input is not None else None,
             scheduledStartTimestamp=helpers.new_timestamp(start_at) if start_at else None,
             version=wrappers_pb2.StringValue(value=""))
 
@@ -144,7 +144,19 @@ def wait_for_orchestration_completion(self, instance_id: str, *,
         try:
             self._logger.info(f"Waiting {timeout}s for instance '{instance_id}' to complete.")
             res: pb.GetInstanceResponse = self._stub.WaitForInstanceCompletion(req, timeout=timeout)
-            return new_orchestration_state(req.instanceId, res)
+            state = new_orchestration_state(req.instanceId, res)
+            if not state:
+                return None
+
+            if state.runtime_status == OrchestrationStatus.FAILED and state.failure_details is not None:
+                details = state.failure_details
+                self._logger.info(f"Instance '{instance_id}' failed: [{details.error_type}] {details.message}")
+            elif state.runtime_status == OrchestrationStatus.TERMINATED:
+                self._logger.info(f"Instance '{instance_id}' was terminated.")
+            elif state.runtime_status == OrchestrationStatus.COMPLETED:
+                self._logger.info(f"Instance '{instance_id}' completed.")
+
+            return state
         except grpc.RpcError as rpc_error:
             if rpc_error.code() == grpc.StatusCode.DEADLINE_EXCEEDED:  # type: ignore
                 # Replace gRPC error with the built-in TimeoutError

durabletask/internal/helpers.py

@@ -3,7 +3,7 @@
 
 import traceback
 from datetime import datetime
-from typing import Union
+from typing import List, Union
 
 from google.protobuf import timestamp_pb2, wrappers_pb2
 
@@ -161,14 +161,13 @@ def new_complete_orchestration_action(
         id: int,
         status: pb.OrchestrationStatus,
         result: Union[str, None] = None,
-        failure_details: Union[pb.TaskFailureDetails, None] = None) -> pb.OrchestratorAction:
-
+        failure_details: Union[pb.TaskFailureDetails, None] = None,
+        carryover_events: Union[List[pb.HistoryEvent], None] = None) -> pb.OrchestratorAction:
     completeOrchestrationAction = pb.CompleteOrchestrationAction(
         orchestrationStatus=status,
         result=get_string_value(result),
-        failureDetails=failure_details)
-
-    # TODO: CarryoverEvents
+        failureDetails=failure_details,
+        carryoverEvents=carryover_events)
 
     return pb.OrchestratorAction(id=id, completeOrchestration=completeOrchestrationAction)
 

durabletask/internal/shared.py

@@ -26,9 +26,10 @@ def get_grpc_channel(host_address: Union[str, None]) -> grpc.Channel:
 
 
 def get_logger(
+        name_suffix: str,
         log_handler: Union[logging.Handler, None] = None,
         log_formatter: Union[logging.Formatter, None] = None) -> logging.Logger:
-    logger = logging.Logger("durabletask")
+    logger = logging.Logger(f"durabletask-{name_suffix}")
 
     # Add a default log handler if none is provided
     if log_handler is None:

durabletask/task.py

@@ -147,6 +147,19 @@ def wait_for_external_event(self, name: str) -> Task:
         """
         pass
 
+    @abstractmethod
+    def continue_as_new(self, new_input: Any, *, save_events: bool = False) -> None:
+        """Continue the orchestration execution as a new instance.
+
+        Parameters
+        ----------
+        new_input : Any
+            The new input to use for the new orchestration instance.
+        save_events : bool
+            A flag indicating whether to add any unprocessed external events in the new orchestration history.
+        """
+        pass
+
 
 class FailureDetails:
     def __init__(self, message: str, error_type: str, stack_trace: Union[str, None]):

durabletask/worker.py

@@ -3,7 +3,6 @@
 
 import concurrent.futures
 import logging
-from dataclasses import dataclass
 from datetime import datetime, timedelta
 from threading import Event, Thread
 from types import GeneratorType
@@ -90,7 +89,7 @@ def __init__(self, *,
                  log_formatter: Union[logging.Formatter, None] = None):
         self._registry = _Registry()
         self._host_address = host_address if host_address else shared.get_default_host_address()
-        self._logger = shared.get_logger(log_handler, log_formatter)
+        self._logger = shared.get_logger("worker", log_handler, log_formatter)
         self._shutdown = Event()
         self._response_stream = None
         self._is_running = False
@@ -149,7 +148,7 @@ def run_loop():
 
                     except grpc.RpcError as rpc_error:
                         if rpc_error.code() == grpc.StatusCode.CANCELLED:  # type: ignore
-                            self._logger.warning(f'Disconnected from {self._host_address}')
+                            self._logger.info(f'Disconnected from {self._host_address}')
                         elif rpc_error.code() == grpc.StatusCode.UNAVAILABLE:  # type: ignore
                             self._logger.warning(
                                 f'The sidecar at address {self._host_address} is unavailable - will continue retrying')
@@ -163,7 +162,7 @@ def run_loop():
                 self._logger.info("No longer listening for work items")
                 return
 
-        self._logger.info(f"starting gRPC worker that connects to {self._host_address}")
+        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
@@ -220,12 +219,6 @@ def _execute_activity(self, req: pb.ActivityRequest, stub: stubs.TaskHubSidecarS
                 f"Failed to deliver activity response for '{req.name}#{req.taskId}' of orchestration ID '{instance_id}' to sidecar: {ex}")
 
 
-@dataclass
-class _ExternalEvent:
-    name: str
-    data: Any
-
-
 class _RuntimeOrchestrationContext(task.OrchestrationContext):
     _generator: Union[Generator[task.Task, Any, Any], None]
     _previous_task: Union[task.Task, None]
@@ -241,8 +234,10 @@ def __init__(self, instance_id: str):
         self._current_utc_datetime = datetime(1000, 1, 1)
         self._instance_id = instance_id
         self._completion_status: Union[pb.OrchestrationStatus, None] = None
-        self._received_events: Dict[str, List[_ExternalEvent]] = {}
+        self._received_events: Dict[str, List[Any]] = {}
         self._pending_events: Dict[str, List[task.CompletableTask]] = {}
+        self._new_input: Union[Any, None] = None
+        self._save_events = False
 
     def run(self, generator: Generator[task.Task, Any, Any]):
         self._generator = generator
@@ -282,6 +277,9 @@ def set_complete(self, result: Any, status: pb.OrchestrationStatus, is_result_en
             return
 
         self._is_complete = True
+        self._completion_status = status
+        self._pending_actions.clear()  # Cancel any pending actions
+
         self._result = result
         result_json: Union[str, None] = None
         if result is not None:
@@ -296,13 +294,44 @@ def set_failed(self, ex: Exception):
 
         self._is_complete = True
         self._pending_actions.clear()  # Cancel any pending actions
+        self._completion_status = pb.ORCHESTRATION_STATUS_FAILED
+
         action = ph.new_complete_orchestration_action(
             self.next_sequence_number(), pb.ORCHESTRATION_STATUS_FAILED, None, ph.new_failure_details(ex)
         )
         self._pending_actions[action.id] = action
 
+    def set_continued_as_new(self, new_input: Any, save_events: bool):
+        if self._is_complete:
+            return
+
+        self._is_complete = True
+        self._pending_actions.clear()  # Cancel any pending actions
+        self._completion_status = pb.ORCHESTRATION_STATUS_CONTINUED_AS_NEW
+        self._new_input = new_input
+        self._save_events = save_events
+
     def get_actions(self) -> List[pb.OrchestratorAction]:
-        return list(self._pending_actions.values())
+        if self._completion_status == pb.ORCHESTRATION_STATUS_CONTINUED_AS_NEW:
+            # When continuing-as-new, we only return a single completion action.
+            carryover_events: Union[List[pb.HistoryEvent], None] = None
+            if self._save_events:
+                carryover_events = []
+                # We need to save the current set of pending events so that they can be
+                # replayed when the new instance starts.
+                for event_name, values in self._received_events.items():
+                    for event_value in values:
+                        encoded_value = shared.to_json(event_value) if event_value else None
+                        carryover_events.append(ph.new_event_raised_event(event_name, encoded_value))
+            action = ph.new_complete_orchestration_action(
+                self.next_sequence_number(),
+                pb.ORCHESTRATION_STATUS_CONTINUED_AS_NEW,
+                result=shared.to_json(self._new_input) if self._new_input is not None else None,
+                failure_details=None,
+                carryover_events=carryover_events)
+            return [action]
+        else:
+            return list(self._pending_actions.values())
 
     def next_sequence_number(self) -> int:
         self._sequence_number += 1
@@ -370,13 +399,13 @@ def wait_for_external_event(self, name: str) -> task.Task:
         # arrives. If there are multiple events with the same name, we return
         # them in the order they were received.
         external_event_task = task.CompletableTask()
-        event_name = name.upper()
+        event_name = name.casefold()
         event_list = self._received_events.get(event_name, None)
         if event_list:
-            event = event_list.pop(0)
+            event_data = event_list.pop(0)
             if not event_list:
                 del self._received_events[event_name]
-            external_event_task.complete(event.data)
+            external_event_task.complete(event_data)
         else:
             task_list = self._pending_events.get(event_name, None)
             if not task_list:
@@ -385,6 +414,12 @@ def wait_for_external_event(self, name: str) -> task.Task:
             task_list.append(external_event_task)
         return external_event_task
 
+    def continue_as_new(self, new_input, *, save_events: bool = False) -> None:
+        if self._is_complete:
+            return
+
+        self.set_continued_as_new(new_input, save_events)
+
 
 class _OrchestrationExecutor:
     _generator: Union[task.Orchestrator, None]
@@ -415,13 +450,16 @@ 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)
 
-        if ctx._completion_status:
+        if not ctx._is_complete:
+            task_count = len(ctx._pending_tasks)
+            event_count = len(ctx._pending_events)
+            self._logger.info(f"{instance_id}: Waiting for {task_count} task(s) and {event_count} event(s).")
+        elif ctx._completion_status and ctx._completion_status is not pb.ORCHESTRATION_STATUS_CONTINUED_AS_NEW:
             completion_status_str = pbh.get_orchestration_status_str(ctx._completion_status)
             self._logger.info(f"{instance_id}: Orchestration completed with status: {completion_status_str}")
 
@@ -570,9 +608,9 @@ def process_event(self, ctx: _RuntimeOrchestrationContext, event: pb.HistoryEven
                 ctx.resume()
             elif event.HasField("eventRaised"):
                 # event names are case-insensitive
-                event_name = event.eventRaised.name.upper()
+                event_name = event.eventRaised.name.casefold()
                 if not ctx.is_replaying:
-                    self._logger.info(f"Event raised: {event_name}")
+                    self._logger.info(f"{ctx.instance_id} Event raised: {event_name}")
                 task_list = ctx._pending_events.get(event_name, None)
                 decoded_result: Union[Any, None] = None
                 if task_list:
@@ -591,7 +629,7 @@ def process_event(self, ctx: _RuntimeOrchestrationContext, event: pb.HistoryEven
                         ctx._received_events[event_name] = event_list
                     if not ph.is_empty(event.eventRaised.input):
                         decoded_result = shared.from_json(event.eventRaised.input.value)
-                    event_list.append(_ExternalEvent(event.eventRaised.name, decoded_result))
+                    event_list.append(decoded_result)
                     if not ctx.is_replaying:
                         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"):

tests/test_activity_executor.py

@@ -2,12 +2,16 @@
 # Licensed under the MIT License.
 
 import json
+import logging
 from typing import Any, Tuple, Union
 
-import durabletask.internal.shared as shared
 from durabletask import task, worker
 
-TEST_LOGGER = shared.get_logger()
+logging.basicConfig(
+    format='%(asctime)s.%(msecs)03d %(name)s %(levelname)s: %(message)s',
+    datefmt='%Y-%m-%d %H:%M:%S',
+    level=logging.DEBUG)
+TEST_LOGGER = logging.getLogger("tests")
 TEST_INSTANCE_ID = 'abc123'
 TEST_TASK_ID = 42
 

tests/test_orchestration_e2e.py

@@ -232,3 +232,39 @@ def orchestrator(ctx: task.OrchestrationContext, _):
         assert state is not None
         assert state.runtime_status == client.OrchestrationStatus.TERMINATED
         assert state.serialized_output == json.dumps("some reason for termination")
+
+
+def test_continue_as_new():
+    all_results = []
+
+    def orchestrator(ctx: task.OrchestrationContext, input: int):
+        result = yield ctx.wait_for_external_event("my_event")
+        if not ctx.is_replaying:
+            # NOTE: Real orchestrations should never interact with nonlocal variables like this.
+            nonlocal all_results
+            all_results.append(result)
+
+        if len(all_results) <= 4:
+            ctx.continue_as_new(max(all_results), save_events=True)
+        else:
+            return all_results
+
+    # 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, input=0)
+        task_hub_client.raise_orchestration_event(id, "my_event", data=1)
+        task_hub_client.raise_orchestration_event(id, "my_event", data=2)
+        task_hub_client.raise_orchestration_event(id, "my_event", data=3)
+        task_hub_client.raise_orchestration_event(id, "my_event", data=4)
+        task_hub_client.raise_orchestration_event(id, "my_event", data=5)
+
+        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(all_results)
+        assert state.serialized_input == json.dumps(4)
+        assert all_results == [1, 2, 3, 4, 5]

tests/test_orchestration_executor.py

@@ -6,6 +6,8 @@
 from datetime import datetime, timedelta
 from typing import List
 
+import pytest
+
 import durabletask.internal.helpers as helpers
 import durabletask.internal.orchestrator_service_pb2 as pb
 from durabletask import task, worker
@@ -585,6 +587,41 @@ def orchestrator(ctx: task.OrchestrationContext, _):
     assert complete_action.result.value == json.dumps("terminated!")
 
 
+@pytest.mark.parametrize("save_events", [True, False])
+def test_continue_as_new(save_events: bool):
+    """Tests the behavior of the continue-as-new API"""
+    def orchestrator(ctx: task.OrchestrationContext, input: int):
+        yield ctx.create_timer(ctx.current_utc_datetime + timedelta(days=1))
+        ctx.continue_as_new(input + 1, save_events=save_events)
+
+    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, encoded_input="1"),
+        helpers.new_event_raised_event("my_event", encoded_input="42"),
+        helpers.new_event_raised_event("my_event", encoded_input="43"),
+        helpers.new_event_raised_event("my_event", encoded_input="44"),
+        helpers.new_timer_created_event(1, datetime.utcnow() + timedelta(days=1))]
+    new_events = [
+        helpers.new_timer_fired_event(1, datetime.utcnow() + timedelta(days=1))]
+
+    # Execute the orchestration. It should be in a running state waiting for the timer to fire
+    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_CONTINUED_AS_NEW
+    assert complete_action.result.value == json.dumps(2)
+    assert len(complete_action.carryoverEvents) == (3 if save_events else 0)
+    for i in range(len(complete_action.carryoverEvents)):
+        event = complete_action.carryoverEvents[i]
+        assert type(event) is pb.HistoryEvent
+        assert event.HasField("eventRaised")
+        assert event.eventRaised.name.casefold() == "my_event".casefold()  # event names are case-insensitive
+        assert event.eventRaised.input.value == json.dumps(42 + i)
+
+
 def test_fan_out():
     """Tests that a fan-out pattern correctly schedules N tasks"""
     def hello(_, name: str):