spider_lib/

engine.rs

use crate::downloader::Downloader;
use crate::error::SpiderError;
use crate::item::{CrawlOutput, ScrapedItem};
use crate::pipeline::ItemPipeline;
use crate::request::Request;
use crate::response::Response;
use crate::scheduler::Scheduler;
use crate::spider::Spider;
use crate::state::EngineState;
use anyhow::Result;
use futures_util::future::join_all;
use kanal::{AsyncReceiver, AsyncSender, bounded_async};
use std::sync::Arc;
use std::sync::atomic::Ordering;
use std::time::Duration;
use tokio::sync::Semaphore;
use tokio::task::JoinSet;
use tracing::{debug, error, info, instrument, warn};

pub struct Engine<S: Spider> {
    scheduler: Arc<Scheduler>,
    req_rx: AsyncReceiver<Request>,
    downloader: Arc<Downloader>,
    spider: Arc<S>,
    item_pipelines: Vec<Box<dyn ItemPipeline<S::Item>>>,
    max_concurrent_downloads: usize,
    #[allow(dead_code)]
    parser_workers: usize,
}

impl<S> Engine<S>
where
    S: Spider + 'static,
    S::Item: ScrapedItem,
{
    pub(crate) fn new(
        scheduler: Arc<Scheduler>,
        req_rx: AsyncReceiver<Request>,
        downloader: Arc<Downloader>,
        spider: S,
        item_pipelines: Vec<Box<dyn ItemPipeline<S::Item>>>,
        max_concurrent_downloads: usize,
        parser_workers: usize,
    ) -> Self {
        Engine {
            scheduler,
            req_rx,
            downloader,
            spider: Arc::new(spider),
            item_pipelines,
            max_concurrent_downloads,
            parser_workers,
        }
    }

    #[instrument(skip_all)]
    pub async fn start_crawl(self) -> Result<(), SpiderError> {
        info!("Engine starting crawl");

        let Engine {
            scheduler,
            req_rx,
            downloader,
            spider,
            item_pipelines,
            max_concurrent_downloads,
            parser_workers,
        } = self;

        let state = EngineState::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 initial_requests_task =
            spawn_initial_requests_task::<S>(scheduler.clone(), spider.clone());

        let downloader_actor = spawn_downloader_actor::<S>(
            req_rx,
            downloader,
            state.clone(),
            res_tx,
            max_concurrent_downloads,
        );

        let parser_actor = spawn_parser_actor::<S>(
            scheduler.clone(),
            spider.clone(),
            state.clone(),
            res_rx,
            item_tx,
            parser_workers,
        );

        let item_processor_actor =
            spawn_item_processor_actor::<S>(state.clone(), item_rx, pipelines.clone());

        monitor_crawl(scheduler.clone(), state.clone()).await;
        info!("Crawl has become idle, initiating shutdown.");

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

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

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

        item_processor_actor.await.map_err(|e| {
            SpiderError::InternalError(format!("Item processor actor task 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<S>,
) -> tokio::task::JoinHandle<()>
where
    S: Spider + 'static,
    S::Item: ScrapedItem,
{
    tokio::spawn(async move {
        for mut url in spider.start_urls() {
            url.set_fragment(None);
            let req = Request::new(url, spider.name(), "parse");
            if let Err(e) = scheduler.enqueue_request(req).await {
                error!("Failed to enqueue initial request: {}", e);
            }
        }
    })
}

async fn monitor_crawl(scheduler: Arc<Scheduler>, state: Arc<EngineState>) {
    loop {
        if scheduler.is_idle() && state.is_idle() {
            // Short delay to ensure no tasks are in the process of being created.
            tokio::time::sleep(Duration::from_millis(50)).await;
            if scheduler.is_idle() && state.is_idle() {
                break;
            }
        }
        tokio::time::sleep(Duration::from_millis(100)).await;
    }
}

fn spawn_downloader_actor<S>(
    req_rx: AsyncReceiver<Request>,
    downloader: Arc<Downloader>,
    state: Arc<EngineState>,
    res_tx: AsyncSender<Response>,
    max_concurrent_downloads: usize,
) -> tokio::task::JoinHandle<()>
where
    S: Spider + 'static,
    S::Item: ScrapedItem,
{
    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(p) => p,
                Err(_) => {
                    warn!("Semaphore closed, shutting down downloader actor.");
                    break;
                }
            };

            state.in_flight_requests.fetch_add(1, Ordering::SeqCst);
            let downloader_clone = downloader.clone();
            let res_tx_clone = res_tx.clone();
            let state_clone = state.clone();

            tasks.spawn(async move {
                debug!("Downloading {}", request.url);
                match downloader_clone.download(request).await {
                    Ok(response) => {
                        if res_tx_clone.send(response).await.is_err() {
                            error!("Response channel closed, cannot send parsed response.");
                        }
                    }
                    Err(e) => error!("Download error: {:?}", e),
                }
                state_clone
                    .in_flight_requests
                    .fetch_sub(1, Ordering::SeqCst);
                drop(permit);
            });
        }
        // Wait for all download tasks to complete
        while let Some(res) = tasks.join_next().await {
            if let Err(e) = res {
                error!("A download task failed: {:?}", e);
            }
        }
    })
}

fn spawn_parser_actor<S>(
    scheduler: Arc<Scheduler>,
    spider: Arc<S>,
    state: Arc<EngineState>,
    res_rx: AsyncReceiver<Response>,
    item_tx: AsyncSender<S::Item>,
    parser_workers: usize,
) -> 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 = spider.clone();
        let scheduler_clone = scheduler.clone();
        let item_tx_clone = item_tx.clone();
        let state_clone = state.clone();

        tasks.spawn(async move {
            while let Ok(response) = internal_parse_rx_clone.recv().await {
                debug!("Parsing response from {}", response.url);
                match spider_clone.parse(response).await {
                    Ok(outputs) => {
                        process_crawl_outputs::<S>(
                            outputs,
                            scheduler_clone.clone(),
                            item_tx_clone.clone(),
                            state_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: CrawlOutput<S::Item>,
    scheduler: Arc<Scheduler>,
    item_tx: AsyncSender<S::Item>,
    state: Arc<EngineState>,
) where
    S: Spider + 'static,
    S::Item: ScrapedItem,
{
    for request in outputs.requests {
        if let Err(e) = scheduler.enqueue_request(request).await {
            error!("Failed to enqueue request from parser: {}", e);
        }
    }
    for item in outputs.items {
        state.processing_items.fetch_add(1, Ordering::SeqCst);
        if item_tx.send(item).await.is_err() {
            error!("Item channel closed, cannot send scraped item.");
            state.processing_items.fetch_sub(1, Ordering::SeqCst); // Decrement if send fails
        }
    }
}

fn spawn_item_processor_actor<S>(
    state: Arc<EngineState>,
    item_rx: AsyncReceiver<S::Item>,
    pipelines: Arc<Vec<Box<dyn ItemPipeline<S::Item>>>>,
) -> tokio::task::JoinHandle<()>
where
    S: Spider + 'static,
    S::Item: ScrapedItem,
{
    let mut tasks = JoinSet::new();
    tokio::spawn(async move {
        while let Ok(item) = item_rx.recv().await {
            let state_clone = state.clone();
            let pipelines_clone = pipelines.clone();
            tasks.spawn(async move {
                debug!("Processing item...");
                let mut item_to_process = Some(item);
                for pipeline in pipelines_clone.iter() {
                    if let Some(it) = item_to_process.take() {
                        match pipeline.process_item(it).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);
            });
        }
        while let Some(res) = tasks.join_next().await {
            if let Err(e) = res {
                error!("An item processing task failed: {:?}", e);
            }
        }
    })
}