1use crate::{
4 crawler::{CrawlerHandle, EnqueueAdmissionReservation},
5 errors::CrawlError,
6 link_extraction::{
7 CrawlPolicy, EnqueueStrategy, ExtractedLink, GlobPattern, LinkExtractor, TransformResult,
8 UrlPattern, compile_globs, strategy_allows,
9 },
10 request::{Request, RequestId, UserData},
11 storage::{AddOptions, ProcessedRequest},
12};
13use futures_util::future::BoxFuture;
14use globset::GlobSet;
15use regex::Regex;
16use std::{collections::HashSet, fmt, sync::Arc};
17use url::Url;
18
19type Transform = dyn for<'r> Fn(&'r mut Request) -> BoxFuture<'r, TransformResult> + Send + Sync;
20
21#[derive(Clone)]
39#[must_use = "enqueue pipelines do nothing unless send is awaited"]
40pub struct EnqueueLinker {
41 crawler: CrawlerHandle,
42 parent_url: Url,
43 parent_depth: u32,
44 extractor: Option<Arc<dyn LinkExtractor>>,
45}
46
47impl EnqueueLinker {
48 pub fn new(crawler: CrawlerHandle, parent: &Request) -> Self {
50 Self {
51 crawler,
52 parent_url: parent.url.clone(),
53 parent_depth: parent.crawl_depth,
54 extractor: None,
55 }
56 }
57
58 pub fn with_extractor(
60 crawler: CrawlerHandle,
61 parent: &Request,
62 extractor: Arc<dyn LinkExtractor>,
63 ) -> Self {
64 Self {
65 crawler,
66 parent_url: parent.url.clone(),
67 parent_depth: parent.crawl_depth,
68 extractor: Some(extractor),
69 }
70 }
71
72 pub fn options(&self) -> EnqueueLinksOptions<'_> {
74 EnqueueLinksOptions::new(self)
75 }
76
77 pub async fn urls(
83 &self,
84 urls: impl IntoIterator<Item = Url>,
85 ) -> Result<EnqueueResult, CrawlError> {
86 self.options().urls(urls).send().await
87 }
88
89 pub async fn all(&self) -> Result<EnqueueResult, CrawlError> {
91 self.options().strategy(EnqueueStrategy::All).send().await
92 }
93
94 pub async fn same_origin(&self) -> Result<EnqueueResult, CrawlError> {
96 self.options()
97 .strategy(EnqueueStrategy::SameOrigin)
98 .send()
99 .await
100 }
101
102 pub async fn same_hostname(&self) -> Result<EnqueueResult, CrawlError> {
104 self.options()
105 .strategy(EnqueueStrategy::SameHostname)
106 .send()
107 .await
108 }
109
110 pub async fn same_domain(&self) -> Result<EnqueueResult, CrawlError> {
112 self.options()
113 .strategy(EnqueueStrategy::SameDomain)
114 .send()
115 .await
116 }
117}
118
119impl fmt::Debug for EnqueueLinker {
120 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
121 formatter
122 .debug_struct("EnqueueLinker")
123 .field("crawler", &self.crawler)
124 .field("parent_url", &self.parent_url)
125 .field("parent_depth", &self.parent_depth)
126 .field(
127 "extractor",
128 &self.extractor.as_ref().map(|_| "<dyn LinkExtractor>"),
129 )
130 .finish()
131 }
132}
133
134#[non_exhaustive]
156#[must_use = "enqueue options do nothing unless passed through the enqueue pipeline"]
157pub struct EnqueueLinksOptions<'a> {
158 linker: &'a EnqueueLinker,
159 candidates: Vec<UrlCandidate>,
160 base_url: Option<Url>,
161 label: Option<String>,
162 user_data: Option<UserData>,
163 selector: Option<String>,
164 strategy: Option<EnqueueStrategy>,
165 globs: Vec<GlobPattern>,
166 regex: Vec<Regex>,
167 exclude: Vec<UrlPattern>,
168 transform: Option<Arc<Transform>>,
169 limit: Option<usize>,
170 forefront: bool,
171}
172
173enum UrlCandidate {
174 Absolute(Url),
175 Raw { url: String, base: Option<Url> },
176}
177
178impl<'a> EnqueueLinksOptions<'a> {
179 fn new(linker: &'a EnqueueLinker) -> Self {
180 Self {
181 linker,
182 candidates: Vec::new(),
183 base_url: None,
184 label: None,
185 user_data: None,
186 selector: None,
187 strategy: None,
188 globs: Vec::new(),
189 regex: Vec::new(),
190 exclude: Vec::new(),
191 transform: None,
192 limit: None,
193 forefront: false,
194 }
195 }
196
197 pub fn urls(mut self, urls: impl IntoIterator<Item = Url>) -> Self {
202 self.candidates
203 .extend(urls.into_iter().map(UrlCandidate::Absolute));
204 self
205 }
206 pub fn raw_urls<S: Into<String>>(mut self, urls: impl IntoIterator<Item = S>) -> Self {
208 self.candidates
209 .extend(urls.into_iter().map(|url| UrlCandidate::Raw {
210 url: url.into(),
211 base: None,
212 }));
213 self
214 }
215 pub fn base_url(mut self, base_url: Url) -> Self {
217 self.base_url = Some(base_url);
218 self
219 }
220 pub fn label(mut self, label: impl Into<String>) -> Self {
223 self.label = Some(label.into());
224 self
225 }
226 pub fn user_data(mut self, user_data: UserData) -> Self {
228 self.user_data = Some(user_data);
229 self
230 }
231 pub fn selector(mut self, selector: impl Into<String>) -> Self {
233 self.selector = Some(selector.into());
234 self
235 }
236 pub fn strategy(mut self, strategy: EnqueueStrategy) -> Self {
238 self.strategy = Some(strategy);
239 self
240 }
241 pub fn globs<G: Into<GlobPattern>>(mut self, globs: impl IntoIterator<Item = G>) -> Self {
243 self.globs.extend(globs.into_iter().map(Into::into));
244 self
245 }
246 pub fn regex(mut self, patterns: impl IntoIterator<Item = Regex>) -> Self {
248 self.regex.extend(patterns);
249 self
250 }
251 pub fn exclude<P: Into<UrlPattern>>(mut self, patterns: impl IntoIterator<Item = P>) -> Self {
253 self.exclude.extend(patterns.into_iter().map(Into::into));
254 self
255 }
256 pub fn transform<F>(mut self, transform: F) -> Self
258 where
259 F: for<'r> Fn(&'r mut Request) -> BoxFuture<'r, TransformResult> + Send + Sync + 'static,
260 {
261 self.transform = Some(Arc::new(transform));
262 self
263 }
264 pub fn limit(mut self, limit: usize) -> Self {
267 self.limit = Some(limit);
268 self
269 }
270 pub fn forefront(mut self, forefront: bool) -> Self {
272 self.forefront = forefront;
273 self
274 }
275
276 pub async fn send(mut self) -> Result<EnqueueResult, CrawlError> {
278 let should_extract = self.selector.is_some()
279 || (self.candidates.is_empty() && self.linker.extractor.is_some());
280 if should_extract {
281 let extractor = self.linker.extractor.as_ref().ok_or_else(|| {
282 CrawlError::non_retryable(anyhow::anyhow!(
283 "selector-based enqueue requires an HTML or browser context; use urls()/raw_urls() on HTTP contexts"
284 ))
285 })?;
286 let extracted = extractor.extract(self.selector.as_deref()).await?;
287 self.candidates.extend(
288 extracted
289 .into_iter()
290 .map(|ExtractedLink { url, base }| UrlCandidate::Raw { url, base }),
291 );
292 }
293
294 let policy = self
295 .linker
296 .crawler
297 .crawl_policy()
298 .unwrap_or_else(|| Arc::new(CrawlPolicy::default()));
299 let strategy = self.strategy.unwrap_or(policy.strategy);
300 let compiled_globs = compile_include_patterns(self.globs)?;
301 let compiled_excludes = compile_patterns(self.exclude)?;
302 let has_globs = !compiled_globs.is_empty();
303 let has_regex = !self.regex.is_empty();
304 let mut skipped = Vec::new();
305 let mut candidates = Vec::with_capacity(self.candidates.len());
306 for candidate in self.candidates {
307 let (resolved, apply_strategy) = match candidate {
308 UrlCandidate::Absolute(url) => (Ok(url), false),
309 UrlCandidate::Raw { url, base } => {
310 let base = base
311 .as_ref()
312 .or(self.base_url.as_ref())
313 .unwrap_or(&self.linker.parent_url);
314 (resolve_raw_url(base, &url).map_err(|_| url), true)
315 }
316 };
317 let url = match resolved {
318 Ok(candidate) => candidate,
319 Err(raw) => {
320 report_skip(&policy, &mut skipped, raw, SkipReason::InvalidUrl);
321 continue;
322 }
323 };
324 let url_text = url.to_string();
325 if apply_strategy && !strategy_allows(strategy, &self.linker.parent_url, &url) {
326 report_skip(
327 &policy,
328 &mut skipped,
329 url_text,
330 SkipReason::StrategyExcluded,
331 );
332 continue;
333 }
334 if let Some(reason) = first_exclusion(&compiled_excludes, &url_text) {
335 report_skip(&policy, &mut skipped, url_text, reason);
336 continue;
337 }
338 let include_override = compiled_globs
339 .iter()
340 .find(|pattern| pattern.matches(&url_text))
341 .map(|pattern| pattern.overrides());
342 if has_globs || has_regex {
343 let regex_matches = self.regex.iter().any(|pattern| pattern.is_match(&url_text));
344 if include_override.is_none() && !regex_matches {
345 let reason = if has_globs {
346 SkipReason::GlobExcluded
347 } else {
348 SkipReason::RegexExcluded
349 };
350 report_skip(&policy, &mut skipped, url_text, reason);
351 continue;
352 }
353 }
354 candidates.push((url, include_override.unwrap_or_default()));
355 }
356
357 let mut seen = HashSet::new();
358 candidates.retain(|(url, _)| seen.insert(url.to_string()));
359 if let Some(limit) = self.limit {
360 candidates.truncate(limit);
361 }
362
363 let mut requests = Vec::with_capacity(candidates.len());
364 let mut admission_reservations = Vec::new();
365 let child_depth = self.linker.parent_depth.saturating_add(1);
366 for (url, overrides) in candidates {
367 let url_text = url.to_string();
368 let mut builder = Request::builder().url(url).crawl_depth(child_depth);
369 if let Some(label) = overrides.label.as_deref().or(self.label.as_deref()) {
370 builder = builder.label(label.to_owned());
371 }
372 if let Some(user_data) = overrides.user_data.as_ref().or(self.user_data.as_ref()) {
373 builder = builder.user_data(user_data.clone());
374 }
375 if let Some(method) = overrides.method {
376 builder = builder.method(method);
377 }
378 if let Some(headers) = overrides.headers {
379 builder = builder.headers(headers);
380 }
381 let mut request = match builder.build() {
382 Ok(request) => request,
383 Err(_) => {
384 report_skip(&policy, &mut skipped, url_text, SkipReason::InvalidUrl);
385 continue;
386 }
387 };
388 if let Some(limit) = policy.max_crawl_depth {
389 if child_depth > limit {
390 report_skip(
391 &policy,
392 &mut skipped,
393 url_text,
394 SkipReason::MaxDepthExceeded {
395 depth: child_depth,
396 limit,
397 },
398 );
399 continue;
400 }
401 }
402 let admission_reservation = if let Some(limit) = policy.max_requests_per_crawl {
403 match reserve_request_slot(&self.linker.crawler, limit).await? {
404 Admission::Reserved(reservation) => Some(reservation),
405 Admission::Untracked => None,
406 Admission::Rejected => {
407 report_skip(
408 &policy,
409 &mut skipped,
410 url_text,
411 SkipReason::MaxRequestsReached { limit },
412 );
413 continue;
414 }
415 }
416 } else {
417 None
418 };
419 let pre_transform_unique_key = request.unique_key.clone();
420 if let Some(transform) = &self.transform {
421 match transform(&mut request).await {
422 TransformResult::Enqueue => {}
423 TransformResult::Skip { reason } => {
424 report_skip(
425 &policy,
426 &mut skipped,
427 url_text,
428 SkipReason::TransformRejected { reason },
429 );
430 continue;
431 }
432 }
433 }
434 if request.unique_key == pre_transform_unique_key {
435 request.unique_key = Request::compute_unique_key(
436 &request.url,
437 &request.method,
438 request.body.as_ref(),
439 );
440 request.id = RequestId::from_unique_key(&request.unique_key);
441 }
442 requests.push(request);
443 admission_reservations.push(admission_reservation);
444 }
445
446 let request_urls: Vec<_> = requests
447 .iter()
448 .map(|request| request.url.to_string())
449 .collect();
450 let batch = self
451 .linker
452 .crawler
453 .add_requests_with_options(
454 requests,
455 AddOptions {
456 forefront: self.forefront,
457 },
458 )
459 .await?;
460 let batch = batch.wait().await?;
461 let mut added = Vec::new();
462 let mut admission_reservations = admission_reservations.into_iter();
463 for (index, request) in batch.processed.into_iter().enumerate() {
464 let admission_reservation = admission_reservations.next().flatten();
465 if request.was_already_present || request.was_already_handled {
466 let url = request_urls
467 .get(index)
468 .cloned()
469 .unwrap_or_else(|| request.unique_key.clone());
470 report_skip(&policy, &mut skipped, url, SkipReason::DuplicateUniqueKey);
471 } else {
472 if let Some(reservation) = admission_reservation {
473 reservation.commit();
474 }
475 added.push(request);
476 }
477 }
478 Ok(EnqueueResult { added, skipped })
479 }
480}
481
482enum Admission {
483 Reserved(EnqueueAdmissionReservation),
484 Untracked,
485 Rejected,
486}
487
488async fn reserve_request_slot(
489 crawler: &CrawlerHandle,
490 limit: u64,
491) -> Result<Admission, CrawlError> {
492 let _admission = crawler.lock_enqueue_admission().await?;
493 let queue = match crawler.request_queue() {
494 Some(queue) => queue,
495 None => return Ok(Admission::Untracked),
496 };
497 let Some(queue_count) = request_count(queue.as_ref()).await else {
498 return Ok(Admission::Untracked);
499 };
500 let admitted = crawler.synchronize_enqueue_admissions(queue_count)?;
501 if admitted >= limit {
502 return Ok(Admission::Rejected);
503 }
504 Ok(Admission::Reserved(crawler.reserve_enqueue_admission()?))
505}
506
507async fn request_count(queue: &dyn crate::storage::RequestQueue) -> Option<u64> {
508 let handled = match queue.handled_count().await {
509 Ok(count) => count,
510 Err(error) => {
511 tracing::debug!(%error, "could not read handled request count; skipping max-request admission check");
512 return None;
513 }
514 };
515 let pending = match queue.pending_count().await {
516 Ok(count) => count,
517 Err(error) => {
518 tracing::debug!(%error, "could not read pending request count; skipping max-request admission check");
519 return None;
520 }
521 };
522 Some(handled.saturating_add(pending))
523}
524
525#[derive(Default)]
526struct PatternOverrides {
527 label: Option<String>,
528 user_data: Option<UserData>,
529 method: Option<crate::request::Method>,
530 headers: Option<crate::request::HeaderMap>,
531}
532
533enum CompiledPattern {
534 Glob(GlobSet),
535 Regex(Regex),
536}
537
538impl CompiledPattern {
539 fn matches(&self, url: &str) -> bool {
540 match self {
541 Self::Glob(pattern) => pattern.is_match(url),
542 Self::Regex(pattern) => pattern.is_match(url),
543 }
544 }
545}
546
547struct CompiledInclude {
548 pattern: CompiledPattern,
549 source: GlobPattern,
550}
551
552impl CompiledInclude {
553 fn matches(&self, url: &str) -> bool {
554 self.pattern.matches(url)
555 }
556
557 fn overrides(&self) -> PatternOverrides {
558 PatternOverrides {
559 label: self.source.label().map(str::to_owned),
560 user_data: self.source.user_data().cloned(),
561 method: self.source.method().cloned(),
562 headers: self.source.headers().cloned(),
563 }
564 }
565}
566
567fn compile_include_patterns(
568 patterns: Vec<GlobPattern>,
569) -> Result<Vec<CompiledInclude>, CrawlError> {
570 patterns
571 .into_iter()
572 .map(|source| {
573 let pattern = compile_pattern(source.pattern().clone())?;
574 Ok(CompiledInclude { pattern, source })
575 })
576 .collect()
577}
578
579fn compile_patterns(patterns: Vec<UrlPattern>) -> Result<Vec<CompiledPattern>, CrawlError> {
580 patterns.into_iter().map(compile_pattern).collect()
581}
582
583fn compile_pattern(pattern: UrlPattern) -> Result<CompiledPattern, CrawlError> {
584 match pattern {
585 UrlPattern::Glob(pattern) => compile_globs(std::slice::from_ref(&pattern))
586 .map(CompiledPattern::Glob)
587 .map_err(CrawlError::non_retryable),
588 UrlPattern::Regex(pattern) => Ok(CompiledPattern::Regex(pattern)),
589 }
590}
591
592fn first_exclusion(patterns: &[CompiledPattern], url: &str) -> Option<SkipReason> {
593 patterns.iter().find_map(|pattern| {
594 if !pattern.matches(url) {
595 return None;
596 }
597 Some(match pattern {
598 CompiledPattern::Glob(_) => SkipReason::GlobExcluded,
599 CompiledPattern::Regex(_) => SkipReason::RegexExcluded,
600 })
601 })
602}
603
604fn report_skip(
605 policy: &CrawlPolicy,
606 skipped: &mut Vec<SkippedUrl>,
607 url: String,
608 reason: SkipReason,
609) {
610 if let Some(handler) = &policy.on_skipped {
611 handler.on_skip(&url, &reason);
612 }
613 skipped.push(SkippedUrl { url, reason });
614}
615
616fn resolve_raw_url(base: &Url, raw: &str) -> Result<Url, url::ParseError> {
617 if let Ok(absolute) = Url::parse(raw) {
618 return Ok(absolute);
619 }
620
621 let first_path_segment = raw
625 .split_once(['?', '#'])
626 .map_or(raw, |(path, _)| path)
627 .split('/')
628 .next()
629 .unwrap_or_default();
630 if first_path_segment.contains(':') {
631 return Err(url::ParseError::RelativeUrlWithoutBase);
632 }
633
634 base.join(raw)
635}
636
637#[derive(Debug)]
649#[non_exhaustive]
650#[must_use = "enqueue results report both accepted and skipped URLs"]
651pub struct EnqueueResult {
652 pub added: Vec<ProcessedRequest>,
654 pub skipped: Vec<SkippedUrl>,
656}
657
658impl EnqueueResult {
659 pub fn added_count(&self) -> usize {
661 self.added.len()
662 }
663 pub fn skipped_count(&self) -> usize {
665 self.skipped.len()
666 }
667}
668
669#[derive(Debug, Clone)]
685pub struct SkippedUrl {
686 pub url: String,
688 pub reason: SkipReason,
690}
691
692#[derive(Debug, Clone, PartialEq, Eq)]
705#[non_exhaustive]
706pub enum SkipReason {
707 MaxDepthExceeded {
709 depth: u32,
711 limit: u32,
713 },
714 MaxRequestsReached {
716 limit: u64,
718 },
719 StrategyExcluded,
721 GlobExcluded,
723 RegexExcluded,
725 TransformRejected {
727 reason: String,
729 },
730 DuplicateUniqueKey,
732 InvalidUrl,
734}
735
736impl fmt::Display for SkipReason {
737 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
738 match self {
739 Self::MaxDepthExceeded { depth, limit } => {
740 write!(formatter, "crawl depth {depth} exceeds limit {limit}")
741 }
742 Self::MaxRequestsReached { limit } => {
743 write!(formatter, "maximum request count {limit} reached")
744 }
745 Self::StrategyExcluded => formatter.write_str("excluded by enqueue strategy"),
746 Self::GlobExcluded => formatter.write_str("excluded by glob patterns"),
747 Self::RegexExcluded => formatter.write_str("excluded by regex patterns"),
748 Self::TransformRejected { reason } => {
749 write!(formatter, "rejected by transform: {reason}")
750 }
751 Self::DuplicateUniqueKey => formatter.write_str("duplicate unique key"),
752 Self::InvalidUrl => formatter.write_str("invalid URL"),
753 }
754 }
755}