From 43e7a3a6cb79093fdbc90dc8d3d25c79b8357f67 Mon Sep 17 00:00:00 2001 From: damocles Date: Sun, 15 Nov 2026 14:33:49 +0100 Subject: [PATCH] swarm-queue-client: mint the next token and reconnect before the queue expires the connection --- swarm-queue-client/src/lib.rs | 164 +++++++++++++++++++++++++++++++++- 1 file changed, 162 insertions(+), 2 deletions(-) diff --git a/swarm-queue-client/src/lib.rs b/swarm-queue-client/src/lib.rs index 83c3de8b3..79af64be4 100644 --- a/swarm-queue-client/src/lib.rs +++ b/swarm-queue-client/src/lib.rs @@ -586,6 +586,30 @@ impl TokenCache { if let Some(minted) = slot.as_ref().filter(|m| m.usable_at((self.now)())) { return Ok(minted.token.token.clone()); } + self.mint_into(&mut slot, mint).await + } + + /// Mint a replacement now and cache it, whatever the cached token has + /// left. The rotation task's way of having the next connection's token + /// ready before the current one expires. + async fn rotate(&self, mint: F) -> Result<(), Error> + where + F: FnOnce() -> Fut, + Fut: Future>, + { + let mut slot = self.slot.lock().await; + self.mint_into(&mut slot, mint).await.map(drop) + } + + async fn mint_into( + &self, + slot: &mut Option, + mint: F, + ) -> Result + where + F: FnOnce() -> Fut, + Fut: Future>, + { // Stamped before the request, so the token's age is never understated. let minted_at = (self.now)(); let token = mint().await?; @@ -594,6 +618,24 @@ impl TokenCache { Ok(value) } + /// How long until the cached token has only [`TOKEN_REFRESH_SKEW`] of its + /// life left: zero once it has, or when the clock reads earlier than the + /// mint (see [`MintedToken::usable_at`]). `None` while nothing is cached. + async fn rotate_in(&self) -> Option { + let slot = self.slot.lock().await; + let minted = slot.as_ref()?; + let Ok(age) = (self.now)().duration_since(minted.minted_at) else { + return Some(std::time::Duration::ZERO); + }; + Some( + minted + .token + .lifetime + .saturating_sub(TOKEN_REFRESH_SKEW) + .saturating_sub(age), + ) + } + /// Drop the cached token when `event` is the server refusing the /// credential. Any other event, a plain disconnect included, keeps it: a /// reconnect that is not an auth refusal must not mint, or the retry loop @@ -640,9 +682,19 @@ const TOKEN_REFRESH_SKEW: std::time::Duration = std::time::Duration::from_mins(2 /// degrades every other client of it. const MAX_RECONNECT_DELAY: std::time::Duration = std::time::Duration::from_mins(1); -/// Exponential from 500ms, capped at [`MAX_RECONNECT_DELAY`]. +/// Immediate for the first attempt after a disconnect, then exponential from +/// 500ms, capped at [`MAX_RECONNECT_DELAY`]. +/// +/// The first attempt is not delayed because most disconnects are the queue +/// closing a connection whose credential expired, which a prompt reconnect +/// with the next token fixes; the backoff only has to protect the identity +/// provider from a connection that keeps failing, and it still does from the +/// second attempt on. The library's own default makes the same exception. fn reconnect_delay(attempts: usize) -> std::time::Duration { - let exp = u32::try_from(attempts.saturating_sub(1)).unwrap_or(u32::MAX); + if attempts <= 1 { + return std::time::Duration::ZERO; + } + let exp = u32::try_from(attempts - 2).unwrap_or(u32::MAX); std::cmp::min( std::time::Duration::from_millis(500u64.saturating_mul(2u64.saturating_pow(exp.min(8)))), MAX_RECONNECT_DELAY, @@ -1071,6 +1123,7 @@ pub async fn connect(cfg: QueueConfig) -> Result { // every retry" are the same code. let cache = std::sync::Arc::new(TokenCache::new(std::time::SystemTime::now)); let events = cache.clone(); + let (rotating, rotation_http, rotation_cfg) = (cache.clone(), http.clone(), cfg.clone()); let mut options = async_nats::ConnectOptions::with_auth_callback(move |_nonce| { let http = http.clone(); @@ -1136,9 +1189,61 @@ pub async fn connect(cfg: QueueConfig) -> Result { source, })?; + tokio::spawn(rotate_before_expiry( + rotating, + client.clone(), + rotation_http, + rotation_cfg, + )); Ok(client) } +/// How long [`rotate_before_expiry`] waits to look again when no token is +/// cached yet (the first connect is still pending) or minting one failed. +const ROTATION_RETRY: std::time::Duration = std::time::Duration::from_secs(30); + +/// Replace the connection's token before the queue expires it. +/// +/// The queue closes every connection when its credential expires +/// (`User Authentication Expired`, by design: `swarm-nats-auth` caps the +/// user JWT at the token's own expiry), and a request that arrives while the +/// client reconnects fails as not connected. Reconnecting on our own schedule +/// moves that drop out of the way of the expiry and, because the next token +/// is minted first, leaves the window one connection round trip long: the +/// auth callback finds a fresh token in the cache and makes no request to the +/// identity provider. +/// +/// A failed mint costs nothing yet, the current token has [`TOKEN_REFRESH_SKEW`] +/// left, so it is retried at a fixed [`ROTATION_RETRY`] interval, not hammered. +/// If every try fails the connection still ends at the expiry and the +/// reconnect path takes over, as it did before this task existed. +/// +/// Runs for the life of the process, like the client it keeps. +async fn rotate_before_expiry( + cache: std::sync::Arc, + client: async_nats::Client, + http: reqwest::Client, + cfg: QueueConfig, +) { + loop { + match cache.rotate_in().await { + None => tokio::time::sleep(ROTATION_RETRY).await, + Some(wait) if !wait.is_zero() => tokio::time::sleep(wait).await, + Some(_) => { + if let Err(e) = cache.rotate(|| mint_token(&http, &cfg, None, None)).await { + tracing::warn!( + error = %chain(&e), + "minting the next queue token ahead of expiry failed; retrying" + ); + tokio::time::sleep(ROTATION_RETRY).await; + } else if let Err(e) = client.force_reconnect().await { + tracing::warn!(error = %e, "reconnecting to the queue with the next token failed"); + } + } + } + } +} + /// Connect presenting the token `token` resolves to, as `agent` presents its /// own queue credential, with reply inboxes under [`agent_inbox_prefix`]. /// @@ -1651,6 +1756,61 @@ mod tests { assert_eq!(take(&cache, &mints).await, "token-2"); } + /// The rotation task sleeps until the cached token is `TOKEN_REFRESH_SKEW` + /// from expiry, and not before. + #[tokio::test] + async fn rotation_is_due_when_the_token_reaches_the_refresh_skew() { + // The stub mint's one-hour life less the two-minute skew. + const LIFE: u64 = 3599 - 120; + let secs = std::time::Duration::from_secs; + let clock = FakeClock::starting_at(START); + let (cache, mints) = (clock.cache(), AtomicUsize::new(0)); + assert_eq!(cache.rotate_in().await, None, "nothing cached yet"); + + take(&cache, &mints).await; + assert_eq!(cache.rotate_in().await, Some(secs(LIFE))); + + clock.advance(60); + assert_eq!(cache.rotate_in().await, Some(secs(LIFE - 60))); + clock.advance(LIFE - 60); + assert_eq!(cache.rotate_in().await, Some(std::time::Duration::ZERO)); + } + + /// A rotation replaces the token even though the cached one is usable, so + /// the next connect presents it without asking the identity provider. + #[tokio::test] + async fn rotating_replaces_a_usable_token_and_the_next_connect_reuses_it() { + let clock = FakeClock::starting_at(START); + let (cache, mints) = (clock.cache(), AtomicUsize::new(0)); + assert_eq!(take(&cache, &mints).await, "token-1"); + + cache + .rotate(|| async { + mints.fetch_add(1, Ordering::SeqCst); + Ok(CachedToken { + token: "token-2".to_owned(), + lifetime: std::time::Duration::from_secs(3599), + }) + }) + .await + .expect("the stub mint succeeds"); + + assert_eq!(take(&cache, &mints).await, "token-2"); + assert_eq!(mints.load(Ordering::SeqCst), 2, "the connect did not mint"); + } + + /// The first attempt after a disconnect is immediate; the backoff, and its + /// cap, protect the identity provider from the second attempt on. + #[test] + fn the_first_reconnect_is_immediate_and_the_backoff_follows() { + let ms = std::time::Duration::from_millis; + assert_eq!(reconnect_delay(0), std::time::Duration::ZERO); + assert_eq!(reconnect_delay(1), std::time::Duration::ZERO); + assert_eq!(reconnect_delay(2), ms(500)); + assert_eq!(reconnect_delay(3), ms(1000)); + assert_eq!(reconnect_delay(100), MAX_RECONNECT_DELAY); + } + /// A suspend moves the wall clock and leaves the monotonic one behind. #[tokio::test] async fn a_token_is_not_reused_after_the_clock_jumps_past_its_life() {