use std::future::Future;
use std::time::Duration;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Retryability {
Retry,
Fatal,
}
pub fn classify(err: &anyhow::Error) -> Retryability {
let msg = err.to_string();
if msg.contains("request failed") {
return Retryability::Retry;
}
if let Some(code) = status_code_in(&msg) {
if is_retryable_status(code) {
return Retryability::Retry;
}
return Retryability::Fatal;
}
Retryability::Fatal
}
fn is_retryable_status(code: u16) -> bool {
code == 408 || code == 429 || (500..600).contains(&code)
}
fn status_code_in(msg: &str) -> Option<u16> {
const PREFIXES: &[&str] = &["returned ", "inference endpoint "];
for prefix in PREFIXES {
if let Some(after) = msg.split_once(prefix).map(|(_, r)| r) {
if let Some(code) = after
.split(|c: char| !c.is_ascii_digit())
.find(|s: &&str| !s.is_empty())
.and_then(|s| s.parse().ok())
{
return Some(code);
}
}
}
None
}
#[derive(Debug, Clone)]
pub struct RetryPolicy {
pub max_retries: u32,
pub base: Duration,
pub max: Duration,
pub jitter: bool,
}
impl Default for RetryPolicy {
fn default() -> Self {
Self {
max_retries: 4,
base: Duration::from_millis(500),
max: Duration::from_secs(8),
jitter: true,
}
}
}
impl RetryPolicy {
pub fn from_env() -> Self {
Self::from_env_or(Self::default())
}
pub fn from_env_or(mut base: Self) -> Self {
if let Some(n) = env_parse::<u32>("NEWT_HTTP_MAX_RETRIES") {
base.max_retries = n;
}
if let Some(ms) = env_parse::<u64>("NEWT_HTTP_BACKOFF_BASE_MS") {
base.base = Duration::from_millis(ms);
}
if let Some(ms) = env_parse::<u64>("NEWT_HTTP_BACKOFF_MAX_MS") {
base.max = Duration::from_millis(ms);
}
if let Ok(v) = std::env::var("NEWT_HTTP_JITTER") {
base.jitter = !matches!(
v.trim().to_ascii_lowercase().as_str(),
"0" | "false" | "off"
);
}
base
}
pub fn for_local_inference() -> Self {
Self::from_env_or(Self {
max_retries: 6,
base: Duration::from_secs(2),
max: Duration::from_secs(30),
jitter: true,
})
}
pub fn immediate(max_retries: u32) -> Self {
Self {
max_retries,
base: Duration::ZERO,
max: Duration::ZERO,
jitter: false,
}
}
pub fn delay_for(&self, attempt: u32) -> Duration {
let capped_ms = self.base_delay_ms(attempt);
if !self.jitter || capped_ms == 0 {
return Duration::from_millis(capped_ms);
}
let half = capped_ms / 2;
let span = capped_ms - half; Duration::from_millis(half + jitter_u64() % (span + 1))
}
fn base_delay_ms(&self, attempt: u32) -> u64 {
let shift = attempt.saturating_sub(1).min(31);
let factor = 1u64 << shift;
let raw = (self.base.as_millis() as u64).saturating_mul(factor);
raw.min(self.max.as_millis() as u64)
}
}
pub async fn with_backoff_notify<T, F, Fut, N>(
policy: &RetryPolicy,
mut op: F,
mut on_retry: N,
) -> anyhow::Result<T>
where
F: FnMut() -> Fut,
Fut: Future<Output = anyhow::Result<T>>,
N: FnMut(u32, Duration),
{
let mut retries = 0u32;
loop {
match op().await {
Ok(value) => return Ok(value),
Err(err) => {
if classify(&err) == Retryability::Fatal || retries >= policy.max_retries {
return Err(err);
}
retries += 1;
let delay = policy.delay_for(retries);
on_retry(retries, delay);
tracing::warn!(
attempt = retries,
delay_ms = delay.as_millis() as u64,
error = %err,
"retrying inference request"
);
tokio::time::sleep(delay).await;
}
}
}
}
pub async fn with_backoff<T, F, Fut>(policy: &RetryPolicy, op: F) -> anyhow::Result<T>
where
F: FnMut() -> Fut,
Fut: Future<Output = anyhow::Result<T>>,
{
with_backoff_notify(policy, op, |_, _| {}).await
}
fn env_parse<T: std::str::FromStr>(key: &str) -> Option<T> {
std::env::var(key).ok()?.trim().parse().ok()
}
fn jitter_u64() -> u64 {
use std::sync::atomic::{AtomicU64, Ordering};
static STATE: AtomicU64 = AtomicU64::new(0);
if STATE.load(Ordering::Relaxed) == 0 {
let seed = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.subsec_nanos() as u64)
.unwrap_or(0)
^ 0x9E37_79B9_7F4A_7C15;
STATE.store(seed | 1, Ordering::Relaxed);
}
let mut x = STATE.fetch_add(0x9E37_79B9_7F4A_7C15, Ordering::Relaxed);
x ^= x >> 12;
x ^= x << 25;
x ^= x >> 27;
x.wrapping_mul(0x2545_F491_4F6C_DD1D)
}
#[cfg(test)]
mod tests {
use super::*;
use std::cell::Cell;
fn err(msg: &str) -> anyhow::Error {
anyhow::anyhow!("{msg}")
}
#[test]
fn classify_transport_failure_is_retry() {
assert_eq!(
classify(&err("vLLM request failed: error sending request for url")),
Retryability::Retry
);
assert_eq!(
classify(&err("Ollama request failed: connection refused")),
Retryability::Retry
);
}
#[test]
fn classify_429_and_408_and_5xx_are_retry() {
assert_eq!(
classify(&err("vLLM returned 429 Too Many Requests: slow down")),
Retryability::Retry
);
assert_eq!(
classify(&err("vLLM returned 408 Request Timeout:")),
Retryability::Retry
);
assert_eq!(
classify(&err("Ollama returned 503 Service Unavailable: down")),
Retryability::Retry
);
assert_eq!(
classify(&err("vLLM returned 500 Internal Server Error: boom")),
Retryability::Retry
);
}
#[test]
fn classify_other_4xx_is_fatal() {
assert_eq!(
classify(&err("vLLM returned 400 Bad Request: nope")),
Retryability::Fatal
);
assert_eq!(
classify(&err("vLLM returned 404 Not Found:")),
Retryability::Fatal
);
}
#[test]
fn classify_unknown_shape_is_fatal() {
assert_eq!(
classify(&err("error decoding response body: expected value")),
Retryability::Fatal
);
}
#[test]
fn status_code_parsing() {
assert_eq!(
status_code_in("vLLM returned 503 Service Unavailable: x"),
Some(503)
);
assert_eq!(status_code_in("Ollama returned 429: y"), Some(429));
assert_eq!(status_code_in("no status here"), None);
}
#[test]
fn classify_inference_endpoint_5xx_is_retry() {
assert_eq!(
classify(&err(
r#"inference endpoint 500 Internal Server Error: {"error":{"message":"instance_id not found"}}"#
)),
Retryability::Retry
);
assert_eq!(
classify(&err("inference endpoint 503 Service Unavailable: down")),
Retryability::Retry
);
assert_eq!(
classify(&err("inference endpoint 429 Too Many Requests: slow")),
Retryability::Retry
);
}
#[test]
fn status_code_parsing_inference_endpoint_format() {
assert_eq!(
status_code_in("inference endpoint 500 Internal Server Error: body"),
Some(500)
);
assert_eq!(
status_code_in("inference endpoint 429 Too Many Requests: body"),
Some(429)
);
assert_eq!(
status_code_in("inference endpoint 400 Bad Request: body"),
Some(400)
);
}
#[test]
fn base_delay_is_exponential_and_capped() {
let p = RetryPolicy {
max_retries: 10,
base: Duration::from_millis(500),
max: Duration::from_secs(8),
jitter: false,
};
assert_eq!(p.base_delay_ms(1), 500);
assert_eq!(p.base_delay_ms(2), 1000);
assert_eq!(p.base_delay_ms(3), 2000);
assert_eq!(p.base_delay_ms(4), 4000);
assert_eq!(p.base_delay_ms(5), 8000);
assert_eq!(p.base_delay_ms(6), 8000);
assert_eq!(p.base_delay_ms(30), 8000);
}
#[test]
fn delay_without_jitter_is_deterministic() {
let p = RetryPolicy {
jitter: false,
..RetryPolicy::default()
};
assert_eq!(p.delay_for(1), Duration::from_millis(500));
assert_eq!(p.delay_for(2), Duration::from_millis(1000));
}
#[test]
fn jittered_delay_stays_within_equal_jitter_band() {
let p = RetryPolicy {
max_retries: 4,
base: Duration::from_millis(1000),
max: Duration::from_secs(8),
jitter: true,
};
for _ in 0..200 {
let ms = p.delay_for(1).as_millis() as u64;
assert!((500..=1000).contains(&ms), "delay {ms} outside [500,1000]");
}
}
#[test]
fn immediate_policy_never_sleeps() {
let p = RetryPolicy::immediate(3);
assert_eq!(p.max_retries, 3);
assert_eq!(p.delay_for(1), Duration::ZERO);
assert_eq!(p.delay_for(3), Duration::ZERO);
}
#[tokio::test]
async fn with_backoff_succeeds_after_transient_failures() {
let calls = Cell::new(0u32);
let result: anyhow::Result<&str> = with_backoff(&RetryPolicy::immediate(5), || {
let n = calls.get() + 1;
calls.set(n);
async move {
if n < 3 {
Err(err("vLLM returned 429 Too Many Requests: slow"))
} else {
Ok("ok")
}
}
})
.await;
assert_eq!(result.unwrap(), "ok");
assert_eq!(calls.get(), 3, "should have retried twice then succeeded");
}
#[tokio::test]
async fn with_backoff_stops_on_fatal() {
let calls = Cell::new(0u32);
let result: anyhow::Result<&str> = with_backoff(&RetryPolicy::immediate(5), || {
calls.set(calls.get() + 1);
async move { Err(err("vLLM returned 400 Bad Request: nope")) }
})
.await;
assert!(result.is_err());
assert_eq!(calls.get(), 1, "fatal error must not retry");
}
#[tokio::test]
async fn with_backoff_gives_up_after_max_retries() {
let calls = Cell::new(0u32);
let result: anyhow::Result<&str> = with_backoff(&RetryPolicy::immediate(3), || {
calls.set(calls.get() + 1);
async move { Err(err("vLLM returned 503 Service Unavailable: down")) }
})
.await;
let e = result.unwrap_err();
assert!(e.to_string().contains("503"), "last error preserved: {e}");
assert_eq!(calls.get(), 4, "1 initial + 3 retries");
}
#[tokio::test]
async fn with_backoff_notify_fires_callback_before_each_sleep() {
let calls = Cell::new(0u32);
let notified = Cell::new(0u32);
let result: anyhow::Result<&str> = with_backoff_notify(
&RetryPolicy::immediate(3),
|| {
let n = calls.get() + 1;
calls.set(n);
async move {
if n <= 2 {
Err(err("vLLM returned 503 Service Unavailable: down"))
} else {
Ok("ok")
}
}
},
|_, _| {
notified.set(notified.get() + 1);
},
)
.await;
assert_eq!(result.unwrap(), "ok");
assert_eq!(calls.get(), 3, "two retries then success");
assert_eq!(
notified.get(),
2,
"callback fired once before each retry sleep"
);
}
#[test]
fn for_local_inference_is_more_patient_than_default() {
let local = RetryPolicy::for_local_inference();
let default = RetryPolicy::default();
assert!(
local.max_retries > default.max_retries,
"local policy must allow more retries"
);
assert!(
local.max > default.max,
"local policy must have a longer backoff ceiling"
);
}
#[test]
fn from_env_or_starts_from_provided_base() {
let base = RetryPolicy {
max_retries: 9,
base: Duration::from_secs(3),
max: Duration::from_secs(60),
jitter: false,
};
let result = RetryPolicy::from_env_or(base.clone());
assert_eq!(result.max_retries, 9);
assert_eq!(result.base, Duration::from_secs(3));
assert_eq!(result.max, Duration::from_secs(60));
assert!(!result.jitter);
}
}