use std::{
sync::atomic::{self, AtomicU32, Ordering},
time,
};
use crate::retries::{ExponentialBackoffDecider, RetryDecider, RetryState};
#[derive(Clone, Debug)]
pub enum SleepyState {
Init,
Ongoing(RetryState),
Expired,
}
pub struct SleepyMan {
max_sleeps: u32,
last_state: SleepyState,
retry_decider: ExponentialBackoffDecider,
}
impl SleepyMan {
pub fn new(max_idles: u32, retry_decider: ExponentialBackoffDecider) -> Self {
Self {
max_sleeps: max_idles,
retry_decider,
last_state: SleepyState::Init,
}
}
pub fn state(&self) -> SleepyState {
self.last_state.clone()
}
pub fn reset(&mut self) {
self.last_state = SleepyState::Init;
}
pub fn next_sleep(&mut self) -> SleepyState {
match &self.last_state {
SleepyState::Init => {
self.last_state = SleepyState::Ongoing(RetryState::new(0, self.max_sleeps, None));
self.last_state.clone()
}
SleepyState::Ongoing(retry_state) => {
self.last_state = match self.retry_decider.decide(retry_state.clone()) {
Some(inner) => SleepyState::Ongoing(inner.clone()),
None => SleepyState::Expired,
};
self.last_state.clone()
}
SleepyState::Expired => SleepyState::Expired,
}
}
}
pub struct IdleMan {
sleepy: SleepyMan,
counter: AtomicU32,
base_sleep: Option<time::Duration>,
max_idle: u32,
}
static DEFAULT_BASIC_SLEEP: time::Duration = time::Duration::from_millis(10);
static DEFAULT_MAX_SLEEP_COUNT: u32 = 5;
impl IdleMan {
pub fn new(max_idle: u32, base_sleep: Option<time::Duration>, sleepy: SleepyMan) -> Self {
Self {
sleepy,
max_idle,
base_sleep,
counter: AtomicU32::new(0),
}
}
#[must_use]
pub fn basic(max_idle: u32) -> Self {
Self::new(
max_idle,
Some(DEFAULT_BASIC_SLEEP),
SleepyMan::new(
DEFAULT_MAX_SLEEP_COUNT,
ExponentialBackoffDecider::default(),
),
)
}
pub fn increment(&mut self) -> Option<time::Duration> {
let current_count = self.counter.load(atomic::Ordering::SeqCst);
if current_count >= self.max_idle {
return match self.sleepy.next_sleep() {
SleepyState::Ongoing(retry_state) => match retry_state.wait {
None => self.base_sleep,
Some(inner) => Some(inner),
},
SleepyState::Expired => {
assert!(self.try_reset(), "Failed to reset counter");
self.sleepy.reset();
self.base_sleep
}
SleepyState::Init => unreachable!(),
};
}
self.counter.fetch_add(1, Ordering::SeqCst);
self.base_sleep
}
pub fn count(&self) -> u32 {
self.counter.load(atomic::Ordering::Acquire)
}
pub fn try_reset(&self) -> bool {
self.counter
.compare_exchange(
self.max_idle,
0,
atomic::Ordering::SeqCst,
atomic::Ordering::Acquire,
)
.is_ok()
}
}