use std::time::Duration;
use tokio::time::Instant;
use crate::retry_policy::RetryPolicy;
pub(crate) struct Budget {
deadline: Instant,
per_attempt: Duration,
}
impl Budget {
pub(crate) fn start(policy: &RetryPolicy) -> Self {
Budget {
deadline: Instant::now() + policy.overall_deadline,
per_attempt: policy.per_attempt_deadline,
}
}
pub(crate) fn next_attempt(&self) -> Option<Duration> {
let now = Instant::now();
if now >= self.deadline {
return None;
}
Some((self.deadline - now).min(self.per_attempt))
}
pub(crate) fn clamp_backoff(&self, backoff: Duration) -> Duration {
backoff.min(self.deadline.saturating_duration_since(Instant::now()))
}
}
pub(crate) struct PairBudget {
deadline: Instant,
}
impl PairBudget {
pub(crate) fn start(budget: Duration) -> Self {
PairBudget {
deadline: Instant::now() + budget,
}
}
pub(crate) fn remaining(&self) -> Duration {
self.deadline.saturating_duration_since(Instant::now())
}
}
#[cfg(test)]
mod tests {
use super::*;
fn policy(per_attempt: Duration, overall: Duration) -> RetryPolicy {
RetryPolicy {
max_attempts: 5,
per_attempt_deadline: per_attempt,
overall_deadline: overall,
base_backoff: Duration::from_millis(1),
leader_ttl: Duration::from_secs(30),
}
}
#[tokio::test]
async fn next_attempt_is_capped_by_per_attempt_deadline() {
let budget = Budget::start(&policy(Duration::from_millis(100), Duration::from_secs(10)));
let attempt = budget.next_attempt().expect("budget remains");
assert!(
attempt <= Duration::from_millis(100),
"attempt must not exceed per_attempt_deadline, got {attempt:?}",
);
assert!(
attempt > Duration::from_millis(90),
"with 10s overall, the attempt should be ~per_attempt_deadline, got {attempt:?}",
);
}
#[tokio::test(start_paused = true)]
async fn next_attempt_returns_none_past_overall_deadline() {
let budget = Budget::start(&policy(
Duration::from_millis(100),
Duration::from_millis(50),
));
tokio::time::sleep(Duration::from_millis(60)).await;
assert_eq!(
budget.next_attempt(),
None,
"no attempt may start after the overall deadline",
);
}
#[tokio::test(start_paused = true)]
async fn next_attempt_is_capped_by_remaining_overall_budget() {
let budget = Budget::start(&policy(Duration::from_secs(10), Duration::from_millis(100)));
tokio::time::sleep(Duration::from_millis(70)).await;
let attempt = budget.next_attempt().expect("budget remains");
assert!(
attempt <= Duration::from_millis(30),
"attempt must be capped by the ~30ms remaining overall budget, got {attempt:?}",
);
}
#[tokio::test(start_paused = true)]
async fn clamp_backoff_caps_at_remaining_overall_budget() {
let budget = Budget::start(&policy(Duration::from_secs(10), Duration::from_millis(100)));
tokio::time::sleep(Duration::from_millis(80)).await;
let sleep = budget.clamp_backoff(Duration::from_secs(5));
assert!(
sleep <= Duration::from_millis(20),
"backoff must clamp to the ~20ms remaining, got {sleep:?}",
);
}
#[tokio::test(start_paused = true)]
async fn clamp_backoff_floors_at_zero_past_deadline() {
let budget = Budget::start(&policy(Duration::from_secs(10), Duration::from_millis(50)));
tokio::time::sleep(Duration::from_millis(60)).await;
assert_eq!(
budget.clamp_backoff(Duration::from_secs(5)),
Duration::ZERO,
"an elapsed deadline must yield a zero-length backoff",
);
}
#[tokio::test(start_paused = true)]
async fn pair_budget_remaining_floors_at_zero() {
let pair = PairBudget::start(Duration::from_millis(50));
tokio::time::sleep(Duration::from_millis(60)).await;
assert_eq!(
pair.remaining(),
Duration::ZERO,
"a fully-consumed pair budget must floor at zero",
);
}
#[tokio::test(start_paused = true)]
async fn pair_budget_remaining_returns_leftover() {
let pair = PairBudget::start(Duration::from_millis(100));
tokio::time::sleep(Duration::from_millis(40)).await;
let left = pair.remaining();
assert!(
left <= Duration::from_millis(60) && left > Duration::from_millis(50),
"≈60ms should remain after consuming 40ms of 100ms, got {left:?}",
);
}
}