[Rust] Harden OAuth token caching

[danilotrninic-db]

What changes are proposed in this pull request?

Parse the OAuth expires_in field from a quoted integer ("3600") in addition
to a plain JSON integer. Previously a token endpoint that returned expires_in
as a JSON string was read as "no lifetime reported," so the token was fetched
fresh on every stream creation instead of being cached.

How this addresses #607

#607 has three asks:

1. Parse quoted expires_in — the code change above.
2. Continue serving a valid cached token when a proactive refresh fails —
   already implemented in TokenCache::get_or_fetch; this PR verifies it and
   adds tests pinning it. The one edge case (a refresh that succeeds with no
   usable expires_in) is deliberately left returning the fresh token uncached
   rather than falling back to the near-expiry cached one — a missing lifetime is
   missing metadata about the token, not evidence it is bad, and the freshly
   minted token is the more likely of the two to still be valid.
3. Add deterministic tests — new TokenCache tests covering concurrency
   invalidation, cancellation, and expiry.

How is this tested?

cargo test --workspace (822 tests, green); cargo fmt --check and
cargo clippy clean. New unit tests cover parse_expires_in (integer, quoted
integer, whitespace, and reject cases) and the TokenCache behaviors above.

No live-server testing was needed: the sole production change is a pure parsing
function with no I/O or state, fully exercised by the unit tests above.

[teodordelibasic-db]

Not all of the comments are directly related to changes in this PR and asks in the issue, but we can broaden the scope a bit. 🙂

[teodordelibasic-db]

Hm okay I can't recall exactly what we discussed offline, but maybe it makes sense to preserve an unexpired cached token after any proactive refresh error. The err.is_retryable() gate classifies the new mint attempt, but it does not say that the access token already in hand is invalid. InvalidUCTokenError also covers malformed successful responses, while an actual rejection from Zerobus reaches HeadersProvider::invalidate separately.

This means stream creation can fail inside the refresh window even though the cached token has not expired. The fallback can be based on the cached token alone:

if let Some(cached) = guard.as_ref().filter(|cached| !cached.is_expired()) {
    warn!(
        table = %table_name,
        retryable = err.is_retryable(),
        "token refresh failed; serving still-valid cached token"
    );
    return Ok(cached.value.clone());
}
return Err(err);

The existing refresh_failure_propagates_non_retryable_error test can use the same setup but expect "valid"; it fails with the current gate and passes once the fallback is independent of retryability.

[teodordelibasic-db]

Could we make invalidation target the cache generation that supplied the rejected token? A stream can receive token A, another stream can refresh the shared slot to token B, and a later rejection of token A then removes token B because invalidate knows only the cache key. There is also a lookup race: get_or_fetch clones the slot before taking its mutex, so invalidation can detach that slot and the pending lookup can still return token A from its Arc.

A generation handle returned with the token would let the provider invalidate only the rejected value:

struct TokenResult {
    value: String,
    generation: TokenGeneration,
}

The provider can retain that handle, and invalidation can remove the map entry only when the handle still matches the slot's current generation. After taking the slot mutex, lookup also needs to confirm that the same Arc is still current in the map.

The private token_cache test module can cover the rollover deterministically: seed a within-buffer token and retain its generation, refresh to a healthy token, invalidate the old generation, then assert that a following lookup returns the healthy token without invoking its mint closure. The current key-only invalidation invokes that closure; generation-aware invalidation does not.

[teodordelibasic-db]

Could we bound proactive refresh before the enclosing stream connection deadline? get_or_fetch waits directly on the mint future, and the default reqwest client has no total request timeout. The gRPC and Arrow connection paths wrap all of setup in recovery_timeout_ms, which defaults to 15 seconds. If the token endpoint stalls, that outer timeout cancels get_or_fetch, so the refresh never returns an error and the valid cached-token fallback never runs.

One approach is to thread a refresh timeout into the cache and convert only proactive refresh expiry into a normal mint error:

let fetch_result = match reason {
    MintReason::Refresh => match tokio::time::timeout(refresh_timeout, fetch(reason)).await {
        Ok(result) => result,
        Err(_) => Err(ZerobusError::TokenFetchError(
            "proactive token refresh timed out".to_string(),
        )),
    },
    _ => fetch(reason).await,
};

That leaves cold misses on the normal connection deadline because they have no cached fallback, while a stalled refresh can still return the unexpired token in time to create the stream.

[teodordelibasic-db]

Could we add a short retry deadline after a failed proactive refresh? Returning the cached token leaves it inside the refresh window, so every caller waiting on the per-key mutex performs the same mint before it can use the fallback. A burst of stream creation can therefore turn one token endpoint failure into many sequential requests and consume each caller's connection budget.

A small cache field can suppress that repeated work without extending token validity:

struct CachedToken {
    value: String,
    expires_at: Instant,
    refresh_retry_at: Option<Instant>,
}

On refresh failure, set refresh_retry_at to a short backoff capped at expires_at. needs_refresh can return false before that deadline, and explicit authentication invalidation still removes the entry immediately.

A focused unit test can seed a 30-second token, make one refresh fail, then perform another lookup with a mint closure that panics. It fails today because the second lookup invokes the closure; after the backoff it returns the cached token without invoking it.

[teodordelibasic-db]

queued_tx.send(()) only proves that each follower task started. It runs before get_or_fetch, so the parent can release the leader before any follower waits on the slot. In that schedule all followers are ordinary cache hits and calls == 1 passes without exercising single-flight contention.

The test already has private access to the map, so it can wait until all followers have cloned the occupied slot before releasing the leader. For example:

let slot = {
    let entries = cache.entries.lock().await;
    Arc::clone(entries.get(&TokenKey::new("id", "secret", "c.s.t")).unwrap())
};

tokio::time::timeout(Duration::from_secs(1), async {
    while Arc::strong_count(&slot) < FOLLOWERS + 3 {
        tokio::task::yield_now().await;
    }
})
.await
.expect("followers did not reach the occupied slot");

gate.notify_one();

The count includes the map, the leader, and this test handle in addition to every follower. With the per-key lock, all followers remain queued until the gate opens; if same-key minting stops being single-flight, the timeout or final mint count fails deterministically.

[teodordelibasic-db]

Could we anchor expires_in before starting the mint? The cache currently adds the lifetime to Instant::now() after the response arrives. A slow response therefore makes the SDK consider the token valid for longer than the issuer does, and a response slower than the reported lifetime can return and cache a token that is already expired.

let fetch_started_at = Instant::now();
let fetched = match fetch(reason).await {
    Ok(fetched) => fetched,
    Err(err) => {
        if let Some(cached) = guard.as_ref().filter(|cached| !cached.is_expired()) {
            return Ok(cached.value.clone());
        }
        return Err(err);
    }
};
let expires_at = fetched
    .expires_in
    .and_then(|ttl| fetch_started_at.checked_add(ttl));

if expires_at.is_some_and(|deadline| deadline <= Instant::now()) {
    if let Some(cached) = guard.as_ref().filter(|cached| !cached.is_expired()) {
        return Ok(cached.value.clone());
    }
    return Err(ZerobusError::TokenFetchError(
        "fetched OAuth token expired before arrival".to_string(),
    ));
}

If that deadline is already past when the response arrives, a refresh can retain an older unexpired token and a cold miss can return a retryable mint error instead of sending an expired token. The new zero-TTL test does not cover this path because parse_expires_in maps zero to None, so production code never installs that cache entry.