Add cancellable tasks

Author: andystaples

durabletask/task.py

@@ -98,7 +98,7 @@ def set_custom_status(self, custom_status: Any) -> None:
         pass
 
     @abstractmethod
-    def create_timer(self, fire_at: Union[datetime, timedelta]) -> Task:
+    def create_timer(self, fire_at: Union[datetime, timedelta]) -> CancellableTask:
         """Create a Timer Task to fire after at the specified deadline.
 
         Parameters
@@ -231,7 +231,7 @@ def call_sub_orchestrator(self, orchestrator: Union[Orchestrator[TInput, TOutput
     # TOOD: Add a timeout parameter, which allows the task to be canceled if the event is
     # not received within the specified timeout. This requires support for task cancellation.
     @abstractmethod
-    def wait_for_external_event(self, name: str) -> CompletableTask:
+    def wait_for_external_event(self, name: str) -> CancellableTask:
         """Wait asynchronously for an event to be raised with the name `name`.
 
         Parameters
@@ -324,6 +324,10 @@ class OrchestrationStateError(Exception):
     pass
 
 
+class TaskCanceledError(Exception):
+    """Exception type for canceled orchestration tasks."""
+
+
 class Task(ABC, Generic[T]):
     """Abstract base class for asynchronous tasks in a durable orchestration."""
     _result: T
@@ -435,6 +439,48 @@ def fail(self, message: str, details: Union[Exception, pb.TaskFailureDetails]):
             self._parent.on_child_completed(self)
 
 
+class CancellableTask(CompletableTask[T]):
+    """A completable task that can be canceled before it finishes."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._is_cancelled = False
+        self._cancel_handler: Optional[Callable[[], None]] = None
+
+    @property
+    def is_cancelled(self) -> bool:
+        """Returns True if the task was canceled, False otherwise."""
+        return self._is_cancelled
+
+    def get_result(self) -> T:
+        if self._is_cancelled:
+            raise TaskCanceledError('The task was canceled.')
+        return super().get_result()
+
+    def set_cancel_handler(self, cancel_handler: Callable[[], None]) -> None:
+        self._cancel_handler = cancel_handler
+
+    def cancel(self) -> bool:
+        """Attempts to cancel this task.
+
+        Returns
+        -------
+        bool
+            True if cancellation was applied, False if the task had already completed.
+        """
+        if self._is_complete:
+            return False
+
+        if self._cancel_handler is not None:
+            self._cancel_handler()
+
+        self._is_cancelled = True
+        self._is_complete = True
+        if self._parent is not None:
+            self._parent.on_child_completed(self)
+        return True
+
+
 class RetryableTask(CompletableTask[T]):
     """A task that can be retried according to a retry policy."""
 
@@ -474,7 +520,7 @@ def compute_next_delay(self) -> Optional[timedelta]:
         return None
 
 
-class TimerTask(CompletableTask[T]):
+class TimerTask(CancellableTask[T]):
 
     def __init__(self) -> None:
         super().__init__()

durabletask/worker.py

@@ -307,7 +307,7 @@ class TaskHubGrpcWorker:
             activity function.
     """
 
