spider_lib/

scheduler.rs

//! Request Scheduler for managing the crawling frontier.
//!
//! This module defines the `Scheduler` component, which is central to coordinating
//! the web crawling process within `spider-lib`. The `Scheduler` is responsible for:
//! - Maintaining a queue of pending `Request`s to be fetched.
//! - Keeping track of `visited_urls` to avoid re-crawling content.
//! - Handling the salvaging of requests that failed to enqueue, ensuring they are not lost.
//! - Providing snapshot capabilities for checkpointing the crawler's state, enabling
//!   fault tolerance and resumption of crawls.
//!
//! It operates asynchronously, processing messages for request enqueuing, URL marking,
//! and state snapshots through an internal actor-like mechanism.

use crate::SchedulerCheckpoint;

use crate::error::SpiderError;
use crate::request::Request;
use crate::utils::BloomFilter;
use crossbeam::queue::SegQueue;
use kanal::{AsyncReceiver, AsyncSender, bounded_async, unbounded_async};
use moka::sync::Cache;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use tracing::{debug, error, info};

enum SchedulerMessage {
    Enqueue(Box<Request>),
    MarkAsVisited(String),
    Shutdown,
}

pub struct Scheduler {
    request_queue: SegQueue<Request>,
    visited_urls: Cache<String, bool>,
    bloom_filter: BloomFilter,
    tx_internal: AsyncSender<SchedulerMessage>,
    pending_requests: AtomicUsize,
    salvaged_requests: SegQueue<Request>,
    pub(crate) is_shutting_down: AtomicBool,
}

impl Scheduler {
    /// Creates a new `Scheduler` and returns a tuple containing the scheduler and a request receiver.
    pub fn new(
        #[cfg(feature = "checkpoint")] initial_state: Option<SchedulerCheckpoint>,
        #[cfg(not(feature = "checkpoint"))] _initial_state: Option<SchedulerCheckpoint>,
    ) -> (Arc<Self>, AsyncReceiver<Request>) {
        let (tx_internal, rx_internal) = unbounded_async();

        let (tx_req_out, rx_req_out) = bounded_async(100);

        let request_queue: SegQueue<Request>;
        let visited_urls: Cache<String, bool>;
        let pending_requests: AtomicUsize;
        let salvaged_requests: SegQueue<Request>;

        #[cfg(feature = "checkpoint")]
        if let Some(state) = initial_state {
            info!(
                "Initializing scheduler from checkpoint with {} requests, {} visited URLs, and {} salvaged requests.",
                state.request_queue.len(),
                state.visited_urls.len(),
                state.salvaged_requests.len(),
            );
            let pending = state.request_queue.len() + state.salvaged_requests.len();
            request_queue = SegQueue::new();
            for request in state.request_queue {
                request_queue.push(request);
            }

            visited_urls = Cache::builder().max_capacity(100000).build();
            for url in state.visited_urls {
                visited_urls.insert(url, true);
            }

            pending_requests = AtomicUsize::new(pending);
            salvaged_requests = SegQueue::new();
            for request in state.salvaged_requests {
                salvaged_requests.push(request);
            }
        } else {
            request_queue = SegQueue::new();
            visited_urls = Cache::builder().max_capacity(100000).build();
            pending_requests = AtomicUsize::new(0);
            salvaged_requests = SegQueue::new();
        }

        #[cfg(not(feature = "checkpoint"))]
        {
            request_queue = SegQueue::new();
            visited_urls = Cache::builder().max_capacity(100000).build();
            pending_requests = AtomicUsize::new(0);
            salvaged_requests = SegQueue::new();
        }

        let scheduler = Arc::new(Scheduler {
            request_queue,
            visited_urls,
            bloom_filter: BloomFilter::new(1000000, 3),
            tx_internal,
            pending_requests,
            salvaged_requests,
            is_shutting_down: AtomicBool::new(false),
        });

        let scheduler_clone = Arc::clone(&scheduler);
        tokio::spawn(async move {
            scheduler_clone.run_loop(rx_internal, tx_req_out).await;
        });

        (scheduler, rx_req_out)
    }

    async fn run_loop(
        &self,
        rx_internal: AsyncReceiver<SchedulerMessage>,
        tx_req_out: AsyncSender<Request>,
    ) {
        info!("Scheduler run_loop started.");
        loop {
            let maybe_request = if !tx_req_out.is_closed() && !self.is_idle() {
                self.request_queue.pop()
            } else {
                None
            };

            if let Some(request) = maybe_request {
                tokio::select! {
                    biased;
                    send_res = tx_req_out.send(request) => {
                        if send_res.is_err() {
                            error!("Crawler receiver dropped. Scheduler can no longer send requests.");
                        }
                        self.pending_requests.fetch_sub(1, Ordering::SeqCst);
                    },
                    recv_res = rx_internal.recv() => {
                        self.pending_requests.fetch_sub(1, Ordering::SeqCst);
                        if !self.handle_message(recv_res).await {
                            break;
                        }
                    }
                }
            } else if !self.handle_message(rx_internal.recv().await).await {
                break;
            }
        }
        info!("Scheduler run_loop finished.");
    }

