Add Asynchronous Token Refresh to RefreshableTokenSource (#455)

## Motivation
Synchronous token refresh introduces unnecessary blocking and latency,
which becomes especially problematic in high QPS (Queries Per Second)
workloads. In such scenarios, blocking token refresh operations can
create bottlenecks, impact throughput, and cause request queuing.

## Solution
This PR introduces asynchronous token refresh to improve responsiveness
and reliability, with particular benefits for high-throughput
applications.

## What changes are proposed in this pull request?
**Async Refresh Option**: Added `withAsyncRefresh(boolean enabled)`
method to enable/disable asynchronous token refresh. When enabled,
tokens are refreshed in the background when they become "stale" (close
to expiry).

### Token State Management
- **Three-State Token System**:
  - `FRESH`: Token is valid and not close to expiry
- `STALE`: Token is valid but will expire soon (triggers async refresh
if enabled)
  - `EXPIRED`: Token has expired (requires blocking refresh)
- `Token` class is now a pure data class, holding only token information
(access token, refresh token, expiry time).
- `TokenSource` implementations are responsible for token state
management and refresh logic.

### Performance Optimizations
- **Non-blocking for Stale Tokens**: When async is enabled and the token
is stale but not expired, API calls continue using the current token
while a background refresh is triggered, reducing latency and avoiding
thread blocking.
- **Blocking Only on Expiry**: If the token is expired, calls will still
block until a new token is fetched, ensuring correctness.
- **Default Stale Duration**: Tokens are considered stale 3 minutes
before expiry, allowing proactive refresh while maintaining validity.

### Thread Safety & Reliability
- **Synchronized Refresh Logic**: All async refresh operations are
synchronized to prevent race conditions and redundant refreshes.
- **Refresh State Tracking**: The implementation tracks:
  - Whether a refresh is already in progress
  - Whether the last refresh succeeded
  - This prevents unnecessary or repeated background refreshes
- **Failure Handling**: If an async refresh fails, subsequent refreshes
will be forced to be synchronous until a successful refresh occurs.

### Configuration Options
**Experimental Features (may change in future releases)**:
- withClockSupplier() - Custom clock supplier for testing
- withAsyncRefresh() - Enable/disable asynchronous token refresh
- withExpiryBuffer() - Configure the expiry buffer duration

## Testing
- **Unit tests** cover async refresh triggering, correct token usage
during background refresh, and thread safety.
- **Time Control for Testing:**  
The `RefreshableTokenSource` class can be configured with a custom clock
supplier via `withClockSupplier(ClockSupplier clockSupplier)`, allowing
tests to precisely control and simulate token expiry and refresh timing.
This enables deterministic testing of token state transitions and
refresh behavior without relying on real system time.

NO_CHANGELOG=true

---------

Co-authored-by: Parth Bansal <[email protected]>

Author: emmyzhou-db

databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java

@@ -5,36 +5,254 @@
 import com.databricks.sdk.core.http.FormRequest;
 import com.databricks.sdk.core.http.HttpClient;
 import com.databricks.sdk.core.http.Request;
+import com.databricks.sdk.core.utils.ClockSupplier;
+import com.databricks.sdk.core.utils.UtcClockSupplier;
+import java.time.Duration;
 import java.time.Instant;
 import java.util.Base64;
 import java.util.Map;
+import java.util.concurrent.CompletableFuture;
 import org.apache.http.HttpHeaders;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /**
  * An OAuth TokenSource which can be refreshed.
  *
- * <p>Calls to getToken() will first check if the token is still valid (currently defined by having
- * at least 10 seconds until expiry). If not, refresh() is called first to refresh the token.
+ * <p>This class supports both synchronous and asynchronous token refresh. When async is enabled,
+ * stale tokens will trigger a background refresh, while expired tokens will block until a new token
+ * is fetched.
  */
 public abstract class RefreshableTokenSource implements TokenSource {
-  protected Token token;
 
+  /**
+   * Enum representing the state of the token. FRESH: Token is valid and not close to expiry. STALE:
+   * Token is valid but will expire soon - an async refresh will be triggered if enabled. EXPIRED:
+   * Token has expired and must be refreshed using a blocking call.
+   */
+  protected enum TokenState {
+    FRESH,
+    STALE,
+    EXPIRED
+  }
+
+  private static final Logger logger = LoggerFactory.getLogger(RefreshableTokenSource.class);
+  // Default duration before expiry to consider a token as 'stale'.
+  private static final Duration DEFAULT_STALE_DURATION = Duration.ofMinutes(3);
+  // Default additional buffer before expiry to consider a token as expired.
+  private static final Duration DEFAULT_EXPIRY_BUFFER = Duration.ofSeconds(40);
+
+  // The current OAuth token. May be null if not yet fetched.
+  protected volatile Token token;
+  // Whether asynchronous refresh is enabled.
+  private boolean asyncEnabled = false;
+  // Duration before expiry to consider a token as 'stale'.
+  private Duration staleDuration = DEFAULT_STALE_DURATION;
+  // Additional buffer before expiry to consider a token as expired.
+  private Duration expiryBuffer = DEFAULT_EXPIRY_BUFFER;
+  // Whether a refresh is currently in progress (for async refresh).
+  private boolean refreshInProgress = false;
+  // Whether the last refresh attempt succeeded.
+  private boolean lastRefreshSucceeded = true;
+  // Clock supplier for current time.
+  private ClockSupplier clockSupplier = new UtcClockSupplier();
+
+  /** Constructs a new {@code RefreshableTokenSource} with no initial token. */
   public RefreshableTokenSource() {}
 
+  /**
+   * Constructor with initial token.
+   *
+   * @param token The initial token to use.
+   */
   public RefreshableTokenSource(Token token) {
     this.token = token;
   }
 
+  /**
+   * Set the clock supplier for current time.
+   *
+   * <p><b>Experimental:</b> This method may change or be removed in future releases.
+   *
+   * @param clockSupplier The clock supplier to use.
+   * @return this instance for chaining
+   */
+  public RefreshableTokenSource withClockSupplier(ClockSupplier clockSupplier) {
+    this.clockSupplier = clockSupplier;
+    return this;
+  }
+
+  /**
+   * Enable or disable asynchronous token refresh.
+   *
+   * <p><b>Experimental:</b> This method may change or be removed in future releases.
+   *
+   * @param enabled true to enable async refresh, false to disable
+   * @return this instance for chaining
+   */
+  public RefreshableTokenSource withAsyncRefresh(boolean enabled) {
+    this.asyncEnabled = enabled;
+    return this;
+  }
+
+  /**
+   * Set the expiry buffer. If the token's lifetime is less than this buffer, it is considered
+   * expired.
+   *
+   * <p><b>Experimental:</b> This method may change or be removed in future releases.
+   *
+   * @param buffer the expiry buffer duration
+   * @return this instance for chaining
+   */
+  public RefreshableTokenSource withExpiryBuffer(Duration buffer) {
+    this.expiryBuffer = buffer;
+    return this;
+  }
+
+  /**
+   * Refresh the OAuth token. Subclasses must implement this to define how the token is refreshed.
+   *
+   * <p>This method may throw an exception if the token cannot be refreshed. The specific exception
+   * type depends on the implementation.
+   *
+   * @return The newly refreshed Token.
+   */
+  protected abstract Token refresh();
+
+  /**
+   * Gets the current token, refreshing if necessary. If async refresh is enabled, may return a
+   * stale token while a refresh is in progress.
+   *
+   * <p>This method may throw an exception if the token cannot be refreshed, depending on the
+   * implementation of {@link #refresh()}.
+   *
+   * @return The current valid token
+   */
+  public Token getToken() {
+    if (asyncEnabled) {
+      return getTokenAsync();
+    }
+    return getTokenBlocking();
+  }
+
+  /**
+   * Determine the state of the current token (fresh, stale, or expired).
+   *
+   * @return The token state
+   */
+  protected TokenState getTokenState(Token t) {
+    if (t == null) {
+      return TokenState.EXPIRED;
+    }
+    Duration lifeTime = Duration.between(Instant.now(clockSupplier.getClock()), t.getExpiry());
+    if (lifeTime.compareTo(expiryBuffer) <= 0) {
+      return TokenState.EXPIRED;
+    }
+    if (lifeTime.compareTo(staleDuration) <= 0) {
+      return TokenState.STALE;
+    }
+    return TokenState.FRESH;
+  }
+
+  /**
+   * Get the current token, blocking to refresh if expired.
+   *
+   * <p>This method may throw an exception if the token cannot be refreshed, depending on the
+   * implementation of {@link #refresh()}.
+   *
+   * @return The current valid token
+   */
+  protected Token getTokenBlocking() {
+    // Use double-checked locking to minimize synchronization overhead on reads:
+    // 1. Check if the token is expired without locking.
+    // 2. If expired, synchronize and check again (another thread may have refreshed it).
+    // 3. If still expired, perform the refresh.
+    if (getTokenState(token) != TokenState.EXPIRED) {
+      return token;
+    }
+    synchronized (this) {
+      if (getTokenState(token) != TokenState.EXPIRED) {
+        return token;
+      }
+      lastRefreshSucceeded = false;
+      try {
+        token = refresh();
+      } catch (Exception e) {
+        logger.error("Failed to refresh token synchronously", e);
+        throw e;
+      }
+      lastRefreshSucceeded = true;
+      return token;
+    }
+  }
+
+  /**
+   * Get the current token, possibly triggering an async refresh if stale. If the token is expired,
+   * blocks to refresh.
+   *
+   * <p>This method may throw an exception if the token cannot be refreshed, depending on the
+   * implementation of {@link #refresh()}.
+   *
+   * @return The current valid or stale token
+   */
+  protected Token getTokenAsync() {
+    Token currentToken = token;
+
+    switch (getTokenState(currentToken)) {
+      case FRESH:
+        return currentToken;
+      case STALE:
+        triggerAsyncRefresh();
+        return currentToken;
+      case EXPIRED:
+        return getTokenBlocking();
+      default:
+        throw new IllegalStateException("Invalid token state.");
+    }
+  }
+
+  /**
+   * Trigger an asynchronous refresh of the token if not already in progress and last refresh
+   * succeeded.
+   */
+  private synchronized void triggerAsyncRefresh() {
+    // Check token state again inside the synchronized block to avoid triggering a refresh if
+    // another thread updated the token in the meantime.
+    if (!refreshInProgress && lastRefreshSucceeded && getTokenState(token) != TokenState.FRESH) {
+      refreshInProgress = true;
+      CompletableFuture.runAsync(
+          () -> {
+            try {
+              // Attempt to refresh the token in the background
+              Token newToken = refresh();
+              synchronized (this) {
+                token = newToken;
+                refreshInProgress = false;
+              }
+            } catch (Exception e) {
+              synchronized (this) {
+                lastRefreshSucceeded = false;
+                refreshInProgress = false;
+                logger.error("Asynchronous token refresh failed", e);
+              }
+            }
+          });
+    }
+  }
+
   /**
    * Helper method implementing OAuth token refresh.
    *
+   * @param hc The HTTP client to use for the request.
    * @param clientId The client ID to authenticate with.
    * @param clientSecret The client secret to authenticate with.
    * @param tokenUrl The authorization URL for fetching tokens.
    * @param params Additional request parameters.
    * @param headers Additional headers.
    * @param position The position of the authentication parameters in the request.
    * @return The newly fetched Token.
+   * @throws DatabricksException if the refresh fails
+   * @throws IllegalArgumentException if the OAuth response contains an error
    */
   protected static Token retrieveToken(
       HttpClient hc,
@@ -75,13 +293,4 @@ protected static Token retrieveToken(
       throw new DatabricksException("Failed to refresh credentials: " + e.getMessage(), e);
     }
   }
-
-  protected abstract Token refresh();
-
-  public synchronized Token getToken() {
-    if (token == null || !token.isValid()) {
-      token = refresh();
-    }
-    return token;
-  }
 }

databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Token.java

@@ -1,7 +1,5 @@
 package com.databricks.sdk.core.oauth;
 
-import com.databricks.sdk.core.utils.ClockSupplier;
-import com.databricks.sdk.core.utils.SystemClockSupplier;
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import java.time.Instant;
@@ -23,16 +21,9 @@ public class Token {
    */
   @JsonProperty private Instant expiry;
 
-  private final ClockSupplier clockSupplier;
-
   /** Constructor for non-refreshable tokens (e.g. M2M). */
   public Token(String accessToken, String tokenType, Instant expiry) {
-    this(accessToken, tokenType, null, expiry, new SystemClockSupplier());
-  }
-
-  /** Constructor for non-refreshable tokens (e.g. M2M) with ClockSupplier */
-  public Token(String accessToken, String tokenType, Instant expiry, ClockSupplier clockSupplier) {
-    this(accessToken, tokenType, null, expiry, clockSupplier);
+    this(accessToken, tokenType, null, expiry);
   }
 
   /** Constructor for refreshable tokens. */
@@ -42,51 +33,48 @@ public Token(
       @JsonProperty("tokenType") String tokenType,
       @JsonProperty("refreshToken") String refreshToken,
       @JsonProperty("expiry") Instant expiry) {
-    this(accessToken, tokenType, refreshToken, expiry, new SystemClockSupplier());
-  }
-
-  /** Constructor for refreshable tokens with ClockSupplier. */
-  public Token(
-      String accessToken,
-      String tokenType,
-      String refreshToken,
-      Instant expiry,
-      ClockSupplier clockSupplier) {
     Objects.requireNonNull(accessToken, "accessToken must be defined");
     Objects.requireNonNull(tokenType, "tokenType must be defined");
     Objects.requireNonNull(expiry, "expiry must be defined");
-    Objects.requireNonNull(clockSupplier, "clockSupplier must be defined");
     this.accessToken = accessToken;
     this.tokenType = tokenType;
     this.refreshToken = refreshToken;
     this.expiry = expiry;
-    this.clockSupplier = clockSupplier;
-  }
-
-  public boolean isExpired() {
-    if (expiry == null) {
-      return false;
-    }
-    // Azure Databricks rejects tokens that expire in 30 seconds or less,
-    // so we refresh the token 40 seconds before it expires.
-    Instant potentiallyExpired = expiry.minusSeconds(40);
-    Instant now = Instant.now(clockSupplier.getClock());
-    return potentiallyExpired.isBefore(now);
-  }
-
-  public boolean isValid() {
-    return accessToken != null && !isExpired();
   }
 
+  /**
+   * Returns the type of the token (e.g., "Bearer").
+   *
+   * @return the token type
+   */
   public String getTokenType() {
     return tokenType;
   }
 
+  /**
+   * Returns the refresh token, if available. May be null for non-refreshable tokens.
+   *
+   * @return the refresh token or null
+   */
   public String getRefreshToken() {
     return refreshToken;
   }
 
+  /**
+   * Returns the access token string.
+   *
+   * @return the access token
+   */
   public String getAccessToken() {
     return accessToken;
   }
+
+  /**
+   * Returns the expiry time of the token as a Instant.
+   *
+   * @return the expiry time
+   */
+  public Instant getExpiry() {
+    return this.expiry;
+  }
 }

databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/SystemClockSupplier.java renamed to databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/UtcClockSupplier.java

@@ -2,7 +2,7 @@
 
 import java.time.Clock;
 
-public class SystemClockSupplier implements ClockSupplier {
+public class UtcClockSupplier implements ClockSupplier {
   @Override
   public Clock getClock() {
     return Clock.systemUTC();

databricks-sdk-java/src/test/java/com/databricks/sdk/core/CliTokenSourceTest.java

@@ -154,7 +154,7 @@ public void testRefreshWithExpiry(
         Token token = tokenSource.refresh();
         assertEquals("Bearer", token.getTokenType());
         assertEquals("test-token", token.getAccessToken());
-        assertEquals(shouldBeExpired, token.isExpired());
+        assertEquals(shouldBeExpired, token.getExpiry().isBefore(Instant.now()));
       }
     }
   }