spider_lib/

crawler.rs

use crate::checkpoint::{Checkpoint, SchedulerCheckpoint};
use crate::downloader::Downloader;
use crate::error::SpiderError;
use crate::item::{ParseOutput, ScrapedItem};
use crate::middleware::{Middleware, MiddlewareAction};
use crate::pipeline::Pipeline;
use crate::request::Request;
use crate::response::Response;
use crate::scheduler::Scheduler;
use crate::spider::Spider;
use crate::state::CrawlerState;
use anyhow::Result;
use futures_util::future::join_all;
use kanal::{AsyncReceiver, AsyncSender, bounded_async};
use std::collections::{HashMap, VecDeque};
use std::fs;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::sync::atomic::Ordering;
use std::time::Duration;
use tokio::sync::Mutex;
use tokio::sync::Semaphore;
use tokio::task::JoinSet;
use tracing::{debug, error, info, warn};

pub struct Crawler<S: Spider, C> {
    scheduler: Arc<Scheduler>,
    req_rx: AsyncReceiver<Request>,
    downloader: Arc<dyn Downloader<Client = C> + Send + Sync>,
    middlewares: Vec<Box<dyn Middleware<C> + Send + Sync>>,
    spider: Arc<Mutex<S>>,
    item_pipelines: Vec<Box<dyn Pipeline<S::Item>>>,
    max_concurrent_downloads: usize,
    parser_workers: usize,
    max_concurrent_pipelines: usize,
    checkpoint_path: Option<PathBuf>,
    checkpoint_interval: Option<Duration>,
}

