spider_lib/

builder.rs

//! Builder for constructing and configuring the `Crawler` instance.
//!
//! This module provides the `CrawlerBuilder`, a fluent API for
//! setting up and customizing a web crawler. It simplifies the process of
//! assembling various `spider-lib` components, including:
//! - Defining concurrency settings for downloads, parsing, and pipelines.
//! - Attaching custom `Downloader` implementations.
//! - Registering `Middleware`s to process requests and responses.
//! - Adding `Pipeline`s to process scraped items.
//! - Configuring checkpointing for persistence and fault tolerance.
//! - Initializing and integrating a `StatCollector` for gathering crawl statistics.
//!
//! The builder handles default configurations (e.g., adding a default User-Agent
//! middleware if none is specified) and loading existing checkpoints.

use crate::ConsoleWriterPipeline;
#[cfg(feature = "checkpoint")]
use crate::checkpoint::Checkpoint;
use crate::{CookieStore, SchedulerCheckpoint};

use crate::downloader::Downloader;
use crate::downloaders::reqwest_client::ReqwestClientDownloader;
use crate::error::SpiderError;
use crate::middleware::Middleware;
use crate::middlewares::user_agent::{UserAgentMiddleware, UserAgentSource};
use crate::pipeline::Pipeline;
use crate::scheduler::Scheduler;
use crate::spider::Spider;
use num_cpus;
#[cfg(feature = "middleware-cookies")]
use std::any::Any;
#[cfg(feature = "checkpoint")]
use std::fs;
use std::marker::PhantomData;
#[cfg(feature = "checkpoint")]
use std::path::{Path, PathBuf};
use std::sync::Arc;
#[cfg(feature = "checkpoint")]
use std::time::Duration;
#[cfg(feature = "middleware-cookies")]
use tracing::info;
#[cfg(feature = "checkpoint")]
use tracing::{debug, warn};

#[cfg(feature = "middleware-cookies")]
use crate::middlewares::cookies::CookieMiddleware;

#[cfg(feature = "middleware-cookies")]
use tokio::sync::Mutex;

use super::Crawler;
use crate::stats::StatCollector;

/// Configuration for the crawler's concurrency settings.
pub struct CrawlerConfig {
    /// The maximum number of concurrent downloads.
    pub max_concurrent_downloads: usize,
    /// The number of workers dedicated to parsing responses.
    pub parser_workers: usize,
    /// The maximum number of concurrent item processing pipelines.
    pub max_concurrent_pipelines: usize,
}

impl Default for CrawlerConfig {
    fn default() -> Self {
        CrawlerConfig {
            max_concurrent_downloads: 5,
            parser_workers: num_cpus::get(),
            max_concurrent_pipelines: 5,
        }
    }
}

pub struct CrawlerBuilder<S: Spider, D = ReqwestClientDownloader>
where
    D: Downloader,
{
    crawler_config: CrawlerConfig,
    downloader: D,
    spider: Option<S>,
    middlewares: Vec<Box<dyn Middleware<D::Client> + Send + Sync>>,
    item_pipelines: Vec<Box<dyn Pipeline<S::Item>>>,
    #[cfg(feature = "checkpoint")]
    checkpoint_path: Option<PathBuf>,
    #[cfg(feature = "checkpoint")]
    checkpoint_interval: Option<Duration>,
    _phantom: PhantomData<S>,
}

impl<S: Spider, D: Default + Downloader> Default for CrawlerBuilder<S, D> {
    fn default() -> Self {
        Self {
            crawler_config: CrawlerConfig::default(),
            downloader: D::default(),
            spider: None,
            middlewares: Vec::new(),
            item_pipelines: Vec::new(),
            #[cfg(feature = "checkpoint")]
            checkpoint_path: None,
            #[cfg(feature = "checkpoint")]
            checkpoint_interval: None,
            _phantom: PhantomData,
        }
    }
}

impl<S: Spider> CrawlerBuilder<S> {
    /// Creates a new `CrawlerBuilder` for a given spider.
    pub fn new(spider: S) -> Self {
        Self {
            spider: Some(spider),
            ..Default::default()
        }
    }
}

