107 lines
2.7 KiB
Rust
107 lines
2.7 KiB
Rust
//! A Session implementation for limiting requests to domains that consistently fail
|
|
//!
|
|
//! ```rust
|
|
//! use apub_breaker_session::BreakerSession;
|
|
//! use std::time::Duration;
|
|
//!
|
|
//! // Create a session that refuses requests for an hour after 10 consecutive failures to a given domain
|
|
//! let breaker_session = BreakerSession::limit(10, Duration::from_secs(60 * 60));
|
|
//! ```
|
|
|
|
#![deny(missing_docs)]
|
|
|
|
use apub_core::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,
|
|
}
|
|
|
|
type BreakerKey = (Host<String>, Option<u16>);
|
|
|
|
/// The BreakerSession type
|
|
#[derive(Clone, Debug)]
|
|
pub struct BreakerSession {
|
|
limit: usize,
|
|
breaker_duration: Duration,
|
|
hosts: Arc<DashMap<BreakerKey, Breaker>>,
|
|
}
|
|
|
|
impl BreakerSession {
|
|
/// Create a new BreakerSession
|
|
///
|
|
/// ```rust
|
|
/// use apub_breaker_session::BreakerSession;
|
|
/// use std::time::Duration;
|
|
///
|
|
/// // Create a session that refuses requests for an hour after 10 consecutive failures to a given domain
|
|
/// let breaker_session = BreakerSession::limit(10, Duration::from_secs(60 * 60));
|
|
/// ```
|
|
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(&mut 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(&mut 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(&mut 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(),
|
|
}
|
|
}
|
|
}
|