use crate::log::StageLogger;
use std::fmt;
use std::time::Duration;
#[derive(Clone, Copy)]
pub struct RetryLog<'a> {
desc: &'a str,
log: &'a StageLogger,
}
impl<'a> RetryLog<'a> {
pub fn new(desc: &'a str, log: &'a StageLogger) -> Self {
Self { desc, log }
}
pub fn desc(&self) -> &str {
self.desc
}
pub(super) fn warn_retry(
&self,
attempt: u32,
max: u32,
cause: &dyn fmt::Display,
delay: Duration,
) {
self.log.warn(&format!(
"{} attempt {}/{} failed ({}); retrying in {}",
self.desc,
attempt,
max,
cause,
crate::progress::format_elapsed(delay)
));
}
pub(super) fn warn_giving_up(&self, attempts: u32) {
self.log.warn(&format!(
"{} failed after {} attempt(s), giving up",
self.desc, attempts
));
}
pub(super) fn note_succeeded(&self, attempts: u32) {
self.log.status(&format!(
"{} succeeded after {} attempt(s)",
self.desc, attempts
)); }
}
#[derive(Debug, Clone, Copy)]
pub struct RetryPolicy {
pub max_attempts: u32,
pub base_delay: Duration,
pub max_delay: Duration,
}
impl RetryPolicy {
pub const UPLOAD: RetryPolicy = RetryPolicy {
max_attempts: 10,
base_delay: Duration::from_millis(50),
max_delay: Duration::from_secs(30),
};
pub const PREFLIGHT: RetryPolicy = RetryPolicy {
max_attempts: 3,
base_delay: Duration::from_millis(200),
max_delay: Duration::from_secs(1),
};
pub const GUARD_PROBE: RetryPolicy = RetryPolicy {
max_attempts: 3,
base_delay: Duration::from_secs(1),
max_delay: Duration::from_secs(30),
};
pub fn delay_for(&self, next_attempt: u32) -> Duration {
let exp = next_attempt.saturating_sub(2);
let mult = 1u64.checked_shl(exp).unwrap_or(u64::MAX);
let ms = (self.base_delay.as_millis() as u64).saturating_mul(mult);
std::cmp::min(Duration::from_millis(ms), self.max_delay)
}
pub fn with_idempotent_floor(self) -> RetryPolicy {
self.with_floor(IDEMPOTENT_PUT_ATTEMPTS)
}
pub fn with_floor(self, min: u32) -> RetryPolicy {
RetryPolicy {
max_attempts: self.max_attempts.max(min),
..self
}
}
pub fn budget_exhausted(&self, next_attempt: u32, deadline: std::time::Instant) -> bool {
match std::time::Instant::now().checked_add(self.delay_for(next_attempt)) {
Some(projected) => projected > deadline,
None => true,
}
}
}
pub const IDEMPOTENT_PUT_ATTEMPTS: u32 = 3;
pub const DEFAULT_MAX_ELAPSED: Duration = Duration::from_secs(15 * 60);
static RETRY_BACKOFF_MILLIS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
static PER_SCOPE_RETRY: std::sync::Mutex<std::collections::BTreeMap<String, ScopeRetry>> =
std::sync::Mutex::new(std::collections::BTreeMap::new());
static CURRENT_SCOPE: std::sync::Mutex<Option<String>> = std::sync::Mutex::new(None);
const UNATTRIBUTED_SCOPE: &str = "(unattributed)";
#[derive(Clone, Copy, Default)]
struct ScopeRetry {
retries: u32,
backoff_ms: u64,
}
#[must_use = "the scope only applies while the guard is alive"]
pub struct RetryScope {
prev: Option<String>,
}
impl RetryScope {
pub fn enter(name: impl Into<String>) -> Self {
let mut cur = CURRENT_SCOPE.lock().unwrap_or_else(|e| e.into_inner());
let prev = cur.replace(name.into());
RetryScope { prev }
}
}
impl Drop for RetryScope {
fn drop(&mut self) {
*CURRENT_SCOPE.lock().unwrap_or_else(|e| e.into_inner()) = self.prev.take();
}
}
thread_local! {
static CURRENT_BUDGET_ANCHOR: std::cell::Cell<Option<std::time::Instant>> =
const { std::cell::Cell::new(None) };
}
pub fn current_budget_anchor() -> Option<std::time::Instant> {
CURRENT_BUDGET_ANCHOR.with(std::cell::Cell::get)
}
#[must_use = "the budget only applies while the guard is alive"]
pub struct PublisherRetryScope {
_label: RetryScope,
prev_anchor: Option<std::time::Instant>,
}
impl PublisherRetryScope {
pub fn enter(name: impl Into<String>) -> Self {
let label = RetryScope::enter(name);
let prev_anchor = CURRENT_BUDGET_ANCHOR.with(|cell| {
let prev = cell.get();
if prev.is_none() {
cell.set(Some(std::time::Instant::now()));
}
prev
});
PublisherRetryScope {
_label: label,
prev_anchor,
}
}
}
impl Drop for PublisherRetryScope {
fn drop(&mut self) {
CURRENT_BUDGET_ANCHOR.with(|cell| cell.set(self.prev_anchor));
}
}
pub fn record_retry_backoff(d: Duration) {
let ms = u64::try_from(d.as_millis()).unwrap_or(u64::MAX);
RETRY_BACKOFF_MILLIS.fetch_add(ms, std::sync::atomic::Ordering::Relaxed);
let key = CURRENT_SCOPE
.lock()
.unwrap_or_else(|e| e.into_inner())
.clone()
.unwrap_or_else(|| UNATTRIBUTED_SCOPE.to_string());
let mut map = PER_SCOPE_RETRY.lock().unwrap_or_else(|e| e.into_inner());
let entry = map.entry(key).or_default();
entry.retries = entry.retries.saturating_add(1);
entry.backoff_ms = entry.backoff_ms.saturating_add(ms);
}
pub fn sleep_backoff_blocking(d: Duration) {
if d.is_zero() {
return;
}
record_retry_backoff(d);
std::thread::sleep(d);
}
pub async fn sleep_backoff_async(d: Duration) {
if d.is_zero() {
return;
}
record_retry_backoff(d);
tokio::time::sleep(d).await;
}
pub fn total_retry_backoff() -> Duration {
Duration::from_millis(RETRY_BACKOFF_MILLIS.load(std::sync::atomic::Ordering::Relaxed))
}
pub fn retry_scope_breakdown() -> Vec<(String, u32, Duration)> {
let map = PER_SCOPE_RETRY.lock().unwrap_or_else(|e| e.into_inner());
let mut rows: Vec<(String, u32, Duration)> = map
.iter()
.map(|(k, v)| (k.clone(), v.retries, Duration::from_millis(v.backoff_ms)))
.collect();
rows.sort_by(|a, b| b.2.cmp(&a.2).then_with(|| a.0.cmp(&b.0)));
rows
}