fix: address pr comments

Signed-off-by: Casper Nielsen <casper@diagrid.io>

Author: CasperGN

durabletask/internal/shared.py

@@ -58,17 +58,6 @@ def get_default_host_address() -> str:
 )
 
 
-def _merge_grpc_options(
-    user_options: Optional[Sequence[tuple[str, Any]]],
-    defaults: Sequence[tuple[str, Any]] = DEFAULT_GRPC_KEEPALIVE_OPTIONS,
-) -> list[tuple[str, Any]]:
-    """Merge user gRPC options with defaults. User options take precedence."""
-    merged = dict(defaults)
-    if user_options:
-        merged.update(dict(user_options))
-    return list(merged.items())
-
-
 def get_grpc_channel(
     host_address: Optional[str],
     secure_channel: bool = False,
@@ -100,7 +89,10 @@ def get_grpc_channel(
             host_address = host_address[len(protocol) :]
             break
 
-    merged_options = _merge_grpc_options(options)
+    merged = dict(DEFAULT_GRPC_KEEPALIVE_OPTIONS)
+    if options:
+        merged.update(dict(options))
+    merged_options = list(merged.items())
     if secure_channel:
         channel = grpc.secure_channel(
             host_address, grpc.ssl_channel_credentials(), options=merged_options

durabletask/worker.py

@@ -307,6 +307,7 @@ def __init__(
         self._channel_options = channel_options
         self._stop_timeout = stop_timeout
         self._current_channel: Optional[grpc.Channel] = None  # Store channel reference for cleanup
+        self._channel_cleanup_threads: list[threading.Thread] = []  # Deferred channel close threads
         self._stream_ready = threading.Event()
         # Use provided concurrency options or create default ones
         self._concurrency_options = (
@@ -388,9 +389,12 @@ async def _async_run_loop(self):
 
         def create_fresh_connection():
             nonlocal current_channel, current_stub, conn_retry_count
-            # Don't call channel.close() — in-flight activity RPCs on the
-            # old stub may still reference the channel from another thread.
-            # The old channel is GC'd once all references are released.
+            # Schedule deferred close of old channel to avoid orphaned TCP
+            # connections. In-flight RPCs on the old stub may still reference
+            # the channel from another thread, so we wait a grace period
+            # before closing instead of closing immediately.
+            if current_channel is not None:
+                self._schedule_deferred_channel_close(current_channel)
             current_channel = None
             current_stub = None
             try:
@@ -415,12 +419,12 @@ def create_fresh_connection():
 
         def invalidate_connection():
             nonlocal current_channel, current_stub, current_reader_thread
-            # Null out references so the next iteration creates a fresh connection.
-            # Do NOT call channel.close() here — in-flight activity RPCs
-            # (CompleteActivityTask) may still be using the stub on another
-            # thread. Closing the channel concurrently causes segfaults in the
-            # gRPC C extension. The old channel is GC'd once all references
-            # (including captured stub refs in activity threads) are released.
+            # Schedule deferred close of old channel to avoid orphaned TCP
+            # connections. In-flight RPCs (e.g. CompleteActivityTask) may still
+            # be using the stub on another thread, so we defer the close by a
+            # grace period instead of closing immediately.
+            if current_channel is not None:
+                self._schedule_deferred_channel_close(current_channel)
             current_channel = None
             self._current_channel = None
             current_stub = None
@@ -717,6 +721,38 @@ def stream_reader():
         except Exception as e:
             self._logger.warning(f"Error while waiting for worker task shutdown: {e}")
 
+    def _schedule_deferred_channel_close(
+        self, old_channel: grpc.Channel, grace_timeout: float = 10.0
+    ):
+        """Schedule a deferred close of an old gRPC channel.
+
+        Waits up to *grace_timeout* seconds for in-flight RPCs to complete
+        before closing the channel.  This prevents orphaned TCP connections
+        while still allowing in-flight work (e.g. ``CompleteActivityTask``
+        calls on another thread) to finish gracefully.
+
+        During ``stop()``, ``_shutdown`` is already set so the wait returns
+        immediately and the channel is closed at once.
+        """
+        # Prune already-finished cleanup threads to avoid unbounded growth
+        self._channel_cleanup_threads = [t for t in self._channel_cleanup_threads if t.is_alive()]
+
+        def _deferred_close():
+            try:
+                # Normal reconnect: wait grace period for RPCs to drain.
+                # Shutdown: _shutdown is already set, returns immediately.
+                self._shutdown.wait(timeout=grace_timeout)
+            finally:
+                try:
+                    old_channel.close()
+                    self._logger.debug("Deferred channel close completed")
+                except Exception as e:
+                    self._logger.debug(f"Error during deferred channel close: {e}")
+
+        thread = threading.Thread(target=_deferred_close, daemon=True, name="ChannelCleanup")
+        thread.start()
+        self._channel_cleanup_threads.append(thread)
+
     def stop(self):
         """Stops the worker and waits for any pending work items to complete."""
         if not self._is_running:
@@ -743,6 +779,11 @@ def stop(self):
             else:
                 self._logger.debug("Worker thread completed successfully")
 
+        # Wait for any deferred channel-cleanup threads to finish
+        for t in self._channel_cleanup_threads:
+            t.join(timeout=5)
+        self._channel_cleanup_threads.clear()
+
         self._async_worker_manager.shutdown()
         self._logger.info("Worker shutdown completed")
         self._is_running = False

tests/durabletask/test_client.py

@@ -2,7 +2,13 @@
 
 from durabletask import client
 from durabletask.internal.grpc_interceptor import DefaultClientInterceptorImpl
-from durabletask.internal.shared import get_default_host_address, get_grpc_channel
+from durabletask.internal.shared import (
+    DEFAULT_GRPC_KEEPALIVE_OPTIONS,
+    get_default_host_address,
+    get_grpc_channel,
+)
+
+EXPECTED_DEFAULT_OPTIONS = list(DEFAULT_GRPC_KEEPALIVE_OPTIONS)
 
 HOST_ADDRESS = "localhost:50051"
 METADATA = [("key1", "value1"), ("key2", "value2")]
@@ -14,7 +20,7 @@ def test_get_grpc_channel_insecure():
         get_grpc_channel(HOST_ADDRESS, False, interceptors=INTERCEPTORS)
         args, kwargs = mock_channel.call_args
         assert args[0] == HOST_ADDRESS
-        assert "options" in kwargs and kwargs["options"] is None
+        assert "options" in kwargs and kwargs["options"] == EXPECTED_DEFAULT_OPTIONS
 
 
 def test_get_grpc_channel_secure():
@@ -26,15 +32,15 @@ def test_get_grpc_channel_secure():
         args, kwargs = mock_channel.call_args
         assert args[0] == HOST_ADDRESS
         assert args[1] == mock_credentials.return_value
-        assert "options" in kwargs and kwargs["options"] is None
+        assert "options" in kwargs and kwargs["options"] == EXPECTED_DEFAULT_OPTIONS
 
 
 def test_get_grpc_channel_default_host_address():
     with patch("grpc.insecure_channel") as mock_channel:
         get_grpc_channel(None, False, interceptors=INTERCEPTORS)
         args, kwargs = mock_channel.call_args
         assert args[0] == get_default_host_address()
-        assert "options" in kwargs and kwargs["options"] is None
+        assert "options" in kwargs and kwargs["options"] == EXPECTED_DEFAULT_OPTIONS
 
 
 def test_get_grpc_channel_with_metadata():
@@ -45,7 +51,7 @@ def test_get_grpc_channel_with_metadata():
         get_grpc_channel(HOST_ADDRESS, False, interceptors=INTERCEPTORS)
         args, kwargs = mock_channel.call_args
         assert args[0] == HOST_ADDRESS
-        assert "options" in kwargs and kwargs["options"] is None
+        assert "options" in kwargs and kwargs["options"] == EXPECTED_DEFAULT_OPTIONS
         mock_intercept_channel.assert_called_once()
 
         # Capture and check the arguments passed to intercept_channel()
@@ -66,61 +72,61 @@ def test_grpc_channel_with_host_name_protocol_stripping():
         get_grpc_channel(prefix + host_name, interceptors=INTERCEPTORS)
         args, kwargs = mock_insecure_channel.call_args
         assert args[0] == host_name
-        assert "options" in kwargs and kwargs["options"] is None
+        assert "options" in kwargs and kwargs["options"] == EXPECTED_DEFAULT_OPTIONS
 
         prefix = "http://"
         get_grpc_channel(prefix + host_name, interceptors=INTERCEPTORS)
         args, kwargs = mock_insecure_channel.call_args
         assert args[0] == host_name
-        assert "options" in kwargs and kwargs["options"] is None
+        assert "options" in kwargs and kwargs["options"] == EXPECTED_DEFAULT_OPTIONS
 
         prefix = "HTTP://"
         get_grpc_channel(prefix + host_name, interceptors=INTERCEPTORS)
         args, kwargs = mock_insecure_channel.call_args
         assert args[0] == host_name
-        assert "options" in kwargs and kwargs["options"] is None
+        assert "options" in kwargs and kwargs["options"] == EXPECTED_DEFAULT_OPTIONS
 
         prefix = "GRPC://"
         get_grpc_channel(prefix + host_name, interceptors=INTERCEPTORS)
         args, kwargs = mock_insecure_channel.call_args
         assert args[0] == host_name
-        assert "options" in kwargs and kwargs["options"] is None
+        assert "options" in kwargs and kwargs["options"] == EXPECTED_DEFAULT_OPTIONS
 
         prefix = ""
         get_grpc_channel(prefix + host_name, interceptors=INTERCEPTORS)
         args, kwargs = mock_insecure_channel.call_args
         assert args[0] == host_name
-        assert "options" in kwargs and kwargs["options"] is None
+        assert "options" in kwargs and kwargs["options"] == EXPECTED_DEFAULT_OPTIONS
 
         prefix = "grpcs://"
         get_grpc_channel(prefix + host_name, interceptors=INTERCEPTORS)
         args, kwargs = mock_secure_channel.call_args
         assert args[0] == host_name
-        assert "options" in kwargs and kwargs["options"] is None
+        assert "options" in kwargs and kwargs["options"] == EXPECTED_DEFAULT_OPTIONS
 
         prefix = "https://"
         get_grpc_channel(prefix + host_name, interceptors=INTERCEPTORS)
         args, kwargs = mock_secure_channel.call_args
         assert args[0] == host_name
-        assert "options" in kwargs and kwargs["options"] is None
+        assert "options" in kwargs and kwargs["options"] == EXPECTED_DEFAULT_OPTIONS
 
         prefix = "HTTPS://"
         get_grpc_channel(prefix + host_name, interceptors=INTERCEPTORS)
         args, kwargs = mock_secure_channel.call_args
         assert args[0] == host_name
-        assert "options" in kwargs and kwargs["options"] is None
+        assert "options" in kwargs and kwargs["options"] == EXPECTED_DEFAULT_OPTIONS
 
         prefix = "GRPCS://"
         get_grpc_channel(prefix + host_name, interceptors=INTERCEPTORS)
         args, kwargs = mock_secure_channel.call_args
         assert args[0] == host_name
-        assert "options" in kwargs and kwargs["options"] is None
+        assert "options" in kwargs and kwargs["options"] == EXPECTED_DEFAULT_OPTIONS
 
         prefix = ""
         get_grpc_channel(prefix + host_name, True, interceptors=INTERCEPTORS)
         args, kwargs = mock_secure_channel.call_args
         assert args[0] == host_name
-        assert "options" in kwargs and kwargs["options"] is None
+        assert "options" in kwargs and kwargs["options"] == EXPECTED_DEFAULT_OPTIONS
 
 
 def test_sync_channel_passes_base_options_and_max_lengths():