use std::collections::HashMap;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::atomic::{AtomicI64, AtomicU32, Ordering};
use std::time::Duration;
use chrono::Utc;
use tokio::sync::{OwnedSemaphorePermit, Semaphore};
use crate::error::RssError;
const HOST_MAX_CONCURRENCY: usize = 1;
const HOST_BASE_COOLDOWN: Duration = Duration::from_secs(2);
const HOST_MAX_COOLDOWN: Duration = Duration::from_secs(60);
const STICKY_SPACING: Duration = Duration::from_secs(1);
const WARM_WINDOW: Duration = Duration::from_secs(120);
const MAX_GATE_WAIT: Duration = Duration::from_secs(60);
fn now_ms() -> i64 {
Utc::now().timestamp_millis()
}
fn authority_of(url: &str) -> String {
match url::Url::parse(url) {
Ok(u) => {
let host = u.host_str().unwrap_or_default().to_ascii_lowercase();
let port = u.port_or_known_default().unwrap_or(0);
format!("{host}:{port}")
}
Err(_) => url.to_string(),
}
}
struct HostSlot {
sem: Arc<Semaphore>,
next_allowed_ms: AtomicI64,
warm_until_ms: AtomicI64,
consecutive_throttles: AtomicU32,
}
impl HostSlot {
fn new(per_host: usize) -> Self {
Self {
sem: Arc::new(Semaphore::new(per_host.max(1))),
next_allowed_ms: AtomicI64::new(0),
warm_until_ms: AtomicI64::new(0),
consecutive_throttles: AtomicU32::new(0),
}
}
}
pub struct HostGate {
slots: Mutex<HashMap<String, Arc<HostSlot>>>,
per_host: usize,
base_cooldown: Duration,
max_cooldown: Duration,
sticky_spacing: Duration,
warm_window: Duration,
max_gate_wait: Duration,
}
impl Default for HostGate {
fn default() -> Self {
Self {
slots: Mutex::new(HashMap::new()),
per_host: HOST_MAX_CONCURRENCY,
base_cooldown: HOST_BASE_COOLDOWN,
max_cooldown: HOST_MAX_COOLDOWN,
sticky_spacing: STICKY_SPACING,
warm_window: WARM_WINDOW,
max_gate_wait: MAX_GATE_WAIT,
}
}
}
impl HostGate {
pub fn from_env() -> Self {
let mut gate = Self::default();
if let Some(n) = env_parse::<usize>("RSS_HOST_CONCURRENCY") {
gate.per_host = n.max(1);
}
if let Some(secs) = env_parse::<u64>("RSS_MAX_COOLDOWN_SECS") {
gate.max_cooldown = Duration::from_secs(secs);
}
if let Some(secs) = env_parse::<u64>("RSS_MAX_GATE_WAIT_SECS") {
gate.max_gate_wait = Duration::from_secs(secs);
}
gate
}
fn slot_for(&self, authority: &str) -> Arc<HostSlot> {
let mut slots = self.slots.lock().expect("host-gate mutex poisoned");
slots
.entry(authority.to_string())
.or_insert_with(|| Arc::new(HostSlot::new(self.per_host)))
.clone()
}
fn slot_for_url(&self, url: &str) -> Arc<HostSlot> {
self.slot_for(&authority_of(url))
}
pub async fn acquire(&self, url: &str) -> Result<OwnedSemaphorePermit, RssError> {
let slot = self.slot_for_url(url);
let permit = slot
.sem
.clone()
.acquire_owned()
.await
.map_err(|_| RssError::Other("rate-limiter semaphore closed".to_string()))?;
let now = now_ms();
let target = slot.next_allowed_ms.load(Ordering::Relaxed);
let wait = Duration::from_millis(u64::try_from((target - now).max(0)).unwrap_or(0));
if wait > self.max_gate_wait {
return Err(RssError::RateLimited {
url: url.to_string(),
retry_after: wait,
});
}
if !wait.is_zero() {
tokio::time::sleep(wait).await;
}
let now = now_ms();
if slot.warm_until_ms.load(Ordering::Relaxed) > now {
let spaced = now + ms(self.sticky_spacing);
slot.next_allowed_ms.fetch_max(spaced, Ordering::Relaxed);
}
Ok(permit)
}
fn cooldown_for(&self, n: u32, retry_after: Option<Duration>) -> Duration {
match retry_after {
Some(d) => d.min(self.max_cooldown),
None => {
let shift = n.saturating_sub(1).min(30);
self.base_cooldown
.saturating_mul(1u32 << shift)
.min(self.max_cooldown)
}
}
}
pub fn note_throttled(&self, url: &str, retry_after: Option<Duration>) {
let slot = self.slot_for_url(url);
let n = slot.consecutive_throttles.fetch_add(1, Ordering::Relaxed) + 1;
let cooldown = self.cooldown_for(n, retry_after);
let now = now_ms();
slot.next_allowed_ms
.fetch_max(now + ms(cooldown), Ordering::Relaxed);
slot.warm_until_ms
.fetch_max(now + ms(self.warm_window), Ordering::Relaxed);
}
pub fn note_success(&self, url: &str) {
self.slot_for_url(url)
.consecutive_throttles
.store(0, Ordering::Relaxed);
}
}
fn ms(d: Duration) -> i64 {
i64::try_from(d.as_millis()).unwrap_or(i64::MAX)
}
fn env_parse<T: std::str::FromStr>(key: &str) -> Option<T> {
std::env::var(key).ok()?.trim().parse().ok()
}
#[cfg(test)]
mod tests {
use super::*;
fn gate() -> HostGate {
HostGate::default()
}
#[test]
fn authority_normalizes_host_and_port() {
assert_eq!(
authority_of("https://Example.com/feed.xml"),
"example.com:443"
);
assert_eq!(authority_of("http://example.com/feed"), "example.com:80");
assert_eq!(
authority_of("https://example.com:8443/x"),
"example.com:8443"
);
assert_ne!(
authority_of("https://a.example.com/x"),
authority_of("https://b.example.com/x")
);
assert_eq!(
authority_of("https://example.com/x"),
authority_of("https://example.com/y")
);
}
#[tokio::test]
async fn cold_host_acquires_without_waiting() {
let g = gate();
let t0 = now_ms();
let _p = g.acquire("https://example.com/feed").await.unwrap();
assert!(now_ms() - t0 < 200, "cold acquire must not sleep");
}
#[tokio::test]
async fn cap_one_serializes_same_authority() {
let g = gate();
let p1 = g.acquire("https://example.com/a").await.unwrap();
let fut = g.acquire("https://example.com/b");
tokio::pin!(fut);
assert!(
tokio::time::timeout(Duration::from_millis(100), &mut fut)
.await
.is_err(),
"second same-host acquire should block while the first permit is held"
);
drop(p1);
assert!(fut.await.is_ok());
}
#[tokio::test]
async fn distinct_authorities_do_not_block_each_other() {
let g = gate();
let _a = g.acquire("https://a.example.com/x").await.unwrap();
let b = tokio::time::timeout(
Duration::from_millis(100),
g.acquire("https://b.example.com/x"),
)
.await;
assert!(
b.is_ok(),
"a distinct host must not be gated behind another"
);
}
#[tokio::test]
async fn throttle_makes_sibling_wait_but_not_a_different_host() {
let g = gate();
let t_throttled = now_ms();
g.note_throttled(
"https://slow.example.com/x",
Some(Duration::from_millis(300)),
);
let t0 = now_ms();
let _other = g.acquire("https://fast.example.com/y").await.unwrap();
assert!(now_ms() - t0 < 300, "unrelated host must not wait");
let _same = g.acquire("https://slow.example.com/z").await.unwrap();
let waited = now_ms() - t_throttled;
assert!(
(250..3000).contains(&waited),
"throttled host should wait ~the cooldown, waited {waited}ms"
);
}
#[test]
fn cooldown_for_escalates_doubling_and_caps() {
let g = HostGate {
base_cooldown: Duration::from_secs(2),
max_cooldown: Duration::from_secs(60),
..HostGate::default()
};
assert_eq!(g.cooldown_for(1, None), Duration::from_secs(2));
assert_eq!(g.cooldown_for(2, None), Duration::from_secs(4));
assert_eq!(g.cooldown_for(3, None), Duration::from_secs(8));
assert_eq!(g.cooldown_for(4, None), Duration::from_secs(16));
assert_eq!(g.cooldown_for(5, None), Duration::from_secs(32));
assert_eq!(g.cooldown_for(6, None), Duration::from_secs(60)); assert_eq!(g.cooldown_for(7, None), Duration::from_secs(60));
assert_eq!(g.cooldown_for(1000, None), Duration::from_secs(60));
assert_eq!(
g.cooldown_for(1, Some(Duration::from_millis(300))),
Duration::from_millis(300)
);
assert_eq!(
g.cooldown_for(9, Some(Duration::from_secs(600))),
Duration::from_secs(60)
);
}
#[test]
fn note_success_resets_escalation_counter() {
let g = gate();
let url = "https://esc.example.com/x";
g.note_throttled(url, None);
g.note_throttled(url, None);
assert_eq!(
g.slot_for("esc.example.com:443")
.consecutive_throttles
.load(Ordering::Relaxed),
2
);
g.note_success(url);
assert_eq!(
g.slot_for("esc.example.com:443")
.consecutive_throttles
.load(Ordering::Relaxed),
0,
"a success must reset the escalation counter so the next throttle starts from base"
);
}
#[tokio::test]
async fn warm_host_spaces_the_next_sibling() {
let g = HostGate {
sticky_spacing: Duration::from_millis(300),
..HostGate::default()
};
let url = "https://warm.example.com/x";
g.note_throttled(url, Some(Duration::from_millis(1)));
let _first = g.acquire(url).await.unwrap(); let t = now_ms();
drop(_first);
let _second = g.acquire(url).await.unwrap();
let waited = now_ms() - t;
assert!(
(200..2000).contains(&waited),
"a warm host should space the next sibling by ~sticky_spacing, waited {waited}ms"
);
}
#[tokio::test]
async fn zero_host_concurrency_is_floored_and_does_not_deadlock() {
let g = HostGate {
per_host: 0,
..HostGate::default()
};
let acquired = tokio::time::timeout(
Duration::from_millis(500),
g.acquire("https://z.example.com/x"),
)
.await;
assert!(
acquired.is_ok(),
"per_host must be floored to 1 so acquire never deadlocks"
);
}
#[tokio::test]
async fn wait_beyond_ceiling_fails_fast_with_rate_limited() {
let g = HostGate {
max_gate_wait: Duration::from_millis(100),
max_cooldown: Duration::from_secs(60),
..HostGate::default()
};
let url = "https://busy.example.com/x";
g.note_throttled(url, Some(Duration::from_secs(30)));
let err = g.acquire(url).await.unwrap_err();
match err {
RssError::RateLimited { retry_after, .. } => {
assert!(
retry_after >= Duration::from_secs(1),
"should report a real wait"
);
}
other => panic!("expected RateLimited, got {other:?}"),
}
}
#[tokio::test]
async fn cancelled_acquire_releases_the_permit() {
let g = gate();
let p1 = g.acquire("https://cx.example.com/a").await.unwrap();
{
let fut = g.acquire("https://cx.example.com/b");
tokio::pin!(fut);
let _ = tokio::time::timeout(Duration::from_millis(50), &mut fut).await;
}
drop(p1);
assert!(
tokio::time::timeout(
Duration::from_millis(200),
g.acquire("https://cx.example.com/c")
)
.await
.is_ok(),
"host must be acquirable after a cancelled waiter and a released permit"
);
}
}