impl<S, C> Crawler<S, C>
where
    S: Spider + 'static,
    S::Item: ScrapedItem,
    C: Send + Sync + 'static,
{
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn new(
        scheduler: Arc<Scheduler>,
        req_rx: AsyncReceiver<Request>,
        downloader: Arc<dyn Downloader<Client = C> + Send + Sync>,
        middlewares: Vec<Box<dyn Middleware<C> + Send + Sync>>,
        spider: S,
        item_pipelines: Vec<Box<dyn Pipeline<S::Item>>>,
        max_concurrent_downloads: usize,
        parser_workers: usize,
        max_concurrent_pipelines: usize,
        checkpoint_path: Option<PathBuf>,
        checkpoint_interval: Option<Duration>,
    ) -> Self {
        Crawler {
            scheduler,
            req_rx,
            downloader,
            middlewares,
            spider: Arc::new(Mutex::new(spider)),
            item_pipelines,
            max_concurrent_downloads,
            parser_workers,
            max_concurrent_pipelines,
            checkpoint_path,
            checkpoint_interval,
        }
    }

    /// Starts the crawl.
    pub async fn start_crawl(self) -> Result<(), SpiderError> {
        info!("Crawler starting crawl");

        let Crawler {
            scheduler,
            req_rx,
            downloader,
            middlewares,
            spider,
            item_pipelines,
            max_concurrent_downloads,
            parser_workers,
            max_concurrent_pipelines,
            checkpoint_path,
            checkpoint_interval,
        } = self;

        let state = CrawlerState::new();
        let pipelines = Arc::new(item_pipelines);
        let channel_capacity = max_concurrent_downloads * 2;

        let (res_tx, res_rx) = bounded_async(channel_capacity);
        let (item_tx, item_rx) = bounded_async(channel_capacity);

        let (salvaged_requests_tx, salvaged_requests_rx) = bounded_async(channel_capacity);

        let initial_requests_task = spawn_initial_requests_task::<S>(
            scheduler.clone(),
            spider.clone(),
            salvaged_requests_tx.clone(),
        );

        let downloader_task = spawn_downloader_task::<S, C>(
            scheduler.clone(),
            req_rx,
            downloader,
            Arc::new(Mutex::new(middlewares)),
            state.clone(),
            res_tx.clone(),
            max_concurrent_downloads,
            salvaged_requests_tx.clone(),
        );

        let parser_task = spawn_parser_task::<S>(
            scheduler.clone(),
            spider.clone(),
            state.clone(),
            res_rx,
            item_tx.clone(),
            parser_workers,
            salvaged_requests_tx.clone(),
        );

        let item_processor_task = spawn_item_processor_task::<S>(
            state.clone(),
            item_rx,
            pipelines.clone(),
            max_concurrent_pipelines,
        );

        if let (Some(path), Some(interval)) = (&checkpoint_path, checkpoint_interval) {
            let scheduler_clone = scheduler.clone();
            let pipelines_clone = pipelines.clone();
            let path_clone = path.clone();
            let salvaged_requests_rx_clone = salvaged_requests_rx.clone();
            tokio::spawn(async move {
                let mut interval_timer = tokio::time::interval(interval);
                interval_timer.tick().await;
                loop {
                    tokio::select! {
                                    _ = interval_timer.tick() => {
                                        if let Ok(scheduler_checkpoint) = scheduler_clone.snapshot().await &&
                                            let Err(e) = save_checkpoint::<S>(&path_clone, scheduler_checkpoint, &pipelines_clone, salvaged_requests_rx_clone.clone()).await {
                                                error!("Periodic checkpoint save failed: {}", e);
                                        }
                                    }
                    }
                }
            });
        }

        tokio::select! {
            _ = tokio::signal::ctrl_c() => {
                info!("Ctrl-C received, initiating graceful shutdown.");
            }
            _ = async {
                loop {
                    if scheduler.is_idle() && state.is_idle() {
                        tokio::time::sleep(Duration::from_millis(50)).await;
                        if scheduler.is_idle() && state.is_idle() {
                            break;
                        }
                    }
                    tokio::time::sleep(Duration::from_millis(100)).await;
                }
            } => {
                info!("Crawl has become idle, initiating shutdown.");
            }
        }

        info!("Initiating actor shutdowns.");

        let scheduler_checkpoint = scheduler.snapshot().await?;

        drop(res_tx);
        drop(item_tx);
        drop(salvaged_requests_tx);

        scheduler.shutdown().await?;

        item_processor_task
            .await
            .map_err(|e| SpiderError::GeneralError(format!("Item processor task failed: {}", e)))?;

        parser_task
            .await
            .map_err(|e| SpiderError::GeneralError(format!("Parser task failed: {}", e)))?;

        downloader_task
            .await
            .map_err(|e| SpiderError::GeneralError(format!("Downloader task failed: {}", e)))?;

        initial_requests_task.await.map_err(|e| {
            SpiderError::GeneralError(format!("Initial requests task failed: {}", e))
        })?;

        // Save final checkpoint BEFORE closing pipelines
        if let Some(path) = &checkpoint_path
            && let Err(e) =
                save_checkpoint::<S>(path, scheduler_checkpoint, &pipelines, salvaged_requests_rx)
                    .await
        {
            error!("Final checkpoint save failed: {}", e);
        }

        // Close all pipelines
        info!("Closing item pipelines...");
        let closing_futures: Vec<_> = pipelines.iter().map(|p| p.close()).collect();
        join_all(closing_futures).await;

        info!("Crawl finished successfully.");
        Ok(())
    }
}

fn spawn_initial_requests_task<S>(
    scheduler: Arc<Scheduler>,
    spider: Arc<Mutex<S>>,
    salvaged_requests_tx: AsyncSender<Request>,
) -> tokio::task::JoinHandle<()>
where
    S: Spider + 'static,
    S::Item: ScrapedItem,
{
    tokio::spawn(async move {
        match spider.lock().await.start_requests() {
            Ok(requests) => {
                for mut req in requests {
                    req.url.set_fragment(None);
                    match scheduler.enqueue_request(req).await {
                        Ok(_) => {}
                        Err((req, e)) => {
                            error!("Failed to enqueue initial request: {}", e);
                            if salvaged_requests_tx.send(req).await.is_err() {
                                error!("Failed to send salvaged request to channel.");
                            }
                        }
                    }
                }
            }
            Err(e) => error!("Failed to create start requests: {}", e),
        }
    })
}

