Implement dynamic token stale period based on token TTL.

Replace the fixed 5-minute stale window with a per-token formula: stale_period = min(TTL / 2, 20 minutes). This ensures that short-lived tokens (e.g. FastPath tokens with a 10-minute TTL) enter the stale window early enough to trigger a proactive async refresh, while long-lived tokens are capped at 20 minutes to maintain a meaningful refresh buffer under a 99.99% uptime SLA.

The stale duration is computed at token fetch time using the remaining TTL as a proxy, since the SDK does not track token issuance timestamps. It is updated after every successful refresh (both synchronous and asynchronous) via computeStaleDuration(), and stored in a volatile field so that the unsynchronized fast-path read in getTokenAsync() sees a consistent value.

Due to the exposed CachedTokenSource builder function setStalePeriod backwards compatibility is maintained through the useDynamicStalePeriod flag. Using the setStalePeriod function disables this flag, reverting the behaviour of the token cache to the legacy one.

Author: mihaimitrea-db

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

@@ -33,6 +33,11 @@ private enum TokenState {
   // Default duration before expiry to consider a token as 'stale'. This value is chosen to cover
   // the maximum monthly downtime allowed by a 99.99% uptime SLA (~4.38 minutes).
   private static final Duration DEFAULT_STALE_DURATION = Duration.ofMinutes(5);
+  // The maximum stale duration that can be achieved before expiry to consider a token as 'stale'
+  // when using the dynamic stale duration method. This value is chosen to cover the maximum
+  // monthly downtime allowed by a 99.99% uptime SLA (~4.38 minutes) while increasing the likelihood
+  // that the token is refreshed asynchronously if the auth server is down.
+  private static final Duration MAX_STALE_DURATION = Duration.ofMinutes(20);
   // Default additional buffer before expiry to consider a token as expired.
   // This is 40 seconds by default since Azure Databricks rejects tokens that are within 30 seconds
   // of expiry.
@@ -42,8 +47,12 @@ private enum TokenState {
   private final TokenSource tokenSource;
   // Whether asynchronous refresh is enabled.
   private boolean asyncDisabled = false;
-  // Duration before expiry to consider a token as 'stale'.
+  // The legacy duration before expiry to consider a token as 'stale'.
   private final Duration staleDuration;
+  // Whether to use the dynamic stale duration computation or defer to the legacy duration.
+  private final boolean useDynamicStaleDuration;
+  // The dynamically computed duration before expiry to consider a token as 'stale'.
+  private volatile Duration dynamicStaleDuration;
   // Additional buffer before expiry to consider a token as expired.
   private final Duration expiryBuffer;
   // Clock supplier for current time.
@@ -60,9 +69,16 @@ private CachedTokenSource(Builder builder) {
     this.tokenSource = builder.tokenSource;
     this.asyncDisabled = builder.asyncDisabled;
     this.staleDuration = builder.staleDuration;
+    this.useDynamicStaleDuration = builder.useDynamicStaleDuration;
     this.expiryBuffer = builder.expiryBuffer;
     this.clockSupplier = builder.clockSupplier;
     this.token = builder.token;
+    
+    if (this.useDynamicStaleDuration && this.token != null) {
+      this.dynamicStaleDuration = computeStaleDuration(this.token);
+    } else {
+      this.dynamicStaleDuration = MAX_STALE_DURATION;
+    }
   }
 
   /**
@@ -75,6 +91,7 @@ public static class Builder {
     private final TokenSource tokenSource;
     private boolean asyncDisabled = false;
     private Duration staleDuration = DEFAULT_STALE_DURATION;
+    private boolean useDynamicStaleDuration = true;
     private Duration expiryBuffer = DEFAULT_EXPIRY_BUFFER;
     private ClockSupplier clockSupplier = new UtcClockSupplier();
     private Token token;
@@ -130,6 +147,7 @@ public Builder setAsyncDisabled(boolean asyncDisabled) {
      */
     public Builder setStaleDuration(Duration staleDuration) {
       this.staleDuration = staleDuration;
+      this.useDynamicStaleDuration = false;
       return this;
     }
 
@@ -188,6 +206,17 @@ public Token getToken() {
     return getTokenAsync();
   }
 
+  private Duration computeStaleDuration(Token t) {
+    Duration ttl = Duration.between(Instant.now(clockSupplier.getClock()), t.getExpiry());
+
+    if (ttl.compareTo(Duration.ZERO) <= 0) {
+      return Duration.ZERO;
+    }
+    
+    Duration halfTtl = ttl.dividedBy(2);
+    return halfTtl.compareTo(MAX_STALE_DURATION) > 0 ? MAX_STALE_DURATION : halfTtl;
+  }
+
   /**
    * Determine the state of the current token (fresh, stale, or expired).
    *
@@ -201,7 +230,8 @@ protected TokenState getTokenState(Token t) {
     if (lifeTime.compareTo(expiryBuffer) <= 0) {
       return TokenState.EXPIRED;
     }
-    if (lifeTime.compareTo(staleDuration) <= 0) {
+    Duration usedStaleDuration = useDynamicStaleDuration ? dynamicStaleDuration : staleDuration;
+    if (lifeTime.compareTo(usedStaleDuration) <= 0) {
       return TokenState.STALE;
     }
     return TokenState.FRESH;
@@ -235,6 +265,11 @@ protected Token getTokenBlocking() {
         throw e;
       }
       lastRefreshSucceeded = true;
+
+      // Update the stale duration only after ensuring that the token is refreshed successfully.
+      if (useDynamicStaleDuration) {
+        dynamicStaleDuration = computeStaleDuration(token);
+      }
       return token;
     }
   }
@@ -281,6 +316,11 @@ private synchronized void triggerAsyncRefresh() {
               synchronized (this) {
                 token = newToken;
                 refreshInProgress = false;
+
+                // Update the stale duration only after ensuring that the token is refreshed successfully.
+                if (useDynamicStaleDuration) {
+                  dynamicStaleDuration = computeStaleDuration(token);
+                }
               }
             } catch (Exception e) {
               synchronized (this) {

databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/CachedTokenSourceTest.java

@@ -17,38 +17,57 @@ public class CachedTokenSourceTest {
   private static final String TOKEN_TYPE = "Bearer";
   private static final String INITIAL_TOKEN = "initial-token";
   private static final String REFRESH_TOKEN = "refreshed-token";
+
   private static final long FRESH_MINUTES = 10;
-  private static final long STALE_MINUTES = 1;
+
+  // Token TTL for the stale scenario: 4 minutes.
+  // dynamicStaleDuration = min(4/2, 20) = 2 min.
+  // After advancing the clock by STALE_ADVANCE_MINUTES = 3, lifeTime = 1 min.
+  // 1 min ≤ 2 min (stale) and 1 min > 40s (not expired) → STALE.
+  private static final long STALE_MINUTES = 4;
+  private static final long STALE_ADVANCE_MINUTES = 3;
+
+  // Token TTL for the capped stale duration scenario: 60 minutes.
+  // dynamicStaleDuration = min(60/2, 20) = 20 min (MAX_STALE_DURATION cap).
+  // After advancing the clock by CAPPED_STALE_ADVANCE_MINUTES = 41, lifeTime = 19 min.
+  // 19 min ≤ 20 min (stale) and 19 min > 40s (not expired) → STALE.
+  private static final long CAPPED_STALE_MINUTES = 60;
+  private static final long CAPPED_STALE_ADVANCE_MINUTES = 41;
+
   private static final long EXPIRED_MINUTES = -1;
 
   private static Stream<Arguments> provideAsyncRefreshScenarios() {
     return Stream.of(
-        Arguments.of("Fresh token, async enabled", FRESH_MINUTES, false, false, INITIAL_TOKEN),
-        Arguments.of("Stale token, async enabled", STALE_MINUTES, false, true, INITIAL_TOKEN),
-        Arguments.of("Expired token, async enabled", EXPIRED_MINUTES, false, true, REFRESH_TOKEN),
-        Arguments.of("Fresh token, async disabled", FRESH_MINUTES, true, false, INITIAL_TOKEN),
-        Arguments.of("Stale token, async disabled", STALE_MINUTES, true, false, INITIAL_TOKEN),
-        Arguments.of("Expired token, async disabled", EXPIRED_MINUTES, true, true, REFRESH_TOKEN));
+        Arguments.of("Fresh token, async enabled", FRESH_MINUTES, 0L, false, false, INITIAL_TOKEN),
+        Arguments.of("Stale token, async enabled", STALE_MINUTES, STALE_ADVANCE_MINUTES, false, true, INITIAL_TOKEN),
+        Arguments.of("Expired token, async enabled", EXPIRED_MINUTES, 0L, false, true, REFRESH_TOKEN),
+        Arguments.of("Fresh token, async disabled", FRESH_MINUTES, 0L, true, false, INITIAL_TOKEN),
+        Arguments.of("Stale token, async disabled", STALE_MINUTES, STALE_ADVANCE_MINUTES, true, false, INITIAL_TOKEN),
+        Arguments.of("Stale token, capped stale duration, async enabled", CAPPED_STALE_MINUTES, CAPPED_STALE_ADVANCE_MINUTES, false, true, INITIAL_TOKEN),
+        Arguments.of("Expired token, async disabled", EXPIRED_MINUTES, 0L, true, true, REFRESH_TOKEN));
   }
 
   @ParameterizedTest(name = "{0}")
   @MethodSource("provideAsyncRefreshScenarios")
   void testAsyncRefreshParametrized(
       String testName,
       long minutesUntilExpiry,
+      long clockAdvanceMinutes,
       boolean asyncDisabled,
       boolean expectRefresh,
       String expectedToken)
       throws Exception {
 
+    TestClockSupplier clockSupplier = new TestClockSupplier(Instant.now());
+
     Token initialToken =
         new Token(
             INITIAL_TOKEN,
             TOKEN_TYPE,
             null,
-            Instant.now().plus(Duration.ofMinutes(minutesUntilExpiry)));
+            Instant.now(clockSupplier.getClock()).plus(Duration.ofMinutes(minutesUntilExpiry)));
     Token refreshedToken =
-        new Token(REFRESH_TOKEN, TOKEN_TYPE, null, Instant.now().plus(Duration.ofMinutes(10)));
+        new Token(REFRESH_TOKEN, TOKEN_TYPE, null, Instant.now(clockSupplier.getClock()).plus(Duration.ofMinutes(10)));
     CountDownLatch refreshCalled = new CountDownLatch(1);
 
     TokenSource tokenSource =
@@ -69,8 +88,12 @@ public Token getToken() {
         new CachedTokenSource.Builder(tokenSource)
             .setAsyncDisabled(asyncDisabled)
             .setToken(initialToken)
+            .setClockSupplier(clockSupplier)
             .build();
 
+    // Advance the clock to put the token in the expected state before calling getToken().
+    clockSupplier.advanceTime(Duration.ofMinutes(clockAdvanceMinutes));
+
     Token token = source.getToken();
 
     boolean refreshed = refreshCalled.await(1, TimeUnit.SECONDS);
@@ -90,13 +113,13 @@ void testAsyncRefreshFailureFallback() throws Exception {
     // Create a mutable clock supplier that we can control
     TestClockSupplier clockSupplier = new TestClockSupplier(Instant.now());
 
-    // Create a token that will be stale (2 minutes until expiry)
+    // Create a token with a TTL of 4 minutes that will be stale in 3 minutes.
     Token staleToken =
         new Token(
             INITIAL_TOKEN,
             TOKEN_TYPE,
             null,
-            Instant.now(clockSupplier.getClock()).plus(Duration.ofMinutes(2)));
+            Instant.now(clockSupplier.getClock()).plus(Duration.ofMinutes(4)));
 
     class TestSource implements TokenSource {
       int refreshCallCount = 0;
@@ -132,6 +155,9 @@ public Token getToken() {
             .setClockSupplier(clockSupplier)
             .build();
 
+    // Advance clock to put the token in the stale window.
+    clockSupplier.advanceTime(Duration.ofMinutes(3));
+
     // First call triggers async refresh, which fails
     // Should return stale token immediately (async refresh)
     Token token = source.getToken();