use rustlavel_core::{Json, Result};
use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use std::sync::{Arc, Mutex};
pub type ProgressFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T>> + Send + 'a>>;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Progress {
pub percent: u8,
pub stage: String,
pub updated_at: u64,
pub finished: bool,
pub failed: Option<String>,
}
impl Progress {
pub fn to_json(&self) -> Json {
Json::object([
("percent", Json::from(self.percent as i64)),
("stage", Json::from(self.stage.as_str())),
("updated_at", Json::from(self.updated_at as i64)),
("finished", Json::from(self.finished)),
("failed", self.failed.as_deref().map_or(Json::Null, Json::from)),
])
}
}
pub trait ProgressStore: Send + Sync + 'static {
fn set<'a>(&'a self, key: &'a str, progress: Progress) -> ProgressFuture<'a, ()>;
fn get<'a>(&'a self, key: &'a str) -> ProgressFuture<'a, Option<Progress>>;
fn purge<'a>(&'a self, before: u64) -> ProgressFuture<'a, usize>;
}
#[derive(Default)]
pub struct MemoryProgress {
inner: Mutex<HashMap<String, Progress>>,
}
impl MemoryProgress {
pub fn new() -> MemoryProgress {
MemoryProgress::default()
}
}
impl ProgressStore for MemoryProgress {
fn set<'a>(&'a self, key: &'a str, progress: Progress) -> ProgressFuture<'a, ()> {
Box::pin(async move {
self.inner.lock().expect("the progress store is poisoned").insert(key.to_string(), progress);
Ok(())
})
}
fn get<'a>(&'a self, key: &'a str) -> ProgressFuture<'a, Option<Progress>> {
Box::pin(async move { Ok(self.inner.lock().expect("the progress store is poisoned").get(key).cloned()) })
}
fn purge<'a>(&'a self, before: u64) -> ProgressFuture<'a, usize> {
Box::pin(async move {
let mut held = self.inner.lock().expect("the progress store is poisoned");
let before_len = held.len();
held.retain(|_, progress| progress.updated_at >= before);
Ok(before_len - held.len())
})
}
}
#[derive(Clone)]
pub struct JobContext {
pub id: String,
pub chain: Option<String>,
store: Option<Arc<dyn ProgressStore>>,
}
impl JobContext {
pub fn new(id: impl Into<String>, chain: Option<String>, store: Option<Arc<dyn ProgressStore>>) -> JobContext {
JobContext { id: id.into(), chain, store }
}
pub fn detached() -> JobContext {
JobContext { id: "detached".into(), chain: None, store: None }
}
pub async fn progress(&self, percent: u8, stage: impl Into<String>) {
let Some(store) = &self.store else { return };
let progress = Progress {
percent: percent.min(100),
stage: stage.into(),
updated_at: crate::time::unix_now().max(0) as u64,
finished: false,
failed: None,
};
for key in self.keys() {
if let Err(error) = store.set(&key, progress.clone()).await {
rustlavel_core::warn!("queue: could not record progress for {key}: {error}");
}
}
}
pub fn keys(&self) -> Vec<String> {
let mut keys = vec![self.id.clone()];
if let Some(chain) = &self.chain {
keys.push(format!("chain:{chain}"));
}
keys
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn progress_is_readable_by_job_id_and_by_chain() {
let store: Arc<dyn ProgressStore> = Arc::new(MemoryProgress::new());
let ctx = JobContext::new("job-1", Some("chain-9".into()), Some(Arc::clone(&store)));
ctx.progress(40, "transcribing").await;
let by_job = store.get("job-1").await.unwrap().expect("recorded");
assert_eq!(by_job.percent, 40);
assert_eq!(by_job.stage, "transcribing");
assert!(!by_job.finished);
let by_chain = store.get("chain:chain-9").await.unwrap().expect("recorded under the chain too");
assert_eq!(by_chain.percent, 40);
}
#[tokio::test]
async fn percent_is_clamped() {
let store: Arc<dyn ProgressStore> = Arc::new(MemoryProgress::new());
let ctx = JobContext::new("j", None, Some(Arc::clone(&store)));
ctx.progress(130, "x").await;
assert_eq!(store.get("j").await.unwrap().unwrap().percent, 100);
}
#[tokio::test]
async fn a_detached_context_swallows_reports() {
JobContext::detached().progress(50, "fine").await;
}
}