From 9cc01a88d86bdf3bd8974246c58603f0c66a6984 Mon Sep 17 00:00:00 2001 From: Bernd Verst Date: Sun, 26 Jul 2026 23:11:58 -0700 Subject: [PATCH 1/2] Reuse sandbox registration transport across retries The sandbox worker registration loop built a brand new SandboxActivitiesGrpcTransport on every attempt, so each transient failure paid for a fresh gRPC channel, stub, interceptor, and Entra token acquisition. That amplified startup and recovery latency exactly when the service was already degraded. The loop now creates the transport once and reuses it across retriable stream failures, replacing it only when a failure leaves the channel unusable (a shut-down channel or a non-gRPC failure such as an error raised while the channel is being created). The transport is closed once when the loop exits, so channels are not leaked. The retry backoff schedule, stop-signal responsiveness, thread join semantics, log messages, and exception types on terminal failure are unchanged. Fixes #188 Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9a361145-b7cc-437d-895d-0bea3d19bf91 --- .../azuremanaged/preview/sandboxes/worker.py | 82 +++++++--- .../test_sandboxes_extension.py | 142 ++++++++++++++++++ 2 files changed, 203 insertions(+), 21 deletions(-) diff --git a/durabletask-azuremanaged/durabletask/azuremanaged/preview/sandboxes/worker.py b/durabletask-azuremanaged/durabletask/azuremanaged/preview/sandboxes/worker.py index 21cb7161..8c73c307 100644 --- a/durabletask-azuremanaged/durabletask/azuremanaged/preview/sandboxes/worker.py +++ b/durabletask-azuremanaged/durabletask/azuremanaged/preview/sandboxes/worker.py @@ -130,30 +130,50 @@ def _stop_sandbox_registration(self) -> None: def _run_sandbox_registration_loop(self) -> None: retry_delay = 1.0 - while not self._sandbox_registration_stop.is_set(): - try: - client = SandboxActivitiesGrpcTransport( - host_address=self._sandbox_host_address, - taskhub=self._sandbox_taskhub, - token_credential=self._sandbox_token_credential, - secure_channel=self._sandbox_secure_channel) + client: Optional[SandboxActivitiesGrpcTransport] = None + try: + while not self._sandbox_registration_stop.is_set(): try: + if client is None: + client = SandboxActivitiesGrpcTransport( + host_address=self._sandbox_host_address, + taskhub=self._sandbox_taskhub, + token_credential=self._sandbox_token_credential, + secure_channel=self._sandbox_secure_channel) client.connect_sandbox_activity_worker(self._registration_messages()) retry_delay = 1.0 - finally: - client.close() - except Exception as ex: - if self._sandbox_registration_stop.is_set(): - break - if not _is_retriable_registration_failure(ex): - self._sandbox_logger.error( - "Sandbox activity worker registration failed permanently: %s", ex) - self._sandbox_registration_stop.set() - break - self._sandbox_logger.warning("Sandbox activity worker registration failed: %s", ex) - delay = random.uniform(0, retry_delay) - self._sandbox_registration_stop.wait(delay) - retry_delay = min(retry_delay * 2, 30.0) + except Exception as ex: + # gRPC channels reconnect on their own after transient stream + # failures, so the transport (and its channel, stub, and cached + # access token) is reused instead of being rebuilt per attempt. + if _requires_new_registration_transport(ex): + self._close_sandbox_registration_transport(client) + client = None + if self._sandbox_registration_stop.is_set(): + break + if not _is_retriable_registration_failure(ex): + self._sandbox_logger.error( + "Sandbox activity worker registration failed permanently: %s", ex) + self._sandbox_registration_stop.set() + break + self._sandbox_logger.warning( + "Sandbox activity worker registration failed: %s", ex) + delay = random.uniform(0, retry_delay) + self._sandbox_registration_stop.wait(delay) + retry_delay = min(retry_delay * 2, 30.0) + finally: + self._close_sandbox_registration_transport(client) + + def _close_sandbox_registration_transport( + self, + client: Optional[SandboxActivitiesGrpcTransport]) -> None: + if client is None: + return + try: + client.close() + except Exception as ex: + self._sandbox_logger.debug( + "Sandbox activity worker registration transport failed to close: %s", ex) def _registration_messages(self) -> Iterator[pb.SandboxActivityWorkerMessage]: yield build_sandbox_worker_start( @@ -264,6 +284,26 @@ def _is_retriable_registration_failure(ex: Exception) -> bool: return isinstance(ex, OSError) +def _requires_new_registration_transport(ex: Exception) -> bool: + """Returns whether a failure leaves the registration transport unusable. + + A gRPC channel transparently reconnects after a stream fails, so the same + transport can serve the next registration attempt. A channel that has been + shut down rejects every later RPC, and a failure raised outside of gRPC (for + example while the channel itself is being created) can leave the transport + only partially initialized, so both cases need a replacement transport. + """ + if not isinstance(ex, grpc.RpcError): + return True + + if ex.code() is not grpc.StatusCode.CANCELLED: + return False + + details = getattr(ex, "details", None) + resolved_details = details() if callable(details) else None + return isinstance(resolved_details, str) and "channel closed" in resolved_details.lower() + + def _resolve_max_concurrent_activities() -> int: value = os.getenv("DTS_SANDBOX_MAX_ACTIVITIES") if value is None: diff --git a/tests/durabletask-azuremanaged/test_sandboxes_extension.py b/tests/durabletask-azuremanaged/test_sandboxes_extension.py index caf96571..15991eeb 100644 --- a/tests/durabletask-azuremanaged/test_sandboxes_extension.py +++ b/tests/durabletask-azuremanaged/test_sandboxes_extension.py @@ -2,6 +2,7 @@ # Licensed under the MIT License. import inspect +import threading import grpc from azure.core.credentials import AccessToken @@ -766,6 +767,147 @@ def test_sandbox_registration_retries_failed_precondition_until_structured_reaso _FakeRpcError(grpc.StatusCode.FAILED_PRECONDITION, "worker profile does not match")) +def test_sandbox_registration_recreates_transport_only_when_channel_is_unusable() -> None: + assert not sandbox_worker._requires_new_registration_transport( + _FakeRpcError(grpc.StatusCode.UNAVAILABLE)) + assert not sandbox_worker._requires_new_registration_transport( + _FakeRpcError(grpc.StatusCode.FAILED_PRECONDITION, "sandbox not ready")) + assert not sandbox_worker._requires_new_registration_transport( + _FakeRpcError(grpc.StatusCode.CANCELLED, "Locally cancelled by application!")) + assert sandbox_worker._requires_new_registration_transport( + _FakeRpcError(grpc.StatusCode.CANCELLED, "Channel closed!")) + assert sandbox_worker._requires_new_registration_transport(OSError("connection reset")) + + +def test_sandbox_registration_reuses_one_transport_across_retriable_failures(monkeypatch) -> None: + worker = _build_registration_test_worker(monkeypatch) + attempts = 0 + + def connect(_messages): + nonlocal attempts + attempts += 1 + if attempts <= 3: + raise _FakeRpcError(grpc.StatusCode.UNAVAILABLE) + worker._sandbox_registration_stop.set() + return object() + + transports = _install_fake_registration_transport(monkeypatch, connect) + backoff = _StubBackoff() + monkeypatch.setattr(sandbox_worker, "random", backoff) + + worker._run_sandbox_registration_loop() + + assert attempts == 4 + assert len(transports) == 1 + assert transports[0].close_count == 1 + # The existing exponential backoff schedule must be preserved. + assert backoff.upper_bounds == [1.0, 2.0, 4.0] + + +def test_sandbox_registration_rebuilds_transport_after_channel_shutdown(monkeypatch) -> None: + worker = _build_registration_test_worker(monkeypatch) + attempts = 0 + + def connect(_messages): + nonlocal attempts + attempts += 1 + if attempts == 1: + raise _FakeRpcError(grpc.StatusCode.CANCELLED, "Channel closed!") + worker._sandbox_registration_stop.set() + return object() + + transports = _install_fake_registration_transport(monkeypatch, connect) + monkeypatch.setattr(sandbox_worker, "random", _StubBackoff()) + + worker._run_sandbox_registration_loop() + + assert attempts == 2 + assert len(transports) == 2 + assert [transport.close_count for transport in transports] == [1, 1] + + +def test_sandbox_registration_loop_exits_promptly_on_stop_signal(monkeypatch) -> None: + worker = _build_registration_test_worker(monkeypatch) + worker._sandbox_heartbeat_interval_seconds = 0.05 + connected = threading.Event() + + def connect(messages): + connected.set() + for _ in messages: + pass + return object() + + transports = _install_fake_registration_transport(monkeypatch, connect) + + thread = threading.Thread(target=worker._run_sandbox_registration_loop, daemon=True) + thread.start() + try: + assert connected.wait(10) + worker._sandbox_registration_stop.set() + thread.join(timeout=10) + assert not thread.is_alive() + finally: + worker._sandbox_registration_stop.set() + thread.join(timeout=10) + + assert len(transports) == 1 + assert transports[0].close_count == 1 + + +class _StubBackoff: + """Records the backoff window used for each retry and never actually sleeps.""" + + def __init__(self) -> None: + self.upper_bounds: list[float] = [] + + def uniform(self, lower: float, upper: float) -> float: + self.upper_bounds.append(upper) + return 0.0 + + +class _FakeRegistrationTransport: + def __init__(self, connect, kwargs) -> None: + self.kwargs = kwargs + self.close_count = 0 + self._connect = connect + + def connect_sandbox_activity_worker(self, messages): + return self._connect(messages) + + def close(self) -> None: + self.close_count += 1 + + +def _install_fake_registration_transport( + monkeypatch, connect) -> list[_FakeRegistrationTransport]: + transports: list[_FakeRegistrationTransport] = [] + + def factory(**kwargs) -> _FakeRegistrationTransport: + transport = _FakeRegistrationTransport(connect, kwargs) + transports.append(transport) + return transport + + monkeypatch.setattr(sandbox_worker, "SandboxActivitiesGrpcTransport", factory) + return transports + + +def _build_registration_test_worker(monkeypatch) -> SandboxWorker: + monkeypatch.setenv("DTS_ENDPOINT", "http://localhost:8080") + monkeypatch.setenv("DTS_TASK_HUB", "env-hub") + monkeypatch.setenv("DTS_WORKER_PROFILE_ID", "env-profile") + monkeypatch.setenv("DTS_SANDBOX_PROVIDER", "Sandbox") + _configure_sandbox_worker_auth(monkeypatch) + + def RegistrationActivity(_ctx, value): + return value + + worker = SandboxWorker() + worker.add_activity(RegistrationActivity) + worker._configure_sandbox_activity_filters() + worker._sandbox_registration_stop.clear() + return worker + + class _RecordingChannel: def __init__(self) -> None: self.methods: list[str] = [] From f9753b74c9cbd1d949b9382d8b148cf6f659e544 Mon Sep 17 00:00:00 2001 From: Bernd Verst Date: Mon, 27 Jul 2026 09:38:18 -0700 Subject: [PATCH 2/2] Compare gRPC status codes with != for consistency grpc.StatusCode members are enum singletons, so `is not` and `!=` behave identically here. Every other grpc.StatusCode comparison in the repository uses ==/!= (durabletask/client.py, durabletask/worker.py, durabletask/internal/grpc_resiliency.py, durabletask/extensions/history_export/client.py), so match that style instead of introducing the only identity comparison. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9a361145-b7cc-437d-895d-0bea3d19bf91 --- .../durabletask/azuremanaged/preview/sandboxes/worker.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/durabletask-azuremanaged/durabletask/azuremanaged/preview/sandboxes/worker.py b/durabletask-azuremanaged/durabletask/azuremanaged/preview/sandboxes/worker.py index 8c73c307..c87d12da 100644 --- a/durabletask-azuremanaged/durabletask/azuremanaged/preview/sandboxes/worker.py +++ b/durabletask-azuremanaged/durabletask/azuremanaged/preview/sandboxes/worker.py @@ -296,7 +296,7 @@ def _requires_new_registration_transport(ex: Exception) -> bool: if not isinstance(ex, grpc.RpcError): return True - if ex.code() is not grpc.StatusCode.CANCELLED: + if ex.code() != grpc.StatusCode.CANCELLED: return False details = getattr(ex, "details", None)