diff --git a/durabletask-azuremanaged/CHANGELOG.md b/durabletask-azuremanaged/CHANGELOG.md index 2a5b7b22..27ae8a9e 100644 --- a/durabletask-azuremanaged/CHANGELOG.md +++ b/durabletask-azuremanaged/CHANGELOG.md @@ -14,6 +14,9 @@ adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). `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. +- 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 63daf86a..fedda9d3 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 @@ -71,9 +72,37 @@ 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. + 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 + 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 20162a5d..14807f8a 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 @@ -13,7 +14,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 @@ -230,5 +234,108 @@ 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_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) + 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()