spider_core/

scheduler.rs

//! # Scheduler Module
//!
//! Implements the request scheduler for managing the crawling frontier and duplicate detection.
//!
//! ## Overview
//!
//! The `Scheduler` is a central component that coordinates the web crawling process
//! by managing the queue of pending requests and tracking visited URLs to prevent
//! duplicate processing. It uses an actor-like design pattern with internal message
//! processing for thread-safe operations.
//!
//! ## Key Responsibilities
//!
//! - **Request Queue Management**: Maintains a queue of pending requests to be processed
//! - **Duplicate Detection**: Tracks visited URLs using Bloom Filter and LRU cache for efficiency
//! - **Request Salvaging**: Handles failed enqueuing attempts to prevent request loss
//! - **State Snapshots**: Provides checkpointing capabilities for crawl resumption
//! - **Concurrent Access**: Thread-safe operations for multi-threaded crawling
//!
//! ## Architecture
//!
//! The scheduler operates asynchronously using an internal message queue to handle
//! operations like request enqueuing, URL marking, and state snapshots. It combines
//! a Bloom Filter for fast preliminary duplicate checks with an LRU cache for
//! definitive tracking, optimizing performance when handling millions of URLs.
//!
//! ## Example
//!
//! ```rust,ignore
//! use spider_core::Scheduler;
//! use spider_util::request::Request;
//! use url::Url;
//!
//! let (scheduler, request_receiver) = Scheduler::new(None);
//!
//! // Enqueue a request
//! let request = Request::new(Url::parse("https://example.com").unwrap());
//! scheduler.enqueue_request(request).await?;
//!
//! // Mark a URL as visited
//! scheduler.send_mark_as_visited("unique_fingerprint".to_string()).await?;
//! ```

#[cfg(feature = "checkpoint")]
use spider_util::constants::DEFAULT_VISITED_CACHE_SIZE;
#[cfg(feature = "checkpoint")]
use crate::SchedulerCheckpoint;

#[cfg(not(feature = "checkpoint"))]
use spider_util::constants::MAX_PENDING_REQUESTS;

use spider_util::constants::{
    BLOOM_FILTER_CAPACITY, BLOOM_FILTER_HASH_FUNCTIONS,
    VISITED_URL_CACHE_CAPACITY, VISITED_URL_CACHE_TTL_SECS,
};
use spider_util::error::SpiderError;
use spider_util::request::Request;
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 log::{debug, error, info, trace, warn};

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

use spider_util::bloom::BloomFilter;

use tokio::sync::Notify;

