use std::future::Future;
use std::time::Duration;
use derive_new::new;
use derive_setters::Setters;
use tokio::time::sleep;
use tracing::warn;
#[derive(Debug, Clone, new, Setters)]
#[setters(prefix = "with_")]
pub struct Retry {
#[setters(skip)]
pub max_attempts: u32,
#[new(value = "Duration::from_millis(100)")]
pub initial_backoff: Duration,
#[new(value = "Duration::from_secs(2)")]
pub max_backoff: Duration,
#[new(value = "true")]
pub retry_on_no_reply: bool,
#[new(value = "true")]
pub retry_on_transport_error: bool,
}
impl Retry {
fn validate(&self) -> crate::bus::Result<()> {
if self.max_attempts == 0 {
return Err(crate::bus::Error::InvalidRetry(
"max_attempts must be greater than zero".to_string(),
));
}
if self.initial_backoff.is_zero() {
return Err(crate::bus::Error::InvalidRetry(
"initial_backoff must be greater than zero".to_string(),
));
}
if self.max_backoff.is_zero() {
return Err(crate::bus::Error::InvalidRetry(
"max_backoff must be greater than zero".to_string(),
));
}
if self.initial_backoff > self.max_backoff {
return Err(crate::bus::Error::InvalidRetry(
"initial_backoff must be less than or equal to max_backoff".to_string(),
));
}
Ok(())
}
}
pub async fn retry_query<T, F, Fut>(
name: &str,
retry: &Retry,
mut attempt: F,
) -> crate::bus::Result<Option<T>>
where
F: FnMut() -> Fut,
Fut: Future<Output = crate::bus::Result<Option<T>>>,
{
retry.validate()?;
let mut backoff = retry.initial_backoff;
for attempt_index in 1..=retry.max_attempts {
match attempt().await {
Ok(Some(response)) => return Ok(Some(response)),
Ok(None) => {
if !retry.retry_on_no_reply || attempt_index == retry.max_attempts {
return Ok(None);
}
warn!(
query = name,
attempt = attempt_index,
max_attempts = retry.max_attempts,
backoff_ms = backoff.as_millis(),
"query returned no reply, retrying"
);
}
Err(crate::bus::Error::Transport(error)) => {
if !retry.retry_on_transport_error || attempt_index == retry.max_attempts {
return Err(crate::bus::Error::Transport(error));
}
warn!(
query = name,
attempt = attempt_index,
max_attempts = retry.max_attempts,
backoff_ms = backoff.as_millis(),
error = %error,
"query failed with transport error, retrying"
);
}
Err(error) => return Err(error),
}
sleep(backoff).await;
backoff = std::cmp::min(backoff.saturating_mul(2), retry.max_backoff);
}
Ok(None)
}
#[cfg(test)]
mod tests {
use super::{Retry, retry_query};
use crate::bus::{Error, Result};
use std::sync::{
Arc,
atomic::{AtomicU32, Ordering},
};
use std::time::Duration;
#[tokio::test]
async fn retry_query_retries_no_reply_until_success() -> Result<()> {
let attempts = Arc::new(AtomicU32::new(0));
let retry = Retry::new(3).with_initial_backoff(Duration::from_millis(1));
let result = retry_query("robot/r1/asset/bundle", &retry, {
let attempts = attempts.clone();
move || {
let attempts = attempts.clone();
async move {
let current = attempts.fetch_add(1, Ordering::SeqCst);
if current < 2 {
Ok(None)
} else {
Ok(Some("model"))
}
}
}
})
.await?;
assert_eq!(result, Some("model"));
assert_eq!(attempts.load(Ordering::SeqCst), 3);
Ok(())
}
#[tokio::test]
async fn retry_query_retries_transport_errors_until_success() -> Result<()> {
let attempts = Arc::new(AtomicU32::new(0));
let retry = Retry::new(2).with_initial_backoff(Duration::from_millis(1));
let result = retry_query("robot/r1/asset/bundle", &retry, {
let attempts = attempts.clone();
move || {
let attempts = attempts.clone();
async move {
let current = attempts.fetch_add(1, Ordering::SeqCst);
if current == 0 {
Err(Error::Transport(std::io::Error::other("transient").into()))
} else {
Ok(Some("model"))
}
}
}
})
.await?;
assert_eq!(result, Some("model"));
assert_eq!(attempts.load(Ordering::SeqCst), 2);
Ok(())
}
#[tokio::test]
async fn retry_query_rejects_invalid_retry() {
let error = retry_query::<(), _, _>("robot/r1/asset/bundle", &Retry::new(0), || async {
Ok(Some(()))
})
.await
.expect_err("invalid retry should fail");
assert!(matches!(error, Error::InvalidRetry(_)));
}
}