retry: share one backon helper for the client secret reads and convert the remaining doubling tables

hive-sock-client/Cargo.toml

@@ -16,7 +16,13 @@ workspace = true
 
 [dependencies]
 anyhow.workspace = true
+backon.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 tokio.workspace = true
 tracing.workspace = true
+
+[dev-dependencies]
+# `test-util` for `#[tokio::test(start_paused = true)]`: the ride-out test
+# waits out the whole schedule on paused time.
+tokio = { workspace = true, features = ["test-util"] }

hive-sock-client/src/lib.rs

@@ -38,17 +38,20 @@ use std::path::Path;
 use std::time::Duration;
 
 use anyhow::{Context, Result, anyhow};
+use backon::{BackoffBuilder as _, ExponentialBuilder, Retryable as _};
 use serde::Serialize;
 use serde::de::DeserializeOwned;
 use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
 use tokio::net::UnixStream;
 
-/// Backoff schedule for [`Retry::RideOutRestart`]. Five entries → up to 5
-/// retries on top of the initial attempt; total wall-clock cap =
-/// 2+4+8+16+30 = 60s. Sized to ride out a service restart (systemd usually
-/// has the unix socket back inside ~5s) without the caller having to
-/// handle the transient itself.
-const RIDE_OUT_RESTART_BACKOFFS_MS: &[u64] = &[2_000, 4_000, 8_000, 16_000, 30_000];
+/// Backoff for [`Retry::RideOutRestart`]: waits of 2, 4, 8, 16 and 30 seconds,
+/// so up to 5 retries on top of the initial attempt and 60s of wall-clock in
+/// all. Sized to ride out a service restart (systemd usually has the unix
+/// socket back inside ~5s) without the caller having to handle the transient
+/// itself.
+const RIDE_OUT_RESTART_FIRST_WAIT: Duration = Duration::from_secs(2);
+const RIDE_OUT_RESTART_MAX_WAIT: Duration = Duration::from_secs(30);
+const RIDE_OUT_RESTART_RETRIES: usize = 5;
 
 /// Bound on connecting. The listener is on the same host, so a connect
 /// still pending after this long is not going to complete.
@@ -73,7 +76,7 @@ pub enum Retry {
     /// backoff would stack sleeps on top of it and delay the rest of the
     /// batch.
     None,
-    /// Back off on `RIDE_OUT_RESTART_BACKOFFS_MS` (~60s total). For
+    /// Back off on the `RIDE_OUT_RESTART_*` schedule (~60s total). For
     /// callers with no natural retry of their own, where a surfaced
     /// transient costs more than the wait.
     RideOutRestart,
@@ -81,10 +84,14 @@ pub enum Retry {
 
 impl Retry {
     /// Sleep schedule between attempts; its length is the retry budget.
-    fn backoffs(self) -> &'static [u64] {
+    fn backoff(self) -> ExponentialBuilder {
+        let schedule = ExponentialBuilder::default();
         match self {
-            Self::None => &[],
-            Self::RideOutRestart => RIDE_OUT_RESTART_BACKOFFS_MS,
+            Self::None => schedule.with_max_times(0),
+            Self::RideOutRestart => schedule
+                .with_min_delay(RIDE_OUT_RESTART_FIRST_WAIT)
+                .with_max_delay(RIDE_OUT_RESTART_MAX_WAIT)
+                .with_max_times(RIDE_OUT_RESTART_RETRIES),
         }
     }
 }
@@ -254,40 +261,42 @@ async fn with_retry(
     retry: Retry,
     response_timeout: Duration,
 ) -> Result<(String, u32)> {
-    let backoffs = retry.backoffs();
-    let mut attempt: usize = 0;
-    loop {
-        match try_once(socket, payload, mode, response_timeout).await {
-            Ok(line) => return Ok((line, u32::try_from(attempt).unwrap_or(u32::MAX))),
-            Err(e) if e.is::<ResponseTimeout>() => return Err(e),
-            Err(e) => {
-                let Some(&sleep_ms) = backoffs.get(attempt) else {
-                    return Err(exhausted(e, backoffs));
-                };
-                tracing::warn!(
-                    attempt = attempt + 1,
-                    sleep_ms,
-                    socket = %socket.display(),
-                    error = %e,
-                    "hive socket attempt failed; retrying"
-                );
-                tokio::time::sleep(Duration::from_millis(sleep_ms)).await;
-                attempt += 1;
+    let mut retries: u32 = 0;
+    let line = (|| try_once(socket, payload, mode, response_timeout))
+        .retry(retry.backoff())
+        .when(|e| !e.is::<ResponseTimeout>())
+        .notify(|e, wait| {
+            retries += 1;
+            tracing::warn!(
+                attempt = retries,
+                sleep_ms = u64::try_from(wait.as_millis()).unwrap_or(u64::MAX),
+                socket = %socket.display(),
+                error = %e,
+                "hive socket attempt failed; retrying"
+            );
+        })
+        .await
+        .map_err(|e| {
+            if e.is::<ResponseTimeout>() {
+                e
+            } else {
+                exhausted(e, retry)
             }
-        }
-    }
+        })?;
+    Ok((line, retries))
 }
 
 /// Annotate the final failure with how long the retry schedule tried, so a
 /// surfaced error says whether it was one shot or a full ride-out.
-fn exhausted(err: anyhow::Error, backoffs: &[u64]) -> anyhow::Error {
-    if backoffs.is_empty() {
+fn exhausted(err: anyhow::Error, retry: Retry) -> anyhow::Error {
+    let waits: Vec<Duration> = retry.backoff().build().collect();
+    if waits.is_empty() {
         return err;
     }
-    let total_s = backoffs.iter().sum::<u64>() / 1_000;
+    let total_s = waits.iter().sum::<Duration>().as_secs();
     err.context(format!(
         "gave up after {} retries over ~{total_s}s",
-        backoffs.len()
+        waits.len()
     ))
 }
 
@@ -395,6 +404,8 @@ mod tests {
     use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
     use tokio::net::UnixListener;
 
+    use backon::BackoffBuilder as _;
+
     use super::{Retry, notify, request, request_within};
 
     /// A fresh socket path under the test's temp dir, unique per test.
@@ -513,8 +524,25 @@ mod tests {
     /// harness's tool callers depend on to not surface a restart.
     #[test]
     fn ride_out_restart_budget_is_a_minute() {
-        let total_ms: u64 = Retry::RideOutRestart.backoffs().iter().sum();
-        assert_eq!(total_ms, 60_000);
-        assert!(Retry::None.backoffs().is_empty());
+        let waits: Vec<Duration> = Retry::RideOutRestart.backoff().build().collect();
+        assert_eq!(waits, [2, 4, 8, 16, 30].map(Duration::from_secs));
+        assert_eq!(waits.iter().sum::<Duration>(), Duration::from_mins(1));
+        assert_eq!(Retry::None.backoff().build().count(), 0);
+    }
+
+    /// A socket that never comes back is waited on for the whole schedule,
+    /// and the surfaced error says so.
+    #[tokio::test(start_paused = true)]
+    async fn ride_out_restart_gives_up_after_the_whole_schedule() {
+        let bogus = std::path::Path::new("/nonexistent/hive/agent.sock");
+        let started = tokio::time::Instant::now();
+        let e = notify(bogus, &(), Retry::RideOutRestart)
+            .await
+            .expect_err("connect to a non-existent socket must fail");
+        assert_eq!(started.elapsed(), Duration::from_mins(1));
+        assert!(
+            format!("{e:#}").contains("gave up after 5 retries over ~60s"),
+            "{e:#}"
+        );
     }
 }