Serialize U2M token refreshes across processes

.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))

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

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

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",
+]

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))
 	}

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))

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

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:
+		}
+	}
+}

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))
+}

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)
+}

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)
+}