Merge remote-tracking branch 'origin/main' into berndverst-index-sandbox-activity-overlap-validatio

# Conflicts:
#	tests/durabletask-azuremanaged/test_sandboxes_extension.py

Author: Bernd Verst

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:

durabletask-azuremanaged/durabletask/azuremanaged/preview/sandboxes/worker.py

@@ -130,30 +130,50 @@ def _stop_sandbox_registration(self) -> None:
 
     def _run_sandbox_registration_loop(self) -> None:
         retry_delay = 1.0
-        while not self._sandbox_registration_stop.is_set():
-            try:
-                client = SandboxActivitiesGrpcTransport(
-                    host_address=self._sandbox_host_address,
-                    taskhub=self._sandbox_taskhub,
-                    token_credential=self._sandbox_token_credential,
-                    secure_channel=self._sandbox_secure_channel)
+        client: Optional[SandboxActivitiesGrpcTransport] = None
+        try:
+            while not self._sandbox_registration_stop.is_set():
                 try:
+                    if client is None:
+                        client = SandboxActivitiesGrpcTransport(
+                            host_address=self._sandbox_host_address,
+                            taskhub=self._sandbox_taskhub,
+                            token_credential=self._sandbox_token_credential,
+                            secure_channel=self._sandbox_secure_channel)
                     client.connect_sandbox_activity_worker(self._registration_messages())
                     retry_delay = 1.0
-                finally:
-                    client.close()
-            except Exception as ex:
-                if self._sandbox_registration_stop.is_set():
-                    break
-                if not _is_retriable_registration_failure(ex):
-                    self._sandbox_logger.error(
-                        "Sandbox activity worker registration failed permanently: %s", ex)
-                    self._sandbox_registration_stop.set()
-                    break
-                self._sandbox_logger.warning("Sandbox activity worker registration failed: %s", ex)
-                delay = random.uniform(0, retry_delay)
-                self._sandbox_registration_stop.wait(delay)
-                retry_delay = min(retry_delay * 2, 30.0)
+                except Exception as ex:
+                    # gRPC channels reconnect on their own after transient stream
+                    # failures, so the transport (and its channel, stub, and cached
+                    # access token) is reused instead of being rebuilt per attempt.
+                    if _requires_new_registration_transport(ex):
+                        self._close_sandbox_registration_transport(client)
+                        client = None
+                    if self._sandbox_registration_stop.is_set():
+                        break
+                    if not _is_retriable_registration_failure(ex):
+                        self._sandbox_logger.error(
+                            "Sandbox activity worker registration failed permanently: %s", ex)
+                        self._sandbox_registration_stop.set()
+                        break
+                    self._sandbox_logger.warning(
+                        "Sandbox activity worker registration failed: %s", ex)
+                    delay = random.uniform(0, retry_delay)
+                    self._sandbox_registration_stop.wait(delay)
+                    retry_delay = min(retry_delay * 2, 30.0)
+        finally:
+            self._close_sandbox_registration_transport(client)
+
+    def _close_sandbox_registration_transport(
+            self,
+            client: Optional[SandboxActivitiesGrpcTransport]) -> None:
+        if client is None:
+            return
+        try:
+            client.close()
+        except Exception as ex:
+            self._sandbox_logger.debug(
+                "Sandbox activity worker registration transport failed to close: %s", ex)
 
     def _registration_messages(self) -> Iterator[pb.SandboxActivityWorkerMessage]:
         yield build_sandbox_worker_start(
@@ -264,6 +284,26 @@ def _is_retriable_registration_failure(ex: Exception) -> bool:
     return isinstance(ex, OSError)
 
 
+def _requires_new_registration_transport(ex: Exception) -> bool:
+    """Returns whether a failure leaves the registration transport unusable.
+
+    A gRPC channel transparently reconnects after a stream fails, so the same
+    transport can serve the next registration attempt. A channel that has been
+    shut down rejects every later RPC, and a failure raised outside of gRPC (for
+    example while the channel itself is being created) can leave the transport
+    only partially initialized, so both cases need a replacement transport.
+    """
+    if not isinstance(ex, grpc.RpcError):
+        return True
+
+    if ex.code() != grpc.StatusCode.CANCELLED:
+        return False
+
+    details = getattr(ex, "details", None)
+    resolved_details = details() if callable(details) else None
+    return isinstance(resolved_details, str) and "channel closed" in resolved_details.lower()
+
+
 def _resolve_max_concurrent_activities() -> int:
     value = os.getenv("DTS_SANDBOX_MAX_ACTIVITIES")
     if value is None:

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