spire_core/frontier/queue/
mod.rs1use 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#[derive(Clone)]
28pub struct Queue {
29 frontier: Arc<dyn DynFrontier>,
30 parent: Option<Arc<Parent>>,
31}
32
33struct Parent {
34 depth: Depth,
35 base: Mutex<Uri>,
38}
39
40impl Queue {
41 #[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 #[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 pub fn rebase(&self, uri: Uri) {
72 if let Some(parent) = &self.parent {
73 *lock(&parent.base) = uri;
74 }
75 }
76
77 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 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 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 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 #[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}