use crate::context::JobContext;
use origin_domain::{AppError, Clock, Job, JobId, JobStatus, Progress, Result};
use origin_events::{EventBus, JobFinished, JobProgress, JobStarted, PlatformEvent};
use std::collections::HashMap;
use std::future::Future;
use std::sync::{Arc, RwLock};
use tokio_util::sync::CancellationToken;
const MAX_FINISHED: usize = 50;
const PROGRESS_STEP: f64 = 0.01;
#[derive(Debug)]
struct Entry {
job: Job,
cancel: CancellationToken,
published_ratio: Option<f64>,
sequence: u64,
}
#[derive(Debug, Default)]
struct Inner {
entries: HashMap<JobId, Entry>,
next_sequence: u64,
}
#[derive(Debug, Clone)]
pub struct Jobs {
inner: Arc<RwLock<Inner>>,
events: EventBus,
clock: Arc<dyn Clock>,
}
#[derive(Debug)]
pub struct JobResult<T> {
id: JobId,
rx: tokio::sync::oneshot::Receiver<Result<T>>,
}
impl<T> JobResult<T> {
pub fn id(&self) -> &JobId {
&self.id
}
pub async fn wait(self) -> Result<T> {
self.rx
.await
.unwrap_or_else(|_| Err(AppError::internal("job ended without producing a result")))
}
}
impl Jobs {
pub fn new(events: EventBus, clock: Arc<dyn Clock>) -> Self {
Self {
inner: Arc::new(RwLock::new(Inner::default())),
events,
clock,
}
}
pub fn spawn<F, Fut>(&self, kind: impl Into<String>, body: F) -> JobId
where
F: FnOnce(JobContext) -> Fut + Send + 'static,
Fut: Future<Output = Result<()>> + Send + 'static,
{
let (id, _result) = self
.spawn_core(kind.into(), false, body)
.expect("spawn without exclusivity never fails");
id
}
pub fn spawn_exclusive<F, Fut>(&self, kind: impl Into<String>, body: F) -> Result<JobId>
where
F: FnOnce(JobContext) -> Fut + Send + 'static,
Fut: Future<Output = Result<()>> + Send + 'static,
{
self.spawn_core(kind.into(), true, body)
.map(|(id, _result)| id)
}
pub fn spawn_awaitable<F, Fut, T>(
&self,
kind: impl Into<String>,
body: F,
) -> (JobId, JobResult<T>)
where
F: FnOnce(JobContext) -> Fut + Send + 'static,
Fut: Future<Output = Result<T>> + Send + 'static,
T: Send + 'static,
{
let (id, result) = self
.spawn_core(kind.into(), false, body)
.expect("spawn_awaitable without exclusivity never fails");
(id, result)
}
pub fn spawn_exclusive_awaitable<F, Fut, T>(
&self,
kind: impl Into<String>,
body: F,
) -> Result<(JobId, JobResult<T>)>
where
F: FnOnce(JobContext) -> Fut + Send + 'static,
Fut: Future<Output = Result<T>> + Send + 'static,
T: Send + 'static,
{
self.spawn_core(kind.into(), true, body)
}
fn spawn_core<F, Fut, T>(
&self,
kind: String,
exclusive: bool,
body: F,
) -> Result<(JobId, JobResult<T>)>
where
F: FnOnce(JobContext) -> Fut + Send + 'static,
Fut: Future<Output = Result<T>> + Send + 'static,
T: Send + 'static,
{
let job = Job::queued(kind.clone(), self.clock.now());
let id = job.id.clone();
let cancel = CancellationToken::new();
{
let mut inner = self
.inner
.write()
.unwrap_or_else(|error| error.into_inner());
if exclusive
&& inner
.entries
.values()
.any(|entry| entry.job.kind == kind && !entry.job.status.is_terminal())
{
return Err(AppError::validation(format!(
"a '{kind}' job is already running"
)));
}
let sequence = inner.next_sequence;
inner.next_sequence += 1;
inner.entries.insert(
id.clone(),
Entry {
job,
cancel: cancel.clone(),
published_ratio: None,
sequence,
},
);
}
let (tx, rx) = tokio::sync::oneshot::channel();
let registry = self.clone();
let context = JobContext::new(id.clone(), registry.clone(), cancel.clone());
let started_id = id.clone();
tokio::spawn(async move {
registry.mark_started(&started_id, kind.clone());
let outcome = tokio::spawn(async move { body(context).await }).await;
let (status, error_message, result): (JobStatus, Option<String>, Result<T>) =
match outcome {
Ok(Ok(value)) if cancel.is_cancelled() => {
(JobStatus::Cancelled, None, Ok(value))
}
Ok(Ok(value)) => (JobStatus::Succeeded, None, Ok(value)),
Ok(Err(error)) => {
let message = error.to_string();
(JobStatus::Failed, Some(message), Err(error))
}
Err(join_error) if join_error.is_cancelled() => (
JobStatus::Cancelled,
None,
Err(AppError::internal("job cancelled")),
),
Err(_) => (
JobStatus::Failed,
Some("the job panicked".to_owned()),
Err(AppError::internal("the job panicked")),
),
};
registry.mark_finished(&started_id, kind, status, error_message);
let _ = tx.send(result);
});
Ok((id.clone(), JobResult { id, rx }))
}
pub async fn get(&self, id: &JobId) -> Option<Job> {
self.inner
.read()
.unwrap_or_else(|error| error.into_inner())
.entries
.get(id)
.map(|e| e.job.clone())
}
pub async fn list(&self) -> Vec<Job> {
let inner = self.inner.read().unwrap_or_else(|error| error.into_inner());
let mut entries: Vec<&Entry> = inner.entries.values().collect();
entries.sort_by_key(|entry| std::cmp::Reverse(entry.sequence));
entries.into_iter().map(|entry| entry.job.clone()).collect()
}
pub async fn running(&self) -> Vec<Job> {
self.list()
.await
.into_iter()
.filter(|job| !job.status.is_terminal())
.collect()
}
pub async fn cancel(&self, id: &JobId) -> Result<()> {
let inner = self.inner.read().unwrap_or_else(|error| error.into_inner());
let entry = inner
.entries
.get(id)
.ok_or_else(|| AppError::validation(format!("unknown job {id}")))?;
if !entry.job.cancelable {
return Err(AppError::validation(format!(
"job {id} cannot be cancelled"
)));
}
entry.cancel.cancel();
tracing::debug!(job_id = %id, "cancellation requested");
Ok(())
}
fn mark_started(&self, id: &JobId, kind: String) {
{
let mut inner = self
.inner
.write()
.unwrap_or_else(|error| error.into_inner());
let Some(entry) = inner.entries.get_mut(id) else {
return;
};
entry.job.status = JobStatus::Running;
entry.job.started_at = self.clock.now();
}
let _ = self.events.publish(PlatformEvent::JobStarted(JobStarted {
job: id.clone(),
kind,
}));
}
pub(crate) async fn report_progress(&self, id: &JobId, current: u64, total: Option<u64>) {
let should_publish = {
let mut inner = self
.inner
.write()
.unwrap_or_else(|error| error.into_inner());
let Some(entry) = inner.entries.get_mut(id) else {
return;
};
entry.job.progress = Progress { current, total };
let ratio = entry.job.progress.ratio();
let publish = match (entry.published_ratio, ratio) {
(None, _) => true,
(Some(_), None) => false,
(Some(previous), Some(now)) => (now - previous).abs() >= PROGRESS_STEP,
};
if publish {
entry.published_ratio = ratio.or(entry.published_ratio).or(Some(0.0));
}
publish
};
if should_publish {
let _ = self.events.publish(PlatformEvent::JobProgress(JobProgress {
job: id.clone(),
current,
total,
}));
}
}
fn mark_finished(&self, id: &JobId, kind: String, status: JobStatus, error: Option<String>) {
{
let mut inner = self
.inner
.write()
.unwrap_or_else(|error| error.into_inner());
if let Some(entry) = inner.entries.get_mut(id) {
entry.job.status = status;
entry.job.finished_at = Some(self.clock.now());
entry.job.error = error.clone();
entry.job.cancelable = false;
}
Self::evict_old_finished(&mut inner);
}
match status {
JobStatus::Failed => tracing::warn!(job_id = %id, kind, ?error, "job failed"),
_ => tracing::debug!(job_id = %id, kind, ?status, "job finished"),
}
let _ = self.events.publish(PlatformEvent::JobFinished(JobFinished {
job: id.clone(),
kind,
status,
error,
}));
}
fn evict_old_finished(inner: &mut Inner) {
let mut finished: Vec<(JobId, u64)> = inner
.entries
.iter()
.filter(|(_, entry)| entry.job.status.is_terminal())
.map(|(id, entry)| (id.clone(), entry.sequence))
.collect();
if finished.len() <= MAX_FINISHED {
return;
}
finished.sort_by_key(|(_, sequence)| *sequence);
for (id, _) in finished.iter().take(finished.len() - MAX_FINISHED) {
inner.entries.remove(id);
}
}
}