use std::sync::Arc;
use std::time::{Duration, Instant};
use dashmap::DashMap;
use uuid::Uuid;
use crate::types::{CrawlResult, ScrapeResult};
#[derive(Debug, Clone)]
pub enum JobState {
Pending {
created_at: Instant,
},
InProgress {
pages_completed: usize,
created_at: Instant,
},
CrawlCompleted {
result: Box<CrawlResult>,
created_at: Instant,
},
BatchCompleted {
results: Vec<(String, Result<ScrapeResult, String>)>,
created_at: Instant,
},
Failed {
message: String,
created_at: Instant,
},
Cancelled {
created_at: Instant,
},
}
impl JobState {
pub fn created_at(&self) -> Instant {
match self {
JobState::Pending { created_at }
| JobState::InProgress { created_at, .. }
| JobState::CrawlCompleted { created_at, .. }
| JobState::BatchCompleted { created_at, .. }
| JobState::Failed { created_at, .. }
| JobState::Cancelled { created_at } => *created_at,
}
}
}
#[derive(Debug)]
pub struct JobRegistry {
jobs: DashMap<Uuid, JobState>,
}
impl JobRegistry {
pub fn new() -> Self {
Self { jobs: DashMap::new() }
}
pub fn create_job(&self) -> Uuid {
let id = Uuid::new_v4();
self.jobs.insert(
id,
JobState::Pending {
created_at: Instant::now(),
},
);
id
}
pub fn get_status(&self, id: &Uuid) -> Option<JobState> {
self.jobs.get(id).map(|entry| entry.value().clone())
}
pub fn update(&self, id: &Uuid, state: JobState) -> bool {
if let Some(mut entry) = self.jobs.get_mut(id) {
*entry = state;
true
} else {
false
}
}
pub fn cancel(&self, id: &Uuid) -> bool {
if let Some(mut entry) = self.jobs.get_mut(id) {
match entry.value() {
JobState::Pending { created_at } | JobState::InProgress { created_at, .. } => {
let created_at = *created_at;
*entry = JobState::Cancelled { created_at };
true
}
_ => false,
}
} else {
false
}
}
pub fn evict_expired(&self, max_age: Duration) {
self.jobs.retain(|_, state| state.created_at().elapsed() < max_age);
}
pub fn spawn_eviction_task(self: &Arc<Self>, max_age: Duration) {
let registry = Arc::clone(self);
tokio::spawn(async move {
loop {
tokio::time::sleep(Duration::from_secs(60)).await;
registry.evict_expired(max_age);
}
});
}
}
impl Default for JobRegistry {
fn default() -> Self {
Self::new()
}
}