Skip to main content

millipede_core/crawler/
builder.rs

1//! Construction of configured crawlers.
2
3use super::{Crawler, CrawlerKind, CrawlerShared, engine::EngineOptions};
4use crate::{
5    autoscale::{AutoscaleMode, AutoscaledPool, AutoscaledPoolOptions},
6    config::{ConfigError, Configuration},
7    handler::{FailedRequestHandler, RequestHandler},
8    link_extraction::CrawlPolicy,
9    storage::{KeyValueStore, RequestQueue, StorageClient, StorageError},
10};
11use std::{sync::Arc, time::Duration};
12
13/// An error produced while building a crawler.
14#[derive(Debug, thiserror::Error)]
15#[non_exhaustive]
16pub enum CrawlerBuildError {
17    /// No request handler was supplied.
18    #[error("a request handler is required")]
19    MissingRequestHandler,
20    /// Neither the builder nor its configuration supplied storage.
21    #[error(
22        "no storage client configured (set .storage_client(...) or Configuration::storage_client)"
23    )]
24    MissingStorage,
25    /// The configured concurrency is zero.
26    #[error("max_concurrency must be at least 1")]
27    ZeroMaxConcurrency,
28    /// The configured dynamic scheduler options are invalid.
29    #[error(
30        "dynamic scheduler options require valid concurrency bounds and a non-zero maybe_run_interval"
31    )]
32    InvalidConcurrencyBounds,
33    /// The configured result channel capacity is zero.
34    #[error("results_capacity must be at least 1")]
35    ZeroResultsCapacity,
36    /// Configuration resolution failed.
37    #[error("configuration: {0}")]
38    Config(#[from] ConfigError),
39    /// Storage initialization failed.
40    #[error("storage: {0}")]
41    Storage(#[from] StorageError),
42}
43
44/// Builds a [`Crawler`] around a crawler kind.
45#[must_use = "builders do nothing unless consumed by build"]
46pub struct CrawlerBuilder<K: CrawlerKind> {
47    kind: K,
48    handler: Option<Arc<dyn RequestHandler<K::Context>>>,
49    failed_handler: Option<Arc<dyn FailedRequestHandler>>,
50    autoscaled_pool: AutoscaledPoolOptions,
51    max_request_retries: u32,
52    max_session_rotations: u32,
53    request_handler_timeout: Duration,
54    internal_operation_timeout: Duration,
55    configuration: Option<Configuration>,
56    storage_client: Option<Arc<dyn StorageClient>>,
57    request_queue: Option<Arc<dyn RequestQueue>>,
58    results_capacity: usize,
59    retry_strategy: Option<Arc<dyn crate::retry_strategy::RetryStrategy>>,
60    crawl_policy: Option<Arc<CrawlPolicy>>,
61}
62
63impl<K: CrawlerKind> CrawlerBuilder<K> {
64    /// Creates a builder with engine defaults.
65    pub fn new(kind: K) -> Self {
66        Self {
67            kind,
68            handler: None,
69            failed_handler: None,
70            autoscaled_pool: AutoscaledPoolOptions {
71                fixed_concurrency: Some(10),
72                max_concurrency: 10,
73                ..Default::default()
74            },
75            max_request_retries: 3,
76            max_session_rotations: 10,
77            request_handler_timeout: Duration::from_secs(60),
78            internal_operation_timeout: Duration::from_secs(30),
79            configuration: None,
80            storage_client: None,
81            request_queue: None,
82            results_capacity: 1024,
83            retry_strategy: None,
84            crawl_policy: None,
85        }
86    }
87
88    /// Sets the request handler.
89    pub fn request_handler<H: RequestHandler<K::Context>>(mut self, handler: H) -> Self {
90        self.handler = Some(Arc::new(handler));
91        self
92    }
93
94    /// Sets the permanent-failure handler.
95    pub fn failed_request_handler<H: FailedRequestHandler>(mut self, handler: H) -> Self {
96        self.failed_handler = Some(Arc::new(handler));
97        self
98    }
99
100    /// Sets the maximum number of concurrent requests and pins concurrency to that value.
101    ///
102    /// Calling this after [`Self::autoscale_mode`] re-pins the crawler to fixed concurrency.
103    /// Calling [`Self::autoscale_mode`] after this method retains `count` as the dynamic ceiling.
104    pub fn max_concurrency(mut self, count: usize) -> Self {
105        self.autoscaled_pool.fixed_concurrency = Some(count);
106        self.autoscaled_pool.max_concurrency = count;
107        self
108    }
109    /// Sets the minimum dynamic concurrency.
110    pub fn min_concurrency(mut self, count: usize) -> Self {
111        self.autoscaled_pool.min_concurrency = count;
112        self
113    }
114    /// Sets the initial desired dynamic concurrency.
115    pub fn desired_concurrency(mut self, count: usize) -> Self {
116        self.autoscaled_pool.desired_concurrency = Some(count);
117        self
118    }
119    /// Selects dynamic autoscaling.
120    ///
121    /// Calling this after [`Self::max_concurrency`] keeps that value as the dynamic ceiling.
122    /// Calling [`Self::max_concurrency`] afterward re-pins concurrency to a fixed value.
123    pub fn autoscale_mode(mut self, mode: AutoscaleMode) -> Self {
124        self.autoscaled_pool.mode = mode;
125        self.autoscaled_pool.fixed_concurrency = None;
126        self
127    }
128    /// Sets the maximum number of task starts permitted per minute.
129    pub fn max_tasks_per_minute(mut self, count: u32) -> Self {
130        self.autoscaled_pool.max_tasks_per_minute = Some(count);
131        self
132    }
133    /// Sets the minimum delay between requests to the same domain.
134    pub fn same_domain_delay(mut self, delay: Duration) -> Self {
135        self.autoscaled_pool.same_domain_delay = delay;
136        self
137    }
138    /// Replaces all autoscaled-pool options for advanced use.
139    ///
140    /// The pool's `task_timeout` bounds request preparation, execution, and handler work after an
141    /// attempt starts. Its `maybe_run_interval` periodically rechecks the queue as a missed-wakeup
142    /// safety net and must be greater than zero.
143    pub fn autoscaled_pool_options(mut self, options: AutoscaledPoolOptions) -> Self {
144        self.autoscaled_pool = options;
145        self
146    }
147    /// Sets the maximum number of ordinary request retries.
148    pub fn max_request_retries(mut self, count: u32) -> Self {
149        self.max_request_retries = count;
150        self
151    }
152    /// Sets the maximum number of session rotations.
153    pub fn max_session_rotations(mut self, count: u32) -> Self {
154        self.max_session_rotations = count;
155        self
156    }
157    /// Sets the request-handler timeout.
158    pub fn request_handler_timeout(mut self, timeout: Duration) -> Self {
159        self.request_handler_timeout = timeout;
160        self
161    }
162    /// Sets the internal storage-operation timeout.
163    pub fn internal_operation_timeout(mut self, timeout: Duration) -> Self {
164        self.internal_operation_timeout = timeout;
165        self
166    }
167    /// Sets the resolved crawler configuration.
168    pub fn configuration(mut self, configuration: Configuration) -> Self {
169        self.configuration = Some(configuration);
170        self
171    }
172    /// Overrides the configuration's storage client.
173    pub fn storage_client(mut self, storage: Arc<dyn StorageClient>) -> Self {
174        self.storage_client = Some(storage);
175        self
176    }
177    /// Overrides the request queue opened from storage.
178    ///
179    /// This hook lets a [`crate::sitemap::RequestQueueWithSitemap`] tandem drive the crawler.
180    /// Callers resuming a pre-populated persistent queue should pair it with
181    /// [`crate::config::ConfigurationBuilder::purge_on_start`] set to `false`. Storage still
182    /// supplies the crawler's key-value store and other storage objects.
183    pub fn request_queue(mut self, queue: Arc<dyn RequestQueue>) -> Self {
184        self.request_queue = Some(queue);
185        self
186    }
187    /// Sets the terminal-result broadcast capacity.
188    pub fn results_capacity(mut self, capacity: usize) -> Self {
189        self.results_capacity = capacity;
190        self
191    }
192
193    /// Installs an attempt-level retry strategy.
194    pub fn retry_strategy<S: crate::retry_strategy::RetryStrategy>(mut self, strategy: S) -> Self {
195        self.retry_strategy = Some(Arc::new(strategy));
196        self
197    }
198
199    /// Sets the long-lived link admission and crawl-limit policy.
200    pub fn crawl_policy(mut self, policy: CrawlPolicy) -> Self {
201        self.crawl_policy = Some(Arc::new(policy));
202        self
203    }
204
205    /// Builds the crawler and opens its configured storage objects.
206    ///
207    /// The default [`Configuration`] purges all data managed by the selected storage client before
208    /// opening the queue and key-value store. Set `purge_on_start(false)` to retain existing data.
209    pub async fn build(self) -> Result<Crawler<K>, CrawlerBuildError> {
210        let handler = self
211            .handler
212            .ok_or(CrawlerBuildError::MissingRequestHandler)?;
213        if self.autoscaled_pool.fixed_concurrency == Some(0) {
214            return Err(CrawlerBuildError::ZeroMaxConcurrency);
215        }
216        if (self.autoscaled_pool.fixed_concurrency.is_none()
217            && (self.autoscaled_pool.min_concurrency < 1
218                || self.autoscaled_pool.max_concurrency < self.autoscaled_pool.min_concurrency))
219            || self.autoscaled_pool.maybe_run_interval.is_zero()
220        {
221            return Err(CrawlerBuildError::InvalidConcurrencyBounds);
222        }
223        if self.results_capacity == 0 {
224            return Err(CrawlerBuildError::ZeroResultsCapacity);
225        }
226        let config = match self.configuration {
227            Some(configuration) => configuration,
228            None => Configuration::builder().build()?,
229        };
230        let storage = self
231            .storage_client
232            .or_else(|| config.storage_client().cloned())
233            .ok_or(CrawlerBuildError::MissingStorage)?;
234        if config.purge_on_start() {
235            storage.purge().await?;
236        }
237        let queue = match self.request_queue {
238            Some(queue) => queue,
239            None => {
240                storage
241                    .open_request_queue(Some(config.default_request_queue_id()))
242                    .await?
243            }
244        };
245        let kvs: Option<Arc<dyn KeyValueStore>> = Some(
246            storage
247                .open_key_value_store(Some(config.default_key_value_store_id()))
248                .await?,
249        );
250        let task_timeout = self.autoscaled_pool.task_timeout;
251        let maybe_run_interval = self.autoscaled_pool.maybe_run_interval;
252        let opts = EngineOptions {
253            max_request_retries: self.max_request_retries,
254            max_session_rotations: self.max_session_rotations,
255            request_handler_timeout: self.request_handler_timeout,
256            internal_operation_timeout: self.internal_operation_timeout,
257            persist_state_interval: config.persist_state_interval(),
258            task_timeout,
259            maybe_run_interval,
260            retry_strategy: self.retry_strategy,
261            max_requests_per_crawl: self
262                .crawl_policy
263                .as_ref()
264                .and_then(|policy| policy.max_requests_per_crawl),
265        };
266        let pool = Arc::new(AutoscaledPool::new(self.autoscaled_pool));
267        let shared = Arc::new(CrawlerShared::new_with_policy(
268            queue,
269            config.events().clone(),
270            self.results_capacity,
271            self.internal_operation_timeout,
272            pool,
273            self.crawl_policy,
274        ));
275        Ok(Crawler {
276            kind: Arc::new(self.kind),
277            shared,
278            config: Arc::new(config),
279            handler,
280            failed_handler: self.failed_handler,
281            kvs,
282            storage: Some(storage),
283            opts,
284            started: std::sync::atomic::AtomicBool::new(false),
285        })
286    }
287}
288
289#[cfg(test)]
290mod tests {
291    use super::*;
292    use crate::{
293        crawler::{BasicContext, BasicKind},
294        memory_client as millipede_storage_memory,
295    };
296
297    fn aimd() -> AutoscaleMode {
298        AutoscaleMode::Aimd {
299            increase_after_successes: 1,
300            decrease_factor: 0.5,
301        }
302    }
303
304    #[tokio::test]
305    async fn invalid_dynamic_bounds_are_rejected_before_storage_resolution() {
306        let zero_min = CrawlerBuilder::new(BasicKind)
307            .request_handler(|_ctx: BasicContext| async { Ok(()) })
308            .autoscale_mode(aimd())
309            .min_concurrency(0)
310            .build()
311            .await;
312        assert!(matches!(
313            zero_min,
314            Err(CrawlerBuildError::InvalidConcurrencyBounds)
315        ));
316
317        let inverted = CrawlerBuilder::new(BasicKind)
318            .request_handler(|_ctx: BasicContext| async { Ok(()) })
319            .min_concurrency(2)
320            .max_concurrency(1)
321            .autoscale_mode(aimd())
322            .build()
323            .await;
324        assert!(matches!(
325            inverted,
326            Err(CrawlerBuildError::InvalidConcurrencyBounds)
327        ));
328    }
329
330    #[tokio::test]
331    async fn zero_maybe_run_interval_is_rejected_before_storage_resolution() {
332        let result = CrawlerBuilder::new(BasicKind)
333            .request_handler(|_ctx: BasicContext| async { Ok(()) })
334            .autoscaled_pool_options(AutoscaledPoolOptions {
335                maybe_run_interval: Duration::ZERO,
336                ..Default::default()
337            })
338            .build()
339            .await;
340
341        assert!(matches!(
342            result,
343            Err(CrawlerBuildError::InvalidConcurrencyBounds)
344        ));
345    }
346
347    #[tokio::test]
348    async fn autoscale_mode_after_max_concurrency_keeps_dynamic_ceiling() {
349        let crawler = CrawlerBuilder::new(BasicKind)
350            .request_handler(|_ctx: BasicContext| async { Ok(()) })
351            .storage_client(Arc::new(
352                millipede_storage_memory::MemoryStorageClient::new(),
353            ))
354            .max_concurrency(5)
355            .autoscale_mode(aimd())
356            .build()
357            .await
358            .unwrap();
359
360        let snapshot = crawler.autoscaler_snapshot();
361        assert!(!snapshot.is_fixed);
362        assert_eq!(snapshot.max_concurrency, 5);
363    }
364
365    #[tokio::test]
366    async fn max_concurrency_after_autoscale_mode_repins_to_fixed() {
367        let crawler = CrawlerBuilder::new(BasicKind)
368            .request_handler(|_ctx: BasicContext| async { Ok(()) })
369            .storage_client(Arc::new(
370                millipede_storage_memory::MemoryStorageClient::new(),
371            ))
372            .autoscale_mode(aimd())
373            .max_concurrency(5)
374            .build()
375            .await
376            .unwrap();
377
378        let snapshot = crawler.autoscaler_snapshot();
379        assert!(snapshot.is_fixed);
380        assert_eq!(snapshot.desired_concurrency, 5);
381    }
382}