[Rust] Harden OAuth token caching

rust/NEXT_CHANGELOG.md

@@ -22,6 +22,39 @@
   only after pending replay succeeds, while initial supervisor handoff and failed or
   cancelled replay promptly drop redundant senders instead of retaining incomplete
   `DoPut` request channels until later teardown.
+- The OAuth `expires_in` field is now parsed from a quoted integer (`"3600"`) in
+  addition to a plain JSON integer. A value that is missing or does not represent
+  a positive integer still yields no token lifetime, as before.
+- OAuth token lifetime is now anchored to when the token request starts rather
+  than when the response arrives, so a slow response no longer makes the SDK treat
+  a token as valid longer than the issuer does. A token already past its
+  (start-anchored) expiry is dead on arrival and, on a cold miss, surfaces a
+  retryable `TokenFetchError`.
+- A stalled proactive token refresh no longer hangs stream creation. For streams
+  built with `.oauth(...)`, the refresh is capped at half the stream's configured
+  setup budget (`recovery_timeout_ms`) — but only when the cached token has more
+  life left than that cap, so a hung endpoint becomes a prompt failure and the
+  still-valid cached token is served in time. When too little of the token's life
+  remains for that fallback to help, the refresh runs unbounded (like a cold miss)
+  so a slow-but-working mint isn't cut off early.
+- When a proactive token refresh fails — for any error, including a non-retryable
+  one such as revoked credentials — the still-valid cached token is served rather
+  than failing the caller. The token was validly issued and hasn't expired, and the
+  server re-validates it on every connection, so the error surfaces only when there
+  is no still-valid token to fall back to. A refresh that returns a dead-on-arrival
+  token falls back the same way.
+- After a proactive refresh falls back to the cached token, a short backoff paces
+  further attempts so a burst of stream creations cannot turn one token-endpoint
+  failure into a mint stampede. The interval scales with the configured token
+  refresh buffer and shrinks as the token nears expiry, down to a floor, and never
+  extends past expiry.
+- OAuth credential invalidation after an auth rejection is more precise and never
+  waits on an in-flight mint. Each cached token carries a monotonic generation, and
+  invalidation records the rejected generation on the cache entry with a lock-free
+  atomic, so it never takes the per-token lock or detaches a mint. A token at or
+  below the recorded generation is dropped before its next use and is never refreshed
+  from, so the next fetch re-mints; a newer token installed by a concurrent refresh
+  has a higher generation and is kept and reused.
 
 ### Documentation
 

rust/README.md

