Add context manager support to `TaskHubGrpcClient` to prevent gRPC channel leaks

[andystaples]

Problem

TaskHubGrpcClient has a close() method but does not implement __enter__/__exit__, so it cannot be used as a context manager. This means callers must remember to call close() manually (typically in a try/finally block), but no existing tests or examples do this — every callsite creates the client and never closes it, leaking the underlying gRPC channel.

AsyncTaskHubGrpcClient already has __aenter__/__aexit__, so there is an asymmetry between the sync and async clients.

Scope

This affects ~40+ callsites across the repository:

File	Unclosed usages
tests/durabletask/test_orchestration_e2e.py	~21
tests/durabletask/entities/test_function_based_entities_e2e.py	~10
tests/durabletask/entities/test_entity_failure_handling.py	~5
tests/durabletask/entities/test_class_based_entities_e2e.py	~4
examples/in_memory_backend_example/test/test_workflows.py	~8
examples/human_interaction.py	1

Proposed Fix

1. Add __enter__/__exit__ to TaskHubGrpcClient (mirrors the existing async pattern):

def __enter__(self):
    return self

def __exit__(self, exc_type, exc_val, exc_tb):
    self.close()

2. Update all test and example callsites to use the with statement:

# 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)

This also improves the getting-started experience and aligns with the TaskHubGrpcWorker pattern already used throughout the codebase.

Notes

• Adding __enter__/__exit__ is a non-breaking additive change
• The channel leak is benign in short-lived test processes but is a real concern in long-running applications or test suites that run many tests
• This issue was surfaced by a reviewer comment on PR Add history retrieval and terminal instance ID paging APIs #133

[berndverst]

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).