Skip to main content

leyline/core/crawl/
mod.rs

1mod clock;
2mod fetch;
3mod frontier;
4mod page;
5mod sitemap;
6mod sites;
7mod visit;
8
9use std::collections::{HashMap, HashSet, VecDeque};
10use std::future::Future;
11use std::pin::Pin;
12
13use futures_util::stream::{FuturesUnordered, Stream};
14use url::Url;
15
16use fetch::{Fetched, fetch};
17pub use frontier::Scope;
18use frontier::{Frontier, Queued, page_key, same_site};
19pub use page::{CrawlPage, CrawlStats, SkipReason};
20use sitemap::SitemapRead;
21use sites::{Site, SitemapFile};
22use visit::Held;
23
24use crate::core::config::{HostLimits, Origin};
25use crate::core::error::{Error, Kind, Result};
26use crate::core::into_url::IntoUrl;
27use crate::core::retry::RetryPolicy;
28use crate::core::robots::{Purpose, Robots};
29use crate::core::session::Session;
30
31const ANY_AGENT: &str = "*";
32const NOFOLLOW: &str = "nofollow";
33const MAX_REDIRECTS: usize = 5;
34const PAGE_RETRIES: u32 = 2;
35const MIN_PRINT_WORDS: usize = 20;
36const MAX_IN_FLIGHT: usize = 16;
37const POLITE_IN_FLIGHT: usize = 1;
38const POLITE_PER_SECOND: f64 = 1.0;
39const PAUSE_STATUSES: [u16; 2] = [429, 503];
40
41type Task = Pin<Box<dyn Future<Output = Event> + Send>>;
42
43enum Event {
44    Robots {
45        origin: Origin,
46        url: Url,
47        robots: Box<Robots>,
48    },
49    Page {
50        origin: Origin,
51        queued: Queued,
52        fetched: Result<Fetched>,
53    },
54    Sitemap {
55        file: SitemapFile,
56        read: SitemapRead,
57    },
58}
59
60#[derive(Debug, Clone, Copy, PartialEq, Eq)]
61#[non_exhaustive]
62pub enum Respect {
63    Robots,
64    AiTraining,
65}
66
67#[must_use = "a crawl is lazy: nothing happens until `.next()` is awaited"]
68pub struct Crawl {
69    session: Session,
70    limits: Option<HostLimits>,
71    pace: HostLimits,
72    seeds: Vec<Result<Url>>,
73    respect: Option<Respect>,
74    max_pages: Option<usize>,
75    started: bool,
76    frontier: Frontier,
77    sites: HashMap<Origin, Site>,
78    busy: HashMap<Origin, usize>,
79    sitemap_files: HashSet<String>,
80    sitemap_queue: VecDeque<SitemapFile>,
81    sitemaps_in_flight: usize,
82    trusted: Option<Url>,
83    page_retry: RetryPolicy,
84    clock: clock::Clock,
85    holding: usize,
86    tasks: FuturesUnordered<Task>,
87    pages_in_flight: usize,
88    ready: VecDeque<Result<CrawlPage>>,
89    kept: HashSet<String>,
90    texts: HashSet<u64>,
91    held: HashMap<String, Vec<Held>>,
92    waiting_on: HashMap<String, String>,
93    settled: HashSet<String>,
94    stats: CrawlStats,
95}
96
97fn polite_limits() -> HostLimits {
98    HostLimits::new()
99        .max_in_flight(POLITE_IN_FLIGHT)
100        .per_second(POLITE_PER_SECOND)
101        .pause_on(PAUSE_STATUSES)
102}
103
104fn inherited(session: &Session) -> HostLimits {
105    session.host_limits().clone().filled_from(&polite_limits())
106}
107
108impl Session {
109    pub fn crawl(&self, seed: impl IntoUrl) -> Crawl {
110        Crawl {
111            session: self.clone(),
112            limits: None,
113            pace: HostLimits::new(),
114            seeds: vec![seed.into_url_with_base(self.base_url())],
115            respect: Some(Respect::Robots),
116            max_pages: None,
117            started: false,
118            frontier: Frontier::default(),
119            sites: HashMap::new(),
120            busy: HashMap::new(),
121            sitemap_files: HashSet::new(),
122            sitemap_queue: VecDeque::new(),
123            sitemaps_in_flight: 0,
124            trusted: None,
125            page_retry: RetryPolicy::none(),
126            clock: clock::Clock::default(),
127            holding: 0,
128            tasks: FuturesUnordered::new(),
129            pages_in_flight: 0,
130            ready: VecDeque::new(),
131            kept: HashSet::new(),
132            texts: HashSet::new(),
133            held: HashMap::new(),
134            waiting_on: HashMap::new(),
135            settled: HashSet::new(),
136            stats: CrawlStats::default(),
137        }
138    }
139}
140
141impl Crawl {
142    pub fn respect(mut self, respect: Respect) -> Self {
143        self.respect = Some(respect);
144        self
145    }
146
147    pub fn ignore_robots(mut self) -> Self {
148        self.respect = None;
149        self
150    }
151
152    pub fn scope(mut self, scope: Scope) -> Self {
153        self.frontier.scope = scope;
154        self
155    }
156
157    pub fn limit(mut self, pages: usize) -> Self {
158        self.max_pages = Some(pages);
159        self
160    }
161
162    pub fn host_limits(mut self, limits: HostLimits) -> Self {
163        self.limits = Some(limits);
164        self
165    }
166
167    #[must_use]
168    pub fn stats(&self) -> &CrawlStats {
169        &self.stats
170    }
171
172    pub async fn next(&mut self) -> Option<Result<CrawlPage>> {
173        std::future::poll_fn(|cx| Pin::new(&mut *self).poll_next(cx)).await
174    }
175
176    fn purpose(&self) -> Option<Purpose> {
177        match self.respect? {
178            Respect::Robots => None,
179            Respect::AiTraining => Some(Purpose::AiTrain),
180        }
181    }
182
183    fn start(&mut self) {
184        if std::mem::replace(&mut self.started, true) {
185            return;
186        }
187        let pace = match self.limits.take() {
188            Some(limits) => match limits.config_error() {
189                Some(message) => {
190                    let error = Error::new(Kind::Config).with_message(message);
191                    self.ready.push_back(Err(error));
192                    self.seeds.clear();
193                    return;
194                }
195                None => limits.armed(),
196            },
197            None => self.session.crawl_limits(|| inherited(&self.session)),
198        };
199        let pace = match self.respect {
200            Some(_) => pace,
201            None => pace.without_floors(),
202        };
203        self.session = self.session.with_crawl_pace(pace.clone());
204        self.page_retry = self
205            .session
206            .default_retry()
207            .clone()
208            .or(RetryPolicy::transient().max_retries(PAGE_RETRIES));
209        self.clock.stopped = Some(Box::pin(self.session.shut_down_signal()));
210        self.pace = pace;
211        for seed in std::mem::take(&mut self.seeds) {
212            match seed {
213                Ok(url) => {
214                    self.trusted.get_or_insert_with(|| url.clone());
215                    self.frontier.add(Queued::seed(url));
216                }
217                Err(err) => self.ready.push_back(Err(err)),
218            }
219        }
220    }
221
222    fn limit_met(&self) -> bool {
223        self.max_pages
224            .is_some_and(|max| self.stats.pages() as usize >= max)
225    }
226
227    fn budget_spent(&self) -> bool {
228        self.max_pages.is_some_and(|max| {
229            self.stats.pages() as usize + self.pages_in_flight + self.holding >= max
230        })
231    }
232
233    fn fetch(&mut self, origin: Origin, queued: Queued) {
234        self.pages_in_flight += 1;
235        *self.busy.entry(origin.clone()).or_default() += 1;
236        self.stats.request();
237        let session = self.session.clone();
238        let (retry, trusted) = (self.page_retry.clone(), self.trusted.clone());
239        let attempt = queued.attempts;
240        self.tasks.push(Box::pin(async move {
241            let fetched = fetch(session, queued.url.clone(), retry, attempt, trusted).await;
242            Event::Page {
243                origin,
244                queued,
245                fetched,
246            }
247        }));
248    }
249
250    fn release(&mut self, origin: &Origin) {
251        self.pages_in_flight -= 1;
252        if let Some(busy) = self.busy.get_mut(origin) {
253            *busy -= 1;
254            if *busy == 0 {
255                self.busy.remove(origin);
256            }
257        }
258    }
259
260    fn redirect(&mut self, queued: Queued, to: Url) {
261        if queued.hops >= MAX_REDIRECTS {
262            self.stats.skip(SkipReason::Redirects);
263            return;
264        }
265        let from = page_key(&queued.url);
266        let target = page_key(&to);
267        if target == from || queued.chain.contains(&target) {
268            self.frontier.redirected(&to);
269        }
270        let widen =
271            queued.seed && same_site(&queued.url, &to) && self.frontier.in_scope(&queued.url);
272        let seed = queued.seed && (queued.hops == 0 || widen || self.frontier.in_scope(&to));
273        let seen = self.frontier.has_seen(&to);
274        let mut chain = queued.chain;
275        chain.push(from);
276        let admitted = self.frontier.add_redirect(Queued {
277            url: to,
278            depth: queued.depth,
279            hops: queued.hops + 1,
280            seed,
281            widen,
282            attempts: 0,
283            chain,
284        });
285        match (admitted, seen) {
286            (false, true) => self.stats.skip(SkipReason::Duplicate),
287            (false, false) => self.stats.skip(SkipReason::Scope),
288            (true, _) => {}
289        }
290    }
291
292    fn handle(&mut self, event: Event) {
293        match event {
294            Event::Robots {
295                origin,
296                url,
297                robots,
298            } => self.robots_read(origin, url, *robots),
299            Event::Sitemap { file, read } => self.sitemap_read(file, read),
300            Event::Page {
301                origin,
302                queued,
303                fetched,
304            } => {
305                self.release(&origin);
306                let key = page_key(&queued.url);
307                match fetched {
308                    Err(err) => {
309                        self.stats.error();
310                        self.ready.push_back(Err(err));
311                    }
312                    Ok(Fetched::Redirect { to }) => {
313                        let target = page_key(&to);
314                        self.redirect(queued, to);
315                        return self.forward(key, target);
316                    }
317                    Ok(Fetched::Later(wait)) => return self.later(queued, wait),
318                    Ok(Fetched::Skip(reason)) => self.stats.skip(reason),
319                    Ok(Fetched::Page(loaded)) => self.page(queued.depth, *loaded),
320                }
321                self.settle(key);
322            }
323        }
324    }
325}
326
327impl std::fmt::Debug for Crawl {
328    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
329        f.debug_struct("Crawl")
330            .field("respect", &self.respect)
331            .field("in_flight", &self.tasks.len())
332            .field("stats", &self.stats)
333            .finish()
334    }
335}