async fn save_checkpoint<S: Spider>(
    path: &Path,
    mut scheduler_checkpoint: SchedulerCheckpoint,
    pipelines: &Arc<Vec<Box<dyn Pipeline<S::Item>>>>,
    salvaged_requests_rx: AsyncReceiver<Request>,
) -> Result<(), SpiderError>
where
    S::Item: ScrapedItem,
{
    info!("Saving checkpoint to {:?}", path);

    let mut pipelines_checkpoint_map = HashMap::new();
    for pipeline in pipelines.iter() {
        if let Some(state) = pipeline.get_state().await? {
            pipelines_checkpoint_map.insert(pipeline.name().to_string(), state);
        }
    }

    let mut salvaged_requests_vec: VecDeque<Request> = VecDeque::new();
    while let Ok(req_option) = salvaged_requests_rx.try_recv() {
        if let Some(req) = req_option {
            salvaged_requests_vec.push_back(req);
        }
    }
    scheduler_checkpoint
        .request_queue
        .extend(salvaged_requests_vec.drain(..));
    if !salvaged_requests_vec.is_empty() {
        warn!(
            "Found {} salvaged requests during checkpoint. These have been added to the request queue.",
            salvaged_requests_vec.len()
        );
    }

    let checkpoint = Checkpoint {
        scheduler: scheduler_checkpoint,
        pipelines: pipelines_checkpoint_map,
    };

    // Write to a temporary file then rename for atomicity
    let tmp_path = path.with_extension("tmp");
    let encoded = rmp_serde::to_vec(&checkpoint)
        .map_err(|e| SpiderError::GeneralError(format!("Failed to serialize checkpoint: {}", e)))?;
    fs::write(&tmp_path, encoded).map_err(|e| {
        SpiderError::GeneralError(format!(
            "Failed to write checkpoint to temporary file: {}",
            e
        ))
    })?;
    fs::rename(&tmp_path, path).map_err(|e| {
        SpiderError::GeneralError(format!("Failed to rename temporary checkpoint file: {}", e))
    })?;

    info!("Checkpoint saved successfully.");
    Ok(())
}

