[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.

[danilotrninic-db]

Done. The fallback now serves the still valid cached token on any proactive refresh error, whether or not it is retryable. On a refresh error serve_valid_cached_fallback returns the cached token as long as it is unexpired and not rejected, and it surfaces the error only when there is no valid token to fall back to. An actual server rejection still reaches HeadersProvider::invalidate separately, as you noted. Covered by refresh_failure_serves_still_valid_token.

[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.

[danilotrninic-db]

Done. Each cached token now carries a monotonic generation, and invalidate records the rejected generation on the slot with an atomic fetch_max. A token whose generation is at or below that watermark is dropped before its next use, so a newer token installed by a concurrent refresh has a higher generation and is kept. Rejecting token A no longer removes token B.

The lookup race is gone because invalidate no longer detaches the slot from the map and no longer takes the token mutex. It only raises the watermark on the slot atomically. A lookup that cloned the slot before the invalidation therefore holds the same slot, observes the raised watermark, and drops the rejected token instead of returning it. Covered by invalidate_only_clears_the_rejected_token and invalidate_does_not_block_on_an_in_flight_mint.

[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.

[danilotrninic-db]

Done. A refresh timeout of half the stream's recovery_timeout_ms is threaded into the cache, and a proactive refresh mint is wrapped in tokio::time::timeout. On expiry it becomes a TokenFetchError and the cached token is served before the outer setup deadline. Cold misses stay unbounded because they have no fallback. One refinement over the snippet is that the bound applies only when the cached token has more life left than the timeout. When too little remains to fall back on, the refresh runs unbounded like a cold miss. Covered by stalled_refresh_serves_cached_token.

[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.

[danilotrninic-db]

Done. I added refresh_retry_at to CachedToken as suggested. On a refresh fallback arm_refresh_backoff sets a short backoff that scales with the refresh buffer, shrinks as the token nears expiry, and is capped at expires_at. While inside that window needs_refresh returns false, so a burst of stream creations serves the cached token instead of each caller minting again. Explicit auth invalidation still takes effect immediately because it raises the rejection watermark independently of the backoff. Covered by failed_refresh_backoff_suppresses_repeat_mint.

[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.

[danilotrninic-db]

Done. The test now clones the occupied slot and waits until Arc::strong_count(&slot) reaches FOLLOWERS + 3, which counts the map, the leader, this test handle, and every follower, before it releases the leader. That way all followers are provably queued on the per key lock rather than being incidental cache hits. A generous timeout guards against a regression hanging the test.

[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.

[danilotrninic-db]

Done. fetch_started_at is captured with Instant::now() before the mint, and expires_at is fetch_started_at + ttl, so a slow response cannot make the token look valid for longer than the issuer allows. A token whose anchored expiry has already passed when it arrives is treated as dead on arrival. On a cold miss it surfaces a retryable error, and during a refresh it falls back to the still valid cached token instead of caching the expired one. Covered by cold_miss_dead_on_arrival_token_surfaces_error and refresh_dead_on_arrival_keeps_cached_token.