spider_lib/

scheduler.rs

use crate::checkpoint::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, oneshot};
use tracing::{debug, error, info};

enum SchedulerMessage {
    Enqueue(Box<Request>),
    Shutdown,
    TakeSnapshot(oneshot::Sender<SchedulerCheckpoint>),
}

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

impl Scheduler {
    /// Creates a new `Scheduler` and returns a tuple containing the scheduler and a request receiver.
    pub fn new(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);

        let (request_queue, visited_urls, pending_requests) = if let Some(state) = initial_state {
            info!(
                "Initializing scheduler from checkpoint with {} requests and {} visited URLs.",
                state.request_queue.len(),
                state.visited_urls.len()
            );
            let pending = state.request_queue.len();
            (
                Arc::new(Mutex::new(state.request_queue)),
                state.visited_urls,
                AtomicUsize::new(pending),
            )
        } else {
            (
                Arc::new(Mutex::new(VecDeque::new())),
                DashSet::new(),
                AtomicUsize::new(0),
            )
        };

        let scheduler = Arc::new(Scheduler {
            request_queue,
            visited_urls,
            tx_internal,
            pending_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;
                let fingerprint = request.fingerprint();

                if self.visited_urls.insert(fingerprint.clone()) {
                    self.request_queue.lock().await.push_back(request);
                    self.pending_requests.fetch_add(1, Ordering::SeqCst);
                } else {
                    debug!(
                        "Skipping already visited URL: {} (fingerprint: {})",
                        request.url, fingerprint
                    );
                }
                true
            }
            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 _ = responder.send(SchedulerCheckpoint {
                    request_queue,
                    visited_urls,
                });
                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.
    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<(), (Request, SpiderError)> {
        let original_request = request.clone();

        self.tx_internal
            .send(SchedulerMessage::Enqueue(Box::new(request)))
            .await
            .map_err(|e| {
                (
                    original_request,
                    SpiderError::GeneralError(format!(
                        "Scheduler: Failed to enqueue request: {}",
                        e
                    )),
                )
            })
    }

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

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