Skip to content

Performance [P2]: Defer synchronous Azure token acquisition - #199

Merged
berndverst merged 1 commit into
mainfrom
berndverst-defer-azure-token-acquisition
Jul 27, 2026
Merged

Performance [P2]: Defer synchronous Azure token acquisition#199
berndverst merged 1 commit into
mainfrom
berndverst-defer-azure-token-acquisition

Conversation

@berndverst

Copy link
Copy Markdown
Member

Summary

AccessTokenManager.__init__() acquired an Azure access token eagerly, so constructing a DurableTaskSchedulerClient or DurableTaskSchedulerWorker performed 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. SandboxWorker goes through the same construction path.

Token acquisition is now deferred to the first intercepted RPC:

  • AccessTokenManager.__init__() no longer calls credential.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 what DTSAsyncDefaultClientInterceptorImpl already 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 in durabletask-azuremanaged/CHANGELOG.md under ## 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:

  • Added test_client_construction_does_not_acquire_token — asserts zero credential calls at client construction, exactly one on the first RPC, that the authorization header is present on that request, and that a second request reuses the cached token.
  • Added test_worker_construction_does_not_acquire_token — asserts zero credential calls at worker construction.
  • Added 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.
  • Reworked 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

  • Provider unit tests: pytest tests/durabletask-azuremanaged/test_durabletask_grpc_interceptor.py tests/durabletask-azuremanaged/test_sandboxes_extension.py -q -> 45 passed (baseline 42; +3 new).
  • Core unit tests: pytest tests/durabletask -q --ignore-glob="*_e2e.py" -> 681 passed, 7 skipped. Baseline on this branch had 1 failure in the timing-sensitive test_sync_client_recreate_cooldown_prevents_immediate_repeated_recreation, which is flaky/environmental and passes here.
  • flake8 clean on all four changed files.
  • pyright (repo pyrightconfig.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

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
Copilot AI review requested due to automatic review settings July 27, 2026 06:09

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@berndverst
berndverst merged commit 0f316cc into main Jul 27, 2026
19 checks passed
@berndverst
berndverst deleted the berndverst-defer-azure-token-acquisition branch July 27, 2026 16:18
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Performance [P2]: Defer synchronous Azure token acquisition

3 participants