Fix sync client channel cleanup

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: Bernd Verst

CHANGELOG.md

@@ -39,6 +39,9 @@ FIXED
 - Improved sync and async gRPC clients so repeated transport failures recreate
   SDK-owned channels, while long-poll deadlines, successful replies, and
   application-level RPC errors do not trigger unnecessary channel replacement.
+- Fixed `TaskHubGrpcClient.close()` so explicit sync client shutdown now closes
+  any previously retired SDK-owned gRPC channels immediately instead of waiting
+  for the delayed cleanup timer.
 
 ## v1.4.0
 

durabletask/client.py

@@ -211,8 +211,10 @@ def __init__(self, *,
         self._client_failure_tracker = FailureTracker(
             self._resiliency_options.channel_recreate_failure_threshold
         )
+        self._closing = False
         self._last_recreate_time = 0.0
         self._recreate_lock = threading.Lock()
+        self._retired_channels: dict[grpc.Channel, threading.Timer] = {}
         self._logger = shared.get_logger("client", log_handler, log_formatter)
         self.default_version = default_version
         self._payload_store = payload_store
@@ -243,9 +245,11 @@ def _invoke_unary(
             return response
 
     def _maybe_recreate_channel(self) -> None:
-        if not self._owns_channel:
+        if not self._owns_channel or self._closing:
             return
         with self._recreate_lock:
+            if self._closing:
+                return
             now = time.monotonic()
             if now - self._last_recreate_time < self._resiliency_options.min_recreate_interval_seconds:
                 return
@@ -259,10 +263,22 @@ def _maybe_recreate_channel(self) -> None:
             self._stub = stubs.TaskHubSidecarServiceStub(self._channel)
             self._last_recreate_time = now
             self._client_failure_tracker.record_success()
-            close_timer = threading.Timer(30.0, old_channel.close)
+            close_timer = threading.Timer(
+                30.0,
+                self._close_retired_channel,
+                args=(old_channel,),
+            )
             close_timer.daemon = True
+            self._retired_channels[old_channel] = close_timer
             close_timer.start()
 
+    def _close_retired_channel(self, channel: grpc.Channel) -> None:
+        with self._recreate_lock:
+            close_timer = self._retired_channels.pop(channel, None)
+            if close_timer is None:
+                return
+        channel.close()
+
     def close(self) -> None:
         """Closes the underlying gRPC channel.
 
@@ -272,7 +288,15 @@ def close(self) -> None:
         it.
         """
         if self._owns_channel:
-            self._channel.close()
+            with self._recreate_lock:
+                self._closing = True
+                retired_channels = list(self._retired_channels.items())
+                self._retired_channels.clear()
+                current_channel = self._channel
+            for retired_channel, close_timer in retired_channels:
+                close_timer.cancel()
+                retired_channel.close()
+            current_channel.close()
 
     def schedule_new_orchestration(self, orchestrator: Union[task.Orchestrator[TInput, TOutput], str], *,
                                    input: Optional[TInput] = None,

tests/durabletask/test_client.py

@@ -382,11 +382,88 @@ def test_sync_client_recreates_sdk_owned_channel_with_original_transport_inputs(
         expected_channel_call,
     ]
     assert client._channel is second_channel
-    mock_timer.assert_called_once_with(30.0, first_channel.close)
+    mock_timer.assert_called_once()
+    timer_call = mock_timer.call_args
+    assert timer_call.args[0] == 30.0
+    assert timer_call.args[1].__self__ is client
+    assert timer_call.args[1].__func__ is TaskHubGrpcClient._close_retired_channel
+    assert timer_call.kwargs == {"args": (first_channel,)}
     assert timer.daemon is True
     timer.start.assert_called_once_with()
 
 
+def test_sync_client_close_closes_retired_channels_immediately():
+    first_channel = MagicMock(name="first-channel")
+    second_channel = MagicMock(name="second-channel")
+    first_stub = MagicMock()
+    first_stub.GetInstance.side_effect = FakeRpcError(grpc.StatusCode.UNAVAILABLE)
+    second_stub = MagicMock()
+    second_stub.GetInstance.return_value = MagicMock(exists=False)
+    close_timer = MagicMock(name="close-timer")
+
+    with patch(
+            "durabletask.client.shared.get_grpc_channel",
+            side_effect=[first_channel, second_channel],
+    ), patch(
+            "durabletask.client.stubs.TaskHubSidecarServiceStub", side_effect=[first_stub, second_stub]
+    ), patch("threading.Timer", return_value=close_timer):
+        client = TaskHubGrpcClient(
+            resiliency_options=GrpcClientResiliencyOptions(
+                channel_recreate_failure_threshold=1,
+                min_recreate_interval_seconds=0.0,
+            )
+        )
+        with pytest.raises(FakeRpcError):
+            client.get_orchestration_state("abc")
+
+        client.close()
+
+    close_timer.cancel.assert_called_once_with()
+    first_channel.close.assert_called_once_with()
+    second_channel.close.assert_called_once_with()
+    assert client._retired_channels == {}
+
+
+def test_sync_client_close_closes_all_retired_sdk_channels_immediately():
+    first_channel = MagicMock(name="first-channel")
+    second_channel = MagicMock(name="second-channel")
+    third_channel = MagicMock(name="third-channel")
+    first_stub = MagicMock()
+    first_stub.GetInstance.side_effect = FakeRpcError(grpc.StatusCode.UNAVAILABLE)
+    second_stub = MagicMock()
+    second_stub.GetInstance.side_effect = FakeRpcError(grpc.StatusCode.UNAVAILABLE)
+    third_stub = MagicMock()
+    timer1 = MagicMock(name="close-timer-1")
+    timer2 = MagicMock(name="close-timer-2")
+
+    with patch(
+            "durabletask.client.shared.get_grpc_channel",
+            side_effect=[first_channel, second_channel, third_channel],
+    ), patch(
+            "durabletask.client.stubs.TaskHubSidecarServiceStub",
+            side_effect=[first_stub, second_stub, third_stub],
+    ), patch("threading.Timer", side_effect=[timer1, timer2]):
+        client = TaskHubGrpcClient(
+            resiliency_options=GrpcClientResiliencyOptions(
+                channel_recreate_failure_threshold=1,
+                min_recreate_interval_seconds=0.0,
+            )
+        )
+        with pytest.raises(FakeRpcError):
+            client.get_orchestration_state("abc")
+        with pytest.raises(FakeRpcError):
+            client.get_orchestration_state("abc")
+
+        client.close()
+
+    timer1.cancel.assert_called_once_with()
+    timer2.cancel.assert_called_once_with()
+    first_channel.close.assert_called_once_with()
+    second_channel.close.assert_called_once_with()
+    third_channel.close.assert_called_once_with()
+    assert client._retired_channels == {}
+
+
 @pytest.mark.parametrize(
     ("stub_method_name", "client_method_name"),
     [
@@ -431,11 +508,13 @@ def test_sync_client_does_not_recreate_caller_owned_channel():
             client.get_orchestration_state("abc")
         with pytest.raises(FakeRpcError):
             client.get_orchestration_state("abc")
+        client.close()
 
     assert client._channel is provided_channel
     mock_get_channel.assert_not_called()
     mock_stub.assert_called_once_with(provided_channel)
     mock_timer.assert_not_called()
+    provided_channel.close.assert_not_called()
 
 
 def test_sync_client_recreate_cooldown_prevents_immediate_repeated_recreation():
@@ -478,7 +557,6 @@ def test_sync_client_recreate_cooldown_prevents_immediate_repeated_recreation():
             client.get_orchestration_state("abc")
         assert client._channel is second_channel
         assert mock_get_channel.call_count == 2
-        mock_timer.assert_called_once_with(30.0, first_channel.close)
 
         with pytest.raises(FakeRpcError):
             client.get_orchestration_state("abc")
@@ -495,10 +573,16 @@ def test_sync_client_recreate_cooldown_prevents_immediate_repeated_recreation():
         expected_channel_call,
         expected_channel_call,
     ]
-    assert mock_timer.call_args_list == [
-        call(30.0, first_channel.close),
-        call(30.0, second_channel.close),
-    ]
+    assert mock_timer.call_count == 2
+    first_timer_call, second_timer_call = mock_timer.call_args_list
+    assert first_timer_call.args[0] == 30.0
+    assert first_timer_call.args[1].__self__ is client
+    assert first_timer_call.args[1].__func__ is TaskHubGrpcClient._close_retired_channel
+    assert first_timer_call.kwargs == {"args": (first_channel,)}
+    assert second_timer_call.args[0] == 30.0
+    assert second_timer_call.args[1].__self__ is client
+    assert second_timer_call.args[1].__func__ is TaskHubGrpcClient._close_retired_channel
+    assert second_timer_call.kwargs == {"args": (second_channel,)}
     assert timer1.daemon is True
     assert timer2.daemon is True
     timer1.start.assert_called_once_with()