use std::collections::{HashMap, VecDeque};
use std::fmt;
use std::ops::Deref;
use std::panic::AssertUnwindSafe;
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex, Weak};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
pub const DEFAULT_FINISHED_RETENTION: Duration = Duration::from_secs(600);
pub const DEFAULT_FINISHED_KEPT: usize = 64;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum JobState {
Queued,
Running,
Done,
Failed,
}
impl JobState {
pub fn is_live(self) -> bool {
matches!(self, JobState::Queued | JobState::Running)
}
}
pub trait JobKind: Send + Sync + Sized + 'static {
const NAME: &'static str;
const FINISHED_RETENTION: Duration = DEFAULT_FINISHED_RETENTION;
const FINISHED_KEPT: usize = DEFAULT_FINISHED_KEPT;
type Order: Clone + fmt::Debug + Send + Sync + 'static;
type Report: Clone + fmt::Debug + Default + Send + Sync + 'static;
fn key(order: &Self::Order) -> &str;
fn run(&self, order: &Self::Order, progress: &Reporter<Self>) -> Result<(), String>;
fn settle(
&self,
_order: &Self::Order,
outcome: Result<(), String>,
_progress: &Reporter<Self>,
) -> (JobState, Option<String>) {
match outcome {
Ok(()) => (JobState::Done, None),
Err(why) => (JobState::Failed, Some(why)),
}
}
fn shutdown_error(&self) -> String {
format!("the {} queue is shutting down", Self::NAME)
}
fn spawn_error(&self, err: &dyn fmt::Display) -> String {
format!("the {} worker could not be started: {err}", Self::NAME)
}
}
pub struct Job<K: JobKind> {
pub job_id: String,
pub key: String,
pub state: JobState,
pub queue_position: u32,
pub progress: Option<(u64, u64)>,
pub error: Option<String>,
pub order: K::Order,
pub report: K::Report,
pub enqueued_at_unix_ms: i64,
pub started_at_unix_ms: i64,
pub finished_at_unix_ms: i64,
}
impl<K: JobKind> Deref for Job<K> {
type Target = K::Report;
fn deref(&self) -> &K::Report {
&self.report
}
}
impl<K: JobKind> Clone for Job<K> {
fn clone(&self) -> Self {
Job {
job_id: self.job_id.clone(),
key: self.key.clone(),
state: self.state,
queue_position: self.queue_position,
progress: self.progress,
error: self.error.clone(),
order: self.order.clone(),
report: self.report.clone(),
enqueued_at_unix_ms: self.enqueued_at_unix_ms,
started_at_unix_ms: self.started_at_unix_ms,
finished_at_unix_ms: self.finished_at_unix_ms,
}
}
}
impl<K: JobKind> fmt::Debug for Job<K> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Job")
.field("kind", &K::NAME)
.field("job_id", &self.job_id)
.field("key", &self.key)
.field("state", &self.state)
.field("queue_position", &self.queue_position)
.field("progress", &self.progress)
.field("error", &self.error)
.field("order", &self.order)
.field("report", &self.report)
.field("enqueued_at_unix_ms", &self.enqueued_at_unix_ms)
.field("started_at_unix_ms", &self.started_at_unix_ms)
.field("finished_at_unix_ms", &self.finished_at_unix_ms)
.finish()
}
}
pub struct Reporter<K: JobKind> {
inner: Arc<Inner<K>>,
key: String,
}
impl<K: JobKind> Clone for Reporter<K> {
fn clone(&self) -> Self {
Reporter {
inner: Arc::clone(&self.inner),
key: self.key.clone(),
}
}
}
impl<K: JobKind> Reporter<K> {
pub fn key(&self) -> &str {
&self.key
}
pub fn progress(&self, done: u64, total: u64) {
self.edit(|job| job.progress = Some((done, total)));
}
pub fn report(&self, apply: impl FnOnce(&mut K::Report)) {
self.edit(|job| apply(&mut job.report));
}
pub fn snapshot(&self) -> Option<Job<K>> {
let state = self.inner.state.lock().unwrap_or_else(|p| p.into_inner());
state.jobs.get(&self.key).cloned()
}
fn edit(&self, apply: impl FnOnce(&mut Job<K>)) {
let mut state = self.inner.state.lock().unwrap_or_else(|p| p.into_inner());
if let Some(job) = state.jobs.get_mut(&self.key) {
apply(job);
}
}
}
pub struct JobQueue<K: JobKind> {
inner: Arc<Inner<K>>,
}
impl<K: JobKind> Clone for JobQueue<K> {
fn clone(&self) -> Self {
JobQueue {
inner: Arc::clone(&self.inner),
}
}
}
struct Inner<K: JobKind> {
state: Mutex<State<K>>,
lane: Mutex<
Option<(
Arc<jera::jobs::chain::ChainQueue>,
tokio::task::JoinHandle<()>,
)>,
>,
kind: K,
in_flight: AtomicUsize,
peak_in_flight: AtomicUsize,
runs: AtomicU64,
seq: AtomicU64,
}
struct State<K: JobKind> {
waiting: VecDeque<String>,
jobs: HashMap<String, Job<K>>,
stopping: bool,
}
impl<K: JobKind> Default for State<K> {
fn default() -> Self {
State {
waiting: VecDeque::new(),
jobs: HashMap::new(),
stopping: false,
}
}
}
pub(crate) fn now_ms() -> i64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_millis() as i64)
.unwrap_or_default()
}
impl<K: JobKind> JobQueue<K> {
pub fn new(kind: K) -> Self {
JobQueue {
inner: Arc::new(Inner {
state: Mutex::new(State::default()),
lane: Mutex::new(None),
kind,
in_flight: AtomicUsize::new(0),
peak_in_flight: AtomicUsize::new(0),
runs: AtomicU64::new(0),
seq: AtomicU64::new(0),
}),
}
}
pub fn next_job_id(&self) -> String {
self.mint()
}
pub fn enqueue(&self, order: K::Order) -> Job<K> {
self.enqueue_inner(order, None)
}
pub fn enqueue_as(&self, order: K::Order, job_id: String) -> Job<K> {
self.enqueue_inner(order, Some(job_id))
}
pub fn status(&self, key: &str) -> Option<Job<K>> {
let state = self.inner.state.lock().unwrap_or_else(|p| p.into_inner());
let job = state.jobs.get(key)?.clone();
Some(with_position(&state, job))
}
pub fn peak_in_flight(&self) -> usize {
self.inner.peak_in_flight.load(Ordering::SeqCst)
}
pub fn runs(&self) -> u64 {
self.inner.runs.load(Ordering::SeqCst)
}
pub fn stop(&self) {
let mut state = self.inner.state.lock().unwrap_or_else(|p| p.into_inner());
state.stopping = true;
}
pub fn jera_stats(&self) -> Option<jera::jobs::chain::ChainStats> {
self.inner
.lane
.lock()
.unwrap_or_else(|p| p.into_inner())
.as_ref()
.map(|(q, _)| q.stats())
}
fn lane(&self) -> Result<Arc<jera::jobs::chain::ChainQueue>, String> {
let mut lane = self.inner.lane.lock().unwrap_or_else(|p| p.into_inner());
if let Some((q, _)) = lane.as_ref() {
return Ok(Arc::clone(q));
}
let rt = gunnar_common::jobs::handle()?;
let queue = jera::jobs::chain::ChainQueue::new(gunnar_common::jobs::ledger_sink());
let weak: Weak<Inner<K>> = Arc::downgrade(&self.inner);
let worker = rt.spawn(Arc::clone(&queue).run(move |spec| {
let key = spec
.payload
.get("key")
.and_then(|v| v.as_str())
.unwrap_or_default()
.to_owned();
let Some(inner) = weak.upgrade() else {
anyhow::bail!("the {} queue is gone", K::NAME)
};
run_one(&inner, &key).map_err(|why| anyhow::anyhow!(why))
}));
*lane = Some((Arc::clone(&queue), worker));
Ok(queue)
}
fn mint(&self) -> String {
let id = self.inner.seq.fetch_add(1, Ordering::SeqCst) + 1;
format!("{}-{id}", K::NAME)
}
fn enqueue_inner(&self, order: K::Order, preminted: Option<String>) -> Job<K> {
let key = K::key(&order).to_owned();
let mut state = self.inner.state.lock().unwrap_or_else(|p| p.into_inner());
if state.stopping {
let now = now_ms();
let refused: Job<K> = Job {
job_id: preminted.unwrap_or_default(),
key,
state: JobState::Failed,
queue_position: 0,
progress: None,
error: Some(self.inner.kind.shutdown_error()),
order,
report: K::Report::default(),
enqueued_at_unix_ms: now,
started_at_unix_ms: 0,
finished_at_unix_ms: now,
};
return refused;
}
if let Some(existing) = state.jobs.get(&key) {
if existing.state.is_live() {
let job = existing.clone();
return with_position(&state, job);
}
}
let job_id = match preminted {
Some(id) => id,
None => self.mint(),
};
let job: Job<K> = Job {
job_id,
key: key.clone(),
state: JobState::Queued,
queue_position: 0,
progress: None,
error: None,
order,
report: K::Report::default(),
enqueued_at_unix_ms: now_ms(),
started_at_unix_ms: 0,
finished_at_unix_ms: 0,
};
state.jobs.insert(key.clone(), job.clone());
state.waiting.push_back(key.clone());
let lane = match self.lane() {
Ok(lane) => lane,
Err(err) => {
state.waiting.pop_back();
let message = self.inner.kind.spawn_error(&err);
let failed = state.jobs.get_mut(&key).expect("just inserted");
failed.state = JobState::Failed;
failed.error = Some(message);
failed.finished_at_unix_ms = now_ms();
return failed.clone();
}
};
let spec = jera::jobs::chain::ChainSpec::new(
K::NAME,
job.job_id.clone(),
K::NAME,
0,
"gunnar",
serde_json::json!({ "key": key }),
)
.with_target(format!("{} {key}", K::NAME))
.with_priority(jera::jobs::chain::Priority::Normal);
let _ = lane.submit(spec);
with_position(&state, job)
}
}
fn with_position<K: JobKind>(state: &State<K>, mut job: Job<K>) -> Job<K> {
job.queue_position = if job.state == JobState::Queued {
state
.waiting
.iter()
.position(|k| *k == job.key)
.map(|n| n as u32)
.unwrap_or(0)
} else {
0
};
job
}
fn run_one<K: JobKind>(inner: &Arc<Inner<K>>, key: &str) -> Result<serde_json::Value, String> {
let (order, job_id) = {
let mut state = inner.state.lock().unwrap_or_else(|p| p.into_inner());
if state.stopping {
return Ok(
serde_json::json!({ "key": key, "ran": false, "why": "the queue is stopping" }),
);
}
if let Some(at) = state.waiting.iter().position(|k| k == key) {
state.waiting.remove(at);
}
match state.jobs.get_mut(key) {
Some(job) if job.state == JobState::Queued => {
job.state = JobState::Running;
job.started_at_unix_ms = now_ms();
(job.order.clone(), job.job_id.clone())
}
_ => return Ok(serde_json::json!({ "key": key, "ran": false })),
}
};
let progress = Reporter {
inner: Arc::clone(inner),
key: key.to_owned(),
};
let live = inner.in_flight.fetch_add(1, Ordering::SeqCst) + 1;
inner.peak_in_flight.fetch_max(live, Ordering::SeqCst);
inner.runs.fetch_add(1, Ordering::SeqCst);
let outcome =
match std::panic::catch_unwind(AssertUnwindSafe(|| inner.kind.run(&order, &progress))) {
Ok(result) => result,
Err(panic) => {
let why = format!("the {} panicked: {}", K::NAME, panic_text(&panic));
tracing::error!("ALERT gunnar-{} job {job_id} ({key}): {why}", K::NAME);
Err(why)
}
};
inner.in_flight.fetch_sub(1, Ordering::SeqCst);
let (state_out, error) = inner.kind.settle(&order, outcome, &progress);
let mut state = inner.state.lock().unwrap_or_else(|p| p.into_inner());
if let Some(job) = state.jobs.get_mut(key) {
job.state = state_out;
job.error = error.clone();
job.finished_at_unix_ms = now_ms();
}
prune::<K>(&mut state);
match state_out {
JobState::Failed => Err(error.unwrap_or_else(|| format!("the {} job failed", K::NAME))),
_ => Ok(serde_json::json!({ "key": key, "job": job_id })),
}
}
impl<K: JobKind> Drop for Inner<K> {
fn drop(&mut self) {
if let Some((_, worker)) = self.lane.lock().unwrap_or_else(|p| p.into_inner()).take() {
worker.abort();
}
}
}
fn prune<K: JobKind>(state: &mut State<K>) {
let now = now_ms();
let cutoff = now - K::FINISHED_RETENTION.as_millis() as i64;
let mut finished: Vec<(String, i64)> = state
.jobs
.iter()
.filter(|(_, job)| !job.state.is_live())
.map(|(key, job)| (key.clone(), job.finished_at_unix_ms))
.collect();
finished.sort_by_key(|job| std::cmp::Reverse(job.1));
for (index, (key, finished_at)) in finished.into_iter().enumerate() {
if index >= K::FINISHED_KEPT || finished_at < cutoff {
state.jobs.remove(&key);
}
}
}
pub(crate) fn panic_text(panic: &Box<dyn std::any::Any + Send>) -> String {
if let Some(text) = panic.downcast_ref::<&'static str>() {
return (*text).to_owned();
}
if let Some(text) = panic.downcast_ref::<String>() {
return text.clone();
}
"a panic with no message".to_owned()
}
#[cfg(test)]
mod tests {
use super::*;
struct Nothing;
impl JobKind for Nothing {
const NAME: &'static str = "nothing";
type Order = String;
type Report = ();
fn key(order: &String) -> &str {
order
}
fn run(&self, _order: &String, _progress: &Reporter<Self>) -> Result<(), String> {
Ok(())
}
}
struct Flaky;
impl JobKind for Flaky {
const NAME: &'static str = "flaky";
type Order = String;
type Report = ();
fn key(order: &String) -> &str {
order
}
fn run(&self, order: &String, _progress: &Reporter<Self>) -> Result<(), String> {
match order.as_str() {
"boom" => panic!("the boom job breaks"),
"bad" => Err("the bad job fails".to_owned()),
_ => Ok(()),
}
}
}
#[test]
fn the_jobs_run_as_jera_jobs_and_a_panic_is_a_failed_one() {
let q = JobQueue::new(Flaky);
for k in ["a", "boom", "b", "bad", "c"] {
q.enqueue(k.to_owned());
}
let t = std::time::Instant::now();
while ["a", "boom", "b", "bad", "c"]
.iter()
.any(|k| q.status(k).is_none_or(|j| j.state.is_live()))
&& t.elapsed() < Duration::from_secs(10)
{
std::thread::sleep(Duration::from_millis(10));
}
let state = |k: &str| q.status(k).map(|j| j.state);
assert_eq!(
[
state("a"),
state("boom"),
state("b"),
state("bad"),
state("c")
],
[
Some(JobState::Done),
Some(JobState::Failed),
Some(JobState::Done),
Some(JobState::Failed),
Some(JobState::Done)
]
);
assert!(
q.status("boom")
.and_then(|j| j.error)
.unwrap_or_default()
.contains("the boom job breaks"),
"the panic's words are on the job"
);
let s = q.jera_stats().expect("the queue runs on a jera lane");
assert_eq!(
(s.done, s.failed),
(3, 2),
"three done and two failed jera jobs: {s:?}"
);
assert_eq!(q.peak_in_flight(), 1, "one at a time");
q.stop();
}
#[test]
fn a_stopped_queue_refuses_and_runs_nothing_more() {
let q = JobQueue::new(Flaky);
q.enqueue("a".to_owned());
let t = std::time::Instant::now();
while q.status("a").is_none_or(|j| j.state.is_live())
&& t.elapsed() < Duration::from_secs(10)
{
std::thread::sleep(Duration::from_millis(10));
}
assert!(
q.jera_stats().is_some_and(|s| s.done == 1),
"the first job is a done jera job"
);
q.stop();
let refused = q.enqueue("b".to_owned());
assert_eq!(refused.state, JobState::Failed);
assert!(refused.error.unwrap_or_default().contains("shutting down"));
std::thread::sleep(Duration::from_millis(100));
assert_eq!(q.runs(), 1, "nothing ran after the stop");
}
fn finished(key: &str, age: i64) -> Job<Nothing> {
Job {
job_id: format!("nothing-{key}"),
key: key.to_owned(),
state: JobState::Done,
queue_position: 0,
progress: None,
error: None,
order: key.to_owned(),
report: (),
enqueued_at_unix_ms: 0,
started_at_unix_ms: 0,
finished_at_unix_ms: now_ms() + age,
}
}
#[test]
fn prune_keeps_the_newest_finished_jobs() {
let mut state = State::<Nothing>::default();
for n in 0..(Nothing::FINISHED_KEPT + 5) {
let key = format!("ns/repo{n}");
state.jobs.insert(key.clone(), finished(&key, n as i64));
}
prune::<Nothing>(&mut state);
assert_eq!(state.jobs.len(), Nothing::FINISHED_KEPT);
assert!(
state
.jobs
.contains_key(&format!("ns/repo{}", Nothing::FINISHED_KEPT + 4)),
"the newest finished job must survive"
);
assert!(
!state.jobs.contains_key("ns/repo0"),
"the oldest finished job must be dropped"
);
}
#[test]
fn prune_never_drops_a_live_job() {
let mut state = State::<Nothing>::default();
let mut running = finished("ns/running", 0);
running.state = JobState::Running;
running.finished_at_unix_ms = 0;
state.jobs.insert("ns/running".into(), running);
for n in 0..(Nothing::FINISHED_KEPT + 5) {
let key = format!("ns/repo{n}");
state.jobs.insert(key.clone(), finished(&key, n as i64));
}
prune::<Nothing>(&mut state);
assert!(
state.jobs.contains_key("ns/running"),
"a running job must survive a prune; forgetting it would leave a \
console button lit for ever"
);
}
#[test]
fn a_panics_message_survives_being_caught() {
let owned: Box<dyn std::any::Any + Send> = Box::new("a static message".to_string());
assert_eq!(panic_text(&owned), "a static message");
let borrowed: Box<dyn std::any::Any + Send> = Box::new("a literal message");
assert_eq!(panic_text(&borrowed), "a literal message");
let odd: Box<dyn std::any::Any + Send> = Box::new(7u32);
assert_eq!(panic_text(&odd), "a panic with no message");
}
}