1use 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#[derive(Debug, thiserror::Error)]
15#[non_exhaustive]
16pub enum CrawlerBuildError {
17 #[error("a request handler is required")]
19 MissingRequestHandler,
20 #[error(
22 "no storage client configured (set .storage_client(...) or Configuration::storage_client)"
23 )]
24 MissingStorage,
25 #[error("max_concurrency must be at least 1")]
27 ZeroMaxConcurrency,
28 #[error(
30 "dynamic scheduler options require valid concurrency bounds and a non-zero maybe_run_interval"
31 )]
32 InvalidConcurrencyBounds,
33 #[error("results_capacity must be at least 1")]
35 ZeroResultsCapacity,
36 #[error("configuration: {0}")]
38 Config(#[from] ConfigError),
39 #[error("storage: {0}")]
41 Storage(#[from] StorageError),
42}
43
44#[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 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 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 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 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 pub fn min_concurrency(mut self, count: usize) -> Self {
111 self.autoscaled_pool.min_concurrency = count;
112 self
113 }
114 pub fn desired_concurrency(mut self, count: usize) -> Self {
116 self.autoscaled_pool.desired_concurrency = Some(count);
117 self
118 }
119 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 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 pub fn same_domain_delay(mut self, delay: Duration) -> Self {
135 self.autoscaled_pool.same_domain_delay = delay;
136 self
137 }
138 pub fn autoscaled_pool_options(mut self, options: AutoscaledPoolOptions) -> Self {
144 self.autoscaled_pool = options;
145 self
146 }
147 pub fn max_request_retries(mut self, count: u32) -> Self {
149 self.max_request_retries = count;
150 self
151 }
152 pub fn max_session_rotations(mut self, count: u32) -> Self {
154 self.max_session_rotations = count;
155 self
156 }
157 pub fn request_handler_timeout(mut self, timeout: Duration) -> Self {
159 self.request_handler_timeout = timeout;
160 self
161 }
162 pub fn internal_operation_timeout(mut self, timeout: Duration) -> Self {
164 self.internal_operation_timeout = timeout;
165 self
166 }
167 pub fn configuration(mut self, configuration: Configuration) -> Self {
169 self.configuration = Some(configuration);
170 self
171 }
172 pub fn storage_client(mut self, storage: Arc<dyn StorageClient>) -> Self {
174 self.storage_client = Some(storage);
175 self
176 }
177 pub fn request_queue(mut self, queue: Arc<dyn RequestQueue>) -> Self {
184 self.request_queue = Some(queue);
185 self
186 }
187 pub fn results_capacity(mut self, capacity: usize) -> Self {
189 self.results_capacity = capacity;
190 self
191 }
192
193 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 pub fn crawl_policy(mut self, policy: CrawlPolicy) -> Self {
201 self.crawl_policy = Some(Arc::new(policy));
202 self
203 }
204
205 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}