Performance [P2]: Defer synchronous Azure token acquisition - #199
Merged
Conversation
AccessTokenManager.__init__ acquired an access token eagerly, which made constructing a DurableTaskSchedulerClient or DurableTaskSchedulerWorker perform a blocking credential round trip to Entra before any RPC was made. DTSDefaultClientInterceptorImpl.__init__ then called get_access_token() immediately, so the cost was paid even for instances that were never used. Token acquisition is now deferred to the first intercepted RPC. The existing double-checked refresh lock in get_access_token() is unchanged, so the deferred first acquisition remains single-flight across threads, and the refresh-before-expiry behavior is unaffected. This changes when credential failures surface: they now raise from the first request instead of from the constructor. The exception type and message are unchanged. This is documented in the azuremanaged changelog. Fixes #187 Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 27f463ce-6eb3-45c0-9bc0-63c68f3660ad
Contributor
There was a problem hiding this comment.
Pull request overview
This PR improves cold-start performance for the Azure Managed provider by deferring synchronous Azure access token acquisition until the first intercepted gRPC call, instead of doing a blocking credential round trip during client/worker construction.
Changes:
- Deferred initial token acquisition in
AccessTokenManager(sync) by initializing with no token/expiry and relying on the existing double-checked lock path on first use. - Removed eager token “priming” in
DTSDefaultClientInterceptorImpl.__init__(); authorization headers are still upserted per-call in_intercept_call(). - Added/updated unit tests to assert zero credential calls during construction and single-flight behavior on the first concurrent acquisition, and documented the timing-only behavior change in the provider changelog.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated no comments.
| File | Description |
|---|---|
| tests/durabletask-azuremanaged/test_durabletask_grpc_interceptor.py | Adds tests ensuring token acquisition is deferred and remains single-flight under concurrency. |
| durabletask-azuremanaged/durabletask/azuremanaged/internal/durabletask_grpc_interceptor.py | Removes eager token priming; relies on per-RPC interception to acquire/upsert authorization. |
| durabletask-azuremanaged/durabletask/azuremanaged/internal/access_token_manager.py | Defers initial token acquisition while preserving existing refresh/locking behavior. |
| durabletask-azuremanaged/CHANGELOG.md | Documents the timing-only behavior change under ## Unreleased. |
andystaples
approved these changes
Jul 27, 2026
berndverst
pushed a commit
that referenced
this pull request
Jul 27, 2026
Resolves conflicts in tests/durabletask-azuremanaged/test_durabletask_grpc_interceptor.py introduced by PR #199 (deferred synchronous token acquisition), keeping both sides: - Imports: kept both main's `from typing import Any` and this branch's `from unittest import mock` / `PackageNotFoundError`. - Test classes: kept main's relocated `_TestTokenCredential` and its new deferred-token tests verbatim; moved this branch's `TestSdkVersionCaching` to sit between `TestDurableTaskGrpcInterceptor` and `TestAccessTokenManagerThreadSafety`. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: e7499a54-f262-49a6-bf91-274af28543ae
berndverst
pushed a commit
that referenced
this pull request
Jul 27, 2026
Address review feedback on the stale-lock cleanup in AsyncAccessTokenManager._get_refresh_lock(). The list comprehension was embedded directly in the `for` statement, which read awkwardly across the wrapped line. Hoist it into a named `stale_loops` local and iterate over that. This is purely a readability change with no behavioural difference: the snapshot is still materialized before mutating the mapping, which is what makes deleting during iteration safe. Note this was not a lint violation - flake8 passes on the file both before and after the change. Also align the async deferred-acquisition test with the naming and assertions that PR #199 introduced for the synchronous manager, and assert explicitly that constructing AsyncAccessTokenManager performs no credential call. The async manager has always deferred its first acquisition, so the cold-start burst flows through the same refresh guard as a later refresh; the test now states that intent directly. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: bb08d0ed-484d-4015-b905-3e164f68678b
berndverst
added a commit
that referenced
this pull request
Jul 27, 2026
* perf: single-flight async Azure token refresh AsyncAccessTokenManager.get_access_token() checked expiry and awaited refresh_token() without any synchronization. Since the async DTS gRPC interceptor awaits it on every RPC, a burst of concurrent grpc.aio calls arriving while the token was missing or inside its renewal window would each independently call the credential, producing a thundering herd of Entra/credential requests, extra latency, and throttling pressure. Mirror the double-checked locking already used by the synchronous AccessTokenManager: re-check token validity after acquiring a refresh guard so only one coroutine contacts the credential per refresh window while the rest reuse the result. The guard is an asyncio.Lock created lazily per running event loop (tracked in a small map behind a threading.Lock), because an asyncio.Lock binds to the loop it is first awaited on and this manager may outlive a loop or be shared across loops. Locks for closed loops are pruned so the map cannot grow unbounded. The cached-token fast path is unchanged and never touches the lock, so there is no added per-RPC overhead. Public signatures, return values, the renewal margin behavior, exception propagation, and logging are all preserved. Adds unit tests covering concurrent cold start, concurrent refresh after expiry, the cached-token path, credential failure propagation without deadlock, and reuse across distinct event loops. Fixes #192 Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: bb08d0ed-484d-4015-b905-3e164f68678b * refactor: hoist stale-loop comprehension into a named local Address review feedback on the stale-lock cleanup in AsyncAccessTokenManager._get_refresh_lock(). The list comprehension was embedded directly in the `for` statement, which read awkwardly across the wrapped line. Hoist it into a named `stale_loops` local and iterate over that. This is purely a readability change with no behavioural difference: the snapshot is still materialized before mutating the mapping, which is what makes deleting during iteration safe. Note this was not a lint violation - flake8 passes on the file both before and after the change. Also align the async deferred-acquisition test with the naming and assertions that PR #199 introduced for the synchronous manager, and assert explicitly that constructing AsyncAccessTokenManager performs no credential call. The async manager has always deferred its first acquisition, so the cold-start burst flows through the same refresh guard as a later refresh; the test now states that intent directly. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: bb08d0ed-484d-4015-b905-3e164f68678b --------- Co-authored-by: Bernd Verst <beverst@microsoft.com> Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: bb08d0ed-484d-4015-b905-3e164f68678b
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
AccessTokenManager.__init__()acquired an Azure access token eagerly, so constructing aDurableTaskSchedulerClientorDurableTaskSchedulerWorkerperformed a blocking credential round trip to Entra before any RPC was made.DTSDefaultClientInterceptorImpl.__init__()then calledget_access_token()immediately, so the cost was paid even for instances that were never used.SandboxWorkergoes through the same construction path.Token acquisition is now deferred to the first intercepted RPC:
AccessTokenManager.__init__()no longer callscredential.get_token(); it starts with_token = None/expiry_time = None.DTSDefaultClientInterceptorImpl.__init__()no longer primes the token or the authorization header.The existing double-checked refresh lock in
get_access_token()is untouched, so the deferred first acquisition is still single-flight across threads, and refresh-before-expiry behavior is unchanged._intercept_call()already upserted the authorization header on every call, so no header is missing on the first request. This mirrors whatDTSAsyncDefaultClientInterceptorImplalready does.Behavior change (timing only)
Credential failures — for example an unavailable managed identity or a misconfigured
DefaultAzureCredential— now surface from the first request instead of from the constructor. The exception type and message are unchanged; only the timing differs. This is called out indurabletask-azuremanaged/CHANGELOG.mdunder## Unreleased.No public API changed: no signature changes, no renames, no removed symbols, no changed import paths. No warm-up hook was added, to avoid expanding public surface — callers that want fail-fast authentication can call
credential.get_token(...)themselves before constructing a client or worker.Tests
tests/durabletask-azuremanaged/test_durabletask_grpc_interceptor.py:test_client_construction_does_not_acquire_token— asserts zero credential calls at client construction, exactly one on the first RPC, that theauthorizationheader is present on that request, and that a second request reuses the cached token.test_worker_construction_does_not_acquire_token— asserts zero credential calls at worker construction.test_deferred_first_acquisition_performs_single_acquisition— 8 concurrent threads on a freshly constructed manager result in exactly one credential call, proving the deferred first acquisition is single-flight.test_concurrent_refresh_performs_single_refresh— the single-flight-on-expiry assertion is preserved, with the token now primed explicitly instead of relying on constructor acquisition.All three new tests were verified to fail against the unfixed source and pass with the fix.
Verification
pytest tests/durabletask-azuremanaged/test_durabletask_grpc_interceptor.py tests/durabletask-azuremanaged/test_sandboxes_extension.py -q-> 45 passed (baseline 42; +3 new).pytest tests/durabletask -q --ignore-glob="*_e2e.py"-> 681 passed, 7 skipped. Baseline on this branch had 1 failure in the timing-sensitivetest_sync_client_recreate_cooldown_prevents_immediate_repeated_recreation, which is flaky/environmental and passes here.flake8clean on all four changed files.pyright(repopyrightconfig.json, strict): 0 diagnostics in the changed source files.pymarkdown -c .pymarkdown.json scan durabletask-azuremanaged/CHANGELOG.md: the two MD013 line-length findings are pre-existing on untouched lines; the added lines introduce none.Scope was kept strictly to this issue — no changes related to #192 (async single-flight refresh) or #195 (SDK version lookup caching), which live in the same files.
Fixes #187