kreuzcrawl 0.2.0

High-performance web crawling engine
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117

//! Per-domain rate limiting layer for the Tower service stack.

use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};

use tower::{Layer, Service};

use super::types::{CrawlRequest, CrawlResponse};
use crate::error::CrawlError;
use crate::traits::RateLimiter;

/// Tower layer that applies per-domain rate limiting using a [`RateLimiter`].
pub struct PerDomainRateLimitLayer {
    rate_limiter: Arc<dyn RateLimiter>,
}

impl PerDomainRateLimitLayer {
    pub fn new(rate_limiter: Arc<dyn RateLimiter>) -> Self {
        Self { rate_limiter }
    }
}

impl<S: Clone> Layer<S> for PerDomainRateLimitLayer {
    type Service = PerDomainRateLimitService<S>;

    fn layer(&self, inner: S) -> Self::Service {
        PerDomainRateLimitService {
            inner,
            rate_limiter: self.rate_limiter.clone(),
        }
    }
}

/// Tower service that enforces per-domain rate limits before forwarding requests.
#[derive(Clone)]
pub struct PerDomainRateLimitService<S> {
    inner: S,
    rate_limiter: Arc<dyn RateLimiter>,
}

impl<S> Service<CrawlRequest> for PerDomainRateLimitService<S>
where
    S: Service<CrawlRequest, Response = CrawlResponse, Error = CrawlError> + Clone + Send + 'static,
    S::Future: Send + 'static,
{
    type Response = CrawlResponse;
    type Error = CrawlError;
    type Future = Pin<Box<dyn Future<Output = Result<CrawlResponse, CrawlError>> + Send>>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx)
    }

    fn call(&mut self, req: CrawlRequest) -> Self::Future {
        let domain = req.domain().unwrap_or_default();
        let rate_limiter = self.rate_limiter.clone();
        let mut inner = self.inner.clone();
        // Swap to preserve readiness (standard Tower pattern)
        std::mem::swap(&mut self.inner, &mut inner);

        Box::pin(async move {
            // Acquire rate limit permit
            if !domain.is_empty() {
                rate_limiter.acquire(&domain).await?;
            }

            // Forward to inner service
            let resp = inner.call(req).await?;

            // Record response for adaptive backoff
            if !domain.is_empty() {
                rate_limiter.record_response(&domain, resp.status).await?;
            }

            Ok(resp)
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::defaults::NoopRateLimiter;
    use tower::Service;

    #[derive(Clone)]
    struct OkService;
    impl Service<CrawlRequest> for OkService {
        type Response = CrawlResponse;
        type Error = CrawlError;
        type Future = std::pin::Pin<Box<dyn std::future::Future<Output = Result<CrawlResponse, CrawlError>> + Send>>;
        fn poll_ready(&mut self, _: &mut std::task::Context<'_>) -> std::task::Poll<Result<(), Self::Error>> {
            std::task::Poll::Ready(Ok(()))
        }
        fn call(&mut self, _: CrawlRequest) -> Self::Future {
            Box::pin(async {
                Ok(CrawlResponse {
                    status: 200,
                    content_type: "text/html".into(),
                    body: "ok".into(),
                    body_bytes: vec![],
                    headers: std::collections::HashMap::new(),
                })
            })
        }
    }

    #[tokio::test]
    async fn test_rate_limit_layer_passes_through() {
        let layer = PerDomainRateLimitLayer::new(std::sync::Arc::new(NoopRateLimiter));
        let mut svc = layer.layer(OkService);
        let resp = svc.call(CrawlRequest::new("http://example.com")).await.unwrap();
        assert_eq!(resp.status, 200);
    }
}