Add context manager support to TaskHubGrpcClient to prevent gRPC channel leaks #134

Activity

1. berndverst commented on May 26, 2026

   

   Bernd Verst (berndverst)

   on May 26, 2026

   Member

   More actions

   Note

   This is a proposal from Copilot generated on behalf of Bernd Verst (@berndverst).

   Validity check

   This issue is still applicable as of main (post-#135). I verified the following on the latest main:

   • TaskHubGrpcClient (sync) in durabletask/client.py still defines close() but does not implement __enter__ / __exit__.
   • AsyncTaskHubGrpcClient continues to expose __aenter__ / __aexit__ that delegate to close(), so the sync/async asymmetry described in the original issue remains.
   • TaskHubGrpcWorker already supports context-manager usage, so the sync client is the only public class in this asymmetric state.
   • PR Add gRPC client and worker connection resiliency #135 (gRPC client and worker connection resiliency) did not add context-manager support; it focused on hardening the recreate/close path. That work makes context-manager support more valuable, because close() now also tears down resiliency state.

   Compatibility with PR #135

   The proposal below is purely additive and does not clash with the resiliency overhaul. Specifically:

   • The current TaskHubGrpcClient.close() already does all the heavy lifting introduced by Add gRPC client and worker connection resiliency #135:
     1. sets self._closing = True so any in-flight _maybe_recreate_channel short-circuits,
     2. joins the in-flight recreate thread (self._recreate_thread) with a bounded timeout,
     3. cancels pending _retired_channels close timers under self._recreate_lock,
     4. closes every retired channel and finally the current channel,
     5. respects self._owns_channel, so caller-owned channels are left untouched.
   • __exit__ simply calls self.close(). No new state, no new locks, no parallel teardown path — the resiliency invariants from Add gRPC client and worker connection resiliency #135 are preserved.
   • The same shape already exists on AsyncTaskHubGrpcClient (__aexit__ -> await self.close()), which Add gRPC client and worker connection resiliency #135 also extended; we are bringing the sync client in line with that established pattern.
   • DurableTaskSchedulerClient (azuremanaged) subclasses TaskHubGrpcClient and does not override close(), so it inherits the new context-manager behavior automatically and continues to flow through the resiliency-aware close().

   Proposed implementation

   Add the two dunder methods immediately after close() in TaskHubGrpcClient:

   def __enter__(self) -> "TaskHubGrpcClient":
       return self
   
   def __exit__(self, exc_type, exc_val, exc_tb) -> None:
       self.close()

   Notes:

   • Return type on __enter__ is annotated to keep type checkers (and the project's Pylance baseline) happy.
   • __exit__ deliberately returns None (not False) so exceptions inside the with block propagate unchanged, matching AsyncTaskHubGrpcClient.__aexit__.
   • No changes to __init__, close(), _maybe_recreate_channel, _schedule_recreate, _close_retired_channel, or any resiliency state introduced by Add gRPC client and worker connection resiliency #135.

   Callsite updates

   Migrate existing callsites to with ...: so they exercise the new path and stop leaking channels in long-running test runs. Scope on current main (the original ~40 figure is now slightly larger because of test growth since the issue was filed):

   File	Unclosed TaskHubGrpcClient(...) usages
   tests/durabletask/test_orchestration_e2e.py	~22
   tests/durabletask/test_batch_actions.py	~12
   tests/durabletask/test_work_item_filters_e2e.py	~11
   tests/durabletask/entities/test_function_based_entities_e2e.py	~11
   tests/durabletask/test_large_payload_e2e.py	~8
   examples/in_memory_backend_example/test/test_workflows.py	~8
   tests/durabletask/entities/test_entity_failure_handling.py	~5
   tests/durabletask/entities/test_class_based_entities_e2e.py	~4
   tests/durabletask/test_orchestration_versioning_e2e.py	~3
   examples/human_interaction.py	1

   The tests/durabletask/test_client.py constructor-focused tests that intentionally do not run RPCs can stay as-is (they assert on construction, not on lifecycle), but the e2e suites above should switch to with.

   The azuremanaged examples (e.g. examples/activity_sequence.py, examples/fanout_fanin.py, examples/distributed-tracing/app.py, examples/large_payload/app.py, the entity examples, etc.) instantiate DurableTaskSchedulerClient and would benefit from the same migration, since they inherit the new context-manager support for free. They are not strictly required for the fix but are worth folding into the same PR to align with the worker-style with pattern those files already use.

   Example migration (unchanged from the original issue, since it still applies):

   # Before
   task_hub_client = client.TaskHubGrpcClient(host_address=HOST)
   id = task_hub_client.schedule_new_orchestration(...)
   state = task_hub_client.wait_for_orchestration_completion(id, timeout=30)
   
   # After
   with client.TaskHubGrpcClient(host_address=HOST) as task_hub_client:
       id = task_hub_client.schedule_new_orchestration(...)
       state = task_hub_client.wait_for_orchestration_completion(id, timeout=30)

   Test plan

   Add focused unit tests in tests/durabletask/test_client.py (next to the existing async-context tests and the _closing/recreate tests added by #135):

   1. __enter__ returns the same instance and __exit__ invokes close() exactly once (assert via a spy on TaskHubGrpcClient.close).
   2. After exiting a with block on an SDK-owned channel, client._closing is True and client._channel.close() has been called (mirrors the existing async assertion).
   3. For a caller-owned channel, __exit__ does not close the channel (verifies _owns_channel is still respected end-to-end through the context-manager path).
   4. Exceptions raised inside the with block propagate, and close() is still invoked on the way out.
   5. A regression test that combines the context manager with the resiliency path from Add gRPC client and worker connection resiliency #135: schedule a recreate (e.g. via _schedule_recreate()), then exit the with block, and assert that any retired-channel timers are cancelled and that the recreate thread is joined — confirming we do not regress the Add gRPC client and worker connection resiliency #135 teardown invariants.

   Run as part of the existing pytest suite plus python -m flake8 on the touched files, per repo conventions.

   Changelog

   Add to CHANGELOG.md under ## Unreleased:

   ADDED
   
   - Added context-manager support (`__enter__` / `__exit__`) to
     `TaskHubGrpcClient` so it can be used with `with` statements, mirroring the
     existing `AsyncTaskHubGrpcClient` async-context-manager support and the
     `TaskHubGrpcWorker` pattern. `DurableTaskSchedulerClient` inherits this
     behavior automatically.

   No durabletask-azuremanaged/CHANGELOG.md entry is required since the azuremanaged client does not override close() and the behavior change is fully described by the core SDK entry; we can optionally add a cross-reference note there if reviewers prefer.

   Summary

   • Issue is still valid; Add gRPC client and worker connection resiliency #135 did not address it.
   • Proposed change is two lines on the sync client, purely additive, and routes through the existing close() — preserving every resiliency invariant added in Add gRPC client and worker connection resiliency #135.
   • Callsite migration list refreshed for current main.
   • Test plan includes a specific regression test that asserts the context-manager path still cleans up resiliency state (_closing, recreate thread, retired channels).
2. Bernd Verst (berndverst)

   closed this as completedin #145on May 26, 2026
3. andystaples

   added a commit that references this issue on Jun 3, 2026

   Add context manager support to TaskHubGrpcClient (microsoft#145)

   ...

   1cb30b0