use crate::compat::sleep;
use crate::runner::RunnerConfig;
use bot_core::ExchangeError;
use std::time::Duration;
pub trait HasItems {
fn has_items(&self) -> bool;
}
impl<T> HasItems for Vec<T> {
fn has_items(&self) -> bool {
!self.is_empty()
}
}
pub enum PollOutcome<T> {
Data(T),
Empty,
Degraded(ExchangeError),
Fatal(ExchangeError),
}
pub struct PollGuard {
label: &'static str,
backoff_ms: u64,
initial_backoff_ms: u64,
max_backoff_ms: u64,
backoff_multiplier: f64,
consecutive_empties: u32,
}
impl PollGuard {
pub fn new(label: &'static str, config: &RunnerConfig) -> Self {
Self {
label,
backoff_ms: 0,
initial_backoff_ms: config.initial_backoff_ms,
max_backoff_ms: config.max_backoff_ms,
backoff_multiplier: config.backoff_multiplier,
consecutive_empties: 0,
}
}
pub async fn execute<T, F, Fut>(&mut self, poll_fn: F) -> PollOutcome<T>
where
T: HasItems,
F: FnOnce() -> Fut,
Fut: std::future::Future<Output = Result<T, ExchangeError>>,
{
if self.backoff_ms > 0 {
sleep(Duration::from_millis(self.backoff_ms)).await;
}
match poll_fn().await {
Ok(result) if result.has_items() => {
self.backoff_ms = 0;
self.consecutive_empties = 0;
PollOutcome::Data(result)
}
Ok(_empty) => {
self.backoff_ms = 0;
self.consecutive_empties += 1;
tracing::debug!(
"[PollGuard:{}] Empty response ({}x consecutive)",
self.label,
self.consecutive_empties
);
PollOutcome::Empty
}
Err(e) if e.is_transient() => {
self.increase_backoff();
if e.is_502() {
tracing::warn!(
"[PollGuard:{}] Exchange unavailable, backoff={}ms: {}",
self.label,
self.backoff_ms,
e
);
PollOutcome::Degraded(e)
} else {
tracing::warn!(
"[PollGuard:{}] Transient error, backoff={}ms: {}",
self.label,
self.backoff_ms,
e
);
PollOutcome::Empty
}
}
Err(e) => {
tracing::error!("[PollGuard:{}] Fatal error: {}", self.label, e);
PollOutcome::Fatal(e)
}
}
}
pub fn looks_exhausted(&self, threshold: u32) -> bool {
self.consecutive_empties >= threshold
}
fn increase_backoff(&mut self) {
self.backoff_ms = if self.backoff_ms == 0 {
self.initial_backoff_ms
} else {
((self.backoff_ms as f64) * self.backoff_multiplier) as u64
};
self.backoff_ms = self.backoff_ms.min(self.max_backoff_ms);
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn has_items_vec() {
let empty: Vec<i32> = vec![];
assert!(!empty.has_items());
let non_empty = vec![1, 2, 3];
assert!(non_empty.has_items());
}
#[test]
fn looks_exhausted_threshold() {
let config = RunnerConfig::default();
let mut guard = PollGuard::new("test", &config);
assert!(!guard.looks_exhausted(3));
guard.consecutive_empties = 2;
assert!(!guard.looks_exhausted(3));
guard.consecutive_empties = 3;
assert!(guard.looks_exhausted(3));
}
#[test]
fn increase_backoff_exponential() {
let config = RunnerConfig {
initial_backoff_ms: 100,
max_backoff_ms: 1000,
backoff_multiplier: 2.0,
..RunnerConfig::default()
};
let mut guard = PollGuard::new("test", &config);
guard.increase_backoff();
assert_eq!(guard.backoff_ms, 100);
guard.increase_backoff();
assert_eq!(guard.backoff_ms, 200);
guard.increase_backoff();
assert_eq!(guard.backoff_ms, 400);
guard.increase_backoff();
assert_eq!(guard.backoff_ms, 800);
guard.increase_backoff();
assert_eq!(guard.backoff_ms, 1000);
}
}