-    _response_stream: Optional[grpc.Future] = None
+    _response_stream: Optional[Any] = None
     _interceptors: Optional[list[shared.ClientInterceptor]] = None
 
     def __init__(
@@ -512,7 +512,11 @@ def should_invalidate_connection(rpc_error):
 
                 def stream_reader():
                     try:
-                        for work_item in self._response_stream:
+                        response_stream = self._response_stream
+                        if response_stream is None:
+                            return
+
+                        for work_item in response_stream:
                             work_item_queue.put(work_item)
                     except Exception as e:
                         work_item_queue.put(e)
@@ -843,7 +847,7 @@ def __init__(self, instance_id: str, registry: _Registry):
         self._version: Optional[str] = None
         self._completion_status: Optional[pb.OrchestrationStatus] = None
         self._received_events: dict[str, list[Any]] = {}
-        self._pending_events: dict[str, list[task.CompletableTask]] = {}
+        self._pending_events: dict[str, list[task.CancellableTask]] = {}
         self._new_input: Optional[Any] = None
         self._save_events = False
         self._encoded_custom_status: Optional[str] = None
@@ -1026,7 +1030,13 @@ def create_timer_internal(
         action = ph.new_create_timer_action(id, fire_at)
         self._pending_actions[id] = action
 
-        timer_task: task.TimerTask = task.TimerTask()
+        timer_task = task.TimerTask()
+
+        def _cancel_timer() -> None:
+            self._pending_actions.pop(id, None)
+            self._pending_tasks.pop(id, None)
+
+        timer_task.set_cancel_handler(_cancel_timer)
         if retryable_task is not None:
             timer_task.set_retryable_parent(retryable_task)
         self._pending_tasks[id] = timer_task
@@ -1234,13 +1244,13 @@ def _exit_critical_section(self) -> None:
             action = pb.OrchestratorAction(id=task_id, sendEntityMessage=entity_unlock_message)
             self._pending_actions[task_id] = action
 
-    def wait_for_external_event(self, name: str) -> task.CompletableTask:
+    def wait_for_external_event(self, name: str) -> task.CancellableTask:
         # Check to see if this event has already been received, in which case we
         # can return it immediately. Otherwise, record out intent to receive an
         # event with the given name so that we can resume the generator when it
         # arrives. If there are multiple events with the same name, we return
         # them in the order they were received.
-        external_event_task: task.CompletableTask = task.CompletableTask()
+        external_event_task: task.CancellableTask = task.CancellableTask()
         event_name = name.casefold()
         event_list = self._received_events.get(event_name, None)
         if event_list:
@@ -1254,6 +1264,19 @@ def wait_for_external_event(self, name: str) -> task.CompletableTask:
                 task_list = []
                 self._pending_events[event_name] = task_list
             task_list.append(external_event_task)
+
+            def _cancel_wait() -> None:
+                waiting_tasks = self._pending_events.get(event_name)
+                if waiting_tasks is None:
+                    return
+                try:
+                    waiting_tasks.remove(external_event_task)
+                except ValueError:
+                    return
+                if not waiting_tasks:
+                    del self._pending_events[event_name]
+
+            external_event_task.set_cancel_handler(_cancel_wait)
         return external_event_task
 
     def continue_as_new(self, new_input, *, save_events: bool = False) -> None:
@@ -1450,6 +1473,13 @@ def process_event(
                             f"{ctx.instance_id}: Ignoring unexpected timerFired event with ID = {timer_id}."
                         )
                     return
+                if not isinstance(timer_task, task.TimerTask):
+                    if not ctx._is_replaying:
+                        self._logger.warning(
+                            f"{ctx.instance_id}: Ignoring timerFired event with non-timer task ID = {timer_id}."
+                        )
+                    return
+
                 timer_task.complete(None)
                 if timer_task._retryable_parent is not None:
                     activity_action = timer_task._retryable_parent._action

tests/durabletask/test_orchestration_executor.py

@@ -143,6 +143,87 @@ def delay_orchestrator(ctx: task.OrchestrationContext, _):
     assert complete_action.result.value == '"done"'  # results are JSON-encoded
 
 
+def test_timer_can_be_cancelled_after_when_any_winner():
+    """Tests cancellation of an outstanding timer task after another task wins when_any."""
+
+    def orchestrator(ctx: task.OrchestrationContext, _):
+        approval = ctx.wait_for_external_event("approval")
+        timeout = ctx.create_timer(timedelta(hours=1))
+        winner = yield task.when_any([approval, timeout])
+        if winner == approval:
+            timeout.cancel()
+            return "approved"
+        return "timed out"
+
+    registry = worker._Registry()
+    name = registry.add_orchestrator(orchestrator)
+    executor = worker._OrchestrationExecutor(registry, TEST_LOGGER)
+
+    start_time = datetime(2020, 1, 1, 12, 0, 0)
+    timeout_fire_at = start_time + timedelta(hours=1)
+
+    result = executor.execute(
+        TEST_INSTANCE_ID,
+        [],
+        [
+            helpers.new_orchestrator_started_event(start_time),
+            helpers.new_execution_started_event(name, TEST_INSTANCE_ID, encoded_input=None),
+        ],
+    )
+    assert len(result.actions) == 1
+    assert result.actions[0].HasField("createTimer")
+    assert result.actions[0].createTimer.fireAt.ToDatetime() == timeout_fire_at
+
+    old_events = [
+        helpers.new_orchestrator_started_event(start_time),
+        helpers.new_execution_started_event(name, TEST_INSTANCE_ID, encoded_input=None),
+        helpers.new_timer_created_event(1, timeout_fire_at),
+    ]
+    result = executor.execute(
+        TEST_INSTANCE_ID,
+        old_events,
+        [helpers.new_event_raised_event("approval", json.dumps(True))],
+    )
+    complete_action = get_and_validate_complete_orchestration_action_list(1, result.actions)
+    assert complete_action.orchestrationStatus == pb.ORCHESTRATION_STATUS_COMPLETED
+    assert complete_action.result.value == '"approved"'
+
+
+def test_only_cancellable_tasks_expose_cancel():
+    """Tests that only timer and external-event tasks expose cancellation state and operations."""
+
+    def dummy_activity(ctx, _):
+        pass
+
+    ctx = worker._RuntimeOrchestrationContext(TEST_INSTANCE_ID, worker._Registry())
+
+    timer_task = ctx.create_timer(timedelta(minutes=5))
+    external_event_task = ctx.wait_for_external_event("approval")
+    activity_task = ctx.call_activity(dummy_activity)
+
+    assert isinstance(timer_task, task.CancellableTask)
+    assert isinstance(external_event_task, task.CancellableTask)
+    assert not isinstance(activity_task, task.CancellableTask)
+    assert hasattr(timer_task, "cancel")
+    assert hasattr(external_event_task, "cancel")
+    assert not hasattr(activity_task, "cancel")
+    assert hasattr(timer_task, "is_cancelled")
+    assert hasattr(external_event_task, "is_cancelled")
+    assert not hasattr(activity_task, "is_cancelled")
+
+
+def test_cancelled_task_get_result_raises_task_canceled_error():
+    """Tests that canceled cancellable tasks raise TaskCanceledError from get_result."""
+
+    cancellable_task = task.CancellableTask()
+
+    assert cancellable_task.cancel() is True
+    assert cancellable_task.is_cancelled is True
+
+    with pytest.raises(task.TaskCanceledError):
+        cancellable_task.get_result()
+
+
 def test_schedule_activity_actions():
     """Test the actions output for the call_activity orchestrator method"""
     def dummy_activity(ctx, _):
@@ -1313,7 +1394,7 @@ def orchestrator(ctx: task.OrchestrationContext, _):
     encoded_output = json.dumps(dummy_activity(None, "Seattle"))
     old_events = old_events + new_events
     new_events = [helpers.new_task_completed_event(2, encoded_output),
-                  helpers.new_timer_fired_event(4, current_timestamp)]
+                  helpers.new_timer_fired_event(4, expected_fire_at)]
     executor = worker._OrchestrationExecutor(registry, TEST_LOGGER)
     result = executor.execute(TEST_INSTANCE_ID, old_events, new_events)
     actions = result.actions