From 30847fc1d1389e1efcb4e8d1913a342b4078de07 Mon Sep 17 00:00:00 2001 From: Bernd Verst Date: Sun, 26 Jul 2026 23:08:29 -0700 Subject: [PATCH] Defer synchronous Azure token acquisition to first request AccessTokenManager.__init__ acquired an access token eagerly, which made constructing a DurableTaskSchedulerClient or DurableTaskSchedulerWorker perform a blocking credential round trip to Entra before any RPC was made. DTSDefaultClientInterceptorImpl.__init__ then called get_access_token() immediately, so the cost was paid even for instances that were never used. Token acquisition is now deferred to the first intercepted RPC. The existing double-checked refresh lock in get_access_token() is unchanged, so the deferred first acquisition remains single-flight across threads, and the refresh-before-expiry behavior is unaffected. This changes when credential failures surface: they now raise from the first request instead of from the constructor. The exception type and message are unchanged. This is documented in the azuremanaged changelog. Fixes #187 Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 27f463ce-6eb3-45c0-9bc0-63c68f3660ad --- durabletask-azuremanaged/CHANGELOG.md | 6 ++ .../internal/access_token_manager.py | 13 +-- .../internal/durabletask_grpc_interceptor.py | 6 +- .../test_durabletask_grpc_interceptor.py | 94 +++++++++++++++---- 4 files changed, 92 insertions(+), 27 deletions(-) diff --git a/durabletask-azuremanaged/CHANGELOG.md b/durabletask-azuremanaged/CHANGELOG.md index 300d855b..2a5b7b22 100644 --- a/durabletask-azuremanaged/CHANGELOG.md +++ b/durabletask-azuremanaged/CHANGELOG.md @@ -7,6 +7,12 @@ adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ## Unreleased +- `DurableTaskSchedulerClient` and `DurableTaskSchedulerWorker` no longer block on an Azure + credential round trip while being constructed. The access token is now acquired on the first + request instead, so constructing a client or worker that is never used costs nothing. As a + result, credential failures (for example an unavailable managed identity or a misconfigured + `DefaultAzureCredential`) now surface from the first request rather than from the constructor. + The exception type and message are unchanged; only the timing differs. - `FailureDetails.error_type` now carries the fully-qualified type name (e.g. `durabletask.task.TaskFailedError`) instead of the bare class name, and the new `FailureDetails.is_caused_by()` helper is available (both inherited from durabletask). See the core `durabletask` changelog for details, including the breaking-change notes. ## v1.8.0 diff --git a/durabletask-azuremanaged/durabletask/azuremanaged/internal/access_token_manager.py b/durabletask-azuremanaged/durabletask/azuremanaged/internal/access_token_manager.py index 94fe49be..63daf86a 100644 --- a/durabletask-azuremanaged/durabletask/azuremanaged/internal/access_token_manager.py +++ b/durabletask-azuremanaged/durabletask/azuremanaged/internal/access_token_manager.py @@ -13,6 +13,7 @@ class AccessTokenManager: _token: AccessToken | None + expiry_time: datetime | None def __init__(self, token_credential: TokenCredential | None, refresh_interval_seconds: int = 600): self._scope = "https://durabletask.io/.default" @@ -22,12 +23,12 @@ def __init__(self, token_credential: TokenCredential | None, refresh_interval_se self._credential = token_credential self._refresh_lock = Lock() - if self._credential is not None: - self._token = self._credential.get_token(self._scope) - self.expiry_time = datetime.fromtimestamp(self._token.expires_on, tz=timezone.utc) - else: - self._token = None - self.expiry_time = None + # Token acquisition is deferred to the first get_access_token() call so that + # constructing a client or worker does not perform a blocking credential round + # trip. The deferred first acquisition still goes through the double-checked + # refresh lock below, so it remains single-flight across threads. + self._token = None + self.expiry_time = None def get_access_token(self) -> AccessToken | None: if self._token is None or self.is_token_expired(): diff --git a/durabletask-azuremanaged/durabletask/azuremanaged/internal/durabletask_grpc_interceptor.py b/durabletask-azuremanaged/durabletask/azuremanaged/internal/durabletask_grpc_interceptor.py index e7cbb10a..37afc6d7 100644 --- a/durabletask-azuremanaged/durabletask/azuremanaged/internal/durabletask_grpc_interceptor.py +++ b/durabletask-azuremanaged/durabletask/azuremanaged/internal/durabletask_grpc_interceptor.py @@ -41,13 +41,13 @@ def __init__( self._metadata.append(("workerid", worker_id)) super().__init__(self._metadata) + # Token acquisition is deferred to the first _intercept_call invocation rather + # than happening in __init__, so that constructing a client or worker does not + # block on a credential round trip before any RPC is made. self._token_manager = None if token_credential is not None: self._token_credential = token_credential self._token_manager = AccessTokenManager(token_credential=self._token_credential) - access_token = self._token_manager.get_access_token() - if access_token is not None: - self._upsert_authorization_header(access_token.token) def _upsert_authorization_header(self, token: str) -> None: found = False diff --git a/tests/durabletask-azuremanaged/test_durabletask_grpc_interceptor.py b/tests/durabletask-azuremanaged/test_durabletask_grpc_interceptor.py index 878253a9..20162a5d 100644 --- a/tests/durabletask-azuremanaged/test_durabletask_grpc_interceptor.py +++ b/tests/durabletask-azuremanaged/test_durabletask_grpc_interceptor.py @@ -7,6 +7,7 @@ from importlib.metadata import version import threading import time +from typing import Any import grpc from azure.core.credentials import AccessToken @@ -39,6 +40,21 @@ def GetInstance(self, request, context): return response +class _TestTokenCredential: + """Minimal TokenCredential stub that counts how often a token is requested.""" + + def __init__(self): + self._lock = threading.Lock() + self.calls = 0 + + def get_token(self, *scopes: str, **kwargs: Any) -> AccessToken: + with self._lock: + self.calls += 1 + call_number = self.calls + time.sleep(0.02) + return AccessToken(f"token-{call_number}", int(time.time()) + 3600) + + class TestDurableTaskGrpcInterceptor(unittest.TestCase): """Tests for the DTSDefaultClientInterceptorImpl class.""" @@ -140,34 +156,76 @@ def test_worker_includes_workerid_header(self): self.assertIn("workerid", metadata) self.assertTrue(metadata["workerid"]) + def test_client_construction_does_not_acquire_token(self): + """Token acquisition is deferred from construction to the first request.""" + credential = _TestTokenCredential() -class _TestTokenCredential: - def __init__(self): - self._lock = threading.Lock() - self.calls = 0 + task_hub_client = DurableTaskSchedulerClient( + host_address=self.server_address, + secure_channel=False, + taskhub="test-taskhub", + token_credential=credential, + ) - def get_token(self, _scope): - with self._lock: - self.calls += 1 - call_number = self.calls - time.sleep(0.02) - return AccessToken(f"token-{call_number}", int(time.time()) + 3600) + self.assertEqual(0, credential.calls, "Constructing a client must not acquire a token") + + task_hub_client.get_orchestration_state("test-instance-id") + + self.assertEqual(1, credential.calls, "The first request should acquire exactly one token") + self.assertIn("authorization", self.mock_servicer.captured_metadata) + + task_hub_client.get_orchestration_state("test-instance-id") + + self.assertEqual(1, credential.calls, "A cached, unexpired token should be reused") + self.assertEqual(2, self.mock_servicer.requests_received) + + def test_worker_construction_does_not_acquire_token(self): + """Token acquisition is deferred from worker construction to the first request.""" + credential = _TestTokenCredential() + + DurableTaskSchedulerWorker( + host_address=self.server_address, + secure_channel=False, + taskhub="test-taskhub", + token_credential=credential, + ) + + self.assertEqual(0, credential.calls, "Constructing a worker must not acquire a token") class TestAccessTokenManagerThreadSafety(unittest.TestCase): + + @staticmethod + def _get_access_token_concurrently(manager: AccessTokenManager, thread_count: int = 8) -> None: + threads = [threading.Thread(target=manager.get_access_token) for _ in range(thread_count)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + def test_deferred_first_acquisition_performs_single_acquisition(self): + credential = _TestTokenCredential() + manager = AccessTokenManager(credential) + + self.assertEqual(0, credential.calls, "Constructing the manager must not acquire a token") + + self._get_access_token_concurrently(manager) + + self.assertEqual(1, credential.calls) + token = manager.get_access_token() + self.assertIsNotNone(token) + self.assertEqual("token-1", token.token) # type: ignore[union-attr] + self.assertEqual(1, credential.calls, "A cached, unexpired token should be reused") + def test_concurrent_refresh_performs_single_refresh(self): credential = _TestTokenCredential() manager = AccessTokenManager(credential) - manager.expiry_time = datetime.now(timezone.utc) - timedelta(seconds=1) + manager.get_access_token() + self.assertEqual(1, credential.calls) - threads = [] - for _ in range(8): - thread = threading.Thread(target=manager.get_access_token) - thread.start() - threads.append(thread) + manager.expiry_time = datetime.now(timezone.utc) - timedelta(seconds=1) - for thread in threads: - thread.join() + self._get_access_token_concurrently(manager) self.assertEqual(2, credential.calls)