Skip to content

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

@andystaples

Description

@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()
  1. 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

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions