Skip to content

Commit 30f6e24

Browse files
berndverstBernd VerstCopilot
authored
Performance [P1]: Single-flight async Azure token refresh (#198)
* perf: single-flight async Azure token refresh AsyncAccessTokenManager.get_access_token() checked expiry and awaited refresh_token() without any synchronization. Since the async DTS gRPC interceptor awaits it on every RPC, a burst of concurrent grpc.aio calls arriving while the token was missing or inside its renewal window would each independently call the credential, producing a thundering herd of Entra/credential requests, extra latency, and throttling pressure. Mirror the double-checked locking already used by the synchronous AccessTokenManager: re-check token validity after acquiring a refresh guard so only one coroutine contacts the credential per refresh window while the rest reuse the result. The guard is an asyncio.Lock created lazily per running event loop (tracked in a small map behind a threading.Lock), because an asyncio.Lock binds to the loop it is first awaited on and this manager may outlive a loop or be shared across loops. Locks for closed loops are pruned so the map cannot grow unbounded. The cached-token fast path is unchanged and never touches the lock, so there is no added per-RPC overhead. Public signatures, return values, the renewal margin behavior, exception propagation, and logging are all preserved. Adds unit tests covering concurrent cold start, concurrent refresh after expiry, the cached-token path, credential failure propagation without deadlock, and reuse across distinct event loops. Fixes #192 Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: bb08d0ed-484d-4015-b905-3e164f68678b * refactor: hoist stale-loop comprehension into a named local Address review feedback on the stale-lock cleanup in AsyncAccessTokenManager._get_refresh_lock(). The list comprehension was embedded directly in the `for` statement, which read awkwardly across the wrapped line. Hoist it into a named `stale_loops` local and iterate over that. This is purely a readability change with no behavioural difference: the snapshot is still materialized before mutating the mapping, which is what makes deleting during iteration safe. Note this was not a lint violation - flake8 passes on the file both before and after the change. Also align the async deferred-acquisition test with the naming and assertions that PR #199 introduced for the synchronous manager, and assert explicitly that constructing AsyncAccessTokenManager performs no credential call. The async manager has always deferred its first acquisition, so the cold-start burst flows through the same refresh guard as a later refresh; the test now states that intent directly. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: bb08d0ed-484d-4015-b905-3e164f68678b --------- Co-authored-by: Bernd Verst <beverst@microsoft.com> Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: bb08d0ed-484d-4015-b905-3e164f68678b
1 parent 0f316cc commit 30f6e24

3 files changed

Lines changed: 141 additions & 2 deletions

File tree

‎durabletask-azuremanaged/CHANGELOG.md‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,9 @@ adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
1414
`DefaultAzureCredential`) now surface from the first request rather than from the constructor.
1515
The exception type and message are unchanged; only the timing differs.
1616
- `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.
17+
- Improved async access token refresh concurrency handling to avoid duplicate
18+
refresh operations under concurrent access, matching the existing sync
19+
behavior.
1720

1821
## v1.8.0
1922

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

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
# Copyright (c) Microsoft Corporation.
22
# Licensed under the MIT License.
3+
import asyncio
34
from datetime import datetime, timedelta, timezone
45
from threading import Lock
56

@@ -71,9 +72,37 @@ def __init__(self, token_credential: AsyncTokenCredential | None,
7172
self._token = None
7273
self.expiry_time = None
7374

75+
# An asyncio.Lock binds itself to the event loop it is first used on, and this
76+
# manager may outlive a loop or be shared across loops. Locks are therefore
77+
# created lazily per running loop, guarded by a plain threading lock because
78+
# different loops may run on different threads.
79+
self._refresh_locks: dict[asyncio.AbstractEventLoop, asyncio.Lock] = {}
80+
self._refresh_locks_guard = Lock()
81+
82+
def _get_refresh_lock(self) -> asyncio.Lock:
83+
loop = asyncio.get_running_loop()
84+
with self._refresh_locks_guard:
85+
lock = self._refresh_locks.get(loop)
86+
if lock is None:
87+
# Discard locks belonging to loops that are no longer usable so the
88+
# mapping does not grow without bound.
89+
stale_loops = [
90+
existing for existing in self._refresh_locks if existing.is_closed()
91+
]
92+
for stale_loop in stale_loops:
93+
del self._refresh_locks[stale_loop]
94+
lock = asyncio.Lock()
95+
self._refresh_locks[loop] = lock
96+
return lock
97+
7498
async def get_access_token(self) -> AccessToken | None:
7599
if self._token is None or self.is_token_expired():
76-
await self.refresh_token()
100+
async with self._get_refresh_lock():
101+
# Re-check under the lock: a concurrent caller may have already
102+
# refreshed the token while this one was waiting, so only a single
103+
# credential request is made per refresh window.
104+
if self._token is None or self.is_token_expired():
105+
await self.refresh_token()
77106
return self._token
78107

79108
def is_token_expired(self) -> bool:

‎tests/durabletask-azuremanaged/test_durabletask_grpc_interceptor.py‎

Lines changed: 108 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
# Copyright (c) Microsoft Corporation.
22
# Licensed under the MIT License.
33

4+
import asyncio
45
import unittest
56
from concurrent import futures
67
from datetime import datetime, timedelta, timezone
@@ -13,7 +14,10 @@
1314
from azure.core.credentials import AccessToken
1415

1516
from durabletask.azuremanaged.client import DurableTaskSchedulerClient
16-
from durabletask.azuremanaged.internal.access_token_manager import AccessTokenManager
17+
from durabletask.azuremanaged.internal.access_token_manager import (
18+
AccessTokenManager,
19+
AsyncAccessTokenManager,
20+
)
1721
from durabletask.azuremanaged.worker import DurableTaskSchedulerWorker
1822
from durabletask.internal.grpc_interceptor import DefaultClientInterceptorImpl
1923
from durabletask.internal import orchestrator_service_pb2 as pb
@@ -230,5 +234,108 @@ def test_concurrent_refresh_performs_single_refresh(self):
230234
self.assertEqual(2, credential.calls)
231235

232236

237+
class _CredentialError(Exception):
238+
pass
239+
240+
241+
class _TestAsyncTokenCredential:
242+
def __init__(self, delay: float = 0.02):
243+
self.calls = 0
244+
self._delay = delay
245+
246+
async def get_token(self, _scope):
247+
# Counting before the await means every coroutine that reaches the credential
248+
# is recorded, which is what makes a thundering herd observable.
249+
self.calls += 1
250+
call_number = self.calls
251+
await asyncio.sleep(self._delay)
252+
return AccessToken(f"token-{call_number}", int(time.time()) + 3600)
253+
254+
255+
class _FlakyAsyncTokenCredential:
256+
def __init__(self):
257+
self.calls = 0
258+
self.fail = True
259+
260+
async def get_token(self, _scope):
261+
self.calls += 1
262+
await asyncio.sleep(0)
263+
if self.fail:
264+
raise _CredentialError("credential unavailable")
265+
return AccessToken("token-recovered", int(time.time()) + 3600)
266+
267+
268+
class TestAsyncAccessTokenManagerSingleFlight(unittest.IsolatedAsyncioTestCase):
269+
async def test_deferred_first_acquisition_performs_single_acquisition(self):
270+
credential = _TestAsyncTokenCredential()
271+
manager = AsyncAccessTokenManager(credential)
272+
273+
self.assertEqual(0, credential.calls, "Constructing the manager must not acquire a token")
274+
275+
# The async manager has always deferred the first acquisition (a constructor
276+
# cannot await), so the cold-start burst goes through the same refresh guard
277+
# as a later refresh and must likewise be single-flight.
278+
tokens = await asyncio.gather(*(manager.get_access_token() for _ in range(16)))
279+
280+
self.assertEqual(1, credential.calls)
281+
self.assertEqual({"token-1"}, {token.token for token in tokens})
282+
283+
async def test_concurrent_refresh_performs_single_refresh(self):
284+
credential = _TestAsyncTokenCredential()
285+
manager = AsyncAccessTokenManager(credential)
286+
await manager.get_access_token()
287+
self.assertEqual(1, credential.calls)
288+
289+
manager.expiry_time = datetime.now(timezone.utc) - timedelta(seconds=1)
290+
291+
tokens = await asyncio.gather(*(manager.get_access_token() for _ in range(16)))
292+
293+
self.assertEqual(2, credential.calls)
294+
self.assertEqual({"token-2"}, {token.token for token in tokens})
295+
296+
async def test_valid_token_is_returned_without_contacting_credential(self):
297+
credential = _TestAsyncTokenCredential()
298+
manager = AsyncAccessTokenManager(credential)
299+
300+
first = await manager.get_access_token()
301+
second = await manager.get_access_token()
302+
303+
self.assertEqual(1, credential.calls)
304+
self.assertIs(first, second)
305+
306+
async def test_credential_failure_propagates_and_does_not_deadlock(self):
307+
credential = _FlakyAsyncTokenCredential()
308+
manager = AsyncAccessTokenManager(credential)
309+
310+
results = await asyncio.gather(
311+
*(manager.get_access_token() for _ in range(4)), return_exceptions=True)
312+
313+
self.assertTrue(all(isinstance(result, _CredentialError) for result in results))
314+
315+
# The guard must have been released on failure so a later attempt can succeed.
316+
credential.fail = False
317+
token = await manager.get_access_token()
318+
self.assertIsNotNone(token)
319+
self.assertEqual("token-recovered", token.token)
320+
321+
322+
class TestAsyncAccessTokenManagerEventLoopBinding(unittest.TestCase):
323+
def test_manager_is_reusable_across_event_loops(self):
324+
credential = _TestAsyncTokenCredential(delay=0)
325+
manager = AsyncAccessTokenManager(credential)
326+
327+
first = asyncio.run(manager.get_access_token())
328+
329+
manager.expiry_time = datetime.now(timezone.utc) - timedelta(seconds=1)
330+
331+
# The second run uses a distinct event loop; the refresh guard must not stay
332+
# bound to the first, now-closed loop.
333+
second = asyncio.run(manager.get_access_token())
334+
335+
self.assertIsNotNone(first)
336+
self.assertIsNotNone(second)
337+
self.assertEqual(2, credential.calls)
338+
339+
233340
if __name__ == "__main__":
234341
unittest.main()

0 commit comments

Comments
 (0)