@@ -325,6 +325,8 @@ let sdk = ZerobusSdk::builder()
 # Ok::<(), databricks_zerobus_ingest_sdk::ZerobusError>(())
 ```
 
+When a proactive refresh fails, the SDK keeps serving the still-valid cached token rather than failing stream creation, including on non-retryable failures such as revoked credentials. The cached token is served until it actually expires, after which the next call mints a fresh one and any failure then surfaces. A proactive refresh is also capped at half the stream's [`recovery_timeout_ms`](#configuration-options), but only when the cached token has more life left than that cap, so a stalled endpoint falls back to the cached token in time. When too little of the token's life remains, the refresh runs unbounded instead, like a cold miss.
+
 ### Custom Authentication
 
 For advanced use cases, you can implement the `HeadersProvider` trait to supply your own authentication headers. This is useful for integrating with a different OAuth provider, using a centralized token caching service, or implementing alternative authentication mechanisms.

rust/sdk/Cargo.toml

@@ -51,6 +51,8 @@ futures = { workspace = true, optional = true }
 
 [dev-dependencies]
 tracing-subscriber.workspace = true
+# Deterministic virtual time (start_paused / advance) for timing-sensitive tests.
+tokio = { workspace = true, features = ["test-util"] }
 # zeroparser e2e + bench dev-deps
 criterion = { version = "0.5", default-features = false, features = ["html_reports"] }
 plotters = { version = "0.3", default-features = false, features = ["svg_backend"] }

rust/sdk/src/builder/stream_builder.rs

@@ -231,6 +231,11 @@ impl<'a> StreamBuilder<'a> {
     }
 
     /// Set the timeout in milliseconds for each recovery attempt.
+    ///
+    /// For streams authenticated with [`oauth`](Self::oauth), this also caps a
+    /// proactive OAuth token refresh at half this value, but only when the cached
+    /// token has more life left than that cap, so a stalled endpoint falls back
+    /// before the attempt deadline. Very near expiry the refresh runs unbounded.
     pub fn recovery_timeout_ms(mut self, ms: u64) -> Self {
         self.grpc_config.recovery_timeout_ms = ms;
         #[cfg(feature = "arrow-flight")]
@@ -383,20 +388,32 @@ impl<'a> StreamBuilder<'a> {
         Ok(())
     }
 
-    /// Resolve the headers provider from the stored auth config.
+    /// Resolve the headers provider from the stored auth config. For OAuth, a
+    /// proactive token refresh is bounded by half the stream's recovery timeout.
     fn resolve_headers_provider(&self) -> ZerobusResult<Arc<dyn HeadersProvider>> {
         match self.auth.as_ref() {
             Some(AuthConfig::OAuth {
                 client_id,
                 client_secret,
-            }) => Ok(Arc::new(OAuthHeadersProvider::with_cache(
-                client_id.clone(),
-                client_secret.clone(),
-                self.table_name.clone(),
-                self.sdk.workspace_id.clone(),
-                self.sdk.unity_catalog_url.clone(),
-                Arc::clone(&self.sdk.token_cache),
-            ))),
+            }) => {
+                // Give a proactive refresh half the recovery timeout so a stalled
+                // refresh falls back to the cached token before the setup deadline
+                // cancels the request. Both recovery-timeout setters write the gRPC
+                // and Arrow configs in lockstep, so grpc_config is a valid single
+                // source even for an Arrow stream; revisit if a per-transport
+                // recovery timeout is ever added.
+                let refresh_timeout =
+                    std::time::Duration::from_millis(self.grpc_config.recovery_timeout_ms) / 2;
+                Ok(Arc::new(OAuthHeadersProvider::with_cache(
+                    client_id.clone(),
+                    client_secret.clone(),
+                    self.table_name.clone(),
+                    self.sdk.workspace_id.clone(),
+                    self.sdk.unity_catalog_url.clone(),
+                    Arc::clone(&self.sdk.token_cache),
+                    Some(refresh_timeout),
+                )))
+            }
             Some(AuthConfig::HeadersProvider(p)) => Ok(Arc::clone(p)),
             #[cfg(feature = "testing")]
             Some(AuthConfig::NoAuth) => Ok(Arc::new(NoAuthHeadersProvider)),

rust/sdk/src/default_token_factory.rs

@@ -239,13 +239,23 @@ impl DefaultTokenFactory {
     }
 
     /// Parses the OAuth `expires_in` field (token lifetime in seconds) into a
-    /// `Duration`. It is optional in the OAuth spec; if it is missing or not a
-    /// positive integer the token has no known TTL and must not be cached.
+    /// `Duration`. A plain JSON integer (`3600`) and a quoted one (`"3600"`) are
+    /// both accepted.
+    ///
+    /// `expires_in` is optional in the OAuth spec; a missing value, or one that
+    /// is not a positive integer, yields `None`.
     fn parse_expires_in(body: &serde_json::Value) -> Option<Duration> {
-        body["expires_in"]
-            .as_u64()
-            .filter(|secs| *secs > 0)
-            .map(Duration::from_secs)
+        let secs = match &body["expires_in"] {
+            serde_json::Value::Number(n) => n.as_u64(),
+            serde_json::Value::String(s) => s.trim().parse::<u64>().ok(),
+            _ => None,
+        }?;
+
+        if secs == 0 {
+            return None;
+        }
+
+        Some(Duration::from_secs(secs))
     }
 
     /// Classifies HTTP status codes as retryable or non-retryable errors.
@@ -355,20 +365,43 @@ mod tests {
 
     #[test]
     fn test_parse_expires_in() {
-        let with_ttl = serde_json::json!({ "expires_in": 3600 });
+        // A JSON integer parses to that many seconds.
+        let integer = serde_json::json!({ "expires_in": 3600 });
+        assert_eq!(
+            DefaultTokenFactory::parse_expires_in(&integer),
+            Some(Duration::from_secs(3600))
+        );
+
+        // A quoted integer parses to the same value.
+        let quoted = serde_json::json!({ "expires_in": "3600" });
         assert_eq!(
-            DefaultTokenFactory::parse_expires_in(&with_ttl),
+            DefaultTokenFactory::parse_expires_in(&quoted),
             Some(Duration::from_secs(3600))
         );
 
-        let missing = serde_json::json!({ "access_token": "abc" });
+        // Surrounding whitespace in the string is trimmed.
+        let padded = serde_json::json!({ "expires_in": " 3600 " });
+        assert_eq!(
+            DefaultTokenFactory::parse_expires_in(&padded),
+            Some(Duration::from_secs(3600))
+        );
+
+        // Absent, zero, and negative all yield no TTL.
+        let missing = serde_json::json!({});
         assert_eq!(DefaultTokenFactory::parse_expires_in(&missing), None);
 
         let zero = serde_json::json!({ "expires_in": 0 });
         assert_eq!(DefaultTokenFactory::parse_expires_in(&zero), None);
 
-        // A string value (non-integer) is not usable and yields no TTL.
-        let non_numeric = serde_json::json!({ "expires_in": "3600" });
+        let negative = serde_json::json!({ "expires_in": -1 });
+        assert_eq!(DefaultTokenFactory::parse_expires_in(&negative), None);
+
+        // A fractional value is not a whole number of seconds.
+        let fractional = serde_json::json!({ "expires_in": 3600.9 });
+        assert_eq!(DefaultTokenFactory::parse_expires_in(&fractional), None);
+
+        // A non-numeric string yields no TTL.
+        let non_numeric = serde_json::json!({ "expires_in": "abc" });
         assert_eq!(DefaultTokenFactory::parse_expires_in(&non_numeric), None);
     }
 

rust/sdk/src/headers_provider.rs

@@ -3,7 +3,9 @@ use crate::token_cache::{TokenCache, DEFAULT_REFRESH_BUFFER};
 use crate::ZerobusResult;
 use async_trait::async_trait;
 use std::collections::HashMap;
+use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::Arc;
+use std::time::Duration;
 
 /// A trait for providing custom headers for gRPC requests.
 ///
@@ -47,13 +49,14 @@ pub trait HeadersProvider: Send + Sync {
     /// Returns a `ZerobusError` if header generation fails (e.g., token request fails).
     async fn get_headers(&self) -> ZerobusResult<HashMap<&'static str, String>>;
 
-    /// Invalidates any cached authentication state so the next `get_headers`
-    /// call re-derives it from scratch.
+    /// Invalidates cached authentication state that the server just rejected.
     ///
-    /// The SDK calls this when the server rejects the supplied credentials with
-    /// an authentication error during stream creation. The default is a no-op,
-    /// which is correct for providers that hold no cache; the built-in OAuth
-    /// provider overrides it to drop its cached token so the next call re-mints.
+    /// The SDK calls this when the server rejects the supplied credentials with an
+    /// authentication error during stream creation. The default is a no-op, which is
+    /// correct for providers that hold no cache. The built-in OAuth provider clears
+    /// the rejected token from its cache, so the next `get_headers` re-mints — unless
+    /// a concurrent refresh already replaced it with a newer token, which is kept and
+    /// served without re-minting.
     async fn invalidate(&self) {}
 }
 
@@ -68,6 +71,14 @@ pub struct OAuthHeadersProvider {
     workspace_id: String,
     unity_catalog_url: String,
     token_cache: Arc<TokenCache>,
+    /// How long a proactive refresh may run before it's treated as failed and the
+    /// cached token is served; `None` leaves it unbounded.
+    refresh_timeout: Option<Duration>,
+    /// Generation of the token last returned by `get_headers`, so `invalidate`
+    /// rejects only that token, not a newer one from a concurrent refresh. 0 means
+    /// there is no cached token to reject: nothing has been served yet, or the most
+    /// recently served token was not cacheable (a no-TTL response or a disabled cache).
+    last_served_generation: AtomicU64,
 }
 
 impl OAuthHeadersProvider {
@@ -92,6 +103,7 @@ impl OAuthHeadersProvider {
             workspace_id,
             unity_catalog_url,
             Arc::new(TokenCache::new(true, DEFAULT_REFRESH_BUFFER)),
+            None,
         )
     }
 
@@ -106,6 +118,7 @@ impl OAuthHeadersProvider {
         workspace_id: String,
         unity_catalog_url: String,
         token_cache: Arc<TokenCache>,
+        refresh_timeout: Option<Duration>,
     ) -> Self {
         Self {
             client_id,
@@ -114,40 +127,66 @@ impl OAuthHeadersProvider {
             workspace_id,
             unity_catalog_url,
             token_cache,
+            refresh_timeout,
+            last_served_generation: AtomicU64::new(0),
         }
     }
 }
 
 #[async_trait]
 impl HeadersProvider for OAuthHeadersProvider {
     async fn get_headers(&self) -> ZerobusResult<HashMap<&'static str, String>> {
-        let token = self
-            .token_cache
-            .get_or_fetch(
+        let fetch = |reason| {
+            DefaultTokenFactory::fetch_token(
+                &self.unity_catalog_url,
+                &self.table_name,
                 &self.client_id,
                 &self.client_secret,
-                &self.table_name,
-                |reason| {
-                    DefaultTokenFactory::fetch_token(
-                        &self.unity_catalog_url,
+                &self.workspace_id,
+                reason,
+            )
+        };
+        let (token, generation) = match self.refresh_timeout {
+            Some(refresh_timeout) => {
+                self.token_cache
+                    .get_or_fetch_within(
+                        &self.client_id,
+                        &self.client_secret,
                         &self.table_name,
+                        refresh_timeout,
+                        fetch,
+                    )
+                    .await?
+            }
+            None => {
+                self.token_cache
+                    .get_or_fetch(
                         &self.client_id,
                         &self.client_secret,
-                        &self.workspace_id,
-                        reason,
+                        &self.table_name,
+                        fetch,
                     )
-                },
-            )
-            .await?;
+                    .await?
+            }
+        };
+        // Remember the served token's generation so invalidate() rejects only it.
+        self.last_served_generation
+            .store(generation, Ordering::SeqCst);
         let mut headers = HashMap::new();
         headers.insert("authorization", format!("Bearer {}", token));
         headers.insert("x-databricks-zerobus-table-name", self.table_name.clone());
         Ok(headers)
     }
 
     async fn invalidate(&self) {
+        let rejected_generation = self.last_served_generation.load(Ordering::SeqCst);
         self.token_cache
-            .invalidate(&self.client_id, &self.client_secret, &self.table_name)
+            .invalidate(
+                &self.client_id,
+                &self.client_secret,
+                &self.table_name,
+                rejected_generation,
+            )
             .await;
     }
 }

rust/sdk/src/stream/arrow/options.rs

@@ -50,6 +50,12 @@ pub struct ArrowStreamConfigurationOptions {
     /// Values whose absolute deadline cannot be represented by the platform's
     /// monotonic clock are rejected when the stream is built.
     ///
+    /// For OAuth-authenticated streams this also caps a proactive token refresh at
+    /// half its value, but only when the cached token has more life left than that
+    /// cap, so a stalled endpoint falls back to the cached token before the attempt
+    /// deadline. When too little of the token's life remains, the refresh runs
+    /// unbounded, like a cold miss.
+    ///
     /// Default: 15,000 (15 seconds)
     pub recovery_timeout_ms: u64,
 

rust/sdk/src/stream_configuration.rs

@@ -47,6 +47,12 @@ pub struct StreamConfigurationOptions {
     ///
     /// If a recovery attempt takes longer than this, it will be retried.
     ///
+    /// For OAuth-authenticated streams this also caps a proactive token refresh at
+    /// half its value, but only when the cached token has more life left than that
+    /// cap, so a stalled endpoint falls back to the cached token before the attempt
+    /// deadline. When too little of the token's life remains, the refresh runs
+    /// unbounded, like a cold miss.
+    ///
     /// Default: 15,000 (15 seconds)
     pub recovery_timeout_ms: u64,