Merge pull request #44 from CasperGN/fix(bug)--shutdown-on-iterator

Fix(bug): shutdown on iterator

Author: sicoyle

durabletask/aio/client.py

@@ -27,6 +27,7 @@
 # If `opentelemetry-instrumentation-grpc` is available, enable the gRPC client interceptor
 try:
     from opentelemetry.instrumentation.grpc import GrpcInstrumentorClient
+
     GrpcInstrumentorClient().instrument()
 except ImportError:
     pass

durabletask/client.py

@@ -24,10 +24,12 @@
 # If `opentelemetry-instrumentation-grpc` is available, enable the gRPC client interceptor
 try:
     from opentelemetry.instrumentation.grpc import GrpcInstrumentorClient
+
     GrpcInstrumentorClient().instrument()
 except ImportError:
     pass
 
+
 class OrchestrationStatus(Enum):
     """The status of an orchestration instance."""
 

durabletask/worker.py

@@ -12,7 +12,7 @@
 from datetime import datetime, timedelta
 from threading import Event, Thread
 from types import GeneratorType
-from typing import Any, Generator, Optional, Sequence, TypeVar, Union
+from typing import Any, Generator, Iterator, Optional, Sequence, TypeVar, Union
 
 import grpc
 from google.protobuf import empty_pb2
@@ -38,7 +38,6 @@
     otel_tracer = None
 
 
-
 class VersionNotRegisteredException(Exception):
     pass
 
@@ -283,7 +282,7 @@ class TaskHubGrpcWorker:
             activity function.
     """
 
-    _response_stream: Optional[grpc.Future] = None
+    _response_stream: Optional[Union[Iterator[grpc.Future], grpc.Future]] = None
     _interceptors: Optional[list[shared.ClientInterceptor]] = None
 
     def __init__(
@@ -421,9 +420,12 @@ def invalidate_connection():
             # Cancel the response stream first to signal the reader thread to stop
             if self._response_stream is not None:
                 try:
-                    self._response_stream.cancel()
-                except Exception:
-                    pass
+                    if hasattr(self._response_stream, "call"):
+                        self._response_stream.call.cancel()  # type: ignore
+                    else:
+                        self._response_stream.cancel()  # type: ignore
+                except Exception as e:
+                    self._logger.warning(f"Error cancelling response stream: {e}")
                 self._response_stream = None
 
             # Wait for the reader thread to finish
@@ -740,7 +742,13 @@ def stop(self):
 
         self._logger.info("Stopping gRPC worker...")
         if self._response_stream is not None:
-            self._response_stream.cancel()
+            try:
+                if hasattr(self._response_stream, "call"):
+                    self._response_stream.call.cancel()  # type: ignore
+                else:
+                    self._response_stream.cancel()  # type: ignore
+            except Exception as e:
+                self._logger.warning(f"Error cancelling response stream: {e}")
         self._shutdown.set()
         # Explicitly close the gRPC channel to ensure OTel interceptors and other resources are cleaned up
         if self._current_channel is not None:
@@ -854,13 +862,15 @@ def _execute_activity(
 
         if otel_tracer is not None:
             span_context = otel_tracer.start_as_current_span(
-                name=f'activity: {req.name}',
-                context=otel_propagator.extract(carrier={"traceparent": req.parentTraceContext.traceParent}),
+                name=f"activity: {req.name}",
+                context=otel_propagator.extract(
+                    carrier={"traceparent": req.parentTraceContext.traceParent}
+                ),
                 attributes={
                     "durabletask.task.instance_id": instance_id,
                     "durabletask.task.id": req.taskId,
                     "durabletask.activity.name": req.name,
-                }
+                },
             )
         else:
             span_context = contextlib.nullcontext()

tests/durabletask/test_registry.py

@@ -165,6 +165,7 @@ def activity2(ctx, input):
     assert registry.get_activity(name1) is activity1
     assert registry.get_activity(name2) is activity2
 
+
 def test_registry_add_named_versioned_orchestrators():
     """Test adding versioned orchestrators."""
     registry = worker._Registry()
@@ -179,7 +180,9 @@ def orchestrator3(ctx, input):
         return "two"
 
     registry.add_named_orchestrator(name="orchestrator", fn=orchestrator1, version_name="v1")
-    registry.add_named_orchestrator(name="orchestrator", fn=orchestrator2, version_name="v2", is_latest=True)
+    registry.add_named_orchestrator(
+        name="orchestrator", fn=orchestrator2, version_name="v2", is_latest=True
+    )
     registry.add_named_orchestrator(name="orchestrator", fn=orchestrator3, version_name="v3")
 
     orquestrator, version = registry.get_orchestrator(name="orchestrator")

tests/durabletask/test_worker_stop.py

@@ -0,0 +1,57 @@
+from unittest.mock import MagicMock, patch
+
+import grpc
+
+from durabletask.worker import TaskHubGrpcWorker
+
+
+# Helper to create a running worker with a mocked runLoop
+def _make_running_worker():
+    worker = TaskHubGrpcWorker()
+    worker._is_running = True
+    worker._runLoop = MagicMock()
+    worker._runLoop.is_alive.return_value = False
+    return worker
+
+
+def test_stop_with_grpc_future():
+    worker = _make_running_worker()
+    mock_future = MagicMock(spec=grpc.Future)
+    worker._response_stream = mock_future
+    worker.stop()
+    mock_future.cancel.assert_called_once()
+
+
+def test_stop_with_generator_call():
+    worker = _make_running_worker()
+    mock_call = MagicMock()
+    mock_stream = MagicMock()
+    mock_stream.call = mock_call
+    worker._response_stream = mock_stream
+    worker.stop()
+    mock_call.cancel.assert_called_once()
+
+
+def test_stop_with_unknown_stream_type(caplog):
+    worker = _make_running_worker()
+    # Not a grpc.Future, no 'call' attribute
+    worker._response_stream = object()
+    with caplog.at_level("WARNING"):
+        worker.stop()
+    assert any("Error cancelling response stream: " in m for m in caplog.text.splitlines())
+
+
+def test_stop_with_none_stream():
+    worker = _make_running_worker()
+    worker._response_stream = None
+    # Should not raise
+    worker.stop()
+
+
+def test_stop_when_not_running():
+    worker = TaskHubGrpcWorker()
+    worker._is_running = False
+    # Should return immediately, not set _shutdown
+    with patch.object(worker._shutdown, "set") as shutdown_set:
+        worker.stop()
+        shutdown_set.assert_not_called()