Improvements

- Add rejection
- cleanup comments
- suborch default versioning

Author: andystaples

durabletask/internal/exceptions.py

@@ -0,0 +1,3 @@
+class AbandonOrchestrationError(Exception):
+    def __init__(self, *args: object) -> None:
+        super().__init__(*args)

durabletask/worker.py

@@ -18,6 +18,7 @@
 from google.protobuf import empty_pb2
 
 import durabletask.internal.helpers as ph
+import durabletask.internal.exceptions as pe
 import durabletask.internal.orchestrator_service_pb2 as pb
 import durabletask.internal.orchestrator_service_pb2_grpc as stubs
 import durabletask.internal.shared as shared
@@ -113,10 +114,10 @@ def __init__(self, version: Optional[str] = None,
         """Initialize versioning options.
 
         Args:
-            version: The specific version to use for orchestrators and activities.
-            default_version: The default version to use if no specific version is provided.
-            match_strategy: The strategy to use for matching versions.
-            failure_strategy: The strategy to use if versioning fails.
+            version: The version of orchestrations that the worker can work on.
+            default_version: The default version that will be used for starting new orchestrations.
+            match_strategy: The versioning strategy for the Durable Task worker.
+            failure_strategy: The versioning failure strategy for the Durable Task worker.
         """
         self.version = version
         self.default_version = default_version
@@ -333,7 +334,7 @@ def add_activity(self, fn: task.Activity) -> str:
         return self._registry.add_activity(fn)
 
     def use_versioning(self, version: VersioningOptions) -> None:
-        """Sets the default version for orchestrators and activities."""
+        """Initializes versioning options for sub-orchestrators and activities."""
         if self._is_running:
             raise RuntimeError("Cannot set default version while the worker is running.")
         self._registry.versioning = version
@@ -564,14 +565,24 @@ def _execute_orchestrator(
             completionToken,
     ):
         try:
-            executor = _OrchestrationExecutor(self._registry, self._logger)
+            executor = _OrchestrationExecutor(self._registry, self._logger, stub)
             result = executor.execute(req.instanceId, req.pastEvents, req.newEvents)
             res = pb.OrchestratorResponse(
                 instanceId=req.instanceId,
                 actions=result.actions,
                 customStatus=ph.get_string_value(result.encoded_custom_status),
                 completionToken=completionToken,
             )
+        except pe.AbandonOrchestrationError as ex:
+            self._logger.info(
+                f"Abandoning orchestration. InstanceId = '{req.instanceId}'. Completion token = '{completionToken}'"
+            )
+            stub.AbandonTaskOrchestratorWorkItem(
+                pb.AbandonOrchestrationTaskRequest(
+                    completionToken=completionToken
+                )
+            )
+            return
         except Exception as ex:
             self._logger.exception(
                 f"An error occurred while trying to execute instance '{req.instanceId}': {ex}"
@@ -633,7 +644,7 @@ class _RuntimeOrchestrationContext(task.OrchestrationContext):
     _generator: Optional[Generator[task.Task, Any, Any]]
     _previous_task: Optional[task.Task]
 
-    def __init__(self, instance_id: str):
+    def __init__(self, instance_id: str, registry: _Registry):
         self._generator = None
         self._is_replaying = True
         self._is_complete = False
@@ -643,6 +654,7 @@ def __init__(self, instance_id: str):
         self._sequence_number = 0
         self._current_utc_datetime = datetime(1000, 1, 1)
         self._instance_id = instance_id
+        self._registry = registry
         self._completion_status: Optional[pb.OrchestrationStatus] = None
         self._received_events: dict[str, list[Any]] = {}
         self._pending_events: dict[str, list[task.CompletableTask]] = {}
@@ -831,14 +843,16 @@ def call_sub_orchestrator(
     ) -> task.Task[TOutput]:
         id = self.next_sequence_number()
         orchestrator_name = task.get_name(orchestrator)
+        default_version = self._registry.versioning.default_version if self._registry.versioning else None
+        orchestrator_version = version if version else default_version
         self.call_activity_function_helper(
             id,
             orchestrator_name,
             input=input,
             retry_policy=retry_policy,
             is_sub_orch=True,
             instance_id=instance_id,
-            version=version,
+            version=orchestrator_version
         )
         return self._pending_tasks.get(id, task.CompletableTask())
 
@@ -937,11 +951,12 @@ def __init__(
 class _OrchestrationExecutor:
     _generator: Optional[task.Orchestrator] = None
 
-    def __init__(self, registry: _Registry, logger: logging.Logger):
+    def __init__(self, registry: _Registry, logger: logging.Logger, stub: stubs.TaskHubSidecarServiceStub):
         self._registry = registry
         self._logger = logger
         self._is_suspended = False
         self._suspended_events: list[pb.HistoryEvent] = []
+        self._stub = stub
 
     def execute(
             self,
@@ -954,9 +969,18 @@ def execute(
                 "The new history event list must have at least one event in it."
             )
 
-        ctx = _RuntimeOrchestrationContext(instance_id)
+        ctx = _RuntimeOrchestrationContext(instance_id, self._registry)
         version_failure = None
         try:
+            # Rebuild local state by replaying old history into the orchestrator function
+            self._logger.debug(
+                f"{instance_id}: Rebuilding local state with {len(old_events)} history event..."
+            )
+            ctx._is_replaying = True
+            for old_event in old_events:
+                self.process_event(ctx, old_event)
+
+            # Process versioning if applicable
             execution_started_events = [e.executionStarted for e in old_events if e.HasField("executionStarted")]
             if self._registry.versioning and len(execution_started_events) > 0:
                 execution_started_event = execution_started_events[-1]
@@ -970,19 +994,7 @@ def execute(
                         f"Error action = '{self._registry.versioning.failure_strategy}'. "
                         f"Version error = '{version_failure}'"
                     )
-                    if self._registry.versioning.failure_strategy == VersionFailureStrategy.FAIL:
-                        raise VersionFailureException
-                    elif self._registry.versioning.failure_strategy == VersionFailureStrategy.REJECT:
-                        # TODO: We don't have abandoned orchestrations yet, so we just fail
-                        raise VersionFailureException
-
-            # Rebuild local state by replaying old history into the orchestrator function
-            self._logger.debug(
-                f"{instance_id}: Rebuilding local state with {len(old_events)} history event..."
-            )
-            ctx._is_replaying = True
-            for old_event in old_events:
-                self.process_event(ctx, old_event)
+                    raise VersionFailureException
 
             # Get new actions by executing newly received events into the orchestrator function
             if self._logger.level <= logging.DEBUG:
@@ -995,10 +1007,13 @@ def execute(
                 self.process_event(ctx, new_event)
 
         except VersionFailureException as ex:
-            if version_failure:
-                ctx.set_failed(version_failure)
-            else:
-                ctx.set_failed(ex)
+            if self._registry.versioning and self._registry.versioning.failure_strategy == VersionFailureStrategy.FAIL:
+                if version_failure:
+                    ctx.set_failed(version_failure)
+                else:
+                    ctx.set_failed(ex)
+            elif self._registry.versioning and self._registry.versioning.failure_strategy == VersionFailureStrategy.REJECT:
+                raise pe.AbandonOrchestrationError
 
         except Exception as ex:
             # Unhandled exceptions fail the orchestration

tests/durabletask-azuremanaged/test_dts_orchestration_versioning_e2e.py

@@ -153,6 +153,60 @@ def test_upper_version_worker_succeeds():
 
 
 def test_upper_version_worker_strict_fails():
+    # Start a worker, which will connect to the sidecar in a background thread
+    instance_id: str = ''
+    thrown = False
+    try:
+        with DurableTaskSchedulerWorker(host_address=endpoint, secure_channel=True,
+                                        taskhub=taskhub_name, token_credential=None) as w:
+            w.add_orchestrator(single_activity)
+            w.add_activity(plus_one)
+            w.use_versioning(worker.VersioningOptions(
+                version="1.0.0",
+                default_version="1.1.0",
+                match_strategy=worker.VersionMatchStrategy.STRICT,
+                failure_strategy=worker.VersionFailureStrategy.REJECT
+            ))
+            w.start()
+
+            task_hub_client = DurableTaskSchedulerClient(host_address=endpoint, secure_channel=True,
+                                                         taskhub=taskhub_name, token_credential=None,
+                                                         default_version="1.1.0")
+            instance_id = task_hub_client.schedule_new_orchestration(single_activity, input=1)
+            state = task_hub_client.wait_for_orchestration_completion(
+                instance_id, timeout=5)
+    except TimeoutError as e:
+        thrown = True
+        assert str(e).find("Timed-out waiting for the orchestration to complete") >= 0
+
+    assert thrown is True
+
+    with DurableTaskSchedulerWorker(host_address=endpoint, secure_channel=True,
+                                    taskhub=taskhub_name, token_credential=None) as w:
+        w.add_orchestrator(single_activity)
+        w.add_activity(plus_one)
+        w.use_versioning(worker.VersioningOptions(
+            version="1.1.0",
+            default_version="1.1.0",
+            match_strategy=worker.VersionMatchStrategy.STRICT,
+            failure_strategy=worker.VersionFailureStrategy.REJECT
+        ))
+        w.start()
+
+        task_hub_client = DurableTaskSchedulerClient(host_address=endpoint, secure_channel=True,
+                                                     taskhub=taskhub_name, token_credential=None,
+                                                     default_version="1.1.0")
+        state = task_hub_client.wait_for_orchestration_completion(
+            instance_id, timeout=5)
+
+    assert state is not None
+    assert state.name == task.get_name(single_activity)
+    assert state.instance_id == instance_id
+    assert state.runtime_status == client.OrchestrationStatus.COMPLETED
+    assert state.failure_details is None
+
+
+def test_reject_abandons_and_reprocess():
     # Start a worker, which will connect to the sidecar in a background thread
     with DurableTaskSchedulerWorker(host_address=endpoint, secure_channel=True,
                                     taskhub=taskhub_name, token_credential=None) as w:
@@ -181,6 +235,9 @@ def test_upper_version_worker_strict_fails():
     assert state.failure_details.message.find("The orchestration version '1.0.0' does not match the worker version '1.1.0'.") >= 0
 
 
+# Sub-orchestration tests
+
+
 def sequence_suborchestator(ctx: task.OrchestrationContext, start_val: int):
     numbers = []
     for current in range(start_val, start_val + 5):