Skip to main content

spire_core/frontier/queue/
mod.rs

1//! Queues: [`Queue`], the handle handlers add links through, and [`Link`].
2
3use std::fmt;
4use std::sync::{Arc, Mutex};
5
6use ::http::Uri;
7use url::Url;
8
9mod link;
10
11pub use link::Link;
12use link::Target;
13
14use crate::crawler::Depth;
15use crate::frontier::dynamic::DynFrontier;
16use crate::frontier::seed::Seed;
17use crate::frontier::{Admission, Frontier};
18use crate::http::{Request, RequestExt};
19use crate::sync::lock;
20use crate::{Error, Result};
21
22/// A cloneable handle for adding seeds to a crawl's [`Frontier`].
23///
24/// Handlers get a queue bound to the request they handle, so
25/// [`Queue::follow`] can resolve relative links and derive the depth of the
26/// requests it queues.
27#[derive(Clone)]
28pub struct Queue {
29    frontier: Arc<dyn DynFrontier>,
30    parent: Option<Arc<Parent>>,
31}
32
33struct Parent {
34    depth: Depth,
35    /// The URI relative links resolve against: the request URI, replaced by
36    /// the final URI once the response is fetched.
37    base: Mutex<Uri>,
38}
39
40impl Queue {
41    /// Creates a queue over `frontier`.
42    #[must_use]
43    pub fn new(frontier: impl Frontier) -> Self {
44        Self::from_dyn(Arc::new(frontier))
45    }
46
47    pub(crate) fn from_dyn(frontier: Arc<dyn DynFrontier>) -> Self {
48        Self {
49            frontier,
50            parent: None,
51        }
52    }
53
54    /// Returns a queue whose [`follow`](Self::follow) treats links as found
55    /// on `request`.
56    #[must_use]
57    pub fn child_of<B>(&self, request: &Request<B>) -> Self {
58        let parent = Parent {
59            depth: request.depth(),
60            base: Mutex::new(request.uri().clone()),
61        };
62
63        Self {
64            frontier: self.frontier.clone(),
65            parent: Some(Arc::new(parent)),
66        }
67    }
68
69    /// Resolves relative links against `uri` from now on, e.g. the URI a
70    /// response was served from after redirects. Clones share the base.
71    pub fn rebase(&self, uri: Uri) {
72        if let Some(parent) = &self.parent {
73            *lock(&parent.base) = uri;
74        }
75    }
76
77    /// Adds `seed` as is.
78    ///
79    /// # Errors
80    ///
81    /// Returns an error if `seed` cannot be converted into a [`Seed`] or the
82    /// frontier fails.
83    pub async fn push<S>(&self, seed: S) -> Result<Admission>
84    where
85        S: TryInto<Seed>,
86        S::Error: Into<Error>,
87    {
88        let seed = seed.try_into().map_err(Into::into)?;
89        self.frontier.push(seed).await
90    }
91
92    /// Adds a link found on the current request, one level deeper.
93    ///
94    /// The followed request has the link's tag and priority, or else the
95    /// seed's. Tags are not inherited: an untagged link goes to the router's
96    /// path routes or fallback, so pagination tags its links explicitly.
97    ///
98    /// The tag is not part of the [`Fingerprint`](crate::frontier::seed::Fingerprint):
99    /// with a deduplicating frontier, a URL already followed under one tag is
100    /// a [`Duplicate`](Admission::Duplicate) under any other.
101    ///
102    /// Relative `href`s resolve against the page; fragments are dropped.
103    /// Links that are not `http(s)` URLs, such as `mailto:` or malformed
104    /// ones, are skipped with [`Admission::Unsupported`]. Outside a handler,
105    /// links must be absolute and keep their depth.
106    ///
107    /// # Errors
108    ///
109    /// Returns an error if the frontier fails.
110    pub async fn follow(&self, link: impl Into<Link>) -> Result<Admission> {
111        let link = link.into();
112        let mut seed = match link.target {
113            Target::Seed(seed) => *seed,
114            Target::Href(href) => match self.resolve(&href) {
115                Some(seed) => seed,
116                None => return Ok(Admission::Unsupported),
117            },
118        };
119
120        if let Some(tag) = link.tag {
121            seed = seed.with_tag(tag);
122        }
123        if let Some(priority) = link.priority {
124            seed = seed.with_priority(priority);
125        }
126        if let Some(parent) = &self.parent {
127            seed = seed.with_depth(parent.depth.child());
128        }
129        self.frontier.push(seed).await
130    }
131
132    /// Follows every link, returning how many were queued.
133    ///
134    /// # Errors
135    ///
136    /// Returns an error if the frontier fails.
137    pub async fn follow_all<I>(&self, links: I) -> Result<usize>
138    where
139        I: IntoIterator<Item: Into<Link>>,
140    {
141        let mut queued = 0;
142        for link in links {
143            if self.follow(link).await?.is_queued() {
144                queued += 1;
145            }
146        }
147        Ok(queued)
148    }
149
150    /// Resolves `href` into a `GET` seed, or `None` if it is not a crawlable
151    /// `http(s)` URL.
152    fn resolve(&self, href: &str) -> Option<Seed> {
153        let href = href.trim();
154        let mut url = match &self.parent {
155            Some(parent) => {
156                let base = lock(&parent.base).to_string();
157                Url::parse(&base).ok()?.join(href).ok()?
158            }
159            None => Url::parse(href).ok()?,
160        };
161
162        if !matches!(url.scheme(), "http" | "https") {
163            return None;
164        }
165        url.set_fragment(None);
166        Seed::get(url.as_str()).ok()
167    }
168
169    /// Returns the number of queued seeds, if known.
170    #[must_use]
171    pub fn size_hint(&self) -> Option<usize> {
172        self.frontier.size_hint()
173    }
174}
175
176impl fmt::Debug for Queue {
177    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
178        let base = self
179            .parent
180            .as_ref()
181            .map(|parent| lock(&parent.base).clone());
182        f.debug_struct("Queue")
183            .field("base", &base)
184            .field("size_hint", &self.size_hint())
185            .finish_non_exhaustive()
186    }
187}
188
189#[cfg(test)]
190mod tests {
191    use super::*;
192    use crate::crawler::Tag;
193    use crate::frontier::{FrontierExt, InMemFrontier, NextSeed};
194
195    async fn drain(frontier: &InMemFrontier) -> Vec<(String, Option<String>, u32)> {
196        let mut popped = Vec::new();
197        while let NextSeed::Ready(seed) = Frontier::pop(frontier).await.unwrap() {
198            let tag = seed.tag().map(Tag::to_string);
199            popped.push((seed.uri().to_string(), tag, seed.depth().get()));
200        }
201        popped
202    }
203
204    fn page(uri: &str) -> Request<()> {
205        let seed = Seed::get(uri)
206            .unwrap()
207            .with_tag("list")
208            .with_depth(Depth::new(2));
209        seed.into_request().map(drop)
210    }
211
212    #[tokio::test]
213    async fn follow_derives_depth_not_tag() {
214        let frontier = Arc::new(InMemFrontier::fifo());
215        let queue = Queue::from_dyn(frontier.clone());
216        let child = queue.child_of(&page("https://example.com/"));
217
218        child.follow("https://example.com/untagged").await.unwrap();
219        child
220            .follow(Link::from("/own").with_tag("detail"))
221            .await
222            .unwrap();
223        let seed = Seed::get("https://example.com/seed")
224            .unwrap()
225            .with_tag("seed");
226        child.follow(seed).await.unwrap();
227        queue.follow("https://example.com/root").await.unwrap();
228
229        let tag = |tag: &str| Some(tag.to_owned());
230        assert_eq!(
231            drain(&frontier).await,
232            [
233                ("https://example.com/untagged".to_owned(), None, 3),
234                ("https://example.com/own".to_owned(), tag("detail"), 3),
235                ("https://example.com/seed".to_owned(), tag("seed"), 3),
236                ("https://example.com/root".to_owned(), None, 0),
237            ]
238        );
239    }
240
241    #[tokio::test]
242    async fn follow_sets_the_link_priority() {
243        let frontier = Arc::new(InMemFrontier::priority());
244        let child = Queue::from_dyn(frontier.clone()).child_of(&page("https://example.com/"));
245
246        child.follow("/plain").await.unwrap();
247        child
248            .follow(Link::from("/scored").with_priority(7))
249            .await
250            .unwrap();
251        let seed = Seed::get("https://example.com/seed")
252            .unwrap()
253            .with_priority(3);
254        child.follow(seed).await.unwrap();
255
256        let mut priorities: Vec<_> = Frontier::snapshot(&*frontier)
257            .unwrap()
258            .iter()
259            .map(|seed| (seed.uri().path().to_owned(), seed.priority()))
260            .collect();
261        priorities.sort_by_key(|(_, priority)| -priority);
262        assert_eq!(
263            priorities,
264            [
265                ("/scored".to_owned(), 7),
266                ("/seed".to_owned(), 3),
267                ("/plain".to_owned(), 0),
268            ]
269        );
270    }
271
272    #[tokio::test]
273    async fn follow_resolves_relative_links() {
274        let frontier = Arc::new(InMemFrontier::fifo());
275        let child = Queue::from_dyn(frontier.clone()).child_of(&page("https://example.com/a/b"));
276
277        for href in ["c", "../d", "/e?x=1", "//other.org/f", "g#section", " h "] {
278            assert!(child.follow(href).await.unwrap().is_queued(), "{href}");
279        }
280        let resolved: Vec<_> = drain(&frontier).await.into_iter().map(|x| x.0).collect();
281        assert_eq!(
282            resolved,
283            [
284                "https://example.com/a/c",
285                "https://example.com/d",
286                "https://example.com/e?x=1",
287                "https://other.org/f",
288                "https://example.com/a/g",
289                "https://example.com/a/h",
290            ]
291        );
292    }
293
294    #[tokio::test]
295    async fn rebase_follows_redirects() {
296        let frontier = Arc::new(InMemFrontier::fifo());
297        let child = Queue::from_dyn(frontier.clone()).child_of(&page("https://example.com/old/"));
298        child
299            .clone()
300            .rebase(Uri::from_static("https://example.org/new/"));
301
302        child.follow("page").await.unwrap();
303        assert_eq!(drain(&frontier).await[0].0, "https://example.org/new/page");
304    }
305
306    #[tokio::test]
307    async fn skips_unsupported_links() {
308        let queue = Queue::new(InMemFrontier::fifo());
309        let child = queue.child_of(&page("https://example.com/"));
310        for href in ["mailto:a@example.com", "javascript:void(0)", "http://[::1"] {
311            let admission = child.follow(href).await.unwrap();
312            assert_eq!(admission, Admission::Unsupported, "{href}");
313        }
314        assert_eq!(
315            queue.follow("/relative").await.unwrap(),
316            Admission::Unsupported
317        );
318
319        let deduped = Queue::new(InMemFrontier::fifo().dedup());
320        let child = deduped.child_of(&page("https://example.com/"));
321        assert_eq!(
322            child
323                .follow_all(["/a", "/a", "tel:123", "/b"])
324                .await
325                .unwrap(),
326            2
327        );
328    }
329
330    #[tokio::test]
331    async fn push_reports_admission() {
332        let queue = Queue::new(InMemFrontier::fifo().dedup());
333        assert!(
334            queue
335                .push("https://example.com/")
336                .await
337                .unwrap()
338                .is_queued()
339        );
340        assert_eq!(
341            queue.push("https://example.com/").await.unwrap(),
342            Admission::Duplicate
343        );
344        assert!(queue.push("relative").await.is_err());
345        assert_eq!(queue.size_hint(), Some(1));
346    }
347}