    async fn handle_message(&self, msg: Result<SchedulerMessage, kanal::ReceiveError>) -> bool {
        match msg {
            Ok(SchedulerMessage::Enqueue(boxed_request)) => {
                let request = *boxed_request;
                self.request_queue.push(request);
                self.pending_requests.fetch_add(1, Ordering::SeqCst);
                true
            }
            Ok(SchedulerMessage::MarkAsVisited(fingerprint)) => {
                self.visited_urls.insert(fingerprint.clone(), true);
                // Also add to Bloom Filter for faster preliminary checks
                self.bloom_filter.add(&fingerprint);
                debug!("Marked URL as visited: {}", fingerprint);
                true
            }
            Ok(SchedulerMessage::Shutdown) | Err(_) => {
                info!("Scheduler received shutdown signal or channel closed. Exiting run_loop.");
                self.is_shutting_down.store(true, Ordering::SeqCst);
                false
            }
        }
    }

    /// Takes a snapshot of the current state of the scheduler.
    /// Note: This temporarily moves elements and restores them, which may affect concurrent access.
    #[cfg(feature = "checkpoint")]
    pub async fn snapshot(&self) -> Result<SchedulerCheckpoint, SpiderError> {
        let visited_urls = dashmap::DashSet::new();
        for entry in self.visited_urls.iter() {
            let (key, _) = entry;
            visited_urls.insert(key.as_ref().clone());
        }

        let mut request_queue = std::collections::VecDeque::new();
        let mut temp_requests = Vec::new();
        while let Some(request) = self.request_queue.pop() {
            temp_requests.push(request);
        }

        for request in temp_requests.into_iter().rev() {
            request_queue.push_back(request.clone());
            self.request_queue.push(request);
        }

        let mut salvaged_requests = std::collections::VecDeque::new();
        let mut temp_salvaged = Vec::new();
        while let Some(request) = self.salvaged_requests.pop() {
            temp_salvaged.push(request);
        }

        for request in temp_salvaged.into_iter().rev() {
            salvaged_requests.push_back(request.clone());
            self.salvaged_requests.push(request);
        }

        Ok(SchedulerCheckpoint {
            request_queue,
            visited_urls,
            salvaged_requests,
        })
    }

    /// Enqueues a new request to be processed.
    pub async fn enqueue_request(&self, request: Request) -> Result<(), SpiderError> {
        if !self.should_enqueue_request(&request) {
            debug!("Request already visited, skipping: {}", request.url);
            return Ok(());
        }

        if self
            .tx_internal
            .send(SchedulerMessage::Enqueue(Box::new(request.clone())))
            .await
            .is_err()
        {
            if !self.is_shutting_down.load(Ordering::SeqCst) {
                error!("Scheduler internal message channel is closed. Salvaging request.");
            }
            self.salvaged_requests.push(request);
            return Err(SpiderError::GeneralError(
                "Scheduler internal channel closed, request salvaged.".into(),
            ));
        }
        Ok(())
    }

    /// Sends a shutdown signal to the scheduler.
    pub async fn shutdown(&self) -> Result<(), SpiderError> {
        self.is_shutting_down.store(true, Ordering::SeqCst);

        self.tx_internal
            .send(SchedulerMessage::Shutdown)
            .await
            .map_err(|e| {
                SpiderError::GeneralError(format!(
                    "Scheduler: Failed to send shutdown signal: {}",
                    e
                ))
            })
    }

    /// Sends a message to the scheduler to mark a URL as visited.
    pub async fn send_mark_as_visited(&self, fingerprint: String) -> Result<(), SpiderError> {
        self.tx_internal
            .send(SchedulerMessage::MarkAsVisited(fingerprint))
            .await
            .map_err(|e| {

                if !self.is_shutting_down.load(Ordering::SeqCst) {
                    error!("Scheduler internal message channel is closed. Failed to mark URL as visited: {}", e);
                }
                SpiderError::GeneralError(format!(
                    "Scheduler: Failed to send MarkAsVisited message: {}",
                    e
                ))
            })
    }

    /// Checks if a URL has been visited using Bloom Filter for fast preliminary check
    pub fn has_been_visited(&self, fingerprint: &str) -> bool {
        if !self.bloom_filter.might_contain(&fingerprint) {
            return false;
        }

        self.visited_urls.contains_key(fingerprint)
    }

    /// Checks if a request should be enqueued by checking if its fingerprint has already been visited
    pub fn should_enqueue_request(&self, request: &Request) -> bool {
        let fingerprint = request.fingerprint();
        !self.has_been_visited(&fingerprint)
    }

    /// Returns the number of pending requests in the scheduler.
    #[inline]
    pub fn len(&self) -> usize {
        self.pending_requests.load(Ordering::SeqCst)
    }

    /// Checks if the scheduler has no pending requests.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Checks if the scheduler is idle (has no pending requests).
    #[inline]
    pub fn is_idle(&self) -> bool {
        self.is_empty()
    }
}