use std::{
collections::HashMap,
sync::Arc,
time::{Duration, Instant},
};
use rand::Rng;
use tokio::sync::Mutex;
use crate::{
frontier::Frontier,
request::{CrawlRequest, FrontierRequest},
};
use super::{domain::domain_key, fingerprint::FingerprintPolicy, policy::PolitenessPolicy};
#[derive(Debug, Default)]
struct DomainState {
in_flight: usize,
next_available_at: Option<Instant>,
robots_delay: Option<Duration>,
}
pub struct CrawlScheduler {
frontier: Arc<dyn Frontier>,
policy: PolitenessPolicy,
fingerprint_policy: FingerprintPolicy,
domains: Mutex<HashMap<String, DomainState>>,
}
pub(crate) enum SchedulerPoll {
Ready(Box<FrontierRequest>),
Pending(Duration),
Empty,
}
enum CandidateState {
Ready,
Pending(Duration),
}
fn delay_with_jitter(base: Duration, jitter: Option<Duration>) -> Duration {
let Some(jitter) = jitter else {
return base;
};
if jitter.is_zero() {
return base;
}
let extra = rand::rng().random_range(Duration::ZERO..=jitter);
base.saturating_add(extra)
}
impl CrawlScheduler {
pub fn new(frontier: impl Frontier + 'static, policy: PolitenessPolicy) -> Self {
Self::from_arc(Arc::new(frontier), policy)
}
pub fn from_arc(frontier: Arc<dyn Frontier>, policy: PolitenessPolicy) -> Self {
Self {
frontier,
policy,
fingerprint_policy: FingerprintPolicy::default(),
domains: Mutex::new(HashMap::new()),
}
}
pub fn with_fingerprint_policy(mut self, policy: FingerprintPolicy) -> Self {
self.fingerprint_policy = policy;
self
}
pub async fn push_request(&self, request: CrawlRequest, depth: usize) -> bool {
self.frontier
.push_request(self.apply_fingerprint(request), depth)
.await
}
pub async fn push_request_force(&self, queued: FrontierRequest) {
self.frontier.push_request_force(queued).await;
}
pub async fn is_empty(&self) -> bool {
self.frontier.is_empty().await
}
pub async fn flush(&self) -> Result<(), crate::error::KumoError> {
self.frontier.flush().await
}
pub async fn try_next_ready(&self) -> Option<FrontierRequest> {
match self.poll_next().await {
SchedulerPoll::Ready(queued) => Some(*queued),
SchedulerPoll::Pending(_) | SchedulerPoll::Empty => None,
}
}
pub(crate) async fn poll_ready(&self) -> SchedulerPoll {
self.poll_next().await
}
pub async fn next_ready(&self) -> Option<FrontierRequest> {
loop {
match self.poll_next().await {
SchedulerPoll::Ready(queued) => return Some(*queued),
SchedulerPoll::Pending(wait) => tokio::time::sleep(wait).await,
SchedulerPoll::Empty => return None,
}
}
}
pub async fn finish(&self, queued: &FrontierRequest) {
let Some(domain) = domain_key(queued.request().url()) else {
return;
};
let mut domains = self.domains.lock().await;
let state = domains.entry(domain.clone()).or_default();
state.in_flight = state.in_flight.saturating_sub(1);
let policy_delay = self.policy.policy_for(&domain).delay();
let robots_delay = if self.policy.respects_robots_crawl_delay() {
state.robots_delay
} else {
None
};
if let Some(delay) = [policy_delay, robots_delay].into_iter().flatten().max() {
let delay = delay_with_jitter(delay, self.policy.jitter_range());
state.next_available_at = Some(Instant::now() + delay);
}
}
pub async fn observe_robots_crawl_delay(&self, url: &str, delay: Duration) {
let Some(domain) = domain_key(url) else {
return;
};
let mut domains = self.domains.lock().await;
let state = domains.entry(domain).or_default();
state.robots_delay = Some(delay);
}
async fn poll_next(&self) -> SchedulerPoll {
let queued_len = self.frontier.len().await;
if queued_len == 0 {
return SchedulerPoll::Empty;
};
let mut deferred = Vec::new();
let mut shortest_wait: Option<Duration> = None;
for _ in 0..queued_len {
let Some(queued) = self.frontier.pop_request().await else {
break;
};
match self.classify_candidate(&queued).await {
CandidateState::Ready => {
for item in deferred {
self.frontier.push_request_force(item).await;
}
return SchedulerPoll::Ready(Box::new(queued));
}
CandidateState::Pending(wait) => {
shortest_wait = Some(shortest_wait.map_or(wait, |current| current.min(wait)));
deferred.push(queued);
}
}
}
for item in deferred {
self.frontier.push_request_force(item).await;
}
shortest_wait.map_or(SchedulerPoll::Empty, SchedulerPoll::Pending)
}
async fn classify_candidate(&self, queued: &FrontierRequest) -> CandidateState {
if let Some(scheduled_at) = queued.scheduled_at()
&& let Ok(wait) = scheduled_at.duration_since(std::time::SystemTime::now())
{
return CandidateState::Pending(wait);
}
let Some(domain) = domain_key(queued.request().url()) else {
return CandidateState::Ready;
};
let mut domains = self.domains.lock().await;
let state = domains.entry(domain.clone()).or_default();
let domain_policy = self.policy.policy_for(&domain);
if state.in_flight >= domain_policy.concurrency() {
CandidateState::Pending(Duration::from_millis(10))
} else if let Some(next) = state.next_available_at {
match next.checked_duration_since(Instant::now()) {
Some(wait) => CandidateState::Pending(wait),
None => {
state.in_flight += 1;
CandidateState::Ready
}
}
} else {
state.in_flight += 1;
CandidateState::Ready
}
}
fn apply_fingerprint(&self, request: CrawlRequest) -> CrawlRequest {
if request.dont_filter_enabled() {
return request;
}
let key = self
.fingerprint_policy
.fingerprint(request.url())
.unwrap_or_else(|_| request.url().to_string());
request.with_dedup_key(key)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn delay_with_jitter_keeps_delay_within_configured_range() {
let base = Duration::from_millis(10);
let jitter = Duration::from_millis(50);
for _ in 0..100 {
let delay = delay_with_jitter(base, Some(jitter));
assert!(delay >= base);
assert!(delay <= base + jitter);
}
}
#[test]
fn delay_with_jitter_returns_base_when_jitter_is_disabled() {
let base = Duration::from_millis(10);
assert_eq!(delay_with_jitter(base, None), base);
assert_eq!(delay_with_jitter(base, Some(Duration::ZERO)), base);
}
}