diff --git a/.nextchanges/cli/token-store-cross-process-lock.md b/.nextchanges/cli/token-store-cross-process-lock.md new file mode 100644 index 00000000000..4a9b055a303 --- /dev/null +++ b/.nextchanges/cli/token-store-cross-process-lock.md @@ -0,0 +1 @@ +* Serialize concurrent U2M token refreshes across CLI invocations so parallel `databricks auth token --force-refresh` calls no longer fail on a contended token cache. ([#6759](https://github.com/databricks/cli/pull/6759)) diff --git a/acceptance/cmd/auth/token/force-refresh-concurrent/out.test.toml b/acceptance/cmd/auth/token/force-refresh-concurrent/out.test.toml new file mode 100644 index 00000000000..98ea5040486 --- /dev/null +++ b/acceptance/cmd/auth/token/force-refresh-concurrent/out.test.toml @@ -0,0 +1,2 @@ +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"] diff --git a/acceptance/cmd/auth/token/force-refresh-concurrent/output.txt b/acceptance/cmd/auth/token/force-refresh-concurrent/output.txt new file mode 100644 index 00000000000..81e3ad031f9 --- /dev/null +++ b/acceptance/cmd/auth/token/force-refresh-concurrent/output.txt @@ -0,0 +1,10 @@ + +=== Access tokens returned + 5 oauth-token + +=== +Errors + +=== +Cached token after the refreshes +oauth-token diff --git a/acceptance/cmd/auth/token/force-refresh-concurrent/script b/acceptance/cmd/auth/token/force-refresh-concurrent/script new file mode 100644 index 00000000000..9d39703a3c3 --- /dev/null +++ b/acceptance/cmd/auth/token/force-refresh-concurrent/script @@ -0,0 +1,19 @@ +setup_test_profile +setup_test_token_cache + +# Several processes ask for a forced refresh of the same profile at once. They +# all read the same cached refresh token, so the refresh and the cache write +# have to be serialized across processes for every invocation to succeed. +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 + +title "Access tokens returned\n" +cat token_*.json | jq -r .access_token | sort | uniq -c + +title "\nErrors\n" +cat err_*.txt + +title "\nCached token after the refreshes\n" +jq -r '.tokens["test-profile"].access_token' ./home/.databricks/token-cache.json diff --git a/acceptance/cmd/auth/token/force-refresh-concurrent/test.toml b/acceptance/cmd/auth/token/force-refresh-concurrent/test.toml new file mode 100644 index 00000000000..034cfa2116d --- /dev/null +++ b/acceptance/cmd/auth/token/force-refresh-concurrent/test.toml @@ -0,0 +1,12 @@ +Ignore = [ + "token_1.json", + "token_2.json", + "token_3.json", + "token_4.json", + "token_5.json", + "err_1.txt", + "err_2.txt", + "err_3.txt", + "err_4.txt", + "err_5.txt", +] diff --git a/cmd/auth/token.go b/cmd/auth/token.go index 7bf710de575..736c7fca3ff 100644 --- a/cmd/auth/token.go +++ b/cmd/auth/token.go @@ -287,7 +287,10 @@ func loadToken(ctx context.Context, args loadTokenArgs) (*oauth2.Token, error) { if err != nil { return nil, err } - allArgs := append([]u2m.PersistentAuthOption{u2m.WithTokenStore(storage.OAuthTokenStore(ctx, args.tokenStore, args.mode))}, args.persistentAuthOpts...) + allArgs := append([]u2m.PersistentAuthOption{ + u2m.WithTokenStore(storage.OAuthTokenStore(ctx, args.tokenStore, args.mode)), + u2m.WithStoreLock(storage.LockTokenStore), + }, args.persistentAuthOpts...) if clientID := u2mClientIDFromProfile(existingProfile); clientID != "" { allArgs = append(allArgs, u2m.WithClientID(clientID)) } diff --git a/libs/auth/credentials.go b/libs/auth/credentials.go index 98278fb53ef..f79d86a826e 100644 --- a/libs/auth/credentials.go +++ b/libs/auth/credentials.go @@ -108,6 +108,7 @@ func (c CLICredentials) Configure(ctx context.Context, cfg *config.Config) (cred opts := []u2m.PersistentAuthOption{ u2m.WithOAuthArgument(oauthArg), u2m.WithTokenStore(storage.OAuthTokenStore(ctx, tokenStore, mode)), + u2m.WithStoreLock(storage.LockTokenStore), } if cfg.AuthType == c.Name() && cfg.ClientID != "" { opts = append(opts, u2m.WithClientID(cfg.ClientID)) diff --git a/libs/auth/credentials_test.go b/libs/auth/credentials_test.go index 2802c9b4279..c3ab4fc65d9 100644 --- a/libs/auth/credentials_test.go +++ b/libs/auth/credentials_test.go @@ -225,10 +225,10 @@ func TestCLICredentialsConfigure_ThreadsResolvedTokenStore(t *testing.T) { _, err := c.Configure(t.Context(), &config.Config{Host: "https://x.cloud.databricks.com"}) require.NoError(t, err) - // Two opts expected: WithOAuthArgument and WithTokenStore. The length - // check is the most resilient way to assert both were passed without - // poking at u2m's unexported state. - assert.Len(t, receivedOpts, 2) + // Three opts expected: WithOAuthArgument, WithTokenStore and + // WithStoreLock. The length check is the most resilient way to assert they + // were passed without poking at u2m's unexported state. + assert.Len(t, receivedOpts, 3) } func TestCLICredentialsConfigure_ClientID(t *testing.T) { @@ -242,25 +242,25 @@ func TestCLICredentialsConfigure_ClientID(t *testing.T) { name: "U2M config file client ID", authType: "databricks-cli", source: config.SourceFile, - wantOpts: 3, + wantOpts: 4, }, { name: "U2M environment client ID", authType: "databricks-cli", source: config.SourceEnv, - wantOpts: 3, + wantOpts: 4, }, { name: "U2M dynamic client ID", authType: "databricks-cli", source: config.SourceDynamicConfig, - wantOpts: 3, + wantOpts: 4, }, { name: "non-U2M config file client ID", authType: "oauth-m2m", source: config.SourceFile, - wantOpts: 2, + wantOpts: 3, }, } @@ -341,7 +341,7 @@ func TestCLICredentialsConfigure_HonorsConfigFileSecureMode(t *testing.T) { // The presence of the second opt is verified by the sibling // test; here we just need Configure to succeed end-to-end when // the config file selects secure storage. - assert.Len(t, opts, 2) + assert.Len(t, opts, 3) return auth.TokenSourceFn(func(_ context.Context) (*oauth2.Token, error) { return &oauth2.Token{AccessToken: "tok"}, nil }), nil diff --git a/libs/auth/storage/lock.go b/libs/auth/storage/lock.go new file mode 100644 index 00000000000..a95211cbb54 --- /dev/null +++ b/libs/auth/storage/lock.go @@ -0,0 +1,76 @@ +package storage + +import ( + "context" + "fmt" + "os" + "path/filepath" + "time" + + "github.com/databricks/cli/libs/env" +) + +const ( + // tokenStoreLockFilePath is the path of the lock file guarding token store + // updates, relative to the user's home directory. It sits next to + // token-cache.json and is used by the keyring backend too, which has no + // file of its own to lock. + tokenStoreLockFilePath = ".databricks/token-cache.lock" + + // lockRetryInterval is how long to wait before retrying a contended lock. + // The critical section is a single OAuth token exchange, so a short poll + // keeps the wait close to the holder's actual runtime. + lockRetryInterval = 20 * time.Millisecond +) + +// LockTokenStore acquires the cross-process lock that serializes token store +// refreshes, and returns a function that releases it. +// +// The CLI is stateless, so two invocations for the same profile otherwise load +// the same cached refresh token, both exchange it, and race to write the +// result back. The lock is advisory and held only around the read-refresh-write +// sequence. +// +// 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. +func LockTokenStore(ctx context.Context) (func(), error) { + home, err := env.UserHomeDir(ctx) + if err != nil { + return nil, fmt.Errorf("failed loading home directory: %w", err) + } + path := filepath.Join(home, tokenStoreLockFilePath) + if err := os.MkdirAll(filepath.Dir(path), ownerExecReadWrite); err != nil { + return nil, fmt.Errorf("mkdir: %w", err) + } + + f, err := os.OpenFile(path, os.O_CREATE|os.O_RDWR, ownerReadWrite) + if err != nil { + return nil, fmt.Errorf("open lock file: %w", err) + } + + for { + locked, err := tryLock(f) + if err != nil { + f.Close() + return nil, fmt.Errorf("lock %s: %w", path, err) + } + if locked { + return func() { + // Closing the file releases the lock on both platforms; the + // explicit unlock keeps the two operations from drifting apart. + unlock(f) + f.Close() + }, nil + } + + timer := time.NewTimer(lockRetryInterval) + select { + case <-ctx.Done(): + timer.Stop() + f.Close() + return nil, ctx.Err() + case <-timer.C: + } + } +} diff --git a/libs/auth/storage/lock_test.go b/libs/auth/storage/lock_test.go new file mode 100644 index 00000000000..4120fd1c0c7 --- /dev/null +++ b/libs/auth/storage/lock_test.go @@ -0,0 +1,96 @@ +package storage + +import ( + "bufio" + "context" + "os" + "os/exec" + "path/filepath" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// holdLockEnvVar makes the test binary re-execute itself as a lock holder, so +// the contention is between two processes. An flock is shared by every +// descriptor in the process that took it, so a second in-process acquisition +// would succeed and prove nothing. +const holdLockEnvVar = "DATABRICKS_TEST_HOLD_TOKEN_STORE_LOCK" + +func setHome(t *testing.T, dir string) { + t.Helper() + t.Setenv("HOME", dir) + t.Setenv("USERPROFILE", dir) +} + +func TestLockTokenStoreCreatesLockFile(t *testing.T) { + home := t.TempDir() + setHome(t, home) + + unlock, err := LockTokenStore(t.Context()) + require.NoError(t, err) + defer unlock() + + assert.FileExists(t, filepath.Join(home, tokenStoreLockFilePath)) +} + +func TestLockTokenStoreIsReleasedByUnlock(t *testing.T) { + home := t.TempDir() + setHome(t, home) + + unlock, err := LockTokenStore(t.Context()) + require.NoError(t, err) + unlock() + + unlock, err = LockTokenStore(t.Context()) + require.NoError(t, err) + unlock() +} + +func TestLockTokenStoreWaitsForAnotherProcess(t *testing.T) { + home := t.TempDir() + setHome(t, home) + + cmd := exec.Command(os.Args[0], "-test.run=^TestHoldsTokenStoreLockHelper$") + cmd.Env = append(os.Environ(), holdLockEnvVar+"=1", "HOME="+home, "USERPROFILE="+home) + stdin, err := cmd.StdinPipe() + require.NoError(t, err) + stdout, err := cmd.StdoutPipe() + require.NoError(t, err) + require.NoError(t, cmd.Start()) + defer func() { + stdin.Close() + _ = cmd.Wait() + }() + + // The helper prints this line once it holds the lock. + scanner := bufio.NewScanner(stdout) + require.True(t, scanner.Scan(), "helper process did not report holding the lock") + require.Equal(t, "locked", scanner.Text()) + + ctx, cancel := context.WithTimeout(t.Context(), 100*time.Millisecond) + defer cancel() + _, err = LockTokenStore(ctx) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +// TestHoldsTokenStoreLockHelper is the child half of +// TestLockTokenStoreWaitsForAnotherProcess. It holds the lock until its stdin +// is closed, and is skipped during a normal test run. +func TestHoldsTokenStoreLockHelper(t *testing.T) { + if os.Getenv(holdLockEnvVar) != "1" { + t.Skip("helper process for TestLockTokenStoreWaitsForAnotherProcess") + } + + unlock, err := LockTokenStore(t.Context()) + require.NoError(t, err) + defer unlock() + + _, err = os.Stdout.WriteString("locked\n") + require.NoError(t, err) + + // Block until the parent closes stdin. + _, _ = os.Stdin.Read(make([]byte, 1)) +} diff --git a/libs/auth/storage/lock_unix.go b/libs/auth/storage/lock_unix.go new file mode 100644 index 00000000000..d2d5127d730 --- /dev/null +++ b/libs/auth/storage/lock_unix.go @@ -0,0 +1,29 @@ +//go:build !windows + +package storage + +import ( + "errors" + "os" + + "golang.org/x/sys/unix" +) + +// tryLock takes an exclusive advisory lock on f without blocking. It reports +// whether the lock was taken. +func tryLock(f *os.File) (bool, error) { + err := unix.Flock(int(f.Fd()), unix.LOCK_EX|unix.LOCK_NB) + switch { + case err == nil: + return true, nil + case errors.Is(err, unix.EWOULDBLOCK): + return false, nil + default: + return false, err + } +} + +// unlock releases the advisory lock held on f. +func unlock(f *os.File) { + _ = unix.Flock(int(f.Fd()), unix.LOCK_UN) +} diff --git a/libs/auth/storage/lock_windows.go b/libs/auth/storage/lock_windows.go new file mode 100644 index 00000000000..5a607e8d2c0 --- /dev/null +++ b/libs/auth/storage/lock_windows.go @@ -0,0 +1,40 @@ +package storage + +import ( + "errors" + "os" + + "golang.org/x/sys/windows" +) + +// lockRegionLength is the number of bytes LockFileEx covers. The lock file has +// no contents, so locking a single byte is enough to make holders exclusive. +const lockRegionLength = 1 + +// tryLock takes an exclusive lock on f without blocking. It reports whether the +// lock was taken. +func tryLock(f *os.File) (bool, error) { + var overlapped windows.Overlapped + err := windows.LockFileEx( + windows.Handle(f.Fd()), + windows.LOCKFILE_EXCLUSIVE_LOCK|windows.LOCKFILE_FAIL_IMMEDIATELY, + 0, + lockRegionLength, + 0, + &overlapped, + ) + switch { + case err == nil: + return true, nil + case errors.Is(err, windows.ERROR_LOCK_VIOLATION): + return false, nil + default: + return false, err + } +} + +// unlock releases the lock held on f. +func unlock(f *os.File) { + var overlapped windows.Overlapped + _ = windows.UnlockFileEx(windows.Handle(f.Fd()), 0, lockRegionLength, 0, &overlapped) +} diff --git a/libs/auth/u2m/persistent_auth.go b/libs/auth/u2m/persistent_auth.go index 8f5bd723fe7..6c546aacd96 100644 --- a/libs/auth/u2m/persistent_auth.go +++ b/libs/auth/u2m/persistent_auth.go @@ -74,6 +74,10 @@ type PersistentAuth struct { // store stores and looks up tokens. store storage.Store + // storeLock serializes the read-refresh-write sequence against other + // processes using the same store. Nil means no cross-process coordination. + storeLock StoreLock + // client is the HTTP client to use for OAuth2 requests. client *http.Client @@ -135,6 +139,10 @@ type PersistentAuth struct { type PersistentAuthOption func(*PersistentAuth) +// StoreLock acquires the cross-process lock guarding the token store and +// returns a function that releases it. See storage.LockTokenStore. +type StoreLock func(ctx context.Context) (func(), error) + // WithTokenStore sets the token store for the PersistentAuth. func WithTokenStore(s storage.Store) PersistentAuthOption { return func(a *PersistentAuth) { @@ -142,6 +150,15 @@ func WithTokenStore(s storage.Store) PersistentAuthOption { } } +// WithStoreLock sets the lock serializing token refreshes against other +// processes. Pass it whenever the store is shared with other CLI invocations, +// which is every store except the in-memory one. +func WithStoreLock(l StoreLock) PersistentAuthOption { + return func(a *PersistentAuth) { + a.storeLock = l + } +} + // WithHttpClient sets the HTTP client for the PersistentAuth. func WithHttpClient(c *http.Client) PersistentAuthOption { return func(a *PersistentAuth) { @@ -380,14 +397,28 @@ func (a *PersistentAuth) recoverStoreUpdate(old, candidate *oauth2.Token) *oauth return nil } +// cachedRefreshedToken returns a token that another process stored while this +// one waited for the store lock, or nil if the cache still holds oldToken or +// the cached token itself needs a refresh. +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 + } + return e.Token +} + // refresh refreshes the token for the given OAuthArgument, storing the new // token in the cache. // -// This read-refresh-write sequence is not coordinated across processes. -// Because the CLI is stateless, two separate CLI invocations can load the same -// cached refresh token, both attempt a refresh, and race to update the cache. -// This should be fixed in a follow-up by adding cross-process coordination -// around refresh and cache writes. +// The read-refresh-write sequence is serialized across processes by the store +// lock: the CLI is stateless, so two invocations for the same profile otherwise +// load the same cached refresh token, both exchange it, and race to update the +// cache. Callers that do not set a store lock (an in-memory store has no other +// process to coordinate with) keep the uncoordinated behavior. func (a *PersistentAuth) refresh(oldToken *oauth2.Token) (*oauth2.Token, error) { // Fail fast with ErrMissingRefreshToken instead of letting the oauth2 // library attempt to refresh and return a misleading error (e.g. "token @@ -396,6 +427,20 @@ func (a *PersistentAuth) refresh(oldToken *oauth2.Token) (*oauth2.Token, error) if oldToken.RefreshToken == "" { return nil, ErrMissingRefreshToken } + if a.storeLock != nil { + unlock, err := a.storeLock(a.ctx) + if err != nil { + return nil, fmt.Errorf("token store lock: %w", err) + } + defer unlock() + + // The token this process read may have been refreshed by the process + // that held the lock. Reusing its result keeps the refresh token from + // being exchanged once per waiter. + if t := a.cachedRefreshedToken(oldToken); t != nil { + return t, nil + } + } cfg, err := a.oauth2Config() if err != nil { return nil, err diff --git a/libs/auth/u2m/persistent_auth_test.go b/libs/auth/u2m/persistent_auth_test.go index 1ec02cf1750..2b1cb6556ac 100644 --- a/libs/auth/u2m/persistent_auth_test.go +++ b/libs/auth/u2m/persistent_auth_test.go @@ -733,6 +733,171 @@ func TestForceRefreshToken_RecoversConcurrentCacheUpdate(t *testing.T) { } } +func TestForceRefreshToken_ReusesTokenRefreshedWhileWaitingForLock(t *testing.T) { + now := time.Now() + old := &oauth2.Token{ + AccessToken: "old-access", + RefreshToken: "old-refresh", + Expiry: now.Add(time.Hour), + } + other := &oauth2.Token{ + AccessToken: "other-process-access", + RefreshToken: "other-process-refresh", + Expiry: now.Add(time.Hour), + } + locked := false + unlocked := false + c := &tokenStoreMock{ + lookup: func(key string) (*oauth2.Token, error) { + // The second lookup is the one made under the lock, by which time + // the process that held it has written its own refreshed token. + if locked { + return other, nil + } + return old, nil + }, + store: func(key string, tok *oauth2.Token) error { + t.Fatalf("store(): want no write, got %q", tok.AccessToken) + return nil + }, + } + arg, err := NewBasicAccountOAuthArgument("https://accounts.cloud.databricks.test", "xyz") + if err != nil { + t.Fatalf("NewBasicAccountOAuthArgument(): %v", err) + } + p, err := NewPersistentAuth( + t.Context(), + WithTokenStore(c), + WithStoreLock(func(context.Context) (func(), error) { + locked = true + return func() { unlocked = true }, nil + }), + // An empty transport fails the test if a token exchange is attempted. + WithHttpClient(&http.Client{Transport: fixtures.SliceTransport{}}), + WithOAuthEndpointSupplier(MockOAuthEndpointSupplier{}), + WithOAuthArgument(arg), + ) + if err != nil { + t.Fatalf("NewPersistentAuth(): %v", err) + } + defer p.Close() + + tok, err := p.ForceRefreshToken() + if err != nil { + t.Fatalf("ForceRefreshToken(): want no error, got %v", err) + } + if tok.AccessToken != other.AccessToken { + t.Errorf("ForceRefreshToken(): want access token %q, got %q", other.AccessToken, tok.AccessToken) + } + if !unlocked { + t.Error("ForceRefreshToken(): want the store lock released") + } +} + +func TestForceRefreshToken_RefreshesWhenCacheUnchangedUnderLock(t *testing.T) { + old := &oauth2.Token{ + AccessToken: "old-access", + RefreshToken: "old-refresh", + Expiry: time.Now().Add(time.Hour), + } + stored := "" + unlocked := false + c := &tokenStoreMock{ + lookup: func(key string) (*oauth2.Token, error) { + return old, nil + }, + store: func(key string, tok *oauth2.Token) error { + stored = tok.AccessToken + return nil + }, + } + arg, err := NewBasicAccountOAuthArgument("https://accounts.cloud.databricks.test", "xyz") + if err != nil { + t.Fatalf("NewBasicAccountOAuthArgument(): %v", err) + } + p, err := NewPersistentAuth( + t.Context(), + WithTokenStore(c), + WithStoreLock(func(context.Context) (func(), error) { + return func() { unlocked = true }, nil + }), + WithHttpClient(&http.Client{ + Transport: fixtures.SliceTransport{ + { + Method: "POST", + Resource: "/oidc/accounts/xyz/v1/token", + Response: `access_token=refreshed&refresh_token=refreshed-refresh&expires_in=3600`, + ResponseHeaders: map[string][]string{ + "Content-Type": {"application/x-www-form-urlencoded"}, + }, + }, + }, + }), + WithOAuthEndpointSupplier(MockOAuthEndpointSupplier{}), + WithOAuthArgument(arg), + ) + if err != nil { + t.Fatalf("NewPersistentAuth(): %v", err) + } + defer p.Close() + + tok, err := p.ForceRefreshToken() + if err != nil { + t.Fatalf("ForceRefreshToken(): want no error, got %v", err) + } + if tok.AccessToken != "refreshed" { + t.Errorf("ForceRefreshToken(): want access token 'refreshed', got %q", tok.AccessToken) + } + if stored != "refreshed" { + t.Errorf("store(): want 'refreshed' written, got %q", stored) + } + if !unlocked { + t.Error("ForceRefreshToken(): want the store lock released") + } +} + +func TestForceRefreshToken_FailsWhenStoreLockFails(t *testing.T) { + c := &tokenStoreMock{ + lookup: func(key string) (*oauth2.Token, error) { + return &oauth2.Token{ + AccessToken: "old-access", + RefreshToken: "old-refresh", + Expiry: time.Now().Add(time.Hour), + }, nil + }, + store: func(key string, tok *oauth2.Token) error { + t.Fatalf("store(): want no write, got %q", tok.AccessToken) + return nil + }, + } + arg, err := NewBasicAccountOAuthArgument("https://accounts.cloud.databricks.test", "xyz") + if err != nil { + t.Fatalf("NewBasicAccountOAuthArgument(): %v", err) + } + p, err := NewPersistentAuth( + t.Context(), + WithTokenStore(c), + WithStoreLock(func(context.Context) (func(), error) { + return nil, errors.New("lock unavailable") + }), + WithHttpClient(&http.Client{Transport: fixtures.SliceTransport{}}), + WithOAuthEndpointSupplier(MockOAuthEndpointSupplier{}), + WithOAuthArgument(arg), + ) + if err != nil { + t.Fatalf("NewPersistentAuth(): %v", err) + } + defer p.Close() + + _, err = p.ForceRefreshToken() + if err == nil { + t.Fatal("ForceRefreshToken(): want an error, got none") + } + if !strings.Contains(err.Error(), "lock unavailable") { + t.Errorf("ForceRefreshToken(): want the lock error, got %v", err) + } +} + func TestIsFreshReplacement(t *testing.T) { now := time.Date(2100, time.January, 1, 0, 0, 0, 0, time.UTC) old := &oauth2.Token{AccessToken: "old", Expiry: now.Add(time.Hour)}