kumo 0.5.0

An async web crawling framework for Rust - Scrapy for Rust
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113

use std::{
    sync::{
        Arc,
        atomic::{AtomicBool, Ordering},
    },
    task::{Context, Poll},
};

use crate::{
    error::KumoError,
    logging::{event, target},
    spider::Spider,
};

use super::builder::CrawlEngine;

pub(super) struct ChannelStore {
    tx: tokio::sync::mpsc::Sender<serde_json::Value>,
    cancelled: Arc<AtomicBool>,
}

impl ChannelStore {
    pub(super) fn new(
        tx: tokio::sync::mpsc::Sender<serde_json::Value>,
        cancelled: Arc<AtomicBool>,
    ) -> Self {
        Self { tx, cancelled }
    }
}

#[async_trait::async_trait]
impl crate::store::ItemStore for ChannelStore {
    async fn store(&self, item: &serde_json::Value) -> Result<(), KumoError> {
        if self.tx.send(item.clone()).await.is_err() {
            self.cancelled.store(true, Ordering::Relaxed);
        }
        Ok(())
    }
}

impl CrawlEngine {
    /// Run the crawl in the background and stream items as they are scraped.
    ///
    /// Unlike [`run`](Self::run), this returns immediately with an [`ItemStream`].
    /// The crawl engine runs in a spawned Tokio task; dropping the stream stops
    /// the crawl gracefully.
    ///
    /// # Example
    /// ```rust,ignore
    /// use tokio_stream::StreamExt;
    ///
    /// let mut stream = CrawlEngine::builder()
    ///     .concurrency(4)
    ///     .stream(MySpider)
    ///     .await?;
    ///
    /// while let Some(item) = stream.next().await {
    ///     println!("{}", item);
    /// }
    /// ```
    pub async fn stream<S>(self, spider: S) -> Result<ItemStream, KumoError>
    where
        S: Spider + 'static,
    {
        let buffer = self.stream_buffer;
        let (tx, rx) = tokio::sync::mpsc::channel(buffer);
        let cancelled = Arc::new(AtomicBool::new(false));
        let mut engine = self.store(ChannelStore::new(tx, cancelled.clone()));
        engine.stream_cancelled = Some(cancelled);
        tokio::spawn(async move {
            if let Err(e) = engine.run(spider).await {
                tracing::error!(
                    target: target::CRAWL,
                    event = event::CRAWL_STREAM_ERROR,
                    error = %e,
                    error_kind = e.kind().as_str(),
                    "crawl.stream_error"
                );
            }
        });
        Ok(ItemStream::new(rx))
    }
}

/// An async stream of scraped items returned by [`CrawlEngine::stream`].
///
/// Implements [`tokio_stream::Stream`]`<Item = serde_json::Value>`.
/// Use [`tokio_stream::StreamExt::next`] to consume items one by one.
///
/// Dropping this stream closes the channel, which causes the background
/// crawl engine to stop gracefully on its next attempted send.
pub struct ItemStream {
    inner: tokio_stream::wrappers::ReceiverStream<serde_json::Value>,
}

impl ItemStream {
    pub(super) fn new(rx: tokio::sync::mpsc::Receiver<serde_json::Value>) -> Self {
        Self {
            inner: tokio_stream::wrappers::ReceiverStream::new(rx),
        }
    }
}

impl tokio_stream::Stream for ItemStream {
    type Item = serde_json::Value;

    fn poll_next(
        mut self: std::pin::Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Option<Self::Item>> {
        std::pin::Pin::new(&mut self.inner).poll_next(cx)
    }
}