reifydb_core/

retry.rs

// SPDX-License-Identifier: Apache-2.0
// Copyright (c) 2026 ReifyDB

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"),
		}
	}
}