Skip to content

Commit 30847fc

Browse files
Bernd VerstCopilot
andcommitted
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
1 parent 4d87430 commit 30847fc

4 files changed

Lines changed: 92 additions & 27 deletions

File tree

durabletask-azuremanaged/CHANGELOG.md

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,12 @@ adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
77

88
## Unreleased
99

10+
- `DurableTaskSchedulerClient` and `DurableTaskSchedulerWorker` no longer block on an Azure
11+
credential round trip while being constructed. The access token is now acquired on the first
12+
request instead, so constructing a client or worker that is never used costs nothing. As a
13+
result, credential failures (for example an unavailable managed identity or a misconfigured
14+
`DefaultAzureCredential`) now surface from the first request rather than from the constructor.
15+
The exception type and message are unchanged; only the timing differs.
1016
- `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.
1117

1218
## v1.8.0

durabletask-azuremanaged/durabletask/azuremanaged/internal/access_token_manager.py

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
class AccessTokenManager:
1414

1515
_token: AccessToken | None
16+
expiry_time: datetime | None
1617

1718
def __init__(self, token_credential: TokenCredential | None, refresh_interval_seconds: int = 600):
1819
self._scope = "https://durabletask.io/.default"
@@ -22,12 +23,12 @@ def __init__(self, token_credential: TokenCredential | None, refresh_interval_se
2223
self._credential = token_credential
2324
self._refresh_lock = Lock()
2425

25-
if self._credential is not None:
26-
self._token = self._credential.get_token(self._scope)
27-
self.expiry_time = datetime.fromtimestamp(self._token.expires_on, tz=timezone.utc)
28-
else:
29-
self._token = None
30-
self.expiry_time = None
26+
# Token acquisition is deferred to the first get_access_token() call so that
27+
# constructing a client or worker does not perform a blocking credential round
28+
# trip. The deferred first acquisition still goes through the double-checked
29+
# refresh lock below, so it remains single-flight across threads.
30+
self._token = None
31+
self.expiry_time = None
3132

3233
def get_access_token(self) -> AccessToken | None:
3334
if self._token is None or self.is_token_expired():

durabletask-azuremanaged/durabletask/azuremanaged/internal/durabletask_grpc_interceptor.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -41,13 +41,13 @@ def __init__(
4141
self._metadata.append(("workerid", worker_id))
4242
super().__init__(self._metadata)
4343

44+
# Token acquisition is deferred to the first _intercept_call invocation rather
45+
# than happening in __init__, so that constructing a client or worker does not
46+
# block on a credential round trip before any RPC is made.
4447
self._token_manager = None
4548
if token_credential is not None:
4649
self._token_credential = token_credential
4750
self._token_manager = AccessTokenManager(token_credential=self._token_credential)
48-
access_token = self._token_manager.get_access_token()
49-
if access_token is not None:
50-
self._upsert_authorization_header(access_token.token)
5151

5252
def _upsert_authorization_header(self, token: str) -> None:
5353
found = False

tests/durabletask-azuremanaged/test_durabletask_grpc_interceptor.py

Lines changed: 76 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
from importlib.metadata import version
88
import threading
99
import time
10+
from typing import Any
1011

1112
import grpc
1213
from azure.core.credentials import AccessToken
@@ -39,6 +40,21 @@ def GetInstance(self, request, context):
3940
return response
4041

4142

43+
class _TestTokenCredential:
44+
"""Minimal TokenCredential stub that counts how often a token is requested."""
45+
46+
def __init__(self):
47+
self._lock = threading.Lock()
48+
self.calls = 0
49+
50+
def get_token(self, *scopes: str, **kwargs: Any) -> AccessToken:
51+
with self._lock:
52+
self.calls += 1
53+
call_number = self.calls
54+
time.sleep(0.02)
55+
return AccessToken(f"token-{call_number}", int(time.time()) + 3600)
56+
57+
4258
class TestDurableTaskGrpcInterceptor(unittest.TestCase):
4359
"""Tests for the DTSDefaultClientInterceptorImpl class."""
4460

@@ -140,34 +156,76 @@ def test_worker_includes_workerid_header(self):
140156
self.assertIn("workerid", metadata)
141157
self.assertTrue(metadata["workerid"])
142158

159+
def test_client_construction_does_not_acquire_token(self):
160+
"""Token acquisition is deferred from construction to the first request."""
161+
credential = _TestTokenCredential()
143162

144-
class _TestTokenCredential:
145-
def __init__(self):
146-
self._lock = threading.Lock()
147-
self.calls = 0
163+
task_hub_client = DurableTaskSchedulerClient(
164+
host_address=self.server_address,
165+
secure_channel=False,
166+
taskhub="test-taskhub",
167+
token_credential=credential,
168+
)
148169

149-
def get_token(self, _scope):
150-
with self._lock:
151-
self.calls += 1
152-
call_number = self.calls
153-
time.sleep(0.02)
154-
return AccessToken(f"token-{call_number}", int(time.time()) + 3600)
170+
self.assertEqual(0, credential.calls, "Constructing a client must not acquire a token")
171+
172+
task_hub_client.get_orchestration_state("test-instance-id")
173+
174+
self.assertEqual(1, credential.calls, "The first request should acquire exactly one token")
175+
self.assertIn("authorization", self.mock_servicer.captured_metadata)
176+
177+
task_hub_client.get_orchestration_state("test-instance-id")
178+
179+
self.assertEqual(1, credential.calls, "A cached, unexpired token should be reused")
180+
self.assertEqual(2, self.mock_servicer.requests_received)
181+
182+
def test_worker_construction_does_not_acquire_token(self):
183+
"""Token acquisition is deferred from worker construction to the first request."""
184+
credential = _TestTokenCredential()
185+
186+
DurableTaskSchedulerWorker(
187+
host_address=self.server_address,
188+
secure_channel=False,
189+
taskhub="test-taskhub",
190+
token_credential=credential,
191+
)
192+
193+
self.assertEqual(0, credential.calls, "Constructing a worker must not acquire a token")
155194

156195

157196
class TestAccessTokenManagerThreadSafety(unittest.TestCase):
197+
198+
@staticmethod
199+
def _get_access_token_concurrently(manager: AccessTokenManager, thread_count: int = 8) -> None:
200+
threads = [threading.Thread(target=manager.get_access_token) for _ in range(thread_count)]
201+
for thread in threads:
202+
thread.start()
203+
for thread in threads:
204+
thread.join()
205+
206+
def test_deferred_first_acquisition_performs_single_acquisition(self):
207+
credential = _TestTokenCredential()
208+
manager = AccessTokenManager(credential)
209+
210+
self.assertEqual(0, credential.calls, "Constructing the manager must not acquire a token")
211+
212+
self._get_access_token_concurrently(manager)
213+
214+
self.assertEqual(1, credential.calls)
215+
token = manager.get_access_token()
216+
self.assertIsNotNone(token)
217+
self.assertEqual("token-1", token.token) # type: ignore[union-attr]
218+
self.assertEqual(1, credential.calls, "A cached, unexpired token should be reused")
219+
158220
def test_concurrent_refresh_performs_single_refresh(self):
159221
credential = _TestTokenCredential()
160222
manager = AccessTokenManager(credential)
161-
manager.expiry_time = datetime.now(timezone.utc) - timedelta(seconds=1)
223+
manager.get_access_token()
224+
self.assertEqual(1, credential.calls)
162225

163-
threads = []
164-
for _ in range(8):
165-
thread = threading.Thread(target=manager.get_access_token)
166-
thread.start()
167-
threads.append(thread)
226+
manager.expiry_time = datetime.now(timezone.utc) - timedelta(seconds=1)
168227

169-
for thread in threads:
170-
thread.join()
228+
self._get_access_token_concurrently(manager)
171229

172230
self.assertEqual(2, credential.calls)
173231

0 commit comments

Comments
 (0)