crw-renderer 0.23.0

//! Process-wide per-host rate-limiter and concurrency cap.
//!
//! Keyed by eTLD+1 (via [`crate::preference::normalize_host`]) so subdomains
//! under one registered domain share a single limiter. Each entry pairs a
//! [`RateLimiter`] (interval enforcement) with a [`Semaphore`] (in-flight
//! cap). Used by both [`crate::FallbackRenderer`] and the crawl/discover
//! loops in `crw-crawl` so single /scrape calls and crawl jobs respect the
//! same global per-host budget.
//!
//! First-write-wins for capacity and RPS — a second caller asking for
//! different values silently reuses the existing limiter (warned on RPS
//! mismatch). This matches the historical behaviour of the
//! `DOMAIN_RATE_LIMITERS` map this module replaces.

use dashmap::DashMap;
use std::sync::{Arc, LazyLock};
use std::time::Duration;
use tokio::sync::Mutex;
use tokio::time::Instant;

use crw_core::{LanePermit, ReservedSemaphore};

/// Stale entries are dropped after this idle duration during periodic GC.
const RATE_LIMITER_TTL: Duration = Duration::from_secs(3600);

type RateLimiterEntry = (Arc<Mutex<RateLimiter>>, ReservedSemaphore, Instant);

static DOMAIN_RATE_LIMITERS: LazyLock<DashMap<String, RateLimiterEntry>> =
    LazyLock::new(DashMap::new);

static LIMITER_CALL_COUNT: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);

/// Get or create a (rate limiter, per-host reserved semaphore) pair for the
/// given **already-normalized** host key. Pass the eTLD+1 form produced by
/// [`crate::preference::normalize_host`].
///
/// The semaphore is a [`ReservedSemaphore`] sized `per_host_max_concurrent +
/// interactive_reserve`: batch/crawl are bounded to `per_host_max_concurrent`
/// per host (target politeness), while interactive gets `interactive_reserve`
/// dedicated slots so it's never queued behind a batch — yet total in-flight
/// per host stays bounded (no target hammering).
pub fn get_host_limiter(
    host_key: &str,
    rps: f64,
    per_host_max_concurrent: usize,
    interactive_reserve: usize,
) -> (Arc<Mutex<RateLimiter>>, ReservedSemaphore) {
    let now = Instant::now();
    if LIMITER_CALL_COUNT
        .fetch_add(1, std::sync::atomic::Ordering::Relaxed)
        .is_multiple_of(64)
    {
        DOMAIN_RATE_LIMITERS
            .retain(|_, (_, _, last_used)| now.duration_since(*last_used) < RATE_LIMITER_TTL);
    }
    let cap = per_host_max_concurrent.max(1);
    let entry = DOMAIN_RATE_LIMITERS
        .entry(host_key.to_string())
        .and_modify(|entry| {
            entry.2 = now;
            let existing_interval = entry.0.try_lock().map(|l| l.min_interval).ok();
            let new_interval = if rps > 0.0 {
                Duration::from_secs_f64(1.0 / rps)
            } else {
                Duration::ZERO
            };
            if let Some(existing) = existing_interval
                && existing != new_interval
            {
                tracing::warn!(
                    domain = host_key,
                    existing_rps = ?existing,
                    requested_rps = rps,
                    "Rate limiter RPS mismatch for domain; using existing limiter"
                );
            }
        })
        .or_insert_with(|| {
            (
                Arc::new(Mutex::new(RateLimiter::new(rps))),
                // total = batch cap + interactive reserve; reserved = interactive.
                ReservedSemaphore::new(cap + interactive_reserve, interactive_reserve, "host"),
                now,
            )
        });
    (entry.0.clone(), entry.1.clone())
}

/// Convenience: acquire the per-host concurrency permit and compute how long
/// the caller must sleep to honour the rate limit. The caller is responsible
/// for performing the sleep (so it doesn't happen while holding any other
/// async lock) and for keeping the permit alive across the request.
///
/// Class-aware (the A reserved lane): interactive takes a reserved per-host slot
/// (never queued behind batch), batch stays bounded to `per_host_max_concurrent`
/// per host. Both honour the rate limiter (target RPS). The returned
/// [`LanePermit`] must be held across the request.
pub async fn acquire(
    host_key: &str,
    rps: f64,
    per_host_max_concurrent: usize,
    interactive_reserve: usize,
) -> (LanePermit, Duration) {
    let (rate_limiter, sem) =
        get_host_limiter(host_key, rps, per_host_max_concurrent, interactive_reserve);
    // Read class on the async side; acquire is infallible (semaphore never closed).
    let permit = sem.acquire(crw_core::current_scrape_class()).await;
    let sleep = rate_limiter.lock().await.next_sleep();
    (permit, sleep)
}

