stream_crawler/
lib.rs

1use std::{collections::HashMap, sync::Arc};
2use tokio::sync::mpsc::{self, Receiver, Sender};
3use tokio_stream::wrappers::ReceiverStream;
4use scraper::{Html, Selector};
5
6/// Represents the state of a task.
7enum TaskState {
8    NotStarted(usize),
9    PENDING(usize),
10    PROCESSED(usize),
11    FAILED(usize),
12}
13
14/// Represents an unprocessed URL.
15struct UnprocessedUrl {
16    parent: Option<String>,
17    url: String,
18}
19
20/// Represents a processed URL.
21///
22/// # Examples
23///
24/// ```
25/// use stream_crawler::ProcessedUrl;
26///
27/// let processed_url = ProcessedUrl {
28///     parent: Some(String::from("https://www.example.com")),
29///     url: String::from("https://www.example.com/page"),
30///     content: String::from("<html>...</html>"),
31///     children: vec![String::from("https://www.example.com/page/1")],
32/// };
33///
34/// assert_eq!(processed_url.url, "https://www.example.com/page");
35/// ```
36#[derive(Debug, PartialEq)]
37pub struct ProcessedUrl {
38    pub parent: Option<String>,
39    pub url: String,
40    pub content: String,
41    pub children: Vec<String>,
42}
43
44/// Scrapes the given URLs, returning a stream of processed URLs.
45///
46/// # Examples
47///
48/// ```
49/// use stream_crawler::scrape;
50/// use tokio_stream::StreamExt;
51///
52/// #[tokio::main]
53/// async fn main() {
54///     let urls = vec![
55///         String::from("https://www.google.com"),
56///         String::from("https://www.twitter.com"),
57///     ];
58///
59///     let mut result_stream = scrape(urls, 3, 5, 9).await;
60///     let mut results = vec![];
61///
62///     while let Some(data) = result_stream.next().await {
63///         results.push(data);
64///         break;
65///     }
66///     assert!(!results.is_empty());
67/// }
68/// ```
69pub async fn scrape(urls: Vec<String>, retry_attempt: usize, number_process: usize, channel_size: usize) -> impl tokio_stream::Stream<Item = ProcessedUrl> {
70    let (buffer_tx, mut buffer_rx) = mpsc::channel::<UnprocessedUrl>(channel_size);
71    let (user_tx, user_rx) = mpsc::channel::<ProcessedUrl>(channel_size);
72    let (fail_tx, mut fail_rx) = mpsc::channel::<(UnprocessedUrl, reqwest::Error)>(channel_size);
73    let (url_tx, mut url_rx) = mpsc::channel::<UnprocessedUrl>(channel_size);
74    
75    let task_state: Arc<tokio::sync::Mutex<HashMap<String, TaskState>>> = Arc::new(tokio::sync::Mutex::new(HashMap::new()));
76
77    for url in urls {
78        let unproc = UnprocessedUrl {
79            parent: None,
80            url: url.clone(),
81        };
82        task_state.clone().lock_owned().await.insert(url.clone(), TaskState::NotStarted(0));
83        buffer_tx.send(unproc).await.unwrap();
84    }
85    
86    let url_tx_clone = url_tx.clone();
87    let cloned_task_state = task_state.clone();
88    tokio::spawn(async move {
89        scheduler(&mut buffer_rx, user_tx, url_tx_clone, fail_tx, number_process, cloned_task_state).await;
90    });
91
92    let cloned_task_state = task_state.clone();
93    tokio::spawn(async move {
94        fail_task(&mut fail_rx, url_tx, retry_attempt, cloned_task_state).await;
95    });
96
97    let cloned_task_state = task_state.clone();
98    tokio::spawn(async move {
99        urls_processor(&mut url_rx, buffer_tx, cloned_task_state).await;
100    });
101
102    ReceiverStream::new(user_rx)
103}
104
105async fn scheduler(buffer_rx: &mut Receiver<UnprocessedUrl>, user_tx: Sender<ProcessedUrl>, url_tx: Sender<UnprocessedUrl>, fail_tx: Sender<(UnprocessedUrl, reqwest::Error)>, number_process: usize, task_state: Arc<tokio::sync::Mutex<HashMap<String, TaskState>>>) {
106    let semaphore = Arc::new(tokio::sync::Semaphore::new(number_process));
107    while let Some(unprocessed_url) = buffer_rx.recv().await {
108        let user_tx_clone = user_tx.clone();
109        let url_tx_clone = url_tx.clone();
110        let fail_tx_clone = fail_tx.clone();
111
112        let mut state = task_state.clone().lock_owned().await;
113        let key_value = state.get(&unprocessed_url.url);
114        match key_value {
115            Some(TaskState::FAILED(value)) => {
116                let val = value.clone();
117                state.insert(unprocessed_url.url.clone(), TaskState::PENDING(val));
118            },
119            Some(TaskState::NotStarted(value)) => {
120                let val = value.clone();
121                state.insert(unprocessed_url.url.clone(), TaskState::PENDING(val));
122            },
123            None => {
124                state.insert(unprocessed_url.url.clone(), TaskState::PENDING(0));
125            },
126            _ => continue
127        }
128        let sem = semaphore.clone();
129        tokio::spawn(async move {
130            let permit = sem.acquire_owned().await.unwrap();
131            process_url(user_tx_clone, url_tx_clone, fail_tx_clone, unprocessed_url).await;
132            drop(permit);
133        });
134    }
135}
136
137async fn fail_task(fail_rx: &mut Receiver<(UnprocessedUrl, reqwest::Error)>, url_tx: Sender<UnprocessedUrl>, retry_attempt: usize, task_state: Arc<tokio::sync::Mutex<HashMap<String, TaskState>>>) {
138    while let Some(unprocessed_url) = fail_rx.recv().await {
139
140
141        let mut state = task_state.clone().lock_owned().await;
142        let key_value = state.get(&unprocessed_url.0.url);
143        match key_value {
144            Some(TaskState::PENDING(value)) => {
145                if unprocessed_url.1.is_connect() && retry_attempt > value.clone() {
146
147                    let val = value.clone();
148                    state.insert(unprocessed_url.0.url.clone(), TaskState::FAILED(val + 1));
149                    url_tx.send(unprocessed_url.0).await.unwrap();
150                }
151            },
152            None => continue,
153            _ => continue
154        }
155    }
156}
157
158async fn urls_processor(url_rx: &mut Receiver<UnprocessedUrl>, buffer_tx: Sender<UnprocessedUrl>, task_state: Arc<tokio::sync::Mutex<HashMap<String, TaskState>>>) {
159    while let Some(unprocessed_url) = url_rx.recv().await {
160
161        let mut state = task_state.clone().lock_owned().await;
162        let key_value = state.get(&unprocessed_url.url);
163        match key_value {
164            Some(TaskState::FAILED(value)) => {
165                println!("url tx {:?}", value);
166                let val = value.clone();
167                state.insert(unprocessed_url.url.clone(), TaskState::FAILED(val));
168                buffer_tx.send(unprocessed_url).await.unwrap();
169            },
170            Some(TaskState::PENDING(value)) => {
171                let val = value.clone();
172                state.insert(unprocessed_url.url.clone(), TaskState::PROCESSED(val));
173            },
174            None => {
175                state.insert(unprocessed_url.url.clone(), TaskState::NotStarted(0));
176                buffer_tx.send(unprocessed_url).await.unwrap();
177            },
178            _ => continue
179        }
180    }
181}
182
183async fn process_url(user_tx: Sender<ProcessedUrl>, url_tx: Sender<UnprocessedUrl>, fail_tx: Sender<(UnprocessedUrl, reqwest::Error)>, unprocessed_url: UnprocessedUrl) {
184    println!("Task :: {}", unprocessed_url.url.clone());
185    let resp = reqwest::get(unprocessed_url.url.clone()).await;
186
187    match resp {
188        Ok(response) => {
189            let text = response.text().await;
190            match text {
191                Ok(body) => {
192                    let urls_extracted = extract_urls_from_a_tags(&body);
193                    let mut endpoints_extracted: Vec<String> = [].to_vec();
194                    for endpoint in &urls_extracted {
195                        if endpoint.starts_with('/') {
196                            endpoints_extracted.push(unprocessed_url.url.clone() + endpoint)
197                        } else {
198                            endpoints_extracted.push(endpoint.clone());
199                        }
200                    }
201
202                    let proc = ProcessedUrl {
203                        parent: unprocessed_url.parent.clone(),
204                        url: unprocessed_url.url.clone(),
205                        content: body,
206                        children: endpoints_extracted.clone(),
207                    };
208
209                    user_tx.send(proc).await.unwrap();
210
211                    for endpoint in endpoints_extracted {
212                        let unproc = UnprocessedUrl {
213                            url: endpoint,
214                            parent: Some(unprocessed_url.url.clone()),
215                        };
216                        url_tx.send(unproc).await.unwrap();
217                    }
218                },
219                Err(e) => {
220                    fail_tx.send((unprocessed_url, e)).await.unwrap();
221                }
222            }
223        },
224        Err(e) => {
225            fail_tx.send((unprocessed_url, e)).await.unwrap();
226        }
227    }
228}
229
230/// Extracts URLs from the <a> tags in the given HTML string.
231///
232/// # Examples
233///
234/// ```
235/// use stream_crawler::extract_urls_from_a_tags;
236///
237/// let html = r#"
238/// <html>
239///     <body>
240///         <a href="https://www.example.com/page1">Page 1</a>
241///         <a href="https://www.example.com/page2">Page 2</a>
242///     </body>
243/// </html>
244/// "#;
245///
246/// let urls = extract_urls_from_a_tags(html);
247/// assert_eq!(urls, vec![
248///     "https://www.example.com/page1",
249///     "https://www.example.com/page2"
250/// ]);
251/// ```
252pub fn extract_urls_from_a_tags(html: &str) -> Vec<String> {
253    let document = Html::parse_document(html);
254    let selector = Selector::parse("a").unwrap();
255    let data = document
256        .select(&selector)
257        .filter_map(|element| element.value().attr("href"))
258        .map(|href| href.to_string())
259        .collect();
260    data
261}
stream_crawler/lib.rs

stream_crawler/
lib.rs
