Merge origin/main into berndverst-cache-azure-managed-sdk-version

Resolves the single conflict in
tests/durabletask-azuremanaged/test_durabletask_grpc_interceptor.py introduced by
PR #198 (single-flight async Azure token refresh), keeping both sides.

Both sides widened an adjacent import into parenthesized form, so git tangled
them around the shared closing parenthesis. Kept every import from both:

- main's `access_token_manager` import widened to `AccessTokenManager` plus
  `AsyncAccessTokenManager`.
- this branch's `durabletask_grpc_interceptor` module import and its
  `DTSAsyncDefaultClientInterceptorImpl` / `DTSDefaultClientInterceptorImpl`
  import.

No test class conflicted: #198 appended its async single-flight classes after
TestAccessTokenManagerThreadSafety, while TestSdkVersionCaching sits before it.
Verified by AST inventory that the merged symbol set equals the union of both
parents exactly - 38 symbols, 38 unique, 0 duplicates.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Copilot-Session: e7499a54-f262-49a6-bf91-274af28543ae

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
@@ -15,7 +16,10 @@
 
 from durabletask.azuremanaged.client import DurableTaskSchedulerClient
 from durabletask.azuremanaged.internal import durabletask_grpc_interceptor
-from durabletask.azuremanaged.internal.access_token_manager import AccessTokenManager
+from durabletask.azuremanaged.internal.access_token_manager import (
+    AccessTokenManager,
+    AsyncAccessTokenManager,
+)
 from durabletask.azuremanaged.internal.durabletask_grpc_interceptor import (
     DTSAsyncDefaultClientInterceptorImpl,
     DTSDefaultClientInterceptorImpl,
@@ -288,5 +292,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()