83 lines
1.9 KiB
Rust
83 lines
1.9 KiB
Rust
use apub_session::Session;
|
|
use dashmap::DashMap;
|
|
use std::{
|
|
sync::Arc,
|
|
time::{Duration, Instant},
|
|
};
|
|
use url::{Host, Url};
|
|
|
|
#[derive(Debug)]
|
|
struct Breaker {
|
|
failure_count: usize,
|
|
broken_at: Instant,
|
|
}
|
|
|
|
#[derive(Clone, Debug)]
|
|
pub struct BreakerSession {
|
|
limit: usize,
|
|
breaker_duration: Duration,
|
|
hosts: Arc<DashMap<(Host<String>, Option<u16>), Breaker>>,
|
|
}
|
|
|
|
impl BreakerSession {
|
|
pub fn limit(limit: usize, breaker_duration: Duration) -> Self {
|
|
Self {
|
|
limit,
|
|
breaker_duration,
|
|
hosts: Arc::new(DashMap::new()),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Session for BreakerSession {
|
|
fn should_procede(&self, url: &Url) -> bool {
|
|
if let Some(host) = url.host() {
|
|
let key = (host.to_owned(), url.port());
|
|
|
|
let mut breaker = self.hosts.entry(key).or_default();
|
|
|
|
if breaker.failure_count < self.limit {
|
|
return true;
|
|
}
|
|
|
|
if Instant::now() > breaker.broken_at + self.breaker_duration {
|
|
breaker.failure_count = 0;
|
|
}
|
|
|
|
false
|
|
} else {
|
|
true
|
|
}
|
|
}
|
|
|
|
fn mark_success(&self, url: &Url) {
|
|
if let Some(host) = url.host() {
|
|
let key = (host.to_owned(), url.port());
|
|
let mut breaker = self.hosts.entry(key).or_default();
|
|
|
|
breaker.failure_count = 0;
|
|
}
|
|
}
|
|
|
|
fn mark_failure(&self, url: &Url) {
|
|
if let Some(host) = url.host() {
|
|
let key = (host.to_owned(), url.port());
|
|
let mut breaker = self.hosts.entry(key).or_default();
|
|
|
|
breaker.failure_count += 1;
|
|
if breaker.failure_count >= self.limit {
|
|
breaker.broken_at = Instant::now();
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Default for Breaker {
|
|
fn default() -> Self {
|
|
Self {
|
|
failure_count: 0,
|
|
broken_at: Instant::now(),
|
|
}
|
|
}
|
|
}
|