millipede_core/crawler/
mod.rs

1//! The crawler engine: lifecycle kinds, handles, and shared state.
2
3mod basic;
4mod builder;
5mod engine;
6mod start;
7
8pub use basic::{BasicContext, BasicKind};
9pub use builder::{CrawlerBuildError, CrawlerBuilder};
10pub use start::{IntoStartRequest, IntoStartRequests};
11
12use crate::{
13    autoscale::AutoscaledPool,
14    config::Configuration,
15    errors::CrawlError,
16    events::{EventBus, EventStream, HandledRequest, ResultStream},
17    handler::{FailedRequestHandler, RequestHandler},
18    link_extraction::CrawlPolicy,
19    request::Request,
20    statistics::{FinalStatistics, StatisticsHandle, StatisticsSnapshot},
21    storage::{AddOptions, BatchAddHandle, RequestQueue, RequestSource},
22};
23use futures_util::future::BoxFuture;
24use std::{
25    fmt,
26    sync::{
27        Arc, Weak,
28        atomic::{AtomicBool, AtomicU64, Ordering},
29    },
30    time::Duration,
31};
32
33use engine::{Engine, EngineOptions};
34
35/// A configured crawler using lifecycle behavior supplied by `K`.
36pub struct Crawler<K: CrawlerKind> {
37    kind: Arc<K>,
38    shared: Arc<CrawlerShared>,
39    config: Arc<Configuration>,
40    handler: Arc<dyn RequestHandler<K::Context>>,
41    failed_handler: Option<Arc<dyn FailedRequestHandler>>,
42    kvs: Option<Arc<dyn crate::storage::KeyValueStore>>,
43    storage: Option<Arc<dyn crate::storage::StorageClient>>,
44    opts: EngineOptions,
45    started: AtomicBool,
46}
47
48/// The no-HTTP crawler: drives the queue and hands requests straight to the handler.
49pub type BasicCrawler = Crawler<BasicKind>;
50
51impl<K: CrawlerKind> Crawler<K> {
52    /// Starts building a crawler around the given kind.
53    pub fn builder(kind: K) -> CrawlerBuilder<K> {
54        CrawlerBuilder::new(kind)
55    }
56
57    /// Runs the crawl to completion.
58    ///
59    /// A crawler runs at most once; a second call returns a non-retryable error.
60    pub async fn run(&self, start: impl IntoStartRequests) -> Result<FinalStatistics, CrawlError> {
61        if self.started.swap(true, Ordering::SeqCst) {
62            return Err(CrawlError::non_retryable(anyhow::anyhow!(
63                "this crawler has already been run"
64            )));
65        }
66        let start_requests = start.into_start_requests()?;
67        let env = CrawlerEnv {
68            shared: self.shared.clone(),
69            config: self.config.clone(),
70            storage: self.storage.clone(),
71            kvs: self.kvs.clone(),
72        };
73        self.kind.start(&env).await?;
74        let result = async {
75            let sources = start_requests
76                .into_iter()
77                .map(RequestSource::from)
78                .collect();
79            let batch = tokio::time::timeout(
80                self.opts.internal_operation_timeout,
81                self.shared.queue.add_batch(sources, AddOptions::default()),
82            )
83            .await
84            .map_err(|_| CrawlError::retry(anyhow::anyhow!("queue add timed out")))??;
85            let _ = batch.wait().await?;
86            self.shared.notify.notify_waiters();
87            Engine {
88                kind: self.kind.clone(),
89                handler: self.handler.clone(),
90                failed_handler: self.failed_handler.clone(),
91                shared: self.shared.clone(),
92                kvs: self.kvs.clone(),
93                opts: self.opts.clone(),
94            }
95            .run()
96            .await
97        }
98        .await;
99        if let Err(error) = self.kind.stop(&env).await {
100            tracing::warn!(%error, "crawler kind stop failed");
101        }
102        result
103    }
104
105    /// Creates a weak handle to this crawler.
106    pub fn handle(&self) -> CrawlerHandle {
107        CrawlerHandle::new(Arc::downgrade(&self.shared))
108    }
109    /// Adds requests and waits until the complete batch has been accepted.
110    pub async fn add_requests(
111        &self,
112        reqs: impl IntoIterator<Item = Request> + Send,
113    ) -> Result<(), CrawlError> {
114        let _ = self.handle().add_requests(reqs).await?.wait().await?;
115        Ok(())
116    }
117    /// Subscribes to terminal request snapshots.
118    pub fn results(&self) -> ResultStream {
119        self.shared.results_tx.subscribe()
120    }
121    /// Subscribes to control-plane crawler events.
122    pub fn events(&self) -> EventStream {
123        self.shared.events.subscribe()
124    }
125    /// Returns the live statistics handle.
126    pub fn stats(&self) -> StatisticsHandle {
127        self.shared.stats.clone()
128    }
129    /// Returns a snapshot of the concurrency scaler.
130    pub fn autoscaler_snapshot(&self) -> AutoscalerSnapshot {
131        AutoscalerSnapshot::from_pool(&self.shared.pool)
132    }
133    /// Signals a graceful drain.
134    pub fn stop(&self) {
135        self.handle().stop();
136    }
137    /// Signals immediate cancellation.
138    pub fn abort(&self) {
139        self.handle().abort();
140    }
141}
142
143pub(crate) struct CrawlerShared {
144    pub(crate) queue: Arc<dyn RequestQueue>,
145    pub(crate) stats: StatisticsHandle,
146    pub(crate) events: EventBus,
147    pub(crate) results_tx: tokio::sync::broadcast::Sender<HandledRequest>,
148    pub(crate) drain: tokio_util::sync::CancellationToken,
149    pub(crate) cancel: tokio_util::sync::CancellationToken,
150    pub(crate) notify: tokio::sync::Notify,
151    pub(crate) internal_operation_timeout: Duration,
152    pub(crate) pool: Arc<AutoscaledPool>,
153    enqueue_admission: Arc<tokio::sync::Mutex<()>>,
154    enqueue_admissions: Arc<AtomicU64>,
155    crawl_policy: Option<Arc<CrawlPolicy>>,
156}
157
158impl CrawlerShared {
159    /// Creates shared crawler state with fresh statistics, result, and cancellation channels.
160    ///
161    /// `results_capacity` must be at least one. The crawler builder validates this before
162    /// constructing shared state.
163    #[allow(dead_code)]
164    pub(crate) fn new(
165        queue: Arc<dyn RequestQueue>,
166        events: EventBus,
167        results_capacity: usize,
168        internal_operation_timeout: Duration,
169        pool: Arc<AutoscaledPool>,
170    ) -> Self {
171        debug_assert!(results_capacity >= 1);
172        let (results_tx, _) = tokio::sync::broadcast::channel(results_capacity);
173        Self {
174            queue,
175            stats: StatisticsHandle::new(),
176            events,
177            results_tx,
178            drain: tokio_util::sync::CancellationToken::new(),
179            cancel: tokio_util::sync::CancellationToken::new(),
180            notify: tokio::sync::Notify::new(),
181            internal_operation_timeout,
182            pool,
183            enqueue_admission: Arc::new(tokio::sync::Mutex::new(())),
184            enqueue_admissions: Arc::new(AtomicU64::new(0)),
185            crawl_policy: None,
186        }
187    }
188
189    pub(crate) fn new_with_policy(
190        queue: Arc<dyn RequestQueue>,
191        events: EventBus,
192        results_capacity: usize,
193        internal_operation_timeout: Duration,
194        pool: Arc<AutoscaledPool>,
195        crawl_policy: Option<Arc<CrawlPolicy>>,
196    ) -> Self {
197        let mut shared = Self::new(
198            queue,
199            events,
200            results_capacity,
201            internal_operation_timeout,
202            pool,
203        );
204        shared.crawl_policy = crawl_policy;
205        shared
206    }
207
208    /// Returns the crawler's request queue.
209    pub fn request_queue(&self) -> &Arc<dyn RequestQueue> {
210        &self.queue
211    }
212
213    /// Returns the configured crawl policy, when one was supplied.
214    pub fn crawl_policy(&self) -> Option<&Arc<CrawlPolicy>> {
215        self.crawl_policy.as_ref()
216    }
217}
218
219/// A point-in-time view of a crawler's concurrency scaler.
220#[non_exhaustive]
221#[derive(Debug, Clone, Copy)]
222pub struct AutoscalerSnapshot {
223    /// Concurrency currently requested by the selected scaling mode.
224    pub desired_concurrency: usize,
225    /// Effective minimum concurrency.
226    pub min_concurrency: usize,
227    /// Effective maximum concurrency.
228    pub max_concurrency: usize,
229    /// Whether concurrency is explicitly fixed.
230    pub is_fixed: bool,
231}
232
233impl AutoscalerSnapshot {
234    fn from_pool(pool: &AutoscaledPool) -> Self {
235        Self {
236            desired_concurrency: pool.desired_concurrency(),
237            min_concurrency: pool.min_concurrency(),
238            max_concurrency: pool.max_concurrency(),
239            is_fixed: pool.is_fixed(),
240        }
241    }
242}
243
244/// A cheaply cloned weak back-reference to a running crawler.
245#[derive(Clone)]
246pub struct CrawlerHandle {
247    inner: Weak<CrawlerShared>,
248}
249
250pub(crate) struct EnqueueAdmissionReservation {
251    admissions: Arc<AtomicU64>,
252    committed: bool,
253}
254
255impl EnqueueAdmissionReservation {
256    pub(crate) fn commit(mut self) {
257        self.committed = true;
258    }
259}
260
261impl Drop for EnqueueAdmissionReservation {
262    fn drop(&mut self) {
263        if !self.committed {
264            self.admissions.fetch_sub(1, Ordering::SeqCst);
265        }
266    }
267}
268
269impl fmt::Debug for CrawlerHandle {
270    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
271        formatter
272            .debug_struct("CrawlerHandle")
273            .field("alive", &(self.inner.strong_count() > 0))
274            .finish()
275    }
276}
277
278impl CrawlerHandle {
279    pub(crate) fn new(inner: Weak<CrawlerShared>) -> Self {
280        Self { inner }
281    }
282
283    /// Adds requests to the crawler's queue.
284    pub async fn add_requests(
285        &self,
286        reqs: impl IntoIterator<Item = Request> + Send,
287    ) -> Result<BatchAddHandle, CrawlError> {
288        self.add_requests_with_options(reqs, AddOptions::default())
289            .await
290    }
291
292    /// Adds requests with explicit queue insertion options.
293    pub async fn add_requests_with_options(
294        &self,
295        reqs: impl IntoIterator<Item = Request> + Send,
296        options: AddOptions,
297    ) -> Result<BatchAddHandle, CrawlError> {
298        let shared = self.inner.upgrade().ok_or_else(|| {
299            CrawlError::non_retryable(anyhow::anyhow!("crawler is no longer running"))
300        })?;
301        let sources = reqs.into_iter().map(RequestSource::from).collect();
302        let handle = tokio::time::timeout(
303            shared.internal_operation_timeout,
304            shared.queue.add_batch(sources, options),
305        )
306        .await
307        .map_err(|_| CrawlError::retry(anyhow::anyhow!("queue add timed out")))??;
308        let handle = handle.notify_on_completion({
309            let shared = shared.clone();
310            move || shared.notify.notify_waiters()
311        });
312        Ok(handle)
313    }
314
315    pub(crate) async fn lock_enqueue_admission(
316        &self,
317    ) -> Result<tokio::sync::OwnedMutexGuard<()>, CrawlError> {
318        let shared = self.inner.upgrade().ok_or_else(|| {
319            CrawlError::non_retryable(anyhow::anyhow!("crawler is no longer running"))
320        })?;
321        Ok(shared.enqueue_admission.clone().lock_owned().await)
322    }
323
324    pub(crate) fn synchronize_enqueue_admissions(
325        &self,
326        observed_queue_count: u64,
327    ) -> Result<u64, CrawlError> {
328        let shared = self.inner.upgrade().ok_or_else(|| {
329            CrawlError::non_retryable(anyhow::anyhow!("crawler is no longer running"))
330        })?;
331        let previous = shared
332            .enqueue_admissions
333            .fetch_max(observed_queue_count, Ordering::SeqCst);
334        Ok(previous.max(observed_queue_count))
335    }
336
337    pub(crate) fn reserve_enqueue_admission(
338        &self,
339    ) -> Result<EnqueueAdmissionReservation, CrawlError> {
340        let shared = self.inner.upgrade().ok_or_else(|| {
341            CrawlError::non_retryable(anyhow::anyhow!("crawler is no longer running"))
342        })?;
343        shared.enqueue_admissions.fetch_add(1, Ordering::SeqCst);
344        Ok(EnqueueAdmissionReservation {
345            admissions: shared.enqueue_admissions.clone(),
346            committed: false,
347        })
348    }
349
350    /// Returns a snapshot of live crawl statistics while the crawler exists.
351    pub fn stats(&self) -> Option<StatisticsSnapshot> {
352        self.inner.upgrade().map(|shared| shared.stats.snapshot())
353    }
354
355    /// Returns a snapshot of the concurrency scaler while the crawler exists.
356    pub fn autoscaler_snapshot(&self) -> Option<AutoscalerSnapshot> {
357        self.inner
358            .upgrade()
359            .map(|shared| AutoscalerSnapshot::from_pool(&shared.pool))
360    }
361
362    /// Subscribes to control-plane crawler events while the crawler exists.
363    pub fn events(&self) -> Option<EventStream> {
364        self.inner.upgrade().map(|shared| shared.events.subscribe())
365    }
366
367    /// Subscribes to terminal request snapshots while the crawler exists.
368    pub fn results(&self) -> Option<crate::events::ResultStream> {
369        self.inner
370            .upgrade()
371            .map(|shared| shared.results_tx.subscribe())
372    }
373
374    /// Returns the crawler's request queue while the crawler exists.
375    pub fn request_queue(&self) -> Option<Arc<dyn RequestQueue>> {
376        self.inner
377            .upgrade()
378            .map(|shared| shared.request_queue().clone())
379    }
380
381    /// Returns the configured crawl policy while the crawler exists.
382    pub fn crawl_policy(&self) -> Option<Arc<CrawlPolicy>> {
383        self.inner
384            .upgrade()
385            .and_then(|shared| shared.crawl_policy().cloned())
386    }
387
388    /// Requests a graceful stop that finishes in-flight work and fetches no more requests.
389    pub fn stop(&self) {
390        if let Some(shared) = self.inner.upgrade() {
391            shared.drain.cancel();
392            shared.notify.notify_waiters();
393        }
394    }
395
396    /// Requests immediate cancellation of crawler work.
397    pub fn abort(&self) {
398        if let Some(shared) = self.inner.upgrade() {
399            shared.cancel.cancel();
400            shared.notify.notify_waiters();
401        }
402    }
403}
404
405/// Shared process-level state supplied to crawler lifecycle hooks.
406pub struct CrawlerEnv {
407    pub(crate) shared: Arc<CrawlerShared>,
408    pub(crate) config: Arc<Configuration>,
409    pub(crate) storage: Option<Arc<dyn crate::storage::StorageClient>>,
410    pub(crate) kvs: Option<Arc<dyn crate::storage::KeyValueStore>>,
411}
412
413impl CrawlerEnv {
414    /// Returns the crawler's adopted event bus.
415    pub fn events(&self) -> &EventBus {
416        &self.shared.events
417    }
418
419    /// Returns the live statistics handle.
420    pub fn stats(&self) -> &StatisticsHandle {
421        &self.shared.stats
422    }
423
424    /// Returns the resolved crawler configuration.
425    pub fn config(&self) -> &Configuration {
426        &self.config
427    }
428
429    /// Returns the resolved storage client used by the crawler.
430    pub fn storage_client(&self) -> Option<&Arc<dyn crate::storage::StorageClient>> {
431        self.storage.as_ref()
432    }
433
434    /// Returns the crawler's resolved key-value store.
435    pub fn kvs(&self) -> Option<&Arc<dyn crate::storage::KeyValueStore>> {
436        self.kvs.as_ref()
437    }
438
439    /// Returns the crawler's request queue.
440    pub fn request_queue(&self) -> &Arc<dyn RequestQueue> {
441        &self.shared.queue
442    }
443
444    /// Creates a weak handle to the crawler.
445    pub fn handle(&self) -> CrawlerHandle {
446        CrawlerHandle::new(Arc::downgrade(&self.shared))
447    }
448}
449
450/// Engine-owned scratch space passed to [`CrawlerKind::before_request`].
451#[non_exhaustive]
452pub struct RequestPrep {
453    /// The request being prepared for its next attempt.
454    pub request: Request,
455}
456
457/// Per-attempt inputs supplied to [`CrawlerKind::execute`].
458#[non_exhaustive]
459pub struct RequestEnv<'a> {
460    /// The request being executed.
461    pub request: Arc<Request>,
462    /// A weak back-reference to the running crawler.
463    pub crawler: CrawlerHandle,
464    /// The crawler's event bus.
465    pub events: &'a EventBus,
466    /// Overrides carried from the previous attempt of this request.
467    pub overrides: crate::retry_strategy::AttemptOverrides,
468}
469
470impl<'a> RequestEnv<'a> {
471    /// Clones these per-attempt inputs so one attempt can try a second execution path (e.g. smart
472    /// HTTP-first promotion re-executing through a browser kind). The struct is non-exhaustive, so
473    /// only core can provide this.
474    pub fn duplicate(&self) -> RequestEnv<'a> {
475        RequestEnv {
476            request: Arc::clone(&self.request),
477            crawler: self.crawler.clone(),
478            events: self.events,
479            overrides: self.overrides.clone(),
480        }
481    }
482}
483
484/// Metadata observed after a kind successfully constructs its handler context.
485///
486/// # Examples
487///
488/// ```
489/// use http::StatusCode;
490/// use millipede_core::crawler::AttemptObservation;
491///
492/// let mut observation = AttemptObservation::default();
493/// observation.status = Some(StatusCode::OK);
494/// observation.response_bytes = Some(1_024);
495/// ```
496#[derive(Debug, Clone, Default)]
497#[non_exhaustive]
498pub struct AttemptObservation {
499    /// HTTP response status, when applicable.
500    pub status: Option<http::StatusCode>,
501    /// Final loaded URL after redirects, when applicable.
502    pub loaded_url: Option<url::Url>,
503    /// Session used by the attempt.
504    pub session_id: Option<crate::session::SessionId>,
505    /// Proxy used by the attempt.
506    pub proxy_info: Option<crate::proxy::ProxyInfo>,
507    /// Buffered response size.
508    pub response_bytes: Option<usize>,
509}
510
511/// The outcome supplied to per-attempt cleanup.
512pub enum RequestOutcome<C> {
513    /// Execution and the user handler succeeded.
514    Handled(C),
515    /// The user handler failed after context creation.
516    HandlerFailed {
517        /// The context returned by execution.
518        ctx: C,
519        /// The handler error, shared with the failure handler.
520        error: Arc<CrawlError>,
521    },
522    /// Request preparation or execution failed before a handler completed.
523    ExecuteFailed {
524        /// The request whose attempt failed.
525        request: Arc<Request>,
526        /// The execution error, shared with the failure handler.
527        error: Arc<CrawlError>,
528    },
529}
530
531/// Defines the complete lifecycle for one crawler flavor.
532pub trait CrawlerKind: Send + Sync + 'static {
533    /// The context passed to user handlers and lifecycle hooks.
534    ///
535    /// A context must be a cheap aliasing handle over shared state, typically through `Arc`-backed
536    /// fields. Clones must observe the same underlying resources so mutations through a handler's
537    /// clone remain visible to `after_success` and `cleanup`; plain-value contexts that diverge on
538    /// clone violate this contract. `Clone` is required because the handler consumes an owned
539    /// context while `after_success` and `cleanup` still need it.
540    type Context: Send + Clone + 'static;
541
542    /// Runs once before the crawler fetches any request.
543    fn start<'a>(&'a self, env: &'a CrawlerEnv) -> BoxFuture<'a, Result<(), CrawlError>> {
544        let _ = env;
545        Box::pin(async { Ok(()) })
546    }
547
548    /// Mutates a request before an attempt executes.
549    fn before_request<'a>(
550        &'a self,
551        prep: &'a mut RequestPrep,
552    ) -> BoxFuture<'a, Result<(), CrawlError>> {
553        let _ = prep;
554        Box::pin(async { Ok(()) })
555    }
556
557    /// Executes one request attempt and constructs its handler context.
558    fn execute<'a>(
559        &'a self,
560        env: RequestEnv<'a>,
561    ) -> BoxFuture<'a, Result<Self::Context, CrawlError>>;
562
563    /// Called once after `execute()` succeeds; feeds statistics, `HandledRequest`, and
564    /// `RetryStrategy` metadata.
565    fn observe(&self, ctx: &Self::Context) -> AttemptObservation {
566        let _ = ctx;
567        AttemptObservation::default()
568    }
569
570    /// Runs after the user handler succeeds.
571    fn after_success<'a>(
572        &'a self,
573        ctx: &'a mut Self::Context,
574    ) -> BoxFuture<'a, Result<(), CrawlError>> {
575        let _ = ctx;
576        Box::pin(async { Ok(()) })
577    }
578
579    /// Runs after every attempt concludes, regardless of its outcome.
580    fn cleanup(
581        &self,
582        outcome: RequestOutcome<Self::Context>,
583    ) -> BoxFuture<'_, Result<(), CrawlError>>;
584
585    /// Runs once when crawler shutdown begins.
586    fn stop<'a>(&'a self, env: &'a CrawlerEnv) -> BoxFuture<'a, Result<(), CrawlError>> {
587        let _ = env;
588        Box::pin(async { Ok(()) })
589    }
590}
591
592#[cfg(test)]
593mod tests {
594    use super::*;
595    use crate::storage::{Lease, LeaseId, ProcessedRequest, ReclaimOptions, StorageResult};
596    use std::sync::Mutex;
597
598    #[derive(Default)]
599    struct TestQueue(Mutex<Vec<Request>>);
600
601    #[async_trait::async_trait]
602    impl RequestQueue for TestQueue {
603        async fn add(&self, request: Request, _: AddOptions) -> StorageResult<ProcessedRequest> {
604            let mut requests = self.0.lock().unwrap();
605            let duplicate = requests
606                .iter()
607                .any(|known| known.unique_key == request.unique_key);
608            let info = ProcessedRequest {
609                request_id: request.id.clone(),
610                unique_key: request.unique_key.clone(),
611                was_already_present: duplicate,
612                was_already_handled: false,
613            };
614            if !duplicate {
615                requests.push(request);
616            }
617            Ok(info)
618        }
619
620        async fn add_batch(
621            &self,
622            requests: Vec<RequestSource>,
623            options: AddOptions,
624        ) -> StorageResult<BatchAddHandle> {
625            let mut added = Vec::with_capacity(requests.len());
626            for source in requests {
627                let RequestSource::Request(request) = source;
628                added.push(self.add(request, options.clone()).await?);
629            }
630            Ok(BatchAddHandle::ready(added))
631        }
632
633        async fn fetch_next(&self) -> StorageResult<Option<Lease>> {
634            Ok(None)
635        }
636        async fn mark_handled(&self, _: Lease) -> StorageResult<()> {
637            Ok(())
638        }
639        async fn reclaim(&self, _: Lease, _: ReclaimOptions) -> StorageResult<()> {
640            Ok(())
641        }
642        async fn renew(&self, _: &LeaseId, _: Duration) -> StorageResult<()> {
643            Ok(())
644        }
645        async fn abandon(&self, _: Lease) -> StorageResult<()> {
646            Ok(())
647        }
648        async fn is_empty(&self) -> StorageResult<bool> {
649            Ok(self.0.lock().unwrap().is_empty())
650        }
651        async fn is_finished(&self) -> StorageResult<bool> {
652            self.is_empty().await
653        }
654        async fn handled_count(&self) -> StorageResult<u64> {
655            Ok(0)
656        }
657        async fn pending_count(&self) -> StorageResult<u64> {
658            Ok(self.0.lock().unwrap().len() as u64)
659        }
660    }
661
662    pub(super) fn shared() -> Arc<CrawlerShared> {
663        let queue = Arc::new(TestQueue::default());
664        Arc::new(CrawlerShared::new(
665            queue,
666            EventBus::default(),
667            8,
668            Duration::from_secs(1),
669            Arc::new(AutoscaledPool::new(
670                crate::autoscale::AutoscaledPoolOptions {
671                    fixed_concurrency: Some(8),
672                    ..Default::default()
673                },
674            )),
675        ))
676    }
677
678    #[tokio::test]
679    async fn crawler_handle_adds_deduplicated_requests_and_observes_liveness() {
680        let shared = shared();
681        let queue = shared.queue.clone();
682        let handle = CrawlerHandle::new(Arc::downgrade(&shared));
683        let request = Request::get("https://example.com/item").build().unwrap();
684        let batch = handle
685            .add_requests([request.clone(), request])
686            .await
687            .unwrap();
688        assert_eq!(batch.added.len(), 2);
689        assert!(!batch.added[0].was_already_present);
690        assert!(batch.added[1].was_already_present);
691        assert_eq!(batch.wait().await.unwrap().processed.len(), 2);
692        assert_eq!(queue.pending_count().await.unwrap(), 1);
693        assert!(handle.stats().is_some());
694        let autoscaler = handle.autoscaler_snapshot().unwrap();
695        assert_eq!(autoscaler.desired_concurrency, 8);
696        assert!(autoscaler.is_fixed);
697        assert!(handle.events().is_some());
698        assert!(handle.results().is_some());
699        assert_eq!(format!("{handle:?}"), "CrawlerHandle { alive: true }");
700
701        drop(shared);
702        assert!(handle.add_requests(Vec::new()).await.is_err());
703        assert!(handle.stats().is_none());
704        assert!(handle.autoscaler_snapshot().is_none());
705        assert_eq!(format!("{handle:?}"), "CrawlerHandle { alive: false }");
706    }
707
708    #[tokio::test]
709    async fn crawler_handle_stop_and_abort_cancel_their_tokens() {
710        let shared = shared();
711        let handle = CrawlerHandle::new(Arc::downgrade(&shared));
712        handle.stop();
713        assert!(shared.drain.is_cancelled());
714        handle.abort();
715        assert!(shared.cancel.is_cancelled());
716    }
717}
millipede_core/crawler/mod.rs

millipede_core/crawler/
mod.rs
