use std::future::Future;
use std::time::Duration;
use futures::StreamExt;
use reqwest::Url;
use crate::http::map_send_error;
use crate::{
PluginError,
PluginResult,
};
pub struct RetryPolicy {
pub max_retries: usize,
pub initial_delay: Duration,
pub exponential_backoff: bool,
}
impl Default for RetryPolicy {
fn default() -> Self {
Self {
max_retries: 3,
initial_delay: Duration::from_millis(100),
exponential_backoff: true,
}
}
}
impl RetryPolicy {
pub fn new(max_retries: usize, initial_delay: Duration, exponential_backoff: bool) -> Self {
Self {
max_retries,
initial_delay,
exponential_backoff,
}
}
pub async fn retry<F, Fut, T>(&self, operation: F) -> PluginResult<T>
where
F: Fn() -> Fut,
Fut: Future<Output = PluginResult<T>>,
{
let mut delay = self.initial_delay;
let mut attempt = 1;
loop {
match operation().await {
Ok(result) => return Ok(result),
Err(e) if attempt < self.max_retries && Self::should_retry(&e, delay) => {
tokio::time::sleep(delay).await;
delay = self.next_delay(delay);
attempt += 1;
}
Err(e) => return Err(e),
}
}
}
pub async fn retry_unsent<F, Fut>(&self, url: &Url, send: F) -> PluginResult<reqwest::Response>
where
F: Fn() -> Fut,
Fut: Future<Output = Result<reqwest::Response, reqwest::Error>>,
{
let mut delay = self.initial_delay;
let mut attempt = 1;
loop {
match send().await {
Ok(response) => return Ok(response),
Err(e) if attempt < self.max_retries && e.is_connect() && !e.is_timeout() => {
tokio::time::sleep(delay).await;
delay = self.next_delay(delay);
attempt += 1;
}
Err(e) => return Err(map_send_error(e, url)),
}
}
}
fn should_retry(error: &PluginError, delay: Duration) -> bool {
error.is_transient() && error.retry_after().is_none_or(|hint| hint <= delay)
}
fn next_delay(&self, delay: Duration) -> Duration {
if self.exponential_backoff {
delay.saturating_mul(2)
} else {
delay
}
}
}
pub async fn retain_with_ci<F, Fut>(ids: Vec<String>, has_ci: F) -> Vec<String>
where
F: Fn(String) -> Fut,
Fut: Future<Output = PluginResult<bool>>,
{
futures::stream::iter(ids)
.map(|id| {
let check = has_ci(id.clone());
async move { check.await.unwrap_or(true).then_some(id) }
})
.buffered(8)
.filter_map(std::future::ready)
.collect()
.await
}
#[cfg(test)]
mod tests {
use std::io::{
Read,
Write,
};
use std::net::TcpListener;
use std::sync::atomic::{
AtomicUsize,
Ordering,
};
use std::sync::Arc;
use super::*;
fn fast_policy() -> RetryPolicy {
RetryPolicy::new(3, Duration::from_millis(5), false)
}
fn client(timeout: Duration) -> reqwest::Client {
let _ = rustls::crypto::ring::default_provider().install_default();
reqwest::Client::builder().timeout(timeout).build().unwrap()
}
async fn attempts_for(error: fn() -> PluginError) -> usize {
let attempts = AtomicUsize::new(0);
let result: PluginResult<()> = fast_policy()
.retry(|| async {
attempts.fetch_add(1, Ordering::SeqCst);
Err(error())
})
.await;
assert!(result.is_err());
attempts.load(Ordering::SeqCst)
}
#[tokio::test]
async fn test_retry_success() {
let policy = RetryPolicy::default();
let result = policy.retry(|| async { Ok::<_, PluginError>(42) }).await;
assert_eq!(result.unwrap(), 42);
}
#[tokio::test]
async fn test_retry_eventual_success() {
let policy = RetryPolicy::new(3, Duration::from_millis(10), false);
let attempts = std::cell::Cell::new(0);
let result = policy
.retry(|| async {
let count = attempts.get() + 1;
attempts.set(count);
if count < 2 {
Err(PluginError::NetworkError("Temporary failure".to_string()))
} else {
Ok(42)
}
})
.await;
assert_eq!(result.unwrap(), 42);
}
#[tokio::test]
async fn retries_network_and_short_hint_unavailable() {
assert_eq!(
attempts_for(|| PluginError::NetworkError("x".into())).await,
3
);
assert_eq!(
attempts_for(|| PluginError::ServiceUnavailable {
message: "x".into(),
retry_after: None,
})
.await,
3
);
assert_eq!(
attempts_for(|| PluginError::ServiceUnavailable {
message: "x".into(),
retry_after: Some(Duration::from_millis(1)),
})
.await,
3
);
}
#[tokio::test]
async fn does_not_retry_permanent_errors_or_long_hints() {
assert_eq!(
attempts_for(|| PluginError::AuthenticationFailed("x".into())).await,
1
);
assert_eq!(
attempts_for(|| PluginError::PipelineNotFound("x".into())).await,
1
);
assert_eq!(attempts_for(|| PluginError::ApiError("x".into())).await, 1);
assert_eq!(
attempts_for(|| PluginError::RateLimited {
message: "x".into(),
retry_after: Some(Duration::from_secs(120)),
})
.await,
1
);
assert_eq!(
attempts_for(|| PluginError::ServiceUnavailable {
message: "x".into(),
retry_after: Some(Duration::from_secs(120)),
})
.await,
1
);
}
#[tokio::test]
async fn retry_unsent_retries_connect_errors() {
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
drop(listener);
let url = Url::parse(&format!("http://{addr}/trigger")).unwrap();
let client = client(Duration::from_secs(5));
let attempts = AtomicUsize::new(0);
let result = fast_policy()
.retry_unsent(&url, || {
attempts.fetch_add(1, Ordering::SeqCst);
client.post(url.clone()).send()
})
.await;
assert!(matches!(result, Err(PluginError::NetworkError(_))));
assert_eq!(attempts.load(Ordering::SeqCst), 3);
}
#[tokio::test]
async fn retry_unsent_never_resends_after_timeout() {
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
let received = Arc::new(AtomicUsize::new(0));
let counter = received.clone();
std::thread::spawn(move || {
for mut stream in listener.incoming().flatten() {
let mut buf = [0u8; 1024];
if stream.read(&mut buf).unwrap_or(0) > 0 {
counter.fetch_add(1, Ordering::SeqCst);
}
std::thread::sleep(Duration::from_millis(300));
let _ = stream.write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n");
}
});
let url = Url::parse(&format!("http://{addr}/trigger")).unwrap();
let client = client(Duration::from_millis(50));
let result = fast_policy()
.retry_unsent(&url, || client.post(url.clone()).send())
.await;
assert!(matches!(result, Err(PluginError::NetworkError(_))));
tokio::time::sleep(Duration::from_millis(100)).await;
assert_eq!(received.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn ids_without_ci_are_dropped_and_unchecked_ones_kept_in_order() {
let ids = ["empty", "ci", "unreachable", "ci-too"]
.map(String::from)
.to_vec();
let kept = retain_with_ci(ids, |id| async move {
match id.as_str() {
"empty" => Ok(false),
"unreachable" => Err(PluginError::NetworkError("reset".into())),
_ => Ok(true),
}
})
.await;
assert_eq!(kept, ["ci", "unreachable", "ci-too"]);
}
}