Serialize U2M token refreshes across processes - #6759
Vivek1106-04 wants to merge 3 commits into
Conversation
Approval status: pending
|
1 similar comment
Approval status: pending
|
|
@simonfaltum ! can you review this . Thank You |
simonfaltum
left a comment
simonfaltum
left a comment
There was a problem hiding this comment.
The shared token cache needs coordination across CLI processes. This covers the ordinary refresh path, but there are still gaps around the lifetime of keyring writes and which refresh token we exchange after taking the lock. I'm requesting changes on those points, the acceptance-test portability issue, and the regression coverage described inline.
This belongs in the CLI. #6540 deliberately moved interactive U2M authentication and persistent token handling out of the SDK, and #6619 consolidated it onto the CLI's storage.Store. The Go and Python SDKs' databricks-cli authentication strategies call databricks auth token --force-refresh, so fixing the CLI covers those consumers too. We don't need an SDK change for this fix.
I think the ownership within the CLI should be: storage supplies coordination associated with its backend and guarantees the lifetime of its writes; U2M controls the complete read-refresh-write transaction. Providing WithTokenStore and WithStoreLock independently makes correctness depend on every caller remembering both. Consider coupling them when addressing the write-lifetime issue. Locking individual Put calls alone wouldn't protect the full transaction.
On the two decisions called out in the description:
- Keeping sequential
--force-refreshcalls as separate exchanges is reasonable. Reusing a replacement written after the caller's initial read fits #6051 without introducing a freshness window. - For the unbounded wait, the important distinction is a dead holder versus a live, slow holder. The inline comment covers the global lock's interaction with endpoint discovery and caller deadlines. The claim that the critical section is just one bounded token exchange needs revisiting.
The guarantee should also stay explicit: this coordinates participating refresh callers. Current storeLoginToken and clearTokenStore operations don't acquire this lock, and older CLI/SDK writers won't participate. Those are existing limitations, not newly introduced races, but we shouldn't describe all cache updates as serialized.
Validation at 2e6a03e, using Codex:
- Targeted
LockTokenStoreandForceRefreshTokentests passed, including with-race. go test ./acceptance -run '^TestAccept/cmd/auth/token' -count=1passed on macOS.- Affected auth packages cross-compiled for Linux/amd64 and Windows/amd64. This wasn't native runtime testing on those platforms.
- Running the new acceptance test with GNU
uniqreproduced the snapshot mismatch in both engine variants. - Removing the command's
WithStoreLockoption still left the new concurrent acceptance test passing for all three repetitions. The mutation was then restored. - Temporary diagnostic tests reproduced an earlier timed-out keyring write overwriting a later lock holder's token, and exchanges using the obsolete refresh token despite rereading an updated one. These used a fake keyring backend and mocked OAuth transport; they didn't reproduce the actual macOS
exit status 45. - In the broader
go test ./libs/auth/... ./cmd/auth/...run, the auth, storage, and command packages passed. U2M stalled in the unchangedTestChallengecallback test and was stopped, so this isn't a claim that the full suite passed.
One minor correction: the rationale in lock_test.go:16-19 says every descriptor in a process shares an flock. Independent opens can contend within the same process; I verified that on macOS. The subprocess contention test is still useful, but that explanation should be corrected.
Please add deterministic coverage of the actual command/credentials wiring and the backend lifetime and cache-reread cases before merging.
| if err != nil { | ||
| return nil, fmt.Errorf("token store lock: %w", err) | ||
| } | ||
| defer unlock() |
There was a problem hiding this comment.
[P2] Keep coordination until the underlying keyring write has actually finished.
Returning from store.Put() doesn't currently mean the keyring operation is done. keyringStore.withTimeout returns after three seconds while leaving the backend operation running. If cache-update recovery also fails, this defer releases the lock while that write is still active. A subsequent holder can then write concurrently, and the earlier operation can finish last and overwrite its token.
I reproduced the stale overwrite with the real LockTokenStore and keyring-store timeout wrapper, using a delayed fake backend. This is a pre-existing backend behavior that makes the new critical section incomplete, rather than a new timeout introduced here.
We need to ensure the backend operation has completed or has been cancelled and joined before releasing coordination, or arrange for lock ownership to outlive the caller's timeout. Please cover this interaction in a regression test.
| func (a *PersistentAuth) cachedRefreshedToken(oldToken *oauth2.Token) *oauth2.Token { | ||
| e, err := a.store.Lookup(a.oAuthArgument.GetCacheKey()) | ||
| if err != nil { | ||
| return nil | ||
| } | ||
| if e.Token.AccessToken == oldToken.AccessToken || needsRefresh(e.Token) { | ||
| return nil |
There was a problem hiding this comment.
[P2] Use the token reread under the lock as the input to any further exchange.
This helper discards the cached entry when its access token needs refreshing or equals the original access token. That entry can still contain a newer refresh token. The caller then falls through to expired := *oldToken and exchanges the refresh token loaded before acquiring the lock.
For rotating refresh tokens, a replacement that's already near expiry, or one with the same access-token string, can therefore lead us to exchange an obsolete refresh token. Mock-transport tests for both cases confirmed that the request sends obsolete-refresh even though the lookup under the lock returned current-refresh.
The reread should give us the authoritative token snapshot, with a separate decision about whether its access token can be returned immediately. If another exchange is needed, use that snapshot's refresh token. Please also propagate reread errors explicitly instead of silently falling back to the earlier snapshot.
| wait | ||
|
|
||
| title "Access tokens returned\n" | ||
| cat token_*.json | jq -r .access_token | sort | uniq -c |
There was a problem hiding this comment.
[P2] Make the count output independent of the platform's uniq.
BSD/macOS uniq -c emits three spaces before 5; GNU uniq -c emits six. The golden file records the BSD form, and the acceptance harness preserves that whitespace, so this fails on Linux.
I reran this acceptance test with GNU uniq on macOS and both engine variants failed with exactly that padding difference. Please normalize the formatting or assert the count in a platform-independent way, then regenerate the snapshot.
| for i in 1 2 3 4 5; do | ||
| $CLI auth token --profile test-profile --force-refresh > "token_$i.json" 2> "err_$i.txt" & | ||
| done | ||
| wait |
There was a problem hiding this comment.
We need a regression test that fails when the production caller stops taking the lock.
I removed WithStoreLock(storage.LockTokenStore) from cmd/auth/token.go and ran this test with -count=3; all repetitions still passed. The fake OIDC server accepts repeated exchanges and always returns the same access token, and the plaintext backend doesn't model keyring write collisions. Five successful outputs therefore don't establish that coordination happened.
The separate lock and reuse unit tests are useful, but they don't exercise the production wiring. Please add a deterministic test using a controlled lock holder or a backend that detects overlapping operations, so removing the lock option from the real caller makes it fail. We also need behavioral coverage for the CLICredentials path; counting options alone won't catch passing the wrong option.
| // It blocks until the lock is available or ctx is done. There is no timeout: the | ||
| // operating system releases the lock when a holder exits, including on a crash, | ||
| // so a lock cannot be left behind by a dead process. |
There was a problem hiding this comment.
Design concern: crash cleanup doesn't bound the time spent behind a live holder.
This is one lock for every profile and both storage modes. It's held across oauth2Config(), which can discover workspace/unified OAuth endpoints using the SDK client's default five-minute retry budget, as well as the token exchange and storage operations. One unavailable workspace can therefore delay refreshes for unrelated healthy workspaces.
The auth token path does supply a deadline, but its default is one hour; CLICredentials doesn't add an equivalent lock-wait deadline. So the description's argument that the critical section is a single exchange bounded by the HTTP timeout isn't sufficient.
Can we make the time bounds explicit and keep avoidable discovery work outside the shared critical section? If user-wide serialization is intentional, please document that tradeoff and test cancellation while a different profile holds the lock.
Take an advisory lock on ~/.databricks/token-cache.lock around the read-refresh-write sequence, and re-read the store once it is held so a token another process just exchanged is reused instead of exchanged again.
- Add Lock to storage.Store so every backend supplies its own coordination and callers can no longer pass a store without its lock. File and keyring stores take the token-cache lock, the memory store is a no-op, and the hint and dual-writing wrappers forward to their inner store. WithStoreLock is removed. - Keep the keyring lock held until any timed-out backend call has returned, so a late write cannot overwrite the next holder's token. - Treat the token re-read under the lock as authoritative: reuse its access token when fresh, otherwise exchange its refresh token, and return re-read errors instead of falling back to the earlier snapshot. - Run endpoint discovery before taking the lock and bound the wait behind a live holder to one minute. The lock stays global because the plaintext backend rewrites the whole token-cache.json on every Put. - Make the acceptance count independent of the platform's uniq. - Add regression tests for the command and CLICredentials wiring, the keyring write lifetime, the re-read cases, and lock timeouts.
2e6a03e to
95e5a3b
Compare
|
An authorized user can trigger integration tests manually by following the instructions below: Trigger: Inputs:
Checks will be approved automatically on success. |
|
@simonfaltum thanks for the thorough review and for verifying each point. I've rebased onto Store/lock ownership. I took your suggestion to couple them. Keyring write lifetime ( Refresh token read under the lock (
Regression coverage of the wiring. Since the caller no longer passes the lock, the mutation to guard against moved into Time bounds ( I kept one lock for all profiles on purpose. The plaintext backend rewrites the whole Scope of the guarantee. The description now says this coordinates refresh callers only.
Validation: |
This branch has not been deployed
2 participants
Changes
Adds cross-process coordination around the U2M read-refresh-write sequence, taking up the TODO in
PersistentAuth.refresh.storage.StoregainsLock(ctx), so each backend supplies its own coordination and a caller cannot pass a store without its lock. The file and keyring stores take an advisory lock on~/.databricks/token-cache.lock(flockon Unix,LockFileExon Windows). The memory store returns a no-op. The not-found-hint and dual-writing wrappers forward to their inner store.PersistentAuth.refreshruns endpoint discovery first, then takes the store's lock and re-reads the cache. The re-read token is authoritative: its access token is reused if another process replaced it and it is fresh; otherwise its refresh token (which may have been rotated) is the one exchanged. A re-read error is returned rather than falling back to the earlier snapshot.withTimeoutreturns after 3 seconds but leaves the call running, and a write finishing after release could otherwise overwrite the next holder's token.ctxcancellation still applies. A dead holder cannot leave the lock behind: the OS releases it on exit.Why
Fixes #6051.
The CLI is stateless, so two invocations for the same profile load the same cached refresh token, both exchange it with the IdP, and race to write the result back. Since the SDKs began appending
--force-refresh, the race fails hard:ForceRefreshTokenhas no fallback to the cached token by design, so a losing writer surfaces asforced token refresh: cache update: exit status 45from the macOS keyring.Design notes:
token-cache.jsonon everyPut, guarded only by an in-process mutex. Per-profile locks would let refreshes of different profiles in different processes overwrite each other's entries. The cost is that a slow token endpoint for one workspace delays refreshes for others, bounded by the HTTP timeout. Endpoint discovery, which can retry for minutes, runs outside the lock.--force-refreshcalls remain separate exchanges. The re-read only short-circuits a process whose token was replaced while it waited.auth login(storeLoginToken),auth logout(clearTokenStore), and older CLI or SDK writers do not take the lock; those are existing limitations, andrecoverStoreUpdatestays in place for them./usr/bin/securitychild process can still complete the write afterwards. That needs a >3 s keyring stall plus a failed cache-update recovery, and can't be closed without cancellation support in go-keyring.Tests
libs/auth/storage/lock_test.go: lock file creation, release,ctxcancellation, the wait timeout, contention with a re-executed test binary, that every shared store (file, keyring, and both as wrapped for U2M) takes the lock, and that the keyring lock is held until a timed-out write returns.libs/auth/u2m/persistent_auth_test.go: the refresh token read under the lock is the one exchanged (replacement near expiry, and same access token with a rotated refresh token); a re-read error is returned; discovery runs before the lock; reuse without an exchange; lock errors surface.cmd/auth/token_test.goandlibs/auth/credentials_test.go: through the real command andCLICredentialswiring with a file store, a refresh waits while another holder has the lock.acceptance/cmd/auth/token/force-refresh-concurrent: five concurrentauth token --force-refreshinvocations all succeed and leave a valid cache entry. The fake OIDC server accepts repeated exchanges, so this is a smoke test; the unit tests above cover the lock.store.Lockinrefresh; exchanging the pre-lock refresh token; ignoring a re-read error; running discovery under the lock; releasing the keyring lock immediately; a no-op file-store lock; the hint wrapper dropping the lock.go test -race ./libs/auth/... ./cmd/auth/...,go test ./acceptance -run 'TestAccept/cmd/auth',GOOS=windows/GOOS=linuxvet,./task fmt-q,./task lint-q(0 issues),./task ws.