leyline/core/crawl/
clock.rs

1use std::future::Future;
2use std::pin::Pin;
3use std::task::{Context, Poll};
4use std::time::Duration;
5
6use futures_util::stream::{FuturesUnordered, Stream, StreamExt};
7use tokio::time::{Instant, Sleep};
8use tokio_util::sync::WaitForCancellationFutureOwned;
9
10use super::frontier::{Queued, page_key};
11use super::{Crawl, CrawlPage, SkipReason};
12use crate::core::config::Origin;
13use crate::core::deadline::Deadline;
14use crate::core::error::Result;
15
16const ROOM_RECHECK: Duration = Duration::from_millis(250);
17const ALARM_FLOOR: Duration = Duration::from_millis(10);
18
19#[derive(Default)]
20pub(super) struct Clock {
21    pub(super) stopped: Option<Pin<Box<WaitForCancellationFutureOwned>>>,
22    alarm: Option<Pin<Box<Sleep>>>,
23    alarm_at: Option<Instant>,
24    delayed: Vec<(Instant, Queued)>,
25}
26
27impl Crawl {
28    pub(super) fn later(&mut self, mut queued: Queued, wait: Duration) {
29        let Some(at) = Instant::now().checked_add(wait) else {
30            self.stats.skip(SkipReason::Status);
31            return self.settle(page_key(&queued.url));
32        };
33        queued.attempts += 1;
34        self.clock.delayed.push((at, queued));
35        self.set_alarm(at);
36    }
37
38    pub(super) fn promote_due(&mut self) {
39        let now = Instant::now();
40        let (due, waiting): (Vec<_>, Vec<_>) = std::mem::take(&mut self.clock.delayed)
41            .into_iter()
42            .partition(|(at, _)| *at <= now);
43        self.clock.delayed = waiting;
44        for (_, queued) in due {
45            self.frontier.push_back(queued);
46        }
47    }
48
49    pub(super) fn set_alarm(&mut self, at: Instant) {
50        let at = at.max(Instant::now() + ALARM_FLOOR);
51        if self.clock.alarm_at.is_none_or(|current| at < current) {
52            self.clock.alarm = Some(Deadline::alarm(at));
53            self.clock.alarm_at = Some(at);
54        }
55    }
56
57    pub(super) fn plan_alarm(&mut self) {
58        let now = Instant::now();
59        let sitemap_hosts: Vec<Origin> = self
60            .sitemap_queue
61            .iter()
62            .filter(|file| file.not_before.is_none_or(|at| at <= now))
63            .filter_map(|file| Origin::of(&file.url))
64            .collect();
65        let blocked = self
66            .frontier
67            .origins()
68            .chain(sitemap_hosts.iter())
69            .map(|origin| self.host_wait(origin))
70            .chain(self.clock.delayed.iter().map(|(at, _)| Some(*at)))
71            .chain(
72                self.sitemap_queue
73                    .iter()
74                    .map(|file| file.not_before.filter(|at| *at > now)),
75            )
76            .flatten()
77            .min();
78        if let Some(at) = blocked.filter(|_| !self.budget_spent()) {
79            self.set_alarm(at);
80        }
81    }
82
83    fn host_wait(&self, origin: &Origin) -> Option<Instant> {
84        let own = self.session.host_limits();
85        match self.pace.ready_at(origin).max(own.ready_at(origin)) {
86            Some(at) => Some(at),
87            None if !(self.pace.ready(origin) && own.ready(origin)) => {
88                Some(Instant::now() + ROOM_RECHECK)
89            }
90            None => None,
91        }
92    }
93
94    fn waiting_work(&self) -> bool {
95        !self.budget_spent()
96            && (!self.frontier.is_empty()
97                || !self.clock.delayed.is_empty()
98                || !self.sitemap_queue.is_empty())
99    }
100
101    fn stopped(&mut self, cx: &mut Context<'_>) -> bool {
102        let signalled = self
103            .clock
104            .stopped
105            .as_mut()
106            .is_some_and(|stopped| stopped.as_mut().poll(cx).is_ready());
107        signalled || self.session.is_shut_down()
108    }
109
110    fn halt(&mut self) {
111        self.frontier.clear();
112        self.sitemap_queue.clear();
113        self.clock.delayed.clear();
114        self.clock.alarm = None;
115        self.clock.alarm_at = None;
116        self.tasks = FuturesUnordered::new();
117    }
118}
119
120impl Stream for Crawl {
121    type Item = Result<CrawlPage>;
122
123    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
124        let this = self.get_mut();
125        this.start();
126        loop {
127            if this.stopped(cx) {
128                this.halt();
129                return Poll::Ready(this.ready.pop_front());
130            }
131            this.schedule();
132            if let Some(item) = this.ready.pop_front() {
133                return Poll::Ready(Some(item));
134            }
135            if this.limit_met() && this.pages_in_flight == 0 {
136                this.halt();
137                return Poll::Ready(None);
138            }
139            match this.tasks.poll_next_unpin(cx) {
140                Poll::Ready(Some(event)) => {
141                    this.handle(event);
142                    continue;
143                }
144                Poll::Ready(None) if !this.waiting_work() => {
145                    if this.release_held() {
146                        continue;
147                    }
148                    return Poll::Ready(None);
149                }
150                Poll::Ready(None) | Poll::Pending => {}
151            }
152            let fired = this
153                .clock
154                .alarm
155                .as_mut()
156                .is_some_and(|alarm| alarm.as_mut().poll(cx).is_ready());
157            if !fired {
158                return Poll::Pending;
159            }
160            this.clock.alarm = None;
161            this.clock.alarm_at = None;
162        }
163    }
164}
leyline/core/crawl/clock.rs

leyline/core/crawl/
clock.rs