#[allow(clippy::too_many_arguments)]
fn spawn_downloader_task<S, C>(
    scheduler: Arc<Scheduler>,
    req_rx: AsyncReceiver<Request>,
    downloader: Arc<dyn Downloader<Client = C> + Send + Sync>,
    middlewares: Arc<Mutex<Vec<Box<dyn Middleware<C> + Send + Sync>>>>,
    state: Arc<CrawlerState>,
    res_tx: AsyncSender<Response>,
    max_concurrent_downloads: usize,
    salvaged_requests_tx: AsyncSender<Request>,
) -> tokio::task::JoinHandle<()>
where
    S: Spider + 'static,
    S::Item: ScrapedItem,
    C: Send + Sync + 'static,
{
    let semaphore = Arc::new(Semaphore::new(max_concurrent_downloads));
    let mut tasks = JoinSet::new();

    tokio::spawn(async move {
        while let Ok(request) = req_rx.recv().await {
            let permit = match semaphore.clone().acquire_owned().await {
                Ok(permit) => permit,
                Err(_) => {
                    warn!("Semaphore closed, shutting down downloader actor.");
                    break;
                }
            };

            state.in_flight_requests.fetch_add(1, Ordering::SeqCst);
            let downloader_clone = Arc::clone(&downloader);
            let middlewares_clone = Arc::clone(&middlewares);
            let res_tx_clone = res_tx.clone();
            let state_clone = Arc::clone(&state);
            let scheduler_clone = Arc::clone(&scheduler);
            let _salvaged_requests_tx_clone = salvaged_requests_tx.clone();

            tasks.spawn(async move {
                let mut processed_request = request;
                let mut early_returned_response: Option<Response> = None;

                // Process request middlewares
                for mw in middlewares_clone.lock().await.iter_mut() {
                    match mw.process_request(downloader_clone.client(), processed_request.clone()).await {
                        Ok(MiddlewareAction::Continue(req)) => {
                            processed_request = req;
                        }
                        Ok(MiddlewareAction::Retry(req, delay)) => {
                            tokio::time::sleep(delay).await;
                            if scheduler_clone.enqueue_request(*req).await.is_err() {
                                error!("Failed to re-enqueue retried request.");
                            }
                            return;
                        }
                        Ok(MiddlewareAction::Drop) => {
                            debug!("Request dropped by middleware.");
                            return;
                        }
                        Ok(MiddlewareAction::ReturnResponse(resp)) => {
                            early_returned_response = Some(resp);
                            break;
                        }
                        Err(e) => {
                            error!("Request middleware error: {:?}", e);
                            return;
                        }
                    }
                }

                // Download or use early response
                let mut response = match early_returned_response {
                    Some(resp) => resp,
                    None => match downloader_clone.download(processed_request).await {
                        Ok(resp) => resp,
                        Err(e) => {
                            error!("Download error: {:?}", e);
                            return;
                        }
                    },
                };

                // Process response middlewares
                for mw in middlewares_clone.lock().await.iter_mut().rev() {
                    match mw.process_response(response.clone()).await {
                        Ok(MiddlewareAction::Continue(res)) => {
                            response = res;
                        }
                        Ok(MiddlewareAction::Retry(req, delay)) => {
                            tokio::time::sleep(delay).await;
                            if scheduler_clone.enqueue_request(*req).await.is_err() {
                                error!("Failed to re-enqueue retried request.");
                            }
                            return;
                        }
                        Ok(MiddlewareAction::Drop) => {
                            debug!("Response dropped by middleware.");
                            return;
                        }
                        Ok(MiddlewareAction::ReturnResponse(_)) => {
                            debug!("ReturnResponse action encountered in process_response; this is unexpected.");
                            continue;
                        }
                        Err(e) => {
                            error!("Response middleware error: {:?}", e);
                            return;
                        }
                    }
                }

                if res_tx_clone.send(response).await.is_err() {
                    error!("Response channel closed, cannot send parsed response.");
                }

                state_clone.in_flight_requests.fetch_sub(1, Ordering::SeqCst);
                drop(permit);
            });
        }
        while let Some(res) = tasks.join_next().await {
            if let Err(e) = res {
                error!("A download task failed: {:?}", e);
            }
        }
    })
}

fn spawn_parser_task<S>(
    scheduler: Arc<Scheduler>,
    spider: Arc<Mutex<S>>,
    state: Arc<CrawlerState>,
    res_rx: AsyncReceiver<Response>,
    item_tx: AsyncSender<S::Item>,
    parser_workers: usize,
    salvaged_requests_tx: AsyncSender<Request>,
) -> tokio::task::JoinHandle<()>
where
    S: Spider + 'static,
    S::Item: ScrapedItem,
{
    let mut tasks = JoinSet::new();
    let internal_parse_tx: AsyncSender<Response>;
    let internal_parse_rx: AsyncReceiver<Response>;
    (internal_parse_tx, internal_parse_rx) = bounded_async(parser_workers * 2);

    // Spawn N parsing worker tasks
    for _ in 0..parser_workers {
        let internal_parse_rx_clone = internal_parse_rx.clone();
        let spider_clone = Arc::clone(&spider);
        let scheduler_clone = Arc::clone(&scheduler);
        let item_tx_clone = item_tx.clone();
        let state_clone = Arc::clone(&state);
        let salvaged_requests_tx_clone = salvaged_requests_tx.clone();

        tasks.spawn(async move {
            while let Ok(response) = internal_parse_rx_clone.recv().await {
                debug!("Parsing response from {}", response.url);
                match spider_clone.lock().await.parse(response).await {
                    Ok(outputs) => {
                        process_crawl_outputs::<S>(
                            outputs,
                            scheduler_clone.clone(),
                            item_tx_clone.clone(),
                            state_clone.clone(),
                            salvaged_requests_tx_clone.clone(),
                        )
                        .await;
                    }
                    Err(e) => error!("Spider parsing error: {:?}", e),
                }
                state_clone.parsing_responses.fetch_sub(1, Ordering::SeqCst);
            }
        });
    }

    tokio::spawn(async move {
        while let Ok(response) = res_rx.recv().await {
            state.parsing_responses.fetch_add(1, Ordering::SeqCst);
            if internal_parse_tx.send(response).await.is_err() {
                error!("Internal parse channel closed, cannot send response to parser worker.");
                state.parsing_responses.fetch_sub(1, Ordering::SeqCst); // Decrement if send fails
            }
        }
        // Drop the sender to signal worker tasks to shut down
        drop(internal_parse_tx);

        // Wait for all parsing tasks to complete
        while let Some(res) = tasks.join_next().await {
            if let Err(e) = res {
                error!("A parsing worker task failed: {:?}", e);
            }
        }
    })
}

