use std::collections::hash_map::Entry;
use std::collections::HashMap;
use std::hash::Hash;
use std::sync::Arc;
use crate::time::Clock;
use crate::time::SystemClock;
pub struct BackoffTracker<T> {
field: HashMap<T, BackoffEntry>,
clock: Arc<dyn Clock>,
config: BackoffConfig,
}
impl<T> Default for BackoffTracker<T> {
fn default() -> Self {
Self::new(Arc::new(SystemClock::default()), BackoffConfig::default())
}
}
impl<T> BackoffTracker<T> {
pub fn new(clock: Arc<dyn Clock>, config: BackoffConfig) -> Self {
Self {
field: HashMap::new(),
clock,
config,
}
}
}
impl<'t, T: Eq + Hash + 't> BackoffTracker<T> {
pub fn backoff<F: Fn() -> X, X>(&mut self, id: T, f: F) -> Option<X> {
self.backoff_generic(id, f, false, |_| true)
}
pub fn backoff_singular<F: Fn() -> X, X>(&mut self, id: T, f: F) -> Option<X> {
self.backoff_generic(id, f, true, |_| true)
}
pub fn backoff_errors<F, O, E>(&mut self, id: T, f: F) -> Option<Result<O, E>>
where
F: Fn() -> Result<O, E>,
{
self.backoff_generic(id, f, false, Result::is_err)
}
pub fn backoff_oks<F, O, E>(&mut self, id: T, f: F) -> Option<Result<O, E>>
where
F: Fn() -> Result<O, E>,
{
self.backoff_generic(id, f, false, Result::is_ok)
}
pub fn backoff_generic<F, X, P>(
&mut self,
id: T,
f: F,
singular: bool,
predicate: P,
) -> Option<X>
where
F: Fn() -> X,
P: Fn(&X) -> bool,
{
if self.is_ready(&id) {
let result = f();
if predicate(&result) {
if singular && !self.field.contains_key(&id) {
self.field = HashMap::new();
}
self.event(id);
}
Some(result)
} else {
None
}
}
}
impl<'t, T: Eq + Hash + 't> BackoffTracker<T> {
pub fn event(&mut self, item: T) {
let now = self.clock.current_timestamp();
match self.field.entry(item) {
Entry::Vacant(entry) => {
entry.insert(BackoffEntry {
count: 1,
last: now,
delay: self.config.initial_delay(),
});
}
Entry::Occupied(mut entry) => {
let BackoffEntry { count, last, delay } = entry.get();
if last + delay < now {
let delay = match self.config.strategy {
BackoffStrategy::Exponential(b) => {
std::cmp::min(delay * b, self.config.max_delay)
}
BackoffStrategy::Cliff(c) => {
if count > &c {
self.config.max_delay
} else {
self.config.min_delay
}
}
};
entry.insert(BackoffEntry {
count: count + 1,
last: now,
delay,
});
}
}
}
}
pub fn singular_event(&mut self, item: T) {
if !self.field.contains_key(&item) {
self.field = HashMap::new();
}
self.event(item);
}
pub fn clear(&mut self, item: &T) {
if self.field.contains_key(item) {
self.field.remove(item);
}
}
pub fn clear_many<'a>(&mut self, items: impl IntoIterator<Item = &'a T>)
where
't: 'a,
{
for item in items {
self.field.remove(item);
}
}
pub fn is_ready(&self, item: &T) -> bool {
match self.field.get(item) {
Some(x) => self.clock.current_timestamp() > x.last + x.delay,
None => true,
}
}
}
struct BackoffEntry {
count: u64,
last: u64,
delay: u64,
}
pub struct BackoffConfig {
pub min_delay: u64,
pub max_delay: u64,
pub strategy: BackoffStrategy,
}
impl BackoffConfig {
pub fn constant(period: u64) -> Self {
Self {
min_delay: period,
max_delay: period,
strategy: BackoffStrategy::Cliff(0),
}
}
pub fn once() -> Self {
Self {
min_delay: u64::MAX,
max_delay: u64::MAX,
strategy: BackoffStrategy::Cliff(0),
}
}
pub fn initial_delay(&self) -> u64 {
match self.strategy {
BackoffStrategy::Cliff(0) => self.max_delay,
_ => self.min_delay,
}
}
}
impl Default for BackoffConfig {
fn default() -> Self {
Self {
min_delay: 30,
max_delay: 3600,
strategy: Default::default(),
}
}
}
pub enum BackoffStrategy {
Exponential(u64),
Cliff(u64),
}
impl Default for BackoffStrategy {
fn default() -> Self {
BackoffStrategy::Exponential(3)
}
}