Initial work - not complete

Author: andystaples

durabletask/worker.py

@@ -346,7 +346,7 @@ def __init__(
         else:
             self._interceptors = None
 
-        self._async_worker_manager = _AsyncWorkerManager(self._concurrency_options)
+        self._async_worker_manager = _AsyncWorkerManager(self._concurrency_options, self._logger)
 
     @property
     def concurrency_options(self) -> ConcurrencyOptions:
@@ -533,27 +533,31 @@ def stream_reader():
                         if work_item.HasField("orchestratorRequest"):
                             self._async_worker_manager.submit_orchestration(
                                 self._execute_orchestrator,
+                                self._cancel_orchestrator,
                                 work_item.orchestratorRequest,
                                 stub,
                                 work_item.completionToken,
                             )
                         elif work_item.HasField("activityRequest"):
                             self._async_worker_manager.submit_activity(
                                 self._execute_activity,
+                                self._cancel_activity,
                                 work_item.activityRequest,
                                 stub,
                                 work_item.completionToken,
                             )
                         elif work_item.HasField("entityRequest"):
                             self._async_worker_manager.submit_entity_batch(
                                 self._execute_entity_batch,
+                                self._cancel_entity_batch,
                                 work_item.entityRequest,
                                 stub,
                                 work_item.completionToken,
                             )
                         elif work_item.HasField("entityRequestV2"):
                             self._async_worker_manager.submit_entity_batch(
                                 self._execute_entity_batch,
+                                self._cancel_entity_batch,
                                 work_item.entityRequestV2,
                                 stub,
                                 work_item.completionToken
@@ -670,6 +674,19 @@ def _execute_orchestrator(
                 f"Failed to deliver orchestrator response for '{req.instanceId}' to sidecar: {ex}"
             )
 
+    def _cancel_orchestrator(
+            self,
+            req: pb.OrchestratorRequest,
+            stub: stubs.TaskHubSidecarServiceStub,
+            completionToken,
+    ):
+        stub.AbandonTaskOrchestratorWorkItem(
+            pb.AbandonOrchestrationTaskRequest(
+                completionToken=completionToken
+            )
+        )
+        self._logger.info(f"Cancelled orchestration task for invocation ID: {req.instanceId}")
+
     def _execute_activity(
             self,
             req: pb.ActivityRequest,
@@ -703,6 +720,19 @@ def _execute_activity(
                 f"Failed to deliver activity response for '{req.name}#{req.taskId}' of orchestration ID '{instance_id}' to sidecar: {ex}"
             )
 
+    def _cancel_activity(
+            self,
+            req: pb.ActivityRequest,
+            stub: stubs.TaskHubSidecarServiceStub,
+            completionToken,
+    ):
+        stub.AbandonTaskActivityWorkItem(
+            pb.AbandonActivityTaskRequest(
+                completionToken=completionToken
+            )
+        )
+        self._logger.info(f"Cancelled activity task for task ID: {req.taskId} on orchestration ID: {req.orchestrationInstance.instanceId}")
+
     def _execute_entity_batch(
             self,
             req: Union[pb.EntityBatchRequest, pb.EntityRequest],
@@ -771,6 +801,19 @@ def _execute_entity_batch(
 
         return batch_result
 
+    def _cancel_entity_batch(
+            self,
+            req: pb.EntityBatchRequest,
+            stub: stubs.TaskHubSidecarServiceStub,
+            completionToken,
+    ):
+        stub.AbandonTaskEntityWorkItem(
+            pb.AbandonEntityTaskRequest(
+                completionToken=completionToken
+            )
+        )
+        self._logger.info(f"Cancelled entity batch task for entity instance ID: {req.instanceId}")
+
 
 class _RuntimeOrchestrationContext(task.OrchestrationContext):
     _generator: Optional[Generator[task.Task, Any, Any]]
@@ -1368,7 +1411,7 @@ def process_event(
                 timer_id = event.timerFired.timerId
                 timer_task = ctx._pending_tasks.pop(timer_id, None)
                 if not timer_task:
-                    # TODO: Should this be an error? When would it ever happen?
+                    # TODO: Should this be an error?  would it ever happen?
                     if not ctx._is_replaying:
                         self._logger.warning(
                             f"{ctx.instance_id}: Ignoring unexpected timerFired event with ID = {timer_id}."
@@ -1920,8 +1963,10 @@ def _is_suspendable(event: pb.HistoryEvent) -> bool:
 
 
 class _AsyncWorkerManager:
-    def __init__(self, concurrency_options: ConcurrencyOptions):
+    def __init__(self, concurrency_options: ConcurrencyOptions, logger: logging.Logger):
         self.concurrency_options = concurrency_options
+        self._logger = logger
+
         self.activity_semaphore = None
         self.orchestration_semaphore = None
         self.entity_semaphore = None
@@ -2031,17 +2076,47 @@ async def run(self):
         )
 
         # Start background consumers for each work type
-        if self.activity_queue is not None and self.orchestration_queue is not None \
-                and self.entity_batch_queue is not None:
-            await asyncio.gather(
-                self._consume_queue(self.activity_queue, self.activity_semaphore),
-                self._consume_queue(
-                    self.orchestration_queue, self.orchestration_semaphore
-                ),
-                self._consume_queue(
-                    self.entity_batch_queue, self.entity_semaphore
+        try:
+            if self.activity_queue is not None and self.orchestration_queue is not None \
+                    and self.entity_batch_queue is not None:
+                await asyncio.gather(
+                    self._consume_queue(self.activity_queue, self.activity_semaphore),
+                    self._consume_queue(
+                        self.orchestration_queue, self.orchestration_semaphore
+                    ),
+                    self._consume_queue(
+                        self.entity_batch_queue, self.entity_semaphore
+                    )
                 )
-            )
+        except Exception as queue_exception:
+            self._logger.error(f"Uncaught error in activity manager thread pool: {queue_exception}")
+            while self.activity_queue is not None and not self.activity_queue.empty():
+                try:
+                    func, cancellation_func, args, kwargs = self.activity_queue.get_nowait()
+                    await self._run_func(cancellation_func, *args, **kwargs)
+                    self._logger.error(f"Activity work item args: {args}, kwargs: {kwargs}")
+                except asyncio.QueueEmpty:
+                    pass
+                except Exception as cancellation_exception:
+                    self._logger.error(f"Uncaught error while cancelling activity work item: {cancellation_exception}")
+            while self.orchestration_queue is not None and not self.orchestration_queue.empty():
+                try:
+                    func, cancellation_func, args, kwargs = self.orchestration_queue.get_nowait()
+                    await self._run_func(cancellation_func, *args, **kwargs)
+                    self._logger.error(f"Orchestration work item args: {args}, kwargs: {kwargs}")
+                except asyncio.QueueEmpty:
+                    pass
+                except Exception as cancellation_exception:
+                    self._logger.error(f"Uncaught error while cancelling orchestration work item: {cancellation_exception}")
+            while self.entity_batch_queue is not None and not self.entity_batch_queue.empty():
+                try:
+                    func, cancellation_func, args, kwargs = self.entity_batch_queue.get_nowait()
+                    await self._run_func(cancellation_func, *args, **kwargs)
+                    self._logger.error(f"Entity batch work item args: {args}, kwargs: {kwargs}")
+                except asyncio.QueueEmpty:
+                    pass
+                except Exception as cancellation_exception:
+                    self._logger.error(f"Uncaught error while cancelling entity batch work item: {cancellation_exception}")
 
     async def _consume_queue(self, queue: asyncio.Queue, semaphore: asyncio.Semaphore):
         # List to track running tasks
@@ -2061,7 +2136,7 @@ async def _consume_queue(self, queue: asyncio.Queue, semaphore: asyncio.Semaphor
             except asyncio.TimeoutError:
                 continue
 
-            func, args, kwargs = work
+            func, cancellation_func, args, kwargs = work
             # Create a concurrent task for processing
             task = asyncio.create_task(
                 self._process_work_item(semaphore, queue, func, args, kwargs)
@@ -2092,26 +2167,26 @@ async def _run_func(self, func, *args, **kwargs):
                 self.thread_pool, lambda: func(*args, **kwargs)
             )
 
-    def submit_activity(self, func, *args, **kwargs):
-        work_item = (func, args, kwargs)
+    def submit_activity(self, func, cancellation_func, *args, **kwargs):
+        work_item = (func, cancellation_func, args, kwargs)
         self._ensure_queues_for_current_loop()
         if self.activity_queue is not None:
             self.activity_queue.put_nowait(work_item)
         else:
             # No event loop running, store in pending list
             self._pending_activity_work.append(work_item)
 
-    def submit_orchestration(self, func, *args, **kwargs):
-        work_item = (func, args, kwargs)
+    def submit_orchestration(self, func, cancellation_func, *args, **kwargs):
+        work_item = (func, cancellation_func, args, kwargs)
         self._ensure_queues_for_current_loop()
         if self.orchestration_queue is not None:
             self.orchestration_queue.put_nowait(work_item)
         else:
             # No event loop running, store in pending list
             self._pending_orchestration_work.append(work_item)
 
-    def submit_entity_batch(self, func, *args, **kwargs):
-        work_item = (func, args, kwargs)
+    def submit_entity_batch(self, func, cancellation_func, *args, **kwargs):
+        work_item = (func, cancellation_func, args, kwargs)
         self._ensure_queues_for_current_loop()
         if self.entity_batch_queue is not None:
             self.entity_batch_queue.put_nowait(work_item)
@@ -2123,7 +2198,7 @@ def shutdown(self):
         self._shutdown = True
         self.thread_pool.shutdown(wait=True)
 
-    def reset_for_new_run(self):
+    async def reset_for_new_run(self):
         """Reset the manager state for a new run."""
         self._shutdown = False
         # Clear any existing queues - they'll be recreated when needed
@@ -2132,18 +2207,28 @@ def reset_for_new_run(self):
             # This ensures no items from previous runs remain
             try:
                 while not self.activity_queue.empty():
-                    self.activity_queue.get_nowait()
+                    func, cancellation_func, args, kwargs = self.activity_queue.get_nowait()
+                    await self._run_func(cancellation_func, *args, **kwargs)
             except Exception:
                 pass
         if self.orchestration_queue is not None:
             try:
                 while not self.orchestration_queue.empty():
-                    self.orchestration_queue.get_nowait()
+                    func, cancellation_func, args, kwargs = self.orchestration_queue.get_nowait()
+                    await self._run_func(cancellation_func, *args, **kwargs)
+            except Exception:
+                pass
+        if self.entity_batch_queue is not None:
+            try:
+                while not self.entity_batch_queue.empty():
+                    func, cancellation_func, args, kwargs = self.entity_batch_queue.get_nowait()
+                    await self._run_func(cancellation_func, *args, **kwargs)
             except Exception:
                 pass
         # Clear pending work lists
         self._pending_activity_work.clear()
         self._pending_orchestration_work.clear()
+        self._pending_entity_batch_work.clear()
 
 
 # Export public API

tests/durabletask/test_worker_concurrency_loop.py

@@ -52,13 +52,21 @@ def dummy_orchestrator(req, stub, completionToken):
         time.sleep(0.1)
         stub.CompleteOrchestratorTask('ok')
 
+    def cancel_dummy_orchestrator(req, stub, completionToken):
+        pass
+
     def dummy_activity(req, stub, completionToken):
         time.sleep(0.1)
         stub.CompleteActivityTask('ok')
 
+    def cancel_dummy_activity(req, stub, completionToken):
+        pass
+
     # Patch the worker's _execute_orchestrator and _execute_activity
     worker._execute_orchestrator = dummy_orchestrator
+    worker._cancel_orchestrator = cancel_dummy_orchestrator
     worker._execute_activity = dummy_activity
+    worker._cancel_activity = cancel_dummy_activity
 
     orchestrator_requests = [DummyRequest('orchestrator', f'orch{i}') for i in range(3)]
     activity_requests = [DummyRequest('activity', f'act{i}') for i in range(4)]
@@ -67,9 +75,9 @@ async def run_test():
         # Start the worker manager's run loop in the background
         worker_task = asyncio.create_task(worker._async_worker_manager.run())
         for req in orchestrator_requests:
-            worker._async_worker_manager.submit_orchestration(dummy_orchestrator, req, stub, DummyCompletionToken())
+            worker._async_worker_manager.submit_orchestration(dummy_orchestrator, cancel_dummy_orchestrator, req, stub, DummyCompletionToken())
         for req in activity_requests:
-            worker._async_worker_manager.submit_activity(dummy_activity, req, stub, DummyCompletionToken())
+            worker._async_worker_manager.submit_activity(dummy_activity, cancel_dummy_activity, req, stub, DummyCompletionToken())
         await asyncio.sleep(1.0)
         orchestrator_count = sum(1 for t, _ in stub.completed if t == 'orchestrator')
         activity_count = sum(1 for t, _ in stub.completed if t == 'activity')

tests/durabletask/test_worker_concurrency_loop_async.py

@@ -50,13 +50,21 @@ async def dummy_orchestrator(req, stub, completionToken):
         await asyncio.sleep(0.1)
         stub.CompleteOrchestratorTask('ok')
 
+    async def cancel_dummy_orchestrator(req, stub, completionToken):
+        pass
+
     async def dummy_activity(req, stub, completionToken):
         await asyncio.sleep(0.1)
         stub.CompleteActivityTask('ok')
 
+    async def cancel_dummy_activity(req, stub, completionToken):
+        pass
+
     # Patch the worker's _execute_orchestrator and _execute_activity
     grpc_worker._execute_orchestrator = dummy_orchestrator
+    grpc_worker._cancel_orchestrator = cancel_dummy_orchestrator
     grpc_worker._execute_activity = dummy_activity
+    grpc_worker._cancel_activity = cancel_dummy_activity
 
     orchestrator_requests = [DummyRequest('orchestrator', f'orch{i}') for i in range(3)]
     activity_requests = [DummyRequest('activity', f'act{i}') for i in range(4)]
@@ -66,9 +74,9 @@ async def run_test():
         stub.completed.clear()
         worker_task = asyncio.create_task(grpc_worker._async_worker_manager.run())
         for req in orchestrator_requests:
-            grpc_worker._async_worker_manager.submit_orchestration(dummy_orchestrator, req, stub, DummyCompletionToken())
+            grpc_worker._async_worker_manager.submit_orchestration(dummy_orchestrator, cancel_dummy_orchestrator, req, stub, DummyCompletionToken())
         for req in activity_requests:
-            grpc_worker._async_worker_manager.submit_activity(dummy_activity, req, stub, DummyCompletionToken())
+            grpc_worker._async_worker_manager.submit_activity(dummy_activity, cancel_dummy_activity, req, stub, DummyCompletionToken())
         await asyncio.sleep(1.0)
         orchestrator_count = sum(1 for t, _ in stub.completed if t == 'orchestrator')
         activity_count = sum(1 for t, _ in stub.completed if t == 'activity')