pub struct Scheduler {
    queue: SegQueue<Request>,
    visited: Cache<String, bool>,
    bloom: std::sync::Arc<parking_lot::RwLock<BloomFilter>>,
    buffer: Arc<std::sync::Mutex<Vec<String>>>,
    notify: Arc<Notify>,
    tx: AsyncSender<SchedulerMessage>,
    pending: AtomicUsize,
    salvaged: SegQueue<Request>,
    pub(crate) is_shutting_down: AtomicBool,
    max_pending: usize,
}

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

        let queue: SegQueue<Request>;
        let visited: Cache<String, bool>;
        let pending: AtomicUsize;
        let salvaged: SegQueue<Request>;

        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 pend = state.request_queue.len() + state.salvaged_requests.len();
            queue = SegQueue::new();
            for request in state.request_queue {
                queue.push(request);
            }

            visited = Cache::builder()
                .max_capacity(VISITED_URL_CACHE_CAPACITY)
                .time_to_idle(std::time::Duration::from_secs(VISITED_URL_CACHE_TTL_SECS))
                .eviction_listener(|_key, _value, _cause| {})
                .build();
            for url in state.visited_urls {
                visited.insert(url, true);
            }

            pending = AtomicUsize::new(pend);
            salvaged = SegQueue::new();
            for request in state.salvaged_requests {
                salvaged.push(request);
            }
        } else {
            queue = SegQueue::new();
            visited = Cache::builder().max_capacity(DEFAULT_VISITED_CACHE_SIZE).build();
            pending = AtomicUsize::new(0);
            salvaged = SegQueue::new();
        }

        let buffer = Arc::new(std::sync::Mutex::new(Vec::new()));
        let notify = Arc::new(Notify::new());

        let scheduler = Arc::new(Scheduler {
            queue,
            visited,
            bloom: std::sync::Arc::new(parking_lot::RwLock::new(BloomFilter::new(
                BLOOM_FILTER_CAPACITY,
                BLOOM_FILTER_HASH_FUNCTIONS,
            ))),
            buffer: buffer.clone(),
            notify: notify.clone(),
            tx,
            pending,
            salvaged,
            is_shutting_down: AtomicBool::new(false),
            max_pending: 30000,
        });

        let scheduler_bloom = Arc::clone(&scheduler);
        let buffer_clone = buffer.clone();
        let notify_clone = notify.clone();
        tokio::spawn(async move {
            scheduler_bloom.flush_buffer(buffer_clone, notify_clone).await;
        });

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

        (scheduler, rx_out)
    }

    #[cfg(not(feature = "checkpoint"))]
    pub fn new(
        _initial_state: Option<()>,
    ) -> (Arc<Self>, AsyncReceiver<Request>) {
        let (tx, rx_internal) = unbounded_async();
        let (tx_out, rx_out) = bounded_async(100);

        let queue = SegQueue::new();
        let visited = Cache::builder()
            .max_capacity(VISITED_URL_CACHE_CAPACITY)
            .time_to_idle(std::time::Duration::from_secs(VISITED_URL_CACHE_TTL_SECS))
            .eviction_listener(|_key, _value, _cause| {})
            .build();
        let pending = AtomicUsize::new(0);
        let salvaged = SegQueue::new();

        let buffer = Arc::new(std::sync::Mutex::new(Vec::new()));
        let notify = Arc::new(Notify::new());

        let scheduler = Arc::new(Scheduler {
            queue,
            visited,
            bloom: std::sync::Arc::new(parking_lot::RwLock::new(BloomFilter::new(
                BLOOM_FILTER_CAPACITY,
                BLOOM_FILTER_HASH_FUNCTIONS,
            ))),
            buffer: buffer.clone(),
            notify: notify.clone(),
            tx,
            pending,
            salvaged,
            is_shutting_down: AtomicBool::new(false),
            max_pending: MAX_PENDING_REQUESTS,
        });

        let scheduler_bloom = Arc::clone(&scheduler);
        let buffer_clone = buffer.clone();
        let notify_clone = notify.clone();
        tokio::spawn(async move {
            scheduler_bloom.flush_buffer(buffer_clone, notify_clone).await;
        });

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

        (scheduler, rx_out)
    }

    async fn run_loop(
        &self,
        rx_internal: AsyncReceiver<SchedulerMessage>,
        tx_out: AsyncSender<Request>,
    ) {
        info!(
            "Scheduler run_loop started with max pending: {}",
            self.max_pending
        );
        loop {
            if let Ok(Some(msg)) = rx_internal.try_recv() {
                trace!("Processing pending internal message");
                if !self.handle_message(Ok(msg)).await {
                    break;
                }
                continue;
            }

            let request = if !tx_out.is_closed() && !self.is_idle() {
                self.queue.pop()
            } else {
                None
            };

            if let Some(request) = request {
                trace!("Sending request to crawler: {}", request.url);
                tokio::select! {
                    send_res = tx_out.send(request) => {
                        if send_res.is_err() {
                            error!("Crawler receiver dropped. Scheduler can no longer send requests.");
                        } else {
                            trace!("Successfully sent request to crawler");
                        }
                        self.pending.fetch_sub(1, Ordering::SeqCst);
                    },
                    recv_res = rx_internal.recv() => {
                        trace!("Received internal message while sending request");
                        if !self.handle_message(recv_res).await {
                            break;
                        }
                        continue;
                    }
                }
            } else {
                trace!("No pending requests, waiting for internal message");
                if !self.handle_message(rx_internal.recv().await).await {
                    break;
                }
            }
        }
        info!(
            "Scheduler run_loop finished with {} pending requests remaining.",
            self.pending.load(Ordering::SeqCst)
        );
    }

    async fn handle_message(&self, msg: Result<SchedulerMessage, kanal::ReceiveError>) -> bool {
        match msg {
            Ok(SchedulerMessage::Enqueue(boxed_request)) => {
                let request = *boxed_request;
                trace!("Enqueuing request: {}", request.url);
                self.queue.push(request);
                self.pending.fetch_add(1, Ordering::SeqCst);
                true
            }
            Ok(SchedulerMessage::MarkAsVisited(fingerprint)) => {
                trace!("Marking URL fingerprint as visited: {}", fingerprint);

                // Insert into visited cache first (clone needed for cache)
                self.visited.insert(fingerprint.clone(), true);

                // Log before moving fingerprint
                debug!("Marked URL as visited: {}", fingerprint);

                // Then move fingerprint into buffer (no clone needed)
                {
                    let mut buffer = self.buffer.lock().unwrap();
                    buffer.push(fingerprint);
                    if buffer.len() >= 100 {
                        self.notify.notify_one();
                    }
                }

                true
            }
            Ok(SchedulerMessage::MarkAsVisitedBatch(mut fingerprints)) => {
                let count = fingerprints.len();
                trace!("Marking {} URL fingerprints as visited in batch", count);
                
                // Insert all fingerprints into visited cache
                for fingerprint in &fingerprints {
                    self.visited.insert(fingerprint.clone(), true);
                }

                // Then extend buffer with the fingerprints (no clone needed)
                {
                    let mut buffer = self.buffer.lock().unwrap();
                    buffer.append(&mut fingerprints);
                    if buffer.len() >= 100 {
                        self.notify.notify_one();
                    }
                }

                debug!("Marked {} URLs as visited in batch", count);
                true
            }
            Ok(SchedulerMessage::Shutdown) => {
                info!("Scheduler received shutdown signal. Exiting run_loop.");
                self.is_shutting_down.store(true, Ordering::SeqCst);
                self.flush_buffer_now();
                false
            }
            Err(_) => {
                warn!("Scheduler internal message channel closed. Exiting run_loop.");
                self.is_shutting_down.store(true, Ordering::SeqCst);
                false
            }
        }
    }

    #[cfg(feature = "checkpoint")]
    pub async fn snapshot(&self) -> Result<SchedulerCheckpoint, SpiderError> {
        let visited_urls = dashmap::DashSet::new();
        for entry in self.visited.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.queue.pop() {
            temp_requests.push(request);
        }

        for request in temp_requests.into_iter() {
            request_queue.push_back(request.clone());
            if !self.is_shutting_down.load(Ordering::SeqCst) {
                self.queue.push(request);
            }
        }

        let mut salvaged_requests = std::collections::VecDeque::new();
        let mut temp_salvaged = Vec::new();

        while let Some(request) = self.salvaged.pop() {
            temp_salvaged.push(request);
        }

        for request in temp_salvaged.into_iter() {
            salvaged_requests.push_back(request.clone());
            if !self.is_shutting_down.load(Ordering::SeqCst) {
                self.salvaged.push(request);
            }
        }

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

    #[cfg(not(feature = "checkpoint"))]
    pub async fn snapshot(&self) -> Result<(), SpiderError> {
        Ok(())
    }

    pub async fn enqueue_request(&self, request: Request) -> Result<(), SpiderError> {
        if !self.should_enqueue(&request) {
            trace!("Request already visited, skipping: {}", request.url);
            return Ok(());
        }

        let pending = self.pending.load(Ordering::SeqCst);
        if pending >= self.max_pending {
            warn!(
                "Maximum pending requests reached ({}), request dropped due to backpressure: {}",
                self.max_pending, request.url
            );
            return Err(SpiderError::GeneralError(
                "Scheduler at maximum capacity, request dropped due to backpressure.".into(),
            ));
        }

        trace!("Enqueuing request: {}", request.url);
        if self
            .tx
            .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: {}",
                    request.url
                );
            }
            self.salvaged.push(request);
            return Err(SpiderError::GeneralError(
                "Scheduler internal channel closed, request salvaged.".into(),
            ));
        }

        trace!("Successfully enqueued request: {}", request.url);
        Ok(())
    }

    pub async fn shutdown(&self) -> Result<(), SpiderError> {
        self.is_shutting_down.store(true, Ordering::SeqCst);

        if !self.tx.is_closed() {
            self.tx
                .send(SchedulerMessage::Shutdown)
                .await
                .map_err(|e| {
                    SpiderError::GeneralError(format!(
                        "Scheduler: Failed to send shutdown signal: {}",
                        e
                    ))
                })
        } else {
            debug!("Scheduler internal channel already closed, skipping shutdown signal");
            Ok(())
        }
    }

    pub async fn mark_visited(&self, fingerprint: String) -> Result<(), SpiderError> {
        trace!(
            "Sending MarkAsVisited message for fingerprint: {}",
            fingerprint
        );
        self.tx
            .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
                ))
            })
    }

    pub async fn mark_visited_batch(&self, fingerprints: Vec<String>) -> Result<(), SpiderError> {
        if fingerprints.is_empty() {
            return Ok(());
        }

        trace!(
            "Sending MarkAsVisitedBatch message for {} fingerprints",
            fingerprints.len()
        );
        self.tx
            .send(SchedulerMessage::MarkAsVisitedBatch(fingerprints))
            .await
            .map_err(|e| {
                if !self.is_shutting_down.load(Ordering::SeqCst) {
                    error!("Scheduler internal message channel is closed. Failed to mark URLs as visited in batch: {}", e);
                }
                SpiderError::GeneralError(format!(
                    "Scheduler: Failed to send MarkAsVisitedBatch message: {}",
                    e
                ))
            })
    }

    pub fn is_visited(&self, fingerprint: &str) -> bool {
        if !self.bloom.read().might_contain(fingerprint) {
            return false;
        }

        {
            let buffer = self.buffer.lock().unwrap();
            if buffer.iter().any(|item| item == fingerprint) {
                return true;
            }
        }

        self.visited.contains_key(fingerprint)
    }

    fn flush_buffer_now(&self) {
        let mut buffer = self.buffer.lock().unwrap();
        if !buffer.is_empty() {
            let items: Vec<String> = buffer.drain(..).collect();
            drop(buffer);

            let mut bloom = self.bloom.write();
            for item in items {
                bloom.add(&item);
            }
        }
    }

    async fn flush_buffer(
        &self,
        _buffer: Arc<std::sync::Mutex<Vec<String>>>,
        notify: Arc<Notify>,
    ) {
        loop {
            tokio::select! {
                _ = notify.notified() => {
                    self.flush_buffer_now();
                }
                _ = tokio::time::sleep(tokio::time::Duration::from_millis(100)) => {
                    self.flush_buffer_now();
                }
            }

            tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
        }
    }

    pub fn should_enqueue(&self, request: &Request) -> bool {
        let fingerprint = request.fingerprint();
        !self.is_visited(&fingerprint)
    }

    #[inline]
    pub fn len(&self) -> usize {
        self.pending.load(Ordering::SeqCst)
    }

    #[inline]
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    #[inline]
    pub fn is_idle(&self) -> bool {
        self.is_empty()
    }
}