Performance [P1]: Single-flight async Azure token refresh

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
 

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:

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()