Performance [P2]: Defer synchronous Azure token acquisition

durabletask-azuremanaged/CHANGELOG.md

@@ -7,6 +7,12 @@ adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## Unreleased
 
+- `DurableTaskSchedulerClient` and `DurableTaskSchedulerWorker` no longer block on an Azure
+  credential round trip while being constructed. The access token is now acquired on the first
+  request instead, so constructing a client or worker that is never used costs nothing. As a
+  result, credential failures (for example an unavailable managed identity or a misconfigured
+  `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.
 
 ## v1.8.0

durabletask-azuremanaged/durabletask/azuremanaged/internal/access_token_manager.py

@@ -13,6 +13,7 @@
 class AccessTokenManager:
 
     _token: AccessToken | None
+    expiry_time: datetime | None
 
     def __init__(self, token_credential: TokenCredential | None, refresh_interval_seconds: int = 600):
         self._scope = "https://durabletask.io/.default"
@@ -22,12 +23,12 @@ def __init__(self, token_credential: TokenCredential | None, refresh_interval_se
         self._credential = token_credential
         self._refresh_lock = Lock()
 
-        if self._credential is not None:
-            self._token = self._credential.get_token(self._scope)
-            self.expiry_time = datetime.fromtimestamp(self._token.expires_on, tz=timezone.utc)
-        else:
-            self._token = None
-            self.expiry_time = None
+        # Token acquisition is deferred to the first get_access_token() call so that
+        # constructing a client or worker does not perform a blocking credential round
+        # trip. The deferred first acquisition still goes through the double-checked
+        # refresh lock below, so it remains single-flight across threads.
+        self._token = None
+        self.expiry_time = None
 
     def get_access_token(self) -> AccessToken | None:
         if self._token is None or self.is_token_expired():

durabletask-azuremanaged/durabletask/azuremanaged/internal/durabletask_grpc_interceptor.py

@@ -41,13 +41,13 @@ def __init__(
             self._metadata.append(("workerid", worker_id))
         super().__init__(self._metadata)
 
+        # Token acquisition is deferred to the first _intercept_call invocation rather
+        # than happening in __init__, so that constructing a client or worker does not
+        # block on a credential round trip before any RPC is made.
         self._token_manager = None
         if token_credential is not None:
             self._token_credential = token_credential
             self._token_manager = AccessTokenManager(token_credential=self._token_credential)
-            access_token = self._token_manager.get_access_token()
-            if access_token is not None:
-                self._upsert_authorization_header(access_token.token)
 
     def _upsert_authorization_header(self, token: str) -> None:
         found = False

tests/durabletask-azuremanaged/test_durabletask_grpc_interceptor.py

@@ -7,6 +7,7 @@
 from importlib.metadata import version
 import threading
 import time
+from typing import Any
 
 import grpc
 from azure.core.credentials import AccessToken
@@ -39,6 +40,21 @@ def GetInstance(self, request, context):
         return response
 
 
+class _TestTokenCredential:
+    """Minimal TokenCredential stub that counts how often a token is requested."""
+
+    def __init__(self):
+        self._lock = threading.Lock()
+        self.calls = 0
+
+    def get_token(self, *scopes: str, **kwargs: Any) -> AccessToken:
+        with self._lock:
+            self.calls += 1
+            call_number = self.calls
+        time.sleep(0.02)
+        return AccessToken(f"token-{call_number}", int(time.time()) + 3600)
+
+
 class TestDurableTaskGrpcInterceptor(unittest.TestCase):
     """Tests for the DTSDefaultClientInterceptorImpl class."""
 
@@ -140,34 +156,76 @@ def test_worker_includes_workerid_header(self):
         self.assertIn("workerid", metadata)
         self.assertTrue(metadata["workerid"])
 
+    def test_client_construction_does_not_acquire_token(self):
+        """Token acquisition is deferred from construction to the first request."""
+        credential = _TestTokenCredential()
 
-class _TestTokenCredential:
-    def __init__(self):
-        self._lock = threading.Lock()
-        self.calls = 0
+        task_hub_client = DurableTaskSchedulerClient(
+            host_address=self.server_address,
+            secure_channel=False,
+            taskhub="test-taskhub",
+            token_credential=credential,
+        )
 
-    def get_token(self, _scope):
-        with self._lock:
-            self.calls += 1
-            call_number = self.calls
-        time.sleep(0.02)
-        return AccessToken(f"token-{call_number}", int(time.time()) + 3600)
+        self.assertEqual(0, credential.calls, "Constructing a client must not acquire a token")
+
+        task_hub_client.get_orchestration_state("test-instance-id")
+
+        self.assertEqual(1, credential.calls, "The first request should acquire exactly one token")
+        self.assertIn("authorization", self.mock_servicer.captured_metadata)
+
+        task_hub_client.get_orchestration_state("test-instance-id")
+
+        self.assertEqual(1, credential.calls, "A cached, unexpired token should be reused")
+        self.assertEqual(2, self.mock_servicer.requests_received)
+
+    def test_worker_construction_does_not_acquire_token(self):
+        """Token acquisition is deferred from worker construction to the first request."""
+        credential = _TestTokenCredential()
+
+        DurableTaskSchedulerWorker(
+            host_address=self.server_address,
+            secure_channel=False,
+            taskhub="test-taskhub",
+            token_credential=credential,
+        )
+
+        self.assertEqual(0, credential.calls, "Constructing a worker must not acquire a token")
 
 
 class TestAccessTokenManagerThreadSafety(unittest.TestCase):
+
+    @staticmethod
+    def _get_access_token_concurrently(manager: AccessTokenManager, thread_count: int = 8) -> None:
+        threads = [threading.Thread(target=manager.get_access_token) for _ in range(thread_count)]
+        for thread in threads:
+            thread.start()
+        for thread in threads:
+            thread.join()
+
+    def test_deferred_first_acquisition_performs_single_acquisition(self):
+        credential = _TestTokenCredential()
+        manager = AccessTokenManager(credential)
+
+        self.assertEqual(0, credential.calls, "Constructing the manager must not acquire a token")
+
+        self._get_access_token_concurrently(manager)
+
+        self.assertEqual(1, credential.calls)
+        token = manager.get_access_token()
+        self.assertIsNotNone(token)
+        self.assertEqual("token-1", token.token)  # type: ignore[union-attr]
+        self.assertEqual(1, credential.calls, "A cached, unexpired token should be reused")
+
     def test_concurrent_refresh_performs_single_refresh(self):
         credential = _TestTokenCredential()
         manager = AccessTokenManager(credential)
-        manager.expiry_time = datetime.now(timezone.utc) - timedelta(seconds=1)
+        manager.get_access_token()
+        self.assertEqual(1, credential.calls)
 
-        threads = []
-        for _ in range(8):
-            thread = threading.Thread(target=manager.get_access_token)
-            thread.start()
-            threads.append(thread)
+        manager.expiry_time = datetime.now(timezone.utc) - timedelta(seconds=1)
 
-        for thread in threads:
-            thread.join()
+        self._get_access_token_concurrently(manager)
 
         self.assertEqual(2, credential.calls)