/// Minimum-interval rate limiter. Public so `crw-crawl` can keep its existing
/// jitter wrapper.
pub struct RateLimiter {
    pub min_interval: Duration,
    last_request: Instant,
}

impl RateLimiter {
    pub fn new(requests_per_second: f64) -> Self {
        if requests_per_second < 0.0 {
            tracing::warn!(
                requests_per_second,
                "Negative requests_per_second value, treating as unlimited"
            );
        }
        let min_interval = if requests_per_second > 0.0 {
            Duration::from_secs_f64(1.0 / requests_per_second)
        } else {
            Duration::ZERO
        };
        Self {
            min_interval,
            last_request: Instant::now() - min_interval,
        }
    }

    /// Compute how long to sleep and update last_request. Caller must sleep
    /// outside any held async lock.
    pub fn next_sleep(&mut self) -> Duration {
        let elapsed = self.last_request.elapsed();
        let sleep = if elapsed < self.min_interval {
            self.min_interval - elapsed
        } else {
            Duration::ZERO
        };
        self.last_request = Instant::now() + sleep;
        sleep
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn rate_limiter_zero_rps_is_unlimited() {
        let mut limiter = RateLimiter::new(0.0);
        assert_eq!(limiter.next_sleep(), Duration::ZERO);
        assert_eq!(limiter.next_sleep(), Duration::ZERO);
    }

    #[test]
    fn rate_limiter_negative_rps_is_unlimited() {
        let mut limiter = RateLimiter::new(-1.0);
        assert_eq!(limiter.next_sleep(), Duration::ZERO);
    }

    #[test]
    fn rate_limiter_enforces_interval() {
        let mut limiter = RateLimiter::new(10.0); // 100ms interval
        let _first = limiter.next_sleep();
        let second = limiter.next_sleep();
        assert!(second > Duration::from_millis(50) && second <= Duration::from_millis(100));
    }

    #[tokio::test]
    async fn acquire_returns_permit_and_sleep() {
        // Unscoped = Interactive = takes a reserved per-host slot, no sleep.
        let (_p, sleep) = acquire("test-acquire-1.example", 0.0, 1, 1).await;
        assert_eq!(sleep, Duration::ZERO);
    }

    #[tokio::test]
    async fn per_host_cap_serializes_batch() {
        // Batch traffic is bounded to `per_host_max_concurrent` (=1 here); a
        // second concurrent batch acquire to the same host must wait.
        crw_core::REQUEST_CLASS
            .scope(crw_core::ScrapeClass::Batch, async {
                let (p1, _) = acquire("test-cap.example", 0.0, 1, 1).await;
                // Second batch acquire must wait — verify it doesn't complete in 50ms.
                let acquire_fut = acquire("test-cap.example", 0.0, 1, 1);
                let race = tokio::time::timeout(Duration::from_millis(50), acquire_fut).await;
                assert!(
                    race.is_err(),
                    "second batch acquire should block while p1 held (batch cap=1)"
                );
                drop(p1);
                // Now it should succeed quickly.
                let _p2 = tokio::time::timeout(
                    Duration::from_millis(200),
                    acquire("test-cap.example", 0.0, 1, 1),
                )
                .await
                .expect("second batch acquire should succeed after release");
            })
            .await;
    }

    #[tokio::test]
    async fn interactive_reserved_slot_survives_batch() {
        // Interactive gets a dedicated reserved slot, so it never blocks on the
        // batch-held per-host cap.
        crw_core::REQUEST_CLASS
            .scope(crw_core::ScrapeClass::Batch, async {
                let (_batch, _) = acquire("test-bypass.example", 0.0, 1, 1).await;
                // A concurrent interactive request to the same host gets straight
                // through (its reserved slot).
                let interactive =
                    crw_core::REQUEST_CLASS.scope(crw_core::ScrapeClass::Interactive, async {
                        tokio::time::timeout(
                            Duration::from_millis(100),
                            acquire("test-bypass.example", 0.0, 1, 1),
                        )
                        .await
                    });
                let (_p, _) = interactive
                    .await
                    .expect("interactive should not block on the batch-held host cap");
            })
            .await;
    }
}