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