use std::time::{Duration, Instant};
#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)]
pub enum AutoEjectState {
Reachable,
Ejected,
}
#[derive(Debug, Clone)]
pub struct AutoEject {
enabled: bool,
failure_limit: u32,
retry_after: Duration,
failures: u32,
next_retry: Option<Instant>,
}
impl AutoEject {
#[must_use]
pub fn new(enabled: bool, failure_limit: u32, retry_after: Duration) -> Self {
Self {
enabled,
failure_limit,
retry_after,
failures: 0,
next_retry: None,
}
}
#[must_use]
pub fn is_enabled(&self) -> bool {
self.enabled
}
#[must_use]
pub fn failure_limit(&self) -> u32 {
self.failure_limit
}
#[must_use]
pub fn retry_after(&self) -> Duration {
self.retry_after
}
#[must_use]
pub fn failure_count(&self) -> u32 {
self.failures
}
#[must_use]
pub fn next_retry(&self) -> Option<Instant> {
self.next_retry
}
pub fn record_attempt(&mut self, now: Instant) -> AutoEjectState {
if !self.enabled {
return AutoEjectState::Reachable;
}
match self.next_retry {
Some(eta) if now < eta => AutoEjectState::Ejected,
Some(_) => {
self.next_retry = None;
AutoEjectState::Reachable
}
None => AutoEjectState::Reachable,
}
}
pub fn record_success(&mut self, _now: Instant) {
self.failures = 0;
self.next_retry = None;
}
pub fn record_failure(&mut self, now: Instant) -> AutoEjectState {
self.failures = self.failures.saturating_add(1);
if self.enabled && self.failures >= self.failure_limit {
self.next_retry = Some(now + self.retry_after);
AutoEjectState::Ejected
} else {
AutoEjectState::Reachable
}
}
pub fn reset(&mut self) {
self.failures = 0;
self.next_retry = None;
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn disabled_never_ejects() {
let mut ae = AutoEject::new(false, 1, Duration::from_secs(1));
let now = Instant::now();
for _ in 0..5 {
assert_eq!(ae.record_failure(now), AutoEjectState::Reachable);
}
assert_eq!(ae.record_attempt(now), AutoEjectState::Reachable);
}
#[test]
fn ejects_after_threshold_and_recovers_after_window() {
let mut ae = AutoEject::new(true, 3, Duration::from_millis(50));
let now = Instant::now();
assert_eq!(ae.record_failure(now), AutoEjectState::Reachable);
assert_eq!(ae.record_failure(now), AutoEjectState::Reachable);
assert_eq!(ae.record_failure(now), AutoEjectState::Ejected);
assert_eq!(ae.record_attempt(now), AutoEjectState::Ejected);
let after = now + Duration::from_millis(51);
assert_eq!(ae.record_attempt(after), AutoEjectState::Reachable);
}
#[test]
fn record_success_clears_state() {
let mut ae = AutoEject::new(true, 2, Duration::from_secs(1));
let now = Instant::now();
ae.record_failure(now);
ae.record_failure(now);
assert_eq!(ae.record_attempt(now), AutoEjectState::Ejected);
ae.record_success(now);
assert_eq!(ae.record_attempt(now), AutoEjectState::Reachable);
assert_eq!(ae.failure_count(), 0);
}
}