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 dashmap::DashSet;
use kanal::{AsyncReceiver, AsyncSender, bounded_async, unbounded_async};
use std::collections::VecDeque;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use tokio::sync::Mutex;
#[cfg(feature = "checkpoint")]
use tokio::sync::oneshot;
use tracing::{debug, error, info};

enum SchedulerMessage {
    Enqueue(Box<Request>),
    MarkAsVisited(String),
    Shutdown,
    #[cfg(feature = "checkpoint")]
    TakeSnapshot(oneshot::Sender<SchedulerCheckpoint>),
}

pub struct Scheduler {
    request_queue: Arc<Mutex<VecDeque<Request>>>,
    visited_urls: DashSet<String>,
    tx_internal: AsyncSender<SchedulerMessage>,
    pending_requests: AtomicUsize,
    salvaged_requests: Arc<Mutex<VecDeque<Request>>>,
}

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(1);

        // Define these variables conditionally based on feature
        let request_queue: Arc<Mutex<VecDeque<Request>>>;
        let visited_urls: DashSet<String>;
        let pending_requests: AtomicUsize;
        let salvaged_requests: Arc<Mutex<VecDeque<Request>>>;

        #[cfg(feature = "checkpoint")]
        if let Some(state) = initial_state {
            // This branch is only compiled if feature is on
            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 = Arc::new(Mutex::new(state.request_queue));
            visited_urls = state.visited_urls;
            pending_requests = AtomicUsize::new(pending);
            salvaged_requests = Arc::new(Mutex::new(state.salvaged_requests));
        } else {
            // This branch is also only compiled if feature is on, but initial_state is None
            request_queue = Arc::new(Mutex::new(VecDeque::new()));
            visited_urls = DashSet::new();
            pending_requests = AtomicUsize::new(0);
            salvaged_requests = Arc::new(Mutex::new(VecDeque::new()));
        }

        #[cfg(not(feature = "checkpoint"))] // This entire block is only compiled if feature is off
        {
            // Always initialize empty when checkpoint is disabled.
            // initial_state is always None here, but we don't care about it.
            request_queue = Arc::new(Mutex::new(VecDeque::new()));
            visited_urls = DashSet::new();
            pending_requests = AtomicUsize::new(0);
            salvaged_requests = Arc::new(Mutex::new(VecDeque::new()));
        }

        let scheduler = Arc::new(Scheduler {
            request_queue,
            visited_urls,
            tx_internal,
            pending_requests,
            salvaged_requests,
        });

        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 {
            // Only try to pop if we know there are items
            let maybe_request = if !tx_req_out.is_closed() && !self.is_idle() {
                self.request_queue.lock().await.pop_front()
            } 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 the queue is empty, we must wait for a new message to arrive.
                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.lock().await.push_back(request);
                self.pending_requests.fetch_add(1, Ordering::SeqCst);
                true
            }
            Ok(SchedulerMessage::MarkAsVisited(fingerprint)) => {
                self.visited_urls.insert(fingerprint.clone());
                debug!("Marked URL as visited: {}", fingerprint);
                true
            }
            #[cfg(feature = "checkpoint")]
            Ok(SchedulerMessage::TakeSnapshot(responder)) => {
                let visited_urls = self.visited_urls.iter().map(|item| item.clone()).collect();
                let request_queue = self.request_queue.lock().await.clone();
                let salvaged_requests = self.salvaged_requests.lock().await.clone();

                let _ = responder.send(SchedulerCheckpoint {
                    request_queue,
                    visited_urls,
                    salvaged_requests,
                });
                true
            }
            Ok(SchedulerMessage::Shutdown) | Err(_) => {
                info!("Scheduler received shutdown signal or channel closed. Exiting run_loop.");
                false
            }
        }
    }

    /// Takes a snapshot of the current state of the scheduler.
    #[cfg(feature = "checkpoint")]
    pub async fn snapshot(&self) -> Result<SchedulerCheckpoint, SpiderError> {
        let (tx, rx) = oneshot::channel();
        self.tx_internal
            .send(SchedulerMessage::TakeSnapshot(tx))
            .await
            .map_err(|e| {
                SpiderError::GeneralError(format!(
                    "Scheduler: Failed to send snapshot request: {}",
                    e
                ))
            })?;
        rx.await.map_err(|e| {
            SpiderError::GeneralError(format!("Scheduler: Failed to receive snapshot: {}", e))
        })
    }

    /// Enqueues a new request to be processed.
    pub async fn enqueue_request(&self, request: Request) -> Result<(), SpiderError> {
        if self
            .tx_internal
            .send(SchedulerMessage::Enqueue(Box::new(request.clone()))) // Clone here for potential salvaging
            .await
            .is_err()
        {
            error!("Scheduler internal message channel is closed. Salvaging request.");
            self.salvaged_requests.lock().await.push_back(request); // Salvage internally
            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.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| {
                SpiderError::GeneralError(format!(
                    "Scheduler: Failed to send MarkAsVisited message: {}",
                    e
                ))
            })
    }

    /// 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()
    }
}