swarm-queue-client: mint the next token and reconnect before the queue expires the connection

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<F, Fut>(&self, mint: F) -> Result<(), Error>
+    where
+        F: FnOnce() -> Fut,
+        Fut: Future<Output = Result<CachedToken, Error>>,
+    {
+        let mut slot = self.slot.lock().await;
+        self.mint_into(&mut slot, mint).await.map(drop)
+    }
+
+    async fn mint_into<F, Fut>(
+        &self,
+        slot: &mut Option<MintedToken>,
+        mint: F,
+    ) -> Result<String, Error>
+    where
+        F: FnOnce() -> Fut,
+        Fut: Future<Output = Result<CachedToken, Error>>,
+    {
         // 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<std::time::Duration> {
+        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<async_nats::Client, Error> {
     // 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<async_nats::Client, Error> {
             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<TokenCache>,
+    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() {