async fn process_crawl_outputs<S>(
    outputs: ParseOutput<S::Item>,
    scheduler: Arc<Scheduler>,
    item_tx: AsyncSender<S::Item>,
    state: Arc<CrawlerState>,
    salvaged_requests_tx: AsyncSender<Request>,
) where
    S: Spider + 'static,
    S::Item: ScrapedItem,
{
    let (items, requests) = outputs.into_parts();
    info!(
        "Processed {} requests and {} items from spider output.",
        requests.len(),
        items.len()
    );

    let mut request_error_total = 0;
    for request in requests {
        match scheduler.enqueue_request(request).await {
            Ok(_) => {}
            Err((req, _)) => {
                request_error_total += 1;
                if salvaged_requests_tx.send(req).await.is_err() {
                    error!(
                        "Failed to send salvaged request to channel from process_crawl_outputs."
                    );
                }
            }
        }
    }
    if request_error_total > 0 {
        error!(
            "Failed to enqueue {} requests: scheduler channel closed.",
            request_error_total
        );
    }

    let mut item_error_total = 0;
    for item in items {
        state.processing_items.fetch_add(1, Ordering::SeqCst);
        if item_tx.send(item).await.is_err() {
            item_error_total += 1;
            state.processing_items.fetch_sub(1, Ordering::SeqCst);
        }
    }
    if item_error_total > 0 {
        error!(
            "Failed to send {} scraped items: channel closed.",
            item_error_total
        );
    }
}

fn spawn_item_processor_task<S>(
    state: Arc<CrawlerState>,
    item_rx: AsyncReceiver<S::Item>,
    pipelines: Arc<Vec<Box<dyn Pipeline<S::Item>>>>,
    max_concurrent_pipelines: usize,
) -> tokio::task::JoinHandle<()>
where
    S: Spider + 'static,
    S::Item: ScrapedItem,
{
    let mut tasks = JoinSet::new();
    let semaphore = Arc::new(Semaphore::new(max_concurrent_pipelines));
    tokio::spawn(async move {
        while let Ok(item) = item_rx.recv().await {
            let permit = match semaphore.clone().acquire_owned().await {
                Ok(p) => p,
                Err(_) => {
                    warn!("Semaphore closed, shutting down item processor actor.");
                    break;
                }
            };

            let state_clone = Arc::clone(&state);
            let pipelines_clone = Arc::clone(&pipelines);
            tasks.spawn(async move {
                let mut item_to_process = Some(item);
                for pipeline in pipelines_clone.iter() {
                    if let Some(item) = item_to_process.take() {
                        match pipeline.process_item(item).await {
                            Ok(Some(next_item)) => item_to_process = Some(next_item),
                            Ok(None) => break, // Item was dropped
                            Err(e) => {
                                error!("Pipeline error for {}: {:?}", pipeline.name(), e);
                                break;
                            }
                        }
                    } else {
                        break;
                    }
                }
                state_clone.processing_items.fetch_sub(1, Ordering::SeqCst);
                drop(permit);
            });
        }
        while let Some(res) = tasks.join_next().await {
            if let Err(e) = res {
                error!("An item processing task failed: {:?}", e);
            }
        }
    })
}