diff --git a/nativelink-store/src/azure_blob_store.rs b/nativelink-store/src/azure_blob_store.rs index 0de1fd508b9..2056767eaf2 100644 --- a/nativelink-store/src/azure_blob_store.rs +++ b/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 diff --git a/nativelink-store/tests/azure_blob_store_test.rs b/nativelink-store/tests/azure_blob_store_test.rs index 685786b329a..3d919b3c9d3 100644 --- a/nativelink-store/tests/azure_blob_store_test.rs +++ b/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"; diff --git a/nativelink-util/src/retry.rs b/nativelink-util/src/retry.rs index e1a67e53499..560bff3a4a9 100644 --- a/nativelink-util/src/retry.rs +++ b/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, } 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); diff --git a/nativelink-util/tests/retry_test.rs b/nativelink-util/tests/retry_test.rs index daa3d43010c..947fcc437b2 100644 --- a/nativelink-util/tests/retry_test.rs +++ b/nativelink-util/tests/retry_test.rs @@ -535,3 +535,253 @@ async fn throttles_within_the_budget_report_retries_and_recovery() -> Result<(), ); Ok(()) } + +fn growing_retrier( + max_retries: usize, + budget: Duration, + cap: Duration, + sleeps: &Arc>>, +) -> Retrier { + let sleeps = Arc::clone(sleeps); + Retrier::new( + Arc::new(move |duration| { + sleeps.lock().unwrap().push(duration); + Box::pin(ready(())) + }), + Arc::new(|delay| delay), + Retry { + max_retries, + ..Default::default() + }, + ) + .with_throttle_budget(budget) + .with_throttle_max_wait(cap) +} + +/// With a throttle cap the n-th consecutive throttle waits the ask times +/// `2^(n-1)`, no more than the cap, and the budget counts the grown waits. +#[nativelink_test] +async fn consecutive_throttles_double_their_wait_up_to_the_cap() -> Result<(), Error> { + let sleeps = Arc::new(Mutex::new(Vec::new())); + let retrier = growing_retrier( + 1, + Duration::from_secs(100), + Duration::from_secs(30), + &sleeps, + ); + let run_count = Arc::new(AtomicI32::new(0)); + + let result = Pin::new(&retrier) + .retry(repeat_with(|| { + if run_count.fetch_add(1, Ordering::Relaxed) < 6 { + return RetryResult::RetryAfter( + make_err!(Code::Unavailable, "Throttled"), + Duration::from_secs(2), + ); + } + RetryResult::Ok(true) + })) + .await?; + + assert_eq!(result, true); + assert_eq!(run_count.load(Ordering::Relaxed), 7); + assert_eq!( + *sleeps.lock().unwrap(), + [2, 4, 8, 16, 30, 30] + .into_iter() + .map(Duration::from_secs) + .collect::>(), + "each throttle doubles the last, capped" + ); + Ok(()) +} + +/// The cap is a ceiling, not the wait: a single throttle under a large cap +/// waits only what was asked. +#[nativelink_test] +async fn the_first_throttle_waits_the_ask_not_the_cap() -> Result<(), Error> { + let sleeps = Arc::new(Mutex::new(Vec::new())); + let retrier = growing_retrier(1, Duration::from_secs(100), Duration::from_mins(1), &sleeps); + let run_count = Arc::new(AtomicI32::new(0)); + + let result = Pin::new(&retrier) + .retry(repeat_with(|| { + if run_count.fetch_add(1, Ordering::Relaxed) < 1 { + return RetryResult::RetryAfter( + make_err!(Code::Unavailable, "Throttled"), + Duration::from_secs(2), + ); + } + RetryResult::Ok(true) + })) + .await?; + + assert_eq!(result, true); + assert_eq!( + *sleeps.lock().unwrap(), + vec![Duration::from_secs(2)], + "one throttle waits the ask, not the 60 s cap" + ); + Ok(()) +} + +/// A jitter function that zeroes every delay, as tests install to keep +/// backoffs instant, cannot flatten the growth back to the ask. +#[nativelink_test] +async fn a_zeroing_jitter_does_not_undo_the_growth() -> Result<(), Error> { + let sleeps = Arc::new(Mutex::new(Vec::new())); + let recorded = Arc::clone(&sleeps); + let retrier = Retrier::new( + Arc::new(move |duration| { + recorded.lock().unwrap().push(duration); + Box::pin(ready(())) + }), + Arc::new(|_| Duration::ZERO), + Retry { + max_retries: 1, + ..Default::default() + }, + ) + .with_throttle_budget(Duration::from_secs(100)) + .with_throttle_max_wait(Duration::from_secs(30)); + let run_count = Arc::new(AtomicI32::new(0)); + + let result = Pin::new(&retrier) + .retry(repeat_with(|| { + if run_count.fetch_add(1, Ordering::Relaxed) < 4 { + return RetryResult::RetryAfter( + make_err!(Code::Unavailable, "Throttled"), + Duration::from_secs(2), + ); + } + RetryResult::Ok(true) + })) + .await?; + + assert_eq!(result, true); + assert_eq!( + *sleeps.lock().unwrap(), + [2, 4, 8, 16] + .into_iter() + .map(Duration::from_secs) + .collect::>(), + "the growth stands when the jitter returns nothing" + ); + Ok(()) +} + +/// The grown waits are what the budget is spent on: 2 + 4 + 8 fits a 15 s +/// budget and the fourth throttle (16 s grown) falls through to the backoff +/// ladder, where it spends the one allowed attempt and waits the plain ask, +/// not the grown one, so the ladder's total stays what the config says. +#[nativelink_test] +async fn grown_waits_spend_the_budget() -> Result<(), Error> { + let sleeps = Arc::new(Mutex::new(Vec::new())); + let retrier = growing_retrier(1, Duration::from_secs(15), Duration::from_mins(1), &sleeps); + let run_count = Arc::new(AtomicI32::new(0)); + + let result = Pin::new(&retrier) + .retry(repeat_with(|| { + run_count.fetch_add(1, Ordering::Relaxed); + RetryResult::::RetryAfter( + make_err!(Code::Unavailable, "Throttled"), + Duration::from_secs(2), + ) + })) + .await; + + assert!(result.is_err()); + assert_eq!( + run_count.load(Ordering::Relaxed), + 5, + "three budgeted throttles, then one attempt and its retry" + ); + assert_eq!( + *sleeps.lock().unwrap(), + [2, 4, 8, 2] + .into_iter() + .map(Duration::from_secs) + .collect::>(), + "the fourth wait is the backoff against the plain 2 s ask" + ); + Ok(()) +} + +/// A transient error that is not a throttle ends the streak: the next +/// throttle waits the plain ask again. +#[nativelink_test] +async fn a_plain_retry_resets_the_throttle_streak() -> Result<(), Error> { + let sleeps = Arc::new(Mutex::new(Vec::new())); + let retrier = growing_retrier(5, Duration::from_secs(100), Duration::from_mins(1), &sleeps); + let run_count = Arc::new(AtomicI32::new(0)); + + let result = Pin::new(&retrier) + .retry(repeat_with(|| { + match run_count.fetch_add(1, Ordering::Relaxed) { + 0 | 1 | 3 => RetryResult::RetryAfter( + make_err!(Code::Unavailable, "Throttled"), + Duration::from_secs(2), + ), + 2 => RetryResult::Retry(make_err!(Code::Unavailable, "Blip")), + _ => RetryResult::Ok(true), + } + })) + .await?; + + assert_eq!(result, true); + let sleeps = sleeps.lock().unwrap(); + assert_eq!(sleeps[0], Duration::from_secs(2)); + assert_eq!(sleeps[1], Duration::from_secs(4)); + // sleeps[2] is the backoff ladder's first step (delay 0 by default). + assert_eq!( + sleeps[3], + Duration::from_secs(2), + "the streak restarted after the plain retry" + ); + Ok(()) +} + +/// A jitter that lands low is mirrored above the grown wait, so no throttle +/// ever sits on the floor: a tenth of the wait becomes nine tenths on top. +#[nativelink_test] +async fn a_low_jitter_is_mirrored_above_the_grown_wait() -> Result<(), Error> { + let sleeps = Arc::new(Mutex::new(Vec::new())); + let retrier = Retrier::new( + Arc::new({ + let sleeps = Arc::clone(&sleeps); + move |duration| { + sleeps.lock().unwrap().push(duration); + Box::pin(ready(())) + } + }), + // A jitter that always lands low. + Arc::new(|delay| delay / 10), + Retry { + max_retries: 1, + ..Default::default() + }, + ) + .with_throttle_budget(Duration::from_secs(100)) + .with_throttle_max_wait(Duration::from_mins(1)); + let run_count = Arc::new(AtomicI32::new(0)); + + let result = Pin::new(&retrier) + .retry(repeat_with(|| { + if run_count.fetch_add(1, Ordering::Relaxed) < 2 { + return RetryResult::RetryAfter( + make_err!(Code::Unavailable, "Throttled"), + Duration::from_secs(10), + ); + } + RetryResult::Ok(true) + })) + .await?; + + assert_eq!(result, true); + assert_eq!( + *sleeps.lock().unwrap(), + vec![Duration::from_secs(19), Duration::from_secs(38)], + "10 s asked, jittered to a tenth, mirrored above: 19 s, then 38 s on the doubled wait" + ); + Ok(()) +}