impl<S: Spider, D: Downloader> CrawlerBuilder<S, D> {
    /// Sets the maximum number of concurrent downloads.
    pub fn max_concurrent_downloads(mut self, limit: usize) -> Self {
        self.crawler_config.max_concurrent_downloads = limit;
        self
    }

    /// Sets the maximum number of concurrent parser workers.
    pub fn max_parser_workers(mut self, limit: usize) -> Self {
        self.crawler_config.parser_workers = limit;
        self
    }

    /// Sets the maximum number of concurrent pipelines.
    pub fn max_concurrent_pipelines(mut self, limit: usize) -> Self {
        self.crawler_config.max_concurrent_pipelines = limit;
        self
    }

    /// Sets a custom downloader for the crawler.
    pub fn downloader(mut self, downloader: D) -> Self {
        self.downloader = downloader;
        self
    }

    /// Adds a middleware to the crawler.
    pub fn add_middleware<M>(mut self, middleware: M) -> Self
    where
        D: Downloader,
        M: Middleware<D::Client> + Send + Sync + 'static,
    {
        self.middlewares.push(Box::new(middleware));
        self
    }

    /// Adds an item pipeline to the crawler.
    pub fn add_pipeline<P>(mut self, pipeline: P) -> Self
    where
        P: Pipeline<S::Item> + 'static,
    {
        self.item_pipelines.push(Box::new(pipeline));
        self
    }

    /// Enables checkpointing and sets the path for the checkpoint file.
    #[cfg(feature = "checkpoint")]
    pub fn with_checkpoint_path<P: AsRef<Path>>(mut self, path: P) -> Self {
        self.checkpoint_path = Some(path.as_ref().to_path_buf());
        self
    }

    /// Sets the interval for periodic checkpointing.
    #[cfg(feature = "checkpoint")]
    pub fn with_checkpoint_interval(mut self, interval: Duration) -> Self {
        self.checkpoint_interval = Some(interval);
        self
    }

    /// Builds the `Crawler` instance, initializing and passing the `StatCollector` along with other configured components.
    #[allow(unused_variables)]
    pub async fn build(mut self) -> Result<Crawler<S, D::Client>, SpiderError>
    where
        D: Downloader + Send + Sync + 'static,
        D::Client: Send + Sync + Clone,
    {
        let spider = self.validate_and_get_spider()?;
        self.ensure_default_components().await?;

        #[cfg(feature = "checkpoint")]
        let (initial_scheduler_state, loaded_cookie_store) =
            self.load_and_restore_checkpoint_state().await?;

        #[cfg(not(feature = "checkpoint"))]
        let (initial_scheduler_state, loaded_cookie_store): (
            Option<SchedulerCheckpoint>,
            Option<CookieStore>,
        ) = (None, None);

        let (scheduler_arc, req_rx) = Scheduler::new(initial_scheduler_state);

        #[cfg(feature = "middleware-cookies")]
        let final_cookie_store = self.resolve_final_cookie_store(loaded_cookie_store);

        let downloader_arc = Arc::new(self.downloader);
        let stats = Arc::new(StatCollector::new());

        let crawler = Crawler::new(
            scheduler_arc,
            req_rx,
            downloader_arc,
            self.middlewares,
            spider,
            self.item_pipelines,
            self.crawler_config.max_concurrent_downloads,
            self.crawler_config.parser_workers,
            self.crawler_config.max_concurrent_pipelines,
            #[cfg(feature = "checkpoint")]
            self.checkpoint_path.take(),
            #[cfg(feature = "checkpoint")]
            self.checkpoint_interval,
            stats,
            #[cfg(feature = "middleware-cookies")]
            final_cookie_store,
        );

        Ok(crawler)
    }

