From 5f433de385e4b6a6e079b72af938f8edb9dd9035 Mon Sep 17 00:00:00 2001 From: Bernd Verst Date: Sun, 26 Jul 2026 23:07:11 -0700 Subject: [PATCH 1/2] 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 --- durabletask-azuremanaged/CHANGELOG.md | 3 + .../internal/access_token_manager.py | 29 ++++- .../test_durabletask_grpc_interceptor.py | 104 +++++++++++++++++- 3 files changed, 134 insertions(+), 2 deletions(-) diff --git a/durabletask-azuremanaged/CHANGELOG.md b/durabletask-azuremanaged/CHANGELOG.md index 300d855b..f01b17c1 100644 --- a/durabletask-azuremanaged/CHANGELOG.md +++ b/durabletask-azuremanaged/CHANGELOG.md @@ -8,6 +8,9 @@ adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ## Unreleased - `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. +- Improved async access token refresh concurrency handling to avoid duplicate + refresh operations under concurrent access, matching the existing sync + behavior. ## 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..f3b2945b 100644 --- a/durabletask-azuremanaged/durabletask/azuremanaged/internal/access_token_manager.py +++ b/durabletask-azuremanaged/durabletask/azuremanaged/internal/access_token_manager.py @@ -1,5 +1,6 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. +import asyncio from datetime import datetime, timedelta, timezone from threading import Lock @@ -70,9 +71,35 @@ def __init__(self, token_credential: AsyncTokenCredential | None, self._token = None self.expiry_time = None + # An asyncio.Lock binds itself to the event loop it is first used on, and this + # manager may outlive a loop or be shared across loops. Locks are therefore + # created lazily per running loop, guarded by a plain threading lock because + # different loops may run on different threads. + self._refresh_locks: dict[asyncio.AbstractEventLoop, asyncio.Lock] = {} + self._refresh_locks_guard = Lock() + + def _get_refresh_lock(self) -> asyncio.Lock: + loop = asyncio.get_running_loop() + with self._refresh_locks_guard: + lock = self._refresh_locks.get(loop) + if lock is None: + # Discard locks belonging to loops that are no longer usable so the + # mapping does not grow without bound. + for stale_loop in [ + existing for existing in self._refresh_locks if existing.is_closed()]: + del self._refresh_locks[stale_loop] + lock = asyncio.Lock() + self._refresh_locks[loop] = lock + return lock + async def get_access_token(self) -> AccessToken | None: if self._token is None or self.is_token_expired(): - await self.refresh_token() + async with self._get_refresh_lock(): + # Re-check under the lock: a concurrent caller may have already + # refreshed the token while this one was waiting, so only a single + # credential request is made per refresh window. + if self._token is None or self.is_token_expired(): + await self.refresh_token() return self._token def is_token_expired(self) -> bool: diff --git a/tests/durabletask-azuremanaged/test_durabletask_grpc_interceptor.py b/tests/durabletask-azuremanaged/test_durabletask_grpc_interceptor.py index 878253a9..7655dc99 100644 --- a/tests/durabletask-azuremanaged/test_durabletask_grpc_interceptor.py +++ b/tests/durabletask-azuremanaged/test_durabletask_grpc_interceptor.py @@ -1,6 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. +import asyncio import unittest from concurrent import futures from datetime import datetime, timedelta, timezone @@ -12,7 +13,10 @@ from azure.core.credentials import AccessToken from durabletask.azuremanaged.client import DurableTaskSchedulerClient -from durabletask.azuremanaged.internal.access_token_manager import AccessTokenManager +from durabletask.azuremanaged.internal.access_token_manager import ( + AccessTokenManager, + AsyncAccessTokenManager, +) from durabletask.azuremanaged.worker import DurableTaskSchedulerWorker from durabletask.internal.grpc_interceptor import DefaultClientInterceptorImpl from durabletask.internal import orchestrator_service_pb2 as pb @@ -172,5 +176,103 @@ def test_concurrent_refresh_performs_single_refresh(self): self.assertEqual(2, credential.calls) +class _CredentialError(Exception): + pass + + +class _TestAsyncTokenCredential: + def __init__(self, delay: float = 0.02): + self.calls = 0 + self._delay = delay + + async def get_token(self, _scope): + # Counting before the await means every coroutine that reaches the credential + # is recorded, which is what makes a thundering herd observable. + self.calls += 1 + call_number = self.calls + await asyncio.sleep(self._delay) + return AccessToken(f"token-{call_number}", int(time.time()) + 3600) + + +class _FlakyAsyncTokenCredential: + def __init__(self): + self.calls = 0 + self.fail = True + + async def get_token(self, _scope): + self.calls += 1 + await asyncio.sleep(0) + if self.fail: + raise _CredentialError("credential unavailable") + return AccessToken("token-recovered", int(time.time()) + 3600) + + +class TestAsyncAccessTokenManagerSingleFlight(unittest.IsolatedAsyncioTestCase): + async def test_concurrent_cold_start_performs_single_acquisition(self): + credential = _TestAsyncTokenCredential() + manager = AsyncAccessTokenManager(credential) + + tokens = await asyncio.gather(*(manager.get_access_token() for _ in range(16))) + + self.assertEqual(1, credential.calls) + self.assertEqual({"token-1"}, {token.token for token in tokens}) + + async def test_concurrent_refresh_performs_single_refresh(self): + credential = _TestAsyncTokenCredential() + manager = AsyncAccessTokenManager(credential) + await manager.get_access_token() + self.assertEqual(1, credential.calls) + + manager.expiry_time = datetime.now(timezone.utc) - timedelta(seconds=1) + + tokens = await asyncio.gather(*(manager.get_access_token() for _ in range(16))) + + self.assertEqual(2, credential.calls) + self.assertEqual({"token-2"}, {token.token for token in tokens}) + + async def test_valid_token_is_returned_without_contacting_credential(self): + credential = _TestAsyncTokenCredential() + manager = AsyncAccessTokenManager(credential) + + first = await manager.get_access_token() + second = await manager.get_access_token() + + self.assertEqual(1, credential.calls) + self.assertIs(first, second) + + async def test_credential_failure_propagates_and_does_not_deadlock(self): + credential = _FlakyAsyncTokenCredential() + manager = AsyncAccessTokenManager(credential) + + results = await asyncio.gather( + *(manager.get_access_token() for _ in range(4)), return_exceptions=True) + + self.assertTrue(all(isinstance(result, _CredentialError) for result in results)) + + # The guard must have been released on failure so a later attempt can succeed. + credential.fail = False + token = await manager.get_access_token() + self.assertIsNotNone(token) + self.assertEqual("token-recovered", token.token) + + +class TestAsyncAccessTokenManagerEventLoopBinding(unittest.TestCase): + def test_manager_is_reusable_across_event_loops(self): + credential = _TestAsyncTokenCredential(delay=0) + manager = AsyncAccessTokenManager(credential) + + first = asyncio.run(manager.get_access_token()) + + manager.expiry_time = datetime.now(timezone.utc) - timedelta(seconds=1) + + # The second run uses a distinct event loop; the refresh guard must not stay + # bound to the first, now-closed loop. + second = asyncio.run(manager.get_access_token()) + + self.assertIsNotNone(first) + self.assertIsNotNone(second) + self.assertEqual(2, credential.calls) + + if __name__ == "__main__": unittest.main() From 4b3bbb2219d83ab7a47b3e02f0536a833d53d89a Mon Sep 17 00:00:00 2001 From: Bernd Verst Date: Mon, 27 Jul 2026 09:40:53 -0700 Subject: [PATCH 2/2] 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 --- .../azuremanaged/internal/access_token_manager.py | 6 ++++-- .../test_durabletask_grpc_interceptor.py | 7 ++++++- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/durabletask-azuremanaged/durabletask/azuremanaged/internal/access_token_manager.py b/durabletask-azuremanaged/durabletask/azuremanaged/internal/access_token_manager.py index d3d0f96c..fedda9d3 100644 --- a/durabletask-azuremanaged/durabletask/azuremanaged/internal/access_token_manager.py +++ b/durabletask-azuremanaged/durabletask/azuremanaged/internal/access_token_manager.py @@ -86,8 +86,10 @@ def _get_refresh_lock(self) -> asyncio.Lock: if lock is None: # Discard locks belonging to loops that are no longer usable so the # mapping does not grow without bound. - for stale_loop in [ - existing for existing in self._refresh_locks if existing.is_closed()]: + stale_loops = [ + existing for existing in self._refresh_locks if existing.is_closed() + ] + for stale_loop in stale_loops: del self._refresh_locks[stale_loop] lock = asyncio.Lock() self._refresh_locks[loop] = lock diff --git a/tests/durabletask-azuremanaged/test_durabletask_grpc_interceptor.py b/tests/durabletask-azuremanaged/test_durabletask_grpc_interceptor.py index 5545457a..14807f8a 100644 --- a/tests/durabletask-azuremanaged/test_durabletask_grpc_interceptor.py +++ b/tests/durabletask-azuremanaged/test_durabletask_grpc_interceptor.py @@ -266,10 +266,15 @@ async def get_token(self, _scope): class TestAsyncAccessTokenManagerSingleFlight(unittest.IsolatedAsyncioTestCase): - async def test_concurrent_cold_start_performs_single_acquisition(self): + async def test_deferred_first_acquisition_performs_single_acquisition(self): credential = _TestAsyncTokenCredential() manager = AsyncAccessTokenManager(credential) + self.assertEqual(0, credential.calls, "Constructing the manager must not acquire a token") + + # The async manager has always deferred the first acquisition (a constructor + # cannot await), so the cold-start burst goes through the same refresh guard + # as a later refresh and must likewise be single-flight. tokens = await asyncio.gather(*(manager.get_access_token() for _ in range(16))) self.assertEqual(1, credential.calls)