use std::thread;
use reifydb_runtime::context::rng::Rng;
use reifydb_value::{error::Error, value::duration::Duration};
use tracing::{debug, warn};
use crate::execution::ExecutionResult;
pub enum Backoff {
None,
Fixed(Duration),
Exponential {
base: Duration,
max: Duration,
},
ExponentialJitter {
base: Duration,
max: Duration,
},
}
pub struct RetryStrategy {
pub max_attempts: u32,
pub backoff: Backoff,
}
impl Default for RetryStrategy {
fn default() -> Self {
Self {
max_attempts: 10,
backoff: Backoff::ExponentialJitter {
base: Duration::from_milliseconds(5).unwrap(),
max: Duration::from_milliseconds(200).unwrap(),
},
}
}
}
impl RetryStrategy {
pub fn no_retry() -> Self {
Self {
max_attempts: 1,
backoff: Backoff::None,
}
}
pub fn default_conflict_retry() -> Self {
Self::default()
}
pub fn with_fixed_backoff(max_attempts: u32, delay: Duration) -> Self {
Self {
max_attempts,
backoff: Backoff::Fixed(delay),
}
}
pub fn with_exponential_backoff(max_attempts: u32, base: Duration, max: Duration) -> Self {
Self {
max_attempts,
backoff: Backoff::Exponential {
base,
max,
},
}
}
pub fn with_jittered_backoff(max_attempts: u32, base: Duration, max: Duration) -> Self {
Self {
max_attempts,
backoff: Backoff::ExponentialJitter {
base,
max,
},
}
}
pub fn execute<F>(&self, rng: &Rng, rql: &str, f: F) -> ExecutionResult
where
F: FnMut() -> ExecutionResult,
{
self.run(rng, rql, f, |result| result.error.as_ref().is_some_and(is_conflict))
}
pub fn retry<T, F>(&self, rng: &Rng, operation: &str, f: F) -> Result<T, Error>
where
F: FnMut() -> Result<T, Error>,
{
self.run(rng, operation, f, |result| result.as_ref().is_err_and(is_conflict))
}
fn run<R>(&self, rng: &Rng, rql: &str, mut f: impl FnMut() -> R, conflicted: impl Fn(&R) -> bool) -> R {
let mut last_result = None;
for attempt in 0..self.max_attempts {
let result = f();
if !conflicted(&result) {
return result;
}
last_result = Some(result);
let is_last_attempt = attempt + 1 >= self.max_attempts;
if is_last_attempt {
warn!(
attempt = attempt + 1,
max_attempts = self.max_attempts,
rql = %rql,
"Transaction conflict retries exhausted"
);
} else {
let delay = compute_backoff(&self.backoff, attempt, rng);
debug!(
attempt = attempt + 1,
max_attempts = self.max_attempts,
delay_us = delay.microseconds().unwrap_or(0) as u64,
rql = %rql,
"Transaction conflict detected, retrying after backoff"
);
if !delay.is_zero() {
thread::sleep(delay.to_std());
}
}
}
last_result.unwrap()
}
}
pub fn is_conflict(err: &Error) -> bool {
err.code == "TXN_001"
}
fn compute_backoff(backoff: &Backoff, attempt: u32, rng: &Rng) -> Duration {
match backoff {
Backoff::None => Duration::zero(),
Backoff::Fixed(d) => *d,
Backoff::Exponential {
base,
max,
} => exponential_cap(*base, *max, attempt),
Backoff::ExponentialJitter {
base,
max,
} => {
let cap = exponential_cap(*base, *max, attempt);
let cap_nanos = cap.as_nanos().unwrap_or(0).max(0) as u64;
if cap_nanos == 0 {
return Duration::zero();
}
let sampled = rng.infra_u64_inclusive(cap_nanos);
Duration::from_nanoseconds(sampled as i64).unwrap()
}
}
}
fn exponential_cap(base: Duration, max: Duration, attempt: u32) -> Duration {
let shift = attempt.min(30);
let multiplier = 1i64 << shift;
base.saturating_mul(multiplier).min(max)
}
#[cfg(test)]
mod retry_tests {
use std::cell::Cell;
use reifydb_runtime::context::rng::Rng;
use reifydb_value::{
error::{Diagnostic, Error},
fragment::Fragment,
value::duration::Duration,
};
use super::{Backoff, RetryStrategy, compute_backoff, exponential_cap};
use crate::{execution::ExecutionResult, metrics::execution::ExecutionMetrics};
fn ok() -> ExecutionResult {
ExecutionResult {
frames: vec![],
named_system_columns: vec![],
error: None,
metrics: ExecutionMetrics::default(),
}
}
fn err(code: &str) -> ExecutionResult {
ExecutionResult {
frames: vec![],
named_system_columns: vec![],
error: Some(Error(Box::new(Diagnostic {
code: code.to_string(),
rql: None,
message: format!("{} test", code),
column: None,
fragment: Fragment::None,
label: None,
help: None,
notes: vec![],
cause: None,
operator_chain: None,
}))),
metrics: ExecutionMetrics::default(),
}
}
fn no_sleep_strategy(max_attempts: u32) -> RetryStrategy {
RetryStrategy {
max_attempts,
backoff: Backoff::None,
}
}
#[test]
fn success_first_try_runs_closure_once() {
let strategy = no_sleep_strategy(5);
let rng = Rng::default();
let calls = Cell::new(0u32);
let result = strategy.execute(&rng, "", || {
calls.set(calls.get() + 1);
ok()
});
assert!(result.is_ok());
assert_eq!(calls.get(), 1);
}
#[test]
fn non_conflict_error_is_not_retried() {
let strategy = no_sleep_strategy(5);
let rng = Rng::default();
let calls = Cell::new(0u32);
let result = strategy.execute(&rng, "", || {
calls.set(calls.get() + 1);
err("TXN_002")
});
assert!(result.is_err());
assert_eq!(calls.get(), 1);
}
#[test]
fn conflict_retries_then_succeeds() {
let strategy = no_sleep_strategy(5);
let rng = Rng::default();
let calls = Cell::new(0u32);
let result = strategy.execute(&rng, "", || {
let n = calls.get();
calls.set(n + 1);
if n < 2 {
err("TXN_001")
} else {
ok()
}
});
assert!(result.is_ok());
assert_eq!(calls.get(), 3);
}
#[test]
fn conflict_exhausts_attempts_returns_last_error() {
let strategy = no_sleep_strategy(4);
let rng = Rng::default();
let calls = Cell::new(0u32);
let result = strategy.execute(&rng, "", || {
calls.set(calls.get() + 1);
err("TXN_001")
});
assert!(result.is_err());
assert_eq!(result.error.as_ref().unwrap().code, "TXN_001");
assert_eq!(calls.get(), 4);
}
#[test]
fn jittered_backoff_stays_within_cap() {
let base = Duration::from_milliseconds(10).unwrap();
let max = Duration::from_milliseconds(100).unwrap();
let backoff = Backoff::ExponentialJitter {
base,
max,
};
let rng = Rng::default();
for attempt in 0..8 {
let cap = exponential_cap(base, max, attempt);
for _ in 0..50 {
let d = compute_backoff(&backoff, attempt, &rng);
assert!(d <= cap, "attempt {}: {:?} exceeds cap {:?}", attempt, d, cap);
}
}
}
#[test]
fn seeded_rng_produces_deterministic_jitter() {
let base = Duration::from_milliseconds(5).unwrap();
let max = Duration::from_milliseconds(200).unwrap();
let backoff = Backoff::ExponentialJitter {
base,
max,
};
let sample = |seed: u64| -> Vec<Duration> {
let rng = Rng::seeded(seed);
(0..8).map(|attempt| compute_backoff(&backoff, attempt, &rng)).collect()
};
assert_eq!(sample(42), sample(42));
assert_ne!(sample(42), sample(43));
}
#[test]
fn seeded_rng_produces_exact_pinned_jitter_values() {
let base = Duration::from_milliseconds(5).unwrap();
let max = Duration::from_milliseconds(200).unwrap();
let backoff = Backoff::ExponentialJitter {
base,
max,
};
let nanos = |seed: u64| -> Vec<u64> {
let rng = Rng::seeded(seed);
(0..8).map(|attempt| compute_backoff(&backoff, attempt, &rng).as_nanos().unwrap() as u64)
.collect()
};
let expected_42: Vec<u64> = vec![
3_848_394,
113_809,
2_934_288,
23_292_485,
77_680_508,
31_066_617,
36_519_179,
190_866_841,
];
let expected_43: Vec<u64> = vec![
3_974_671, 4_842_103, 12_057_439, 29_830_325, 72_334_216, 22_229_100, 36_417_439, 81_417_246,
];
assert_eq!(nanos(42), expected_42);
assert_eq!(nanos(43), expected_43);
assert_eq!(nanos(42), expected_42);
assert_eq!(nanos(43), expected_43);
}
#[test]
fn exponential_cap_saturates_at_max() {
let base = Duration::from_milliseconds(5).unwrap();
let max = Duration::from_milliseconds(200).unwrap();
assert_eq!(exponential_cap(base, max, 0), Duration::from_milliseconds(5).unwrap());
assert_eq!(exponential_cap(base, max, 1), Duration::from_milliseconds(10).unwrap());
assert_eq!(exponential_cap(base, max, 5), Duration::from_milliseconds(160).unwrap());
assert_eq!(exponential_cap(base, max, 6), max);
assert_eq!(exponential_cap(base, max, 100), max);
}
#[test]
fn default_uses_jittered_backoff() {
let s = RetryStrategy::default();
assert_eq!(s.max_attempts, 10);
match s.backoff {
Backoff::ExponentialJitter {
base,
max,
} => {
assert_eq!(base, Duration::from_milliseconds(5).unwrap());
assert_eq!(max, Duration::from_milliseconds(200).unwrap());
}
_ => panic!("expected ExponentialJitter default"),
}
}
}