    #[cfg(feature = "middleware-cookies")]
    fn resolve_final_cookie_store(
        &self,
        loaded_cookie_store: Option<CookieStore>,
    ) -> Arc<Mutex<CookieStore>> {
        let mut final_cookie_store = Arc::new(Mutex::new(loaded_cookie_store.unwrap_or_default()));

        // Detect CookieMiddleware and use its store, overriding checkpoint or default
        for mw_box in &self.middlewares {
            if let Some(cookie_mw) =
                (mw_box.as_ref() as &dyn Any).downcast_ref::<CookieMiddleware>()
            {
                info!(
                    "Found CookieMiddleware, using its cookie store for Crawler. This overrides any checkpointed store."
                );
                final_cookie_store = cookie_mw.store.clone();
                break;
            }
        }
        final_cookie_store
    }

    #[cfg(feature = "checkpoint")]
    async fn load_and_restore_checkpoint_state(
        &mut self,
    ) -> Result<(Option<SchedulerCheckpoint>, Option<CookieStore>), SpiderError> {
        let mut initial_scheduler_state = None;
        let mut loaded_pipelines_state = None;
        #[cfg(feature = "middleware-cookies")]
        let mut loaded_cookie_store = None;

        if let Some(path) = &self.checkpoint_path {
            debug!("Attempting to load checkpoint from {:?}", path);
            match fs::read(path) {
                Ok(bytes) => match rmp_serde::from_slice::<Checkpoint>(&bytes) {
                    Ok(checkpoint) => {
                        initial_scheduler_state = Some(checkpoint.scheduler);
                        loaded_pipelines_state = Some(checkpoint.pipelines);

                        #[cfg(feature = "middleware-cookies")]
                        {
                            info!("Checkpoint loaded, restoring cookie store data.");
                            loaded_cookie_store = Some(checkpoint.cookie_store);
                        }
                    }
                    Err(e) => warn!("Failed to deserialize checkpoint from {:?}: {}", path, e),
                },
                Err(e) => warn!("Failed to read checkpoint file {:?}: {}", path, e),
            }
        }

        if let Some(pipeline_states) = loaded_pipelines_state {
            for (name, state) in pipeline_states {
                if let Some(pipeline) = self.item_pipelines.iter().find(|p| p.name() == name) {
                    pipeline.restore_state(state).await?;
                } else {
                    warn!("Checkpoint contains state for unknown pipeline: {}", name);
                }
            }
        }

        #[cfg(not(feature = "middleware-cookies"))]
        return Ok((initial_scheduler_state, None));
        #[cfg(feature = "middleware-cookies")]
        return Ok((initial_scheduler_state, loaded_cookie_store));
    }

    fn validate_and_get_spider(&mut self) -> Result<S, SpiderError> {
        if self.crawler_config.max_concurrent_downloads == 0 {
            return Err(SpiderError::ConfigurationError(
                "max_concurrent_downloads must be greater than 0.".to_string(),
            ));
        }
        if self.crawler_config.parser_workers == 0 {
            return Err(SpiderError::ConfigurationError(
                "parser_workers must be greater than 0.".to_string(),
            ));
        }
        self.spider.take().ok_or_else(|| {
            SpiderError::ConfigurationError("Crawler must have a spider.".to_string())
        })
    }

    async fn ensure_default_components(&mut self) -> Result<(), SpiderError> {
        if self.item_pipelines.is_empty() {
            self.item_pipelines
                .push(Box::new(ConsoleWriterPipeline::new()));
        }

        let has_user_agent_middleware = self
            .middlewares
            .iter()
            .any(|m| m.name() == "UserAgentMiddleware");

        if !has_user_agent_middleware {
            let pkg_name = env!("CARGO_PKG_NAME");
            let pkg_version = env!("CARGO_PKG_VERSION");
            let default_user_agent = format!("{}/{}", pkg_name, pkg_version);

            let default_user_agent_mw = UserAgentMiddleware::builder()
                .source(UserAgentSource::List(vec![default_user_agent.clone()]))
                .fallback_user_agent(default_user_agent)
                .build()?;
            self.middlewares.insert(0, Box::new(default_user_agent_mw));
        }

        Ok(())
    }
}