Grow and jitter consecutive throttle waits in the retrier

nativelink-store/src/azure_blob_store.rs

@@ -227,7 +227,8 @@ where
                 spec.common.retry.clone(),
             )
             .with_client("azure")
-            .with_throttle_budget(throttle_retry_budget),
+            .with_throttle_budget(throttle_retry_budget)
+            .with_throttle_max_wait(throttle_retry_max_delay),
             consider_expired_after_s: i64::from(spec.common.consider_expired_after_s),
             max_retry_buffer_per_request: spec
                 .common

nativelink-store/tests/azure_blob_store_test.rs

@@ -553,6 +553,49 @@ async fn throttles_do_not_spend_the_retry_budget() -> Result<(), Error> {
     Ok(())
 }
 
+/// Consecutive throttles with no header double from the floor up to the
+/// cap: 2, 4, 8, 16, 32 s for five of them under the 60 s default cap, 62 s
+/// in all, where five floors would be 10 s.
+#[nativelink_test(start_paused = true)]
+async fn consecutive_throttles_double_from_the_floor_up_to_the_cap() -> Result<(), Error> {
+    const VALUE: &str = "test_content";
+
+    let mock = MockTransport::new();
+    for _ in 0..5 {
+        mock.push_download(error_response(StatusCode::ServiceUnavailable));
+    }
+    mock.push_download(download_ok(StatusCode::Ok, VALUE.as_bytes().to_vec()));
+
+    let retry = Retry {
+        max_retries: 0,
+        delay: 0.0,
+        jitter: 0.0,
+        ..Default::default()
+    };
+    let store = create_test_store(Arc::clone(&mock), Some(retry), None)?;
+    let floor = Duration::from_secs_f32(CommonObjectSpec::default().throttle_retry_delay_s);
+    let cap = Duration::from_secs_f32(CommonObjectSpec::default().throttle_retry_max_delay_s);
+    let doubled = floor + floor * 2 + floor * 4 + floor * 8 + floor * 16;
+    assert!(
+        floor * 16 < cap,
+        "the fifth wait must still be under the cap"
+    );
+
+    let started = tokio::time::Instant::now();
+    let result = store
+        .get_part_unchunked(create_test_digest()?, 0, None)
+        .await?;
+
+    assert_eq!(result, VALUE.as_bytes());
+    assert_eq!(mock.request_count(), 6);
+    assert!(
+        started.elapsed() >= doubled && started.elapsed() < doubled + floor,
+        "waited {:?}, expected the doubled waits of {doubled:?}, not five floors",
+        started.elapsed()
+    );
+    Ok(())
+}
+
 #[nativelink_test(start_paused = true)]
 async fn a_throttled_download_without_a_header_waits_the_configured_floor() -> Result<(), Error> {
     const VALUE: &str = "test_content";

nativelink-util/src/retry.rs

@@ -99,6 +99,10 @@ pub struct Retrier {
     /// How long one `retry` call may spend on `RetryAfter` waits before
     /// they start spending attempts. Zero means they always do.
     throttle_budget: Duration,
+    /// Longest a throttle wait may grow to. Set, each consecutive throttle
+    /// doubles the wait the store asked for, capped here, and jitters it.
+    /// Unset, the wait is what was asked.
+    throttle_max_wait: Option<Duration>,
 }
 
 impl core::fmt::Debug for Retrier {
@@ -140,6 +144,7 @@ impl Retrier {
             client: "unknown",
             outcome_fn: None,
             throttle_budget: Duration::ZERO,
+            throttle_max_wait: None,
         }
     }
 
@@ -167,6 +172,46 @@ impl Retrier {
         self
     }
 
+    /// Lets consecutive throttles wait longer each time. The first throttle
+    /// in a row waits what the server asked, the second twice that, the
+    /// third four times, and so on; `max_wait` is the ceiling that growth
+    /// stops at, not a wait of its own. The grown wait is spread upward by
+    /// the jitter and never less than what was asked. Growth applies to the
+    /// waits the throttle budget pays for; once the budget is spent the
+    /// backoff ladder takes over and waits at least the plain ask, so the
+    /// ladder's total stays what the retry config says. A store that is
+    /// told to slow down with no `Retry-After` otherwise has every client
+    /// retry in lockstep at the configured floor.
+    #[must_use]
+    pub const fn with_throttle_max_wait(mut self, max_wait: Duration) -> Self {
+        self.throttle_max_wait = Some(max_wait);
+        self
+    }
+
+    /// The wait for the `n`-th consecutive throttle whose server asked for
+    /// `asked`: `asked` itself for the first, doubled once per throttle
+    /// before it, capped at the configured ceiling, then jittered.
+    fn throttle_wait(&self, asked: Duration, throttles_in_a_row: u32) -> Duration {
+        let Some(max_wait) = self.throttle_max_wait else {
+            return asked;
+        };
+        // One doubling per earlier throttle in the row; past 31 the
+        // multiplier would overflow and the ceiling applies anyway.
+        let doublings = throttles_in_a_row.saturating_sub(1).min(31);
+        let grown = asked.saturating_mul(1 << doublings).min(max_wait);
+        // `Retry::make_jitter_fn` spreads a delay either way, and on the
+        // first throttle half of every fleet would land at or below the
+        // ask and retry in lockstep. A low result is mirrored above the
+        // grown wait instead, so every step spreads upward only; a jitter
+        // that zeroes delays (as tests install) leaves the grown wait.
+        let jittered = (self.jitter_fn)(grown);
+        if jittered.is_zero() {
+            grown
+        } else {
+            grown + jittered.abs_diff(grown)
+        }
+    }
+
     fn report(&self, outcome: RetryOutcome) {
         record_retry_outcome(self.client, outcome.as_str());
         if let Some(outcome_fn) = &self.outcome_fn {
@@ -225,6 +270,7 @@ impl Retrier {
             tokio::pin!(operation);
             let mut attempt = 0;
             let mut throttled = Duration::ZERO;
+            let mut throttles_in_a_row: u32 = 0;
             loop {
                 attempt += 1;
                 let (err, min_wait, is_throttle) = match operation.next().await {
@@ -245,8 +291,14 @@ impl Retrier {
                     Some(RetryResult::Err(e)) => {
                         return Err(e.append(format!("On attempt {attempt}")));
                     }
-                    Some(RetryResult::Retry(err)) => (err, Duration::ZERO, false),
-                    Some(RetryResult::RetryAfter(err, min_wait)) => (err, min_wait, true),
+                    Some(RetryResult::Retry(err)) => {
+                        throttles_in_a_row = 0;
+                        (err, Duration::ZERO, false)
+                    }
+                    Some(RetryResult::RetryAfter(err, min_wait)) => {
+                        throttles_in_a_row = throttles_in_a_row.saturating_add(1);
+                        (err, min_wait, true)
+                    }
                 };
                 if !self.should_retry(err.code) {
                     if err.code == Code::NotFound {
@@ -261,15 +313,20 @@ impl Retrier {
                 // not an attempt: the server said when to come back, and
                 // the budget bounds how long that is listened to. A zero
                 // wait would spin, so it goes through the backoff instead.
-                if is_throttle
-                    && !min_wait.is_zero()
-                    && throttled.saturating_add(min_wait) <= self.throttle_budget
-                {
-                    throttled = throttled.saturating_add(min_wait);
-                    attempt -= 1;
-                    self.report(RetryOutcome::Retried);
-                    (self.sleep_fn)(min_wait).await;
-                    continue;
+                // The grown wait is what the budget pays for. Past the
+                // budget the backoff ladder below waits at least the plain
+                // ask, never the grown one, so the ladder's total stays what
+                // the retry config says and the caller's own deadline (a
+                // worker's download window, Bazel's remote timeout) holds.
+                if is_throttle && !min_wait.is_zero() {
+                    let wait = self.throttle_wait(min_wait, throttles_in_a_row);
+                    if throttled.saturating_add(wait) <= self.throttle_budget {
+                        throttled = throttled.saturating_add(wait);
+                        attempt -= 1;
+                        self.report(RetryOutcome::Retried);
+                        (self.sleep_fn)(wait).await;
+                        continue;
+                    }
                 }
                 let Some(backoff) = iter.next() else {
                     self.report(RetryOutcome::Exhausted);