Implement dynamic token stale period based on token TTL (#677)

## Summary

Extends the token refresh buffer from a fixed 5 minutes to a dynamic
period that adapts to token lifetime, improving reliability across
different token types while maintaining backward compatibility.

The TTL is computed based on the remaining time to live at the moment
the token is received and not its real TTL. For example, if the token
has a TTL of 60 minutes but was acquired 20 minutes before being used by
the SDK (e.g. through the CLI), its effective TTL will be 60 − 20 = 40
minutes.

## Why

The previous 5 minute stale period barely covered the allowed monthly
downtime of ~4.32 minutes. Thus, if the auth services were to be down
and a request were to come 4 minute into the stale period it would have
only 1 minute to obtain a new valid token before expiry. With an
extended stale period of 20 minutes, the SDK has an extra ~15 minutes of
stale but valid tokens to use, allowing the auth system to recover.

## Changes

- Increase the maximum stale period from 5 to 20 minutes to support
99.95%
   availability.
- Implement dynamic stale period calculation: min(TTL × 0.5, 20
minutes).
  - Compute stale period per-token at acquisition time.

### Backward compatibility:
- The public Builder method setStaleDuration() is preserved. Calling it
disables dynamic mode via a useDynamicStaleDuration flag, reverting the
behavior to the legacy fixed-duration stale window. This ensures that
any caller already configuring a custom stale period is unaffected.

## Implementation:
- Add computeStaleDuration(Token) that computes min(TTL / 2,
MAX_STALE_DURATION).
- Add useDynamicStaleDuration flag to the Builder, defaulting to true;
setStaleDuration() sets it to false.
- Add volatile dynamicStaleDuration field; initialized to
MAX_STALE_DURATION as a safe default when no token is pre-set, or
computed from the pre-set token via computeStaleDuration() when one is
provided.
- Update getTokenState() to use dynamicStaleDuration or the legacy
staleDuration based on the flag.
- Update getTokenBlocking() to recompute the stale period after a
successful synchronous refresh.
- Update triggerAsyncRefresh() to recompute the stale period after a
successful async refresh.

## Testing

- Update testAsyncRefreshParametrized to use TestClockSupplier for
deterministic time control, adding a clockAdvanceMinutes parameter to
bring tokens into the stale window without relying on wall-clock timing.
- Add a capped stale duration scenario: 60-min TTL token advanced 41
minutes leaves lifeTime = 19 min ≤ 20 min cap → STALE, verifying
MAX_STALE_DURATION is correctly applied.
- Update testAsyncRefreshFailureFallback to use a 4-minute TTL token
advanced by 3 minutes to reliably enter the stale window under the
dynamic formula.

Author: mihaimitrea-db

NEXT_CHANGELOG.md

@@ -11,5 +11,6 @@
 ### Documentation
 
 ### Internal Changes
+* Implement dynamic auth token stale period based on initial token lifetime. Increased up to 20 mins for standard OAuth with proportionally shorter periods for short-lived tokens. Manually setting the stale period using the CachedTokeSource builder reverts the behaviour to the legacy fixed stale duration.
 
 ### API Changes

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'.
-  private final Duration staleDuration;
+  // The legacy duration before expiry to consider a token as 'stale'.
+  private final Duration staticStaleDuration;
+  // 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.
@@ -59,10 +68,17 @@ private enum TokenState {
   private CachedTokenSource(Builder builder) {
     this.tokenSource = builder.tokenSource;
     this.asyncDisabled = builder.asyncDisabled;
-    this.staleDuration = builder.staleDuration;
+    this.staticStaleDuration = 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 = Duration.ofMinutes(0);
+    }
   }
 
   /**
@@ -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,21 @@ public Token getToken() {
     return getTokenAsync();
   }
 
+  private Duration computeStaleDuration(Token t) {
+    if (t.getExpiry() == null) {
+      return Duration.ZERO; // Tokens with no expiry are considered permanent.
+    }
+
+    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).
    *
@@ -197,10 +230,15 @@ protected TokenState getTokenState(Token t) {
     if (t == null) {
       return TokenState.EXPIRED;
     }
+    if (t.getExpiry() == null) {
+      return TokenState.FRESH; // Tokens with no expiry are considered permanent.
+    }
+
     Duration lifeTime = Duration.between(Instant.now(clockSupplier.getClock()), t.getExpiry());
     if (lifeTime.compareTo(expiryBuffer) <= 0) {
       return TokenState.EXPIRED;
     }
+    Duration staleDuration = useDynamicStaleDuration ? dynamicStaleDuration : staticStaleDuration;
     if (lifeTime.compareTo(staleDuration) <= 0) {
       return TokenState.STALE;
     }
@@ -228,13 +266,22 @@ protected Token getTokenBlocking() {
         return token;
       }
       lastRefreshSucceeded = false;
+      Token newToken;
       try {
-        token = tokenSource.getToken();
+        newToken = tokenSource.getToken();
       } catch (Exception e) {
         logger.error("Failed to refresh token synchronously", e);
         throw e;
       }
       lastRefreshSucceeded = true;
+
+      // Write dynamicStaleDuration before publishing the new token via the volatile write,
+      // so unsynchronized readers that see the new token are guaranteed to also see the
+      // updated dynamicStaleDuration.
+      if (useDynamicStaleDuration && newToken != null) {
+        dynamicStaleDuration = computeStaleDuration(newToken);
+      }
+      token = newToken;
       return token;
     }
   }
@@ -279,6 +326,12 @@ private synchronized void triggerAsyncRefresh() {
               // Attempt to refresh the token in the background.
               Token newToken = tokenSource.getToken();
               synchronized (this) {
+                // Write dynamicStaleDuration before publishing the new token via the volatile
+                // write, so unsynchronized readers that see the new token are guaranteed to also
+                // see the updated dynamicStaleDuration.
+                if (useDynamicStaleDuration && newToken != null) {
+                  dynamicStaleDuration = computeStaleDuration(newToken);
+                }
                 token = newToken;
                 refreshInProgress = false;
               }

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

@@ -17,38 +17,81 @@ 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 +112,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 +137,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 +179,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();