use parking_lot::Mutex;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::collections::{BTreeMap, HashMap, HashSet};
use std::io::Read;
use std::path::Path;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use crate::catalog::META_DEK_LEN;
use crate::durable_file::DurableRoot;
use crate::error::MongrelError;
pub const JOBS_FILENAME: &str = "JOBS";
const MAGIC: &[u8; 8] = b"MONGRJOB";
const JOBS_FORMAT_VERSION: u16 = 1;
const MAX_JOBS_BYTES: u64 = 64 * 1024 * 1024;
const MAX_CHECKPOINT_BYTES: usize = 1024 * 1024;
pub const DEFAULT_MAX_CONCURRENT_JOBS: usize = 2;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum JobState {
Pending,
Running,
Paused,
Cancelling,
Succeeded,
Failed,
RollingBack,
}
impl JobState {
#[cfg(test)]
pub(crate) const ALL: [JobState; 7] = [
JobState::Pending,
JobState::Running,
JobState::Paused,
JobState::Cancelling,
JobState::Succeeded,
JobState::Failed,
JobState::RollingBack,
];
pub fn is_terminal(self) -> bool {
matches!(self, JobState::Succeeded | JobState::Failed)
}
pub fn can_transition(self, next: JobState) -> bool {
use JobState::{Cancelling, Failed, Paused, Pending, RollingBack, Running, Succeeded};
matches!(
(self, next),
(Pending, Running)
| (Pending, Cancelling)
| (Running, Paused)
| (Running, Cancelling)
| (Running, RollingBack)
| (Running, Succeeded)
| (Paused, Pending)
| (Paused, Cancelling)
| (Cancelling, RollingBack)
| (Cancelling, Failed)
| (RollingBack, Failed)
)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum JobKind {
IndexBuild,
ColumnBackfill,
SchemaValidation,
MaterializedViewRebuild,
KeyRotation,
LargeImport,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct JobTarget {
pub table: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub index: Option<String>,
}
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct JobProgress {
pub fraction: f64,
pub done: u64,
pub total: u64,
}
impl Default for JobProgress {
fn default() -> Self {
Self {
fraction: 0.0,
done: 0,
total: 0,
}
}
}
impl JobProgress {
pub fn new(fraction: f64, done: u64, total: u64) -> Result<Self, JobError> {
if !(0.0..=1.0).contains(&fraction) {
return Err(JobError::InvalidProgress(format!(
"fraction {fraction} is outside [0.0, 1.0]"
)));
}
if total > 0 && done > total {
return Err(JobError::InvalidProgress(format!(
"done {done} exceeds total {total}"
)));
}
Ok(Self {
fraction,
done,
total,
})
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct JobRecord {
pub job_id: u64,
pub kind: JobKind,
pub state: JobState,
pub target: JobTarget,
pub progress: JobProgress,
pub created_at_micros: u64,
pub updated_at_micros: u64,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub error: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub checkpoint: Option<Vec<u8>>,
}
#[derive(Debug, thiserror::Error)]
pub enum JobError {
#[error("job {job_id} not found")]
NotFound {
job_id: u64,
},
#[error("illegal job state transition from {from:?} to {to:?}")]
IllegalTransition {
from: JobState,
to: JobState,
},
#[error("job {job_id} is {actual:?}, expected {expected:?}")]
UnexpectedState {
job_id: u64,
expected: JobState,
actual: JobState,
},
#[error("concurrent job limit reached: {active} active of {limit} allowed")]
ConcurrencyLimit {
active: usize,
limit: usize,
},
#[error("job {job_id} still has a live drive in this process")]
DriveActive {
job_id: u64,
},
#[error("invalid job progress: {0}")]
InvalidProgress(String),
#[error("job cancelled")]
Cancelled,
#[error("job phase failed: {0}")]
Phase(String),
#[error("injected fault: {0}")]
InjectedFault(String),
#[error("job checkpoint of {bytes} bytes exceeds the {limit}-byte limit")]
CheckpointTooLarge {
bytes: usize,
limit: usize,
},
#[error("job registry storage: {0}")]
Storage(String),
#[error("io error: {0}")]
Io(#[from] std::io::Error),
}
impl From<JobError> for MongrelError {
fn from(error: JobError) -> Self {
match error {
JobError::NotFound { job_id } => MongrelError::NotFound(format!("job {job_id}")),
JobError::Cancelled => MongrelError::Cancelled,
JobError::ConcurrencyLimit { active, limit } => MongrelError::ResourceLimitExceeded {
resource: "concurrent jobs",
requested: active,
limit,
},
JobError::InvalidProgress(message) => MongrelError::InvalidArgument(message),
JobError::Io(error) => MongrelError::Io(error),
other => MongrelError::Other(other.to_string()),
}
}
}
#[derive(Debug, Clone, Default)]
pub struct CancellationToken {
flag: Arc<AtomicBool>,
}
impl CancellationToken {
pub fn is_cancelled(&self) -> bool {
self.flag.load(Ordering::Acquire)
}
pub fn cancel(&self) {
self.flag.store(true, Ordering::Release);
}
pub fn check(&self) -> Result<(), JobError> {
if self.is_cancelled() {
Err(JobError::Cancelled)
} else {
Ok(())
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BuildPhase {
RecordPending,
PinSnapshot,
BuildHidden,
CatchUp,
Validate,
Publish,
ReleaseOld,
}
impl BuildPhase {
pub const ALL: [BuildPhase; 7] = [
BuildPhase::RecordPending,
BuildPhase::PinSnapshot,
BuildPhase::BuildHidden,
BuildPhase::CatchUp,
BuildPhase::Validate,
BuildPhase::Publish,
BuildPhase::ReleaseOld,
];
pub fn label(self) -> &'static str {
match self {
BuildPhase::RecordPending => "record_pending",
BuildPhase::PinSnapshot => "pin_snapshot",
BuildPhase::BuildHidden => "build_hidden",
BuildPhase::CatchUp => "catch_up",
BuildPhase::Validate => "validate",
BuildPhase::Publish => "publish",
BuildPhase::ReleaseOld => "release_old",
}
}
pub fn before_hook(self) -> &'static str {
match self {
BuildPhase::RecordPending => "job.record_pending.before",
BuildPhase::PinSnapshot => "job.pin_snapshot.before",
BuildPhase::BuildHidden => "job.build_hidden.before",
BuildPhase::CatchUp => "job.catch_up.before",
BuildPhase::Validate => "job.validate.before",
BuildPhase::Publish => "job.publish.before",
BuildPhase::ReleaseOld => "job.release_old.before",
}
}
pub fn after_hook(self) -> &'static str {
match self {
BuildPhase::RecordPending => "job.record_pending.after",
BuildPhase::PinSnapshot => "job.pin_snapshot.after",
BuildPhase::BuildHidden => "job.build_hidden.after",
BuildPhase::CatchUp => "job.catch_up.after",
BuildPhase::Validate => "job.validate.after",
BuildPhase::Publish => "job.publish.after",
BuildPhase::ReleaseOld => "job.release_old.after",
}
}
fn invoke<J: BuildPublishJob + ?Sized>(
self,
job: &mut J,
context: &JobContext,
) -> Result<(), JobError> {
match self {
BuildPhase::RecordPending => job.record_pending(context),
BuildPhase::PinSnapshot => job.pin_snapshot(context),
BuildPhase::BuildHidden => job.build_hidden(context),
BuildPhase::CatchUp => job.catch_up(context),
BuildPhase::Validate => job.validate(context),
BuildPhase::Publish => job.publish(context),
BuildPhase::ReleaseOld => job.release_old(context),
}
}
}
pub trait BuildPublishJob {
fn checkpoint_state(&self) -> Vec<u8> {
Vec::new()
}
fn restore_checkpoint(&mut self, _state: &[u8]) -> Result<(), JobError> {
Ok(())
}
fn record_pending(&mut self, _context: &JobContext) -> Result<(), JobError> {
Ok(())
}
fn pin_snapshot(&mut self, _context: &JobContext) -> Result<(), JobError> {
Ok(())
}
fn build_hidden(&mut self, _context: &JobContext) -> Result<(), JobError> {
Ok(())
}
fn catch_up(&mut self, _context: &JobContext) -> Result<(), JobError> {
Ok(())
}
fn validate(&mut self, _context: &JobContext) -> Result<(), JobError> {
Ok(())
}
fn publish(&mut self, _context: &JobContext) -> Result<(), JobError> {
Ok(())
}
fn release_old(&mut self, _context: &JobContext) -> Result<(), JobError> {
Ok(())
}
fn rollback(&mut self) -> Result<(), JobError> {
Ok(())
}
}
pub struct JobContext<'a> {
registry: &'a JobRegistry,
job_id: u64,
token: CancellationToken,
}
impl JobContext<'_> {
pub fn job_id(&self) -> u64 {
self.job_id
}
pub fn token(&self) -> &CancellationToken {
&self.token
}
pub fn check_cancelled(&self) -> Result<(), JobError> {
self.token.check()
}
pub fn report_progress(&self, done: u64, total: u64) -> Result<(), JobError> {
if total == 0 {
return Err(JobError::InvalidProgress(
"unit progress requires a nonzero total".to_string(),
));
}
let fraction = done as f64 / total as f64;
self.registry
.update_progress(self.job_id, JobProgress::new(fraction, done, total)?)
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
struct JobsSnapshot {
next_job_id: u64,
jobs: Vec<JobRecord>,
}
#[derive(Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
struct JobsEnvelope {
format_version: u16,
registry: JobsSnapshot,
}
#[derive(Debug, Clone, Default)]
struct RegistryInner {
next_job_id: u64,
jobs: BTreeMap<u64, JobRecord>,
}
pub struct JobRegistry {
root: DurableRoot,
meta_dek: Option<[u8; META_DEK_LEN]>,
inner: Mutex<RegistryInner>,
tokens: Mutex<HashMap<u64, CancellationToken>>,
active_drives: Mutex<HashSet<u64>>,
max_concurrent_jobs: usize,
}
impl std::fmt::Debug for JobRegistry {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("JobRegistry")
.field("root", &self.root)
.field("max_concurrent_jobs", &self.max_concurrent_jobs)
.finish_non_exhaustive()
}
}
impl JobRegistry {
pub fn open(dir: &Path, meta_dek: Option<&[u8; META_DEK_LEN]>) -> Result<Self, JobError> {
#[cfg(not(feature = "encryption"))]
if meta_dek.is_some() {
return Err(JobError::Storage(
"a metadata key was supplied but the `encryption` feature is disabled".to_string(),
));
}
let root = DurableRoot::open(dir)?;
let mut inner = match read_durable(&root, meta_dek)? {
Some(snapshot) => validate_snapshot(snapshot)?,
None => RegistryInner {
next_job_id: 1,
jobs: BTreeMap::new(),
},
};
let recovered = recover_after_crash(&mut inner);
let registry = Self {
root,
meta_dek: meta_dek.copied(),
inner: Mutex::new(inner),
tokens: Mutex::new(HashMap::new()),
active_drives: Mutex::new(HashSet::new()),
max_concurrent_jobs: DEFAULT_MAX_CONCURRENT_JOBS,
};
if recovered {
registry.persist_locked(®istry.inner.lock())?;
}
Ok(registry)
}
pub fn with_max_concurrent_jobs(mut self, limit: usize) -> Self {
self.max_concurrent_jobs = limit.max(1);
self
}
pub fn submit(&self, kind: JobKind, target: JobTarget) -> Result<u64, JobError> {
if target.table.is_empty() {
return Err(JobError::InvalidProgress(
"job target table must not be empty".to_string(),
));
}
if kind == JobKind::IndexBuild && target.index.is_none() {
return Err(JobError::InvalidProgress(
"an index-build job requires a target index".to_string(),
));
}
let mut next = self.inner.lock().clone();
let now = unix_micros();
let job_id = next.next_job_id;
next.next_job_id = next
.next_job_id
.checked_add(1)
.ok_or_else(|| JobError::Storage("job id space exhausted".to_string()))?;
next.jobs.insert(
job_id,
JobRecord {
job_id,
kind,
state: JobState::Pending,
target,
progress: JobProgress::default(),
created_at_micros: now,
updated_at_micros: now,
error: None,
checkpoint: None,
},
);
self.persist_and_swap(next)?;
self.tokens
.lock()
.insert(job_id, CancellationToken::default());
Ok(job_id)
}
pub fn get(&self, job_id: u64) -> Option<JobRecord> {
self.inner.lock().jobs.get(&job_id).cloned()
}
pub fn list(&self) -> Vec<JobRecord> {
self.inner.lock().jobs.values().cloned().collect()
}
pub fn cancellation_token(&self, job_id: u64) -> Option<CancellationToken> {
if !self.inner.lock().jobs.contains_key(&job_id) {
return None;
}
Some(self.tokens.lock().entry(job_id).or_default().clone())
}
pub fn pause(&self, job_id: u64) -> Result<(), JobError> {
self.transition(job_id, JobState::Paused)
}
pub fn resume(&self, job_id: u64) -> Result<(), JobError> {
if self.active_drives.lock().contains(&job_id) {
return Err(JobError::DriveActive { job_id });
}
self.transition(job_id, JobState::Pending)
}
pub fn cancel(&self, job_id: u64) -> Result<(), JobError> {
let state = self.get(job_id).ok_or(JobError::NotFound { job_id })?.state;
match state {
JobState::Pending | JobState::Paused => {
let mut next = self.inner.lock().clone();
let record = next.jobs.get_mut(&job_id).expect("record checked above");
apply_transition(record, JobState::Cancelling)?;
apply_transition(record, JobState::Failed)?;
record.error = Some(match state {
JobState::Pending => "cancelled before the job started".to_string(),
_ => "cancelled while the job was paused".to_string(),
});
self.persist_and_swap(next)?;
self.tokens.lock().entry(job_id).or_default().cancel();
Ok(())
}
JobState::Running => {
self.transition(job_id, JobState::Cancelling)?;
self.tokens.lock().entry(job_id).or_default().cancel();
Ok(())
}
JobState::Cancelling | JobState::RollingBack => Ok(()),
JobState::Succeeded | JobState::Failed => Err(JobError::IllegalTransition {
from: state,
to: JobState::Cancelling,
}),
}
}
fn update_progress(&self, job_id: u64, progress: JobProgress) -> Result<(), JobError> {
let mut next = self.inner.lock().clone();
let record = next
.jobs
.get_mut(&job_id)
.ok_or(JobError::NotFound { job_id })?;
if record.state != JobState::Running {
return Err(JobError::UnexpectedState {
job_id,
expected: JobState::Running,
actual: record.state,
});
}
record.progress = progress;
record.updated_at_micros = unix_micros();
self.persist_and_swap(next)
}
fn transition(&self, job_id: u64, to: JobState) -> Result<(), JobError> {
let mut next = self.inner.lock().clone();
let record = next
.jobs
.get_mut(&job_id)
.ok_or(JobError::NotFound { job_id })?;
apply_transition(record, to)?;
self.persist_and_swap(next)
}
fn admit(&self, job_id: u64) -> Result<(), JobError> {
let mut next = self.inner.lock().clone();
let active = next
.jobs
.values()
.filter(|record| {
matches!(
record.state,
JobState::Running | JobState::Cancelling | JobState::RollingBack
)
})
.count();
let record = next
.jobs
.get(&job_id)
.ok_or(JobError::NotFound { job_id })?;
if record.state != JobState::Pending {
return Err(JobError::IllegalTransition {
from: record.state,
to: JobState::Running,
});
}
if active >= self.max_concurrent_jobs {
return Err(JobError::ConcurrencyLimit {
active,
limit: self.max_concurrent_jobs,
});
}
let record = next.jobs.get_mut(&job_id).expect("record checked above");
apply_transition(record, JobState::Running)?;
self.persist_and_swap(next)
}
fn save_checkpoint(
&self,
job_id: u64,
checkpoint: Vec<u8>,
progress: JobProgress,
) -> Result<(), JobError> {
if checkpoint.len() > MAX_CHECKPOINT_BYTES {
return Err(JobError::CheckpointTooLarge {
bytes: checkpoint.len(),
limit: MAX_CHECKPOINT_BYTES,
});
}
let mut next = self.inner.lock().clone();
let record = next
.jobs
.get_mut(&job_id)
.ok_or(JobError::NotFound { job_id })?;
if !matches!(record.state, JobState::Running | JobState::Paused) {
return Err(JobError::UnexpectedState {
job_id,
expected: JobState::Running,
actual: record.state,
});
}
record.checkpoint = Some(checkpoint);
record.progress = progress;
record.updated_at_micros = unix_micros();
self.persist_and_swap(next)
}
fn complete(&self, job_id: u64) -> Result<(), JobError> {
let mut next = self.inner.lock().clone();
let record = next
.jobs
.get_mut(&job_id)
.ok_or(JobError::NotFound { job_id })?;
apply_transition(record, JobState::Succeeded)?;
record.checkpoint = None;
record.progress.fraction = 1.0;
self.persist_and_swap(next)
}
fn begin_rollback(&self, job_id: u64) -> Result<(), JobError> {
self.transition(job_id, JobState::RollingBack)
}
fn fail(&self, job_id: u64, error: String) -> Result<(), JobError> {
let mut next = self.inner.lock().clone();
let record = next
.jobs
.get_mut(&job_id)
.ok_or(JobError::NotFound { job_id })?;
apply_transition(record, JobState::Failed)?;
record.error = Some(error);
self.persist_and_swap(next)
}
fn persist_and_swap(&self, next: RegistryInner) -> Result<(), JobError> {
self.persist_locked(&next)?;
*self.inner.lock() = next;
Ok(())
}
fn persist_locked(&self, inner: &RegistryInner) -> Result<(), JobError> {
let snapshot = JobsSnapshot {
next_job_id: inner.next_job_id,
jobs: inner.jobs.values().cloned().collect(),
};
write_durable(&self.root, &snapshot, self.meta_dek.as_ref())
}
}
pub fn run_build_publish<J: BuildPublishJob + ?Sized>(
registry: &JobRegistry,
job_id: u64,
job: &mut J,
) -> Result<(), JobError> {
let record = registry.get(job_id).ok_or(JobError::NotFound { job_id })?;
if record.state != JobState::Pending {
return Err(JobError::IllegalTransition {
from: record.state,
to: JobState::Running,
});
}
let mut completed_phases = 0_usize;
if let Some(checkpoint) = &record.checkpoint {
let decoded = decode_build_checkpoint(checkpoint)?;
job.restore_checkpoint(&decoded.state)?;
completed_phases = usize::from(decoded.completed_phases);
}
registry.admit(job_id)?;
let _drive = DriveGuard { registry, job_id };
let token = registry
.cancellation_token(job_id)
.ok_or(JobError::NotFound { job_id })?;
let context = JobContext {
registry,
job_id,
token,
};
for (index, phase) in BuildPhase::ALL
.iter()
.copied()
.enumerate()
.skip(completed_phases)
{
let state = registry
.get(job_id)
.ok_or(JobError::NotFound { job_id })?
.state;
match state {
JobState::Running => {}
JobState::Paused => return Ok(()),
JobState::Cancelling => {
return rollback_and_fail(registry, job_id, job, JobError::Cancelled);
}
state => {
return Err(JobError::IllegalTransition {
from: state,
to: JobState::Running,
});
}
}
if let Err(fault) = mongreldb_fault::inject(phase.before_hook()) {
return park_on_fault(registry, job_id, job, fault);
}
let outcome = phase.invoke(job, &context);
if let Err(fault) = mongreldb_fault::inject(phase.after_hook()) {
return park_on_fault(registry, job_id, job, fault);
}
if let Err(error) = outcome {
return rollback_and_fail(registry, job_id, job, error);
}
completed_phases = index + 1;
let checkpoint = encode_build_checkpoint(completed_phases as u8, &job.checkpoint_state())?;
let total = BuildPhase::ALL.len() as u64;
let done = completed_phases as u64;
let progress = JobProgress::new(done as f64 / total as f64, done, total)?;
let state = registry
.get(job_id)
.ok_or(JobError::NotFound { job_id })?
.state;
match state {
JobState::Running => registry.save_checkpoint(job_id, checkpoint, progress)?,
JobState::Paused => {
registry.save_checkpoint(job_id, checkpoint, progress)?;
return Ok(());
}
JobState::Cancelling => {
return rollback_and_fail(registry, job_id, job, JobError::Cancelled);
}
state => {
return Err(JobError::IllegalTransition {
from: state,
to: JobState::Running,
});
}
}
}
match registry.complete(job_id) {
Ok(()) => Ok(()),
Err(JobError::IllegalTransition { from, .. }) => match from {
JobState::Paused => Ok(()),
JobState::Cancelling => rollback_and_fail(registry, job_id, job, JobError::Cancelled),
state => Err(JobError::IllegalTransition {
from: state,
to: JobState::Succeeded,
}),
},
Err(error) => Err(error),
}
}
fn park_on_fault<J: BuildPublishJob + ?Sized>(
registry: &JobRegistry,
job_id: u64,
job: &mut J,
fault: mongreldb_fault::Fault,
) -> Result<(), JobError> {
match registry.pause(job_id) {
Ok(()) => Err(JobError::InjectedFault(fault.to_string())),
Err(JobError::IllegalTransition {
from: JobState::Cancelling,
..
}) => rollback_and_fail(registry, job_id, job, JobError::Cancelled),
Err(error) => Err(error),
}
}
fn rollback_and_fail<J: BuildPublishJob + ?Sized>(
registry: &JobRegistry,
job_id: u64,
job: &mut J,
error: JobError,
) -> Result<(), JobError> {
match registry.begin_rollback(job_id) {
Ok(()) => {}
Err(JobError::IllegalTransition { .. }) => return Err(error),
Err(storage) => return Err(storage),
}
let message = match job.rollback() {
Ok(()) => error.to_string(),
Err(rollback_error) => format!("{error}; rollback also failed: {rollback_error}"),
};
registry.fail(job_id, message)?;
Err(error)
}
struct DriveGuard<'a> {
registry: &'a JobRegistry,
job_id: u64,
}
impl Drop for DriveGuard<'_> {
fn drop(&mut self) {
self.registry.active_drives.lock().remove(&self.job_id);
}
}
#[derive(Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
struct BuildPublishCheckpoint {
completed_phases: u8,
state: Vec<u8>,
}
fn encode_build_checkpoint(completed_phases: u8, state: &[u8]) -> Result<Vec<u8>, JobError> {
serde_json::to_vec(&BuildPublishCheckpoint {
completed_phases,
state: state.to_vec(),
})
.map_err(|error| JobError::Storage(format!("job checkpoint serialize: {error}")))
}
fn decode_build_checkpoint(bytes: &[u8]) -> Result<BuildPublishCheckpoint, JobError> {
let checkpoint: BuildPublishCheckpoint = serde_json::from_slice(bytes)
.map_err(|error| JobError::Storage(format!("job checkpoint deserialize: {error}")))?;
if usize::from(checkpoint.completed_phases) > BuildPhase::ALL.len() {
return Err(JobError::Storage(format!(
"job checkpoint claims {} completed phases, only {} exist",
checkpoint.completed_phases,
BuildPhase::ALL.len()
)));
}
Ok(checkpoint)
}
fn apply_transition(record: &mut JobRecord, to: JobState) -> Result<(), JobError> {
if !record.state.can_transition(to) {
return Err(JobError::IllegalTransition {
from: record.state,
to,
});
}
record.state = to;
record.updated_at_micros = unix_micros();
Ok(())
}
fn recover_after_crash(inner: &mut RegistryInner) -> bool {
let now = unix_micros();
let mut changed = false;
for record in inner.jobs.values_mut() {
match record.state {
JobState::Running => {
record.state = JobState::Paused;
record.updated_at_micros = now;
changed = true;
}
JobState::Cancelling => {
record.state = JobState::Failed;
record.error =
Some("cancelled (process restarted while the job was cancelling)".to_string());
record.updated_at_micros = now;
changed = true;
}
JobState::RollingBack => {
record.state = JobState::Failed;
const NOTE: &str = "rollback interrupted by process restart";
record.error = Some(match record.error.take() {
Some(error) => format!("{error}; {NOTE}"),
None => NOTE.to_string(),
});
record.updated_at_micros = now;
changed = true;
}
JobState::Pending | JobState::Paused | JobState::Succeeded | JobState::Failed => {}
}
}
changed
}
fn validate_snapshot(snapshot: JobsSnapshot) -> Result<RegistryInner, JobError> {
let mut jobs = BTreeMap::new();
for record in snapshot.jobs {
if jobs.insert(record.job_id, record).is_some() {
return Err(JobError::Storage(
"duplicate job id in registry file".to_string(),
));
}
}
let max_id = jobs.keys().next_back().copied().unwrap_or(0);
if snapshot.next_job_id <= max_id {
return Err(JobError::Storage(format!(
"registry allocator at {} would reissue job id {max_id}",
snapshot.next_job_id
)));
}
Ok(RegistryInner {
next_job_id: snapshot.next_job_id.max(1),
jobs,
})
}
fn encode(snapshot: &JobsSnapshot) -> Result<Vec<u8>, JobError> {
serde_json::to_vec(&JobsEnvelope {
format_version: JOBS_FORMAT_VERSION,
registry: snapshot.clone(),
})
.map_err(|error| JobError::Storage(format!("job registry serialize: {error}")))
}
fn decode(body: &[u8]) -> Result<JobsSnapshot, JobError> {
let envelope: JobsEnvelope = serde_json::from_slice(body)
.map_err(|error| JobError::Storage(format!("job registry deserialize: {error}")))?;
if envelope.format_version != JOBS_FORMAT_VERSION {
return Err(JobError::Storage(format!(
"unsupported job registry format version {}",
envelope.format_version
)));
}
Ok(envelope.registry)
}
fn plaintext_frame(body: &[u8]) -> Vec<u8> {
let hash = Sha256::digest(body);
let mut out = Vec::with_capacity(body.len() + 8 + 32);
out.extend_from_slice(MAGIC);
out.extend_from_slice(&hash);
out.extend_from_slice(body);
out
}
#[cfg(feature = "encryption")]
fn seal(body: &[u8], meta_dek: Option<&[u8; META_DEK_LEN]>) -> Result<Vec<u8>, JobError> {
match meta_dek {
Some(dek) => crate::encryption::encrypt_blob(dek, body)
.map_err(|error| JobError::Storage(format!("job registry seal: {error}"))),
None => Ok(plaintext_frame(body)),
}
}
#[cfg(not(feature = "encryption"))]
fn seal(body: &[u8], meta_dek: Option<&[u8; META_DEK_LEN]>) -> Result<Vec<u8>, JobError> {
if meta_dek.is_some() {
return Err(JobError::Storage(
"a metadata key was supplied but the `encryption` feature is disabled".to_string(),
));
}
Ok(plaintext_frame(body))
}
#[cfg(feature = "encryption")]
fn open_payload(
bytes: &[u8],
meta_dek: Option<&[u8; META_DEK_LEN]>,
) -> Result<JobsSnapshot, JobError> {
match meta_dek {
Some(dek) => {
let body = crate::encryption::decrypt_blob(dek, bytes).map_err(|_| {
JobError::Storage(
"job registry authentication failed (wrong key or tampered)".to_string(),
)
})?;
decode(&body)
}
None => parse_plaintext(bytes),
}
}
#[cfg(not(feature = "encryption"))]
fn open_payload(
bytes: &[u8],
meta_dek: Option<&[u8; META_DEK_LEN]>,
) -> Result<JobsSnapshot, JobError> {
if meta_dek.is_some() {
return Err(JobError::Storage(
"a metadata key was supplied but the `encryption` feature is disabled".to_string(),
));
}
parse_plaintext(bytes)
}
fn write_durable(
root: &DurableRoot,
snapshot: &JobsSnapshot,
meta_dek: Option<&[u8; META_DEK_LEN]>,
) -> Result<(), JobError> {
let body = encode(snapshot)?;
let payload = seal(&body, meta_dek)?;
root.write_atomic(JOBS_FILENAME, &payload)?;
Ok(())
}
fn read_durable(
root: &DurableRoot,
meta_dek: Option<&[u8; META_DEK_LEN]>,
) -> Result<Option<JobsSnapshot>, JobError> {
let file = match root.open_regular(JOBS_FILENAME) {
Ok(file) => file,
Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None),
Err(error) => return Err(error.into()),
};
let length = file.metadata()?.len();
if length > MAX_JOBS_BYTES {
return Err(JobError::Storage(format!(
"job registry of {length} bytes exceeds the {MAX_JOBS_BYTES}-byte limit"
)));
}
let mut bytes = Vec::with_capacity(length as usize);
file.take(MAX_JOBS_BYTES + 1).read_to_end(&mut bytes)?;
if bytes.len() as u64 != length {
return Err(JobError::Storage(
"job registry length changed while reading".to_string(),
));
}
open_payload(&bytes, meta_dek).map(Some)
}
fn parse_plaintext(bytes: &[u8]) -> Result<JobsSnapshot, JobError> {
if bytes.len() < 8 + 32 || &bytes[..8] != MAGIC {
return Err(JobError::Storage(
"job registry magic mismatch (corrupt or sealed with a key)".to_string(),
));
}
let (tag, body) = bytes[8..].split_at(32);
let calc = Sha256::digest(body);
if tag != calc.as_slice() {
return Err(JobError::Storage(
"job registry checksum mismatch (tampered or torn)".to_string(),
));
}
decode(body)
}
fn unix_micros() -> u64 {
let micros = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|duration| duration.as_micros())
.unwrap_or(0);
u64::try_from(micros).unwrap_or(u64::MAX)
}
#[cfg(test)]
mod tests {
use super::*;
fn open_temp() -> (tempfile::TempDir, JobRegistry) {
let dir = tempfile::tempdir().unwrap();
let registry = JobRegistry::open(dir.path(), None).unwrap();
(dir, registry)
}
fn target() -> JobTarget {
JobTarget {
table: "items".to_string(),
index: Some("items_idx".to_string()),
}
}
#[test]
fn transition_graph_matches_the_documented_edges() {
let legal: [(JobState, JobState); 11] = [
(JobState::Pending, JobState::Running),
(JobState::Pending, JobState::Cancelling),
(JobState::Running, JobState::Paused),
(JobState::Running, JobState::Cancelling),
(JobState::Running, JobState::RollingBack),
(JobState::Running, JobState::Succeeded),
(JobState::Paused, JobState::Pending),
(JobState::Paused, JobState::Cancelling),
(JobState::Cancelling, JobState::RollingBack),
(JobState::Cancelling, JobState::Failed),
(JobState::RollingBack, JobState::Failed),
];
for from in JobState::ALL {
for to in JobState::ALL {
let expected = legal.contains(&(from, to));
assert_eq!(from.can_transition(to), expected, "edge {from:?} -> {to:?}");
}
}
for terminal in [JobState::Succeeded, JobState::Failed] {
assert!(terminal.is_terminal());
assert!(JobState::ALL
.iter()
.all(|&next| !terminal.can_transition(next)));
}
}
#[test]
fn progress_validation_rejects_out_of_range_values() {
assert!(JobProgress::new(0.5, 1, 2).is_ok());
assert!(JobProgress::new(0.0, 0, 0).is_ok());
assert!(JobProgress::new(1.0, 7, 7).is_ok());
assert!(JobProgress::new(f64::NAN, 0, 0).is_err());
assert!(JobProgress::new(-0.1, 0, 0).is_err());
assert!(JobProgress::new(1.1, 0, 0).is_err());
assert!(JobProgress::new(0.5, 3, 2).is_err());
}
#[test]
fn persistence_round_trip_preserves_records_and_allocator() {
let dir = tempfile::tempdir().unwrap();
let first = JobRegistry::open(dir.path(), None).unwrap();
let a = first.submit(JobKind::IndexBuild, target()).unwrap();
let b = first
.submit(
JobKind::LargeImport,
JobTarget {
table: "bulk".to_string(),
index: None,
},
)
.unwrap();
assert_eq!((a, b), (1, 2));
first.admit(a).unwrap();
first
.save_checkpoint(
a,
b"opaque-resume-state".to_vec(),
JobProgress::new(0.5, 4, 8).unwrap(),
)
.unwrap();
drop(first);
let reopened = JobRegistry::open(dir.path(), None).unwrap();
let record = reopened.get(a).unwrap();
assert_eq!(record.state, JobState::Paused);
assert_eq!(record.kind, JobKind::IndexBuild);
assert_eq!(record.target, target());
assert_eq!(record.progress, JobProgress::new(0.5, 4, 8).unwrap());
assert_eq!(
record.checkpoint.as_deref(),
Some(b"opaque-resume-state".as_slice())
);
assert!(record.error.is_none());
assert!(record.updated_at_micros >= record.created_at_micros);
assert_eq!(reopened.get(b).unwrap().state, JobState::Pending);
let c = reopened
.submit(
JobKind::KeyRotation,
JobTarget {
table: "items".to_string(),
index: None,
},
)
.unwrap();
assert_eq!(c, 3);
assert_eq!(reopened.list().len(), 3);
}
#[test]
fn crash_recovery_maps_active_states_to_safe_ones() {
let dir = tempfile::tempdir().unwrap();
let registry = JobRegistry::open(dir.path(), None)
.unwrap()
.with_max_concurrent_jobs(8);
let running = registry.submit(JobKind::IndexBuild, target()).unwrap();
let cancelling = registry.submit(JobKind::IndexBuild, target()).unwrap();
let rolling_back = registry.submit(JobKind::IndexBuild, target()).unwrap();
let pending = registry.submit(JobKind::IndexBuild, target()).unwrap();
let paused = registry.submit(JobKind::IndexBuild, target()).unwrap();
let succeeded = registry.submit(JobKind::IndexBuild, target()).unwrap();
let failed = registry.submit(JobKind::IndexBuild, target()).unwrap();
registry.admit(running).unwrap();
registry
.save_checkpoint(
running,
b"resume-bytes".to_vec(),
JobProgress::new(0.25, 1, 4).unwrap(),
)
.unwrap();
registry.admit(cancelling).unwrap();
registry.cancel(cancelling).unwrap();
registry.admit(rolling_back).unwrap();
registry.begin_rollback(rolling_back).unwrap();
registry.admit(paused).unwrap();
registry.pause(paused).unwrap();
registry.admit(succeeded).unwrap();
registry.complete(succeeded).unwrap();
registry.admit(failed).unwrap();
registry.begin_rollback(failed).unwrap();
registry.fail(failed, "boom".to_string()).unwrap();
drop(registry);
let recovered = JobRegistry::open(dir.path(), None).unwrap();
let running = recovered.get(running).unwrap();
assert_eq!(running.state, JobState::Paused);
assert_eq!(
running.checkpoint.as_deref(),
Some(b"resume-bytes".as_slice())
);
assert_eq!(running.progress, JobProgress::new(0.25, 1, 4).unwrap());
let cancelling = recovered.get(cancelling).unwrap();
assert_eq!(cancelling.state, JobState::Failed);
assert!(cancelling.error.unwrap().contains("cancelled"));
let rolling_back = recovered.get(rolling_back).unwrap();
assert_eq!(rolling_back.state, JobState::Failed);
assert!(rolling_back.error.unwrap().contains("rollback interrupted"));
assert_eq!(recovered.get(pending).unwrap().state, JobState::Pending);
assert_eq!(recovered.get(paused).unwrap().state, JobState::Paused);
assert_eq!(recovered.get(succeeded).unwrap().state, JobState::Succeeded);
let failed = recovered.get(failed).unwrap();
assert_eq!(failed.state, JobState::Failed);
assert_eq!(failed.error.as_deref(), Some("boom"));
drop(recovered);
let again = JobRegistry::open(dir.path(), None).unwrap();
assert_eq!(again.list().len(), 7);
assert!(again.list().iter().all(|record| !matches!(
record.state,
JobState::Running | JobState::Cancelling | JobState::RollingBack
)));
}
#[test]
fn admission_enforces_the_concurrency_bound() {
let (_dir, registry) = open_temp();
let registry = registry.with_max_concurrent_jobs(1);
let a = registry.submit(JobKind::IndexBuild, target()).unwrap();
let b = registry.submit(JobKind::IndexBuild, target()).unwrap();
registry.admit(a).unwrap();
let error = registry.admit(b).unwrap_err();
assert!(
matches!(
error,
JobError::ConcurrencyLimit {
active: 1,
limit: 1
}
),
"expected ConcurrencyLimit, got {error:?}"
);
assert_eq!(registry.get(b).unwrap().state, JobState::Pending);
registry.complete(a).unwrap();
registry.admit(b).unwrap();
}
#[test]
fn pause_resume_cancel_follow_the_graph() {
let (_dir, registry) = open_temp();
let job = registry.submit(JobKind::IndexBuild, target()).unwrap();
assert!(matches!(
registry.pause(job),
Err(JobError::IllegalTransition {
from: JobState::Pending,
to: JobState::Paused
})
));
registry.admit(job).unwrap();
registry.pause(job).unwrap();
assert_eq!(registry.get(job).unwrap().state, JobState::Paused);
registry.resume(job).unwrap();
assert_eq!(registry.get(job).unwrap().state, JobState::Pending);
registry.admit(job).unwrap();
registry.cancel(job).unwrap();
assert_eq!(registry.get(job).unwrap().state, JobState::Cancelling);
assert!(registry.cancellation_token(job).unwrap().is_cancelled());
registry.cancel(job).unwrap();
registry.begin_rollback(job).unwrap();
registry.fail(job, "cancelled".to_string()).unwrap();
assert!(matches!(
registry.cancel(job),
Err(JobError::IllegalTransition {
from: JobState::Failed,
to: JobState::Cancelling
})
));
assert!(matches!(
registry.resume(job),
Err(JobError::IllegalTransition {
from: JobState::Failed,
to: JobState::Pending
})
));
}
#[test]
fn cancel_without_a_worker_completes_synchronously() {
let (_dir, registry) = open_temp();
let queued = registry.submit(JobKind::IndexBuild, target()).unwrap();
registry.cancel(queued).unwrap();
let record = registry.get(queued).unwrap();
assert_eq!(record.state, JobState::Failed);
assert_eq!(
record.error.as_deref(),
Some("cancelled before the job started")
);
assert!(registry.cancellation_token(queued).unwrap().is_cancelled());
let parked = registry.submit(JobKind::IndexBuild, target()).unwrap();
registry.admit(parked).unwrap();
registry.pause(parked).unwrap();
registry.cancel(parked).unwrap();
let record = registry.get(parked).unwrap();
assert_eq!(record.state, JobState::Failed);
assert_eq!(
record.error.as_deref(),
Some("cancelled while the job was paused")
);
}
#[test]
fn missing_file_opens_empty_and_file_is_created_on_first_mutation() {
let dir = tempfile::tempdir().unwrap();
let registry = JobRegistry::open(dir.path(), None).unwrap();
assert!(registry.list().is_empty());
assert!(!dir.path().join(JOBS_FILENAME).exists());
registry.submit(JobKind::IndexBuild, target()).unwrap();
assert!(dir.path().join(JOBS_FILENAME).exists());
}
#[test]
fn tampered_file_fails_closed() {
let dir = tempfile::tempdir().unwrap();
let registry = JobRegistry::open(dir.path(), None).unwrap();
registry.submit(JobKind::IndexBuild, target()).unwrap();
drop(registry);
let path = dir.path().join(JOBS_FILENAME);
let mut bytes = std::fs::read(&path).unwrap();
let last = bytes.len() - 1;
bytes[last] ^= 0x01;
std::fs::write(&path, &bytes).unwrap();
let error = JobRegistry::open(dir.path(), None).unwrap_err();
assert!(
matches!(&error, JobError::Storage(message) if message.contains("checksum")),
"expected checksum failure, got {error:?}"
);
bytes[..8].copy_from_slice(b"NOTAJOB!");
std::fs::write(&path, &bytes).unwrap();
let error = JobRegistry::open(dir.path(), None).unwrap_err();
assert!(
matches!(&error, JobError::Storage(message) if message.contains("magic")),
"expected magic failure, got {error:?}"
);
std::fs::write(&path, b"MON").unwrap();
assert!(matches!(
JobRegistry::open(dir.path(), None),
Err(JobError::Storage(_))
));
}
#[test]
fn unsupported_format_version_fails_closed() {
let dir = tempfile::tempdir().unwrap();
let body = serde_json::to_vec(&serde_json::json!({
"format_version": 99,
"registry": { "next_job_id": 1, "jobs": [] }
}))
.unwrap();
let payload = plaintext_frame(&body);
std::fs::write(dir.path().join(JOBS_FILENAME), payload).unwrap();
let error = JobRegistry::open(dir.path(), None).unwrap_err();
assert!(
matches!(&error, JobError::Storage(message) if message.contains("version 99")),
"expected version failure, got {error:?}"
);
}
#[test]
fn allocator_inconsistency_fails_closed() {
let dir = tempfile::tempdir().unwrap();
let body = serde_json::to_vec(&serde_json::json!({
"format_version": 1,
"registry": {
"next_job_id": 1,
"jobs": [{
"job_id": 1,
"kind": "IndexBuild",
"state": "Pending",
"target": { "table": "items" },
"progress": { "fraction": 0.0, "done": 0, "total": 0 },
"created_at_micros": 1,
"updated_at_micros": 1
}]
}
}))
.unwrap();
std::fs::write(dir.path().join(JOBS_FILENAME), plaintext_frame(&body)).unwrap();
assert!(matches!(
JobRegistry::open(dir.path(), None),
Err(JobError::Storage(_))
));
}
#[test]
fn resume_rejects_a_job_with_a_live_drive() {
let (_dir, registry) = open_temp();
let job = registry.submit(JobKind::IndexBuild, target()).unwrap();
registry.admit(job).unwrap();
registry.pause(job).unwrap();
registry.active_drives.lock().insert(job);
assert!(matches!(
registry.resume(job),
Err(JobError::DriveActive { job_id }) if job_id == job
));
assert_eq!(registry.get(job).unwrap().state, JobState::Paused);
registry.active_drives.lock().remove(&job);
registry.resume(job).unwrap();
assert_eq!(registry.get(job).unwrap().state, JobState::Pending);
}
#[test]
fn checkpoint_size_is_bounded() {
let (_dir, registry) = open_temp();
let job = registry.submit(JobKind::IndexBuild, target()).unwrap();
registry.admit(job).unwrap();
let oversized = vec![0_u8; MAX_CHECKPOINT_BYTES + 1];
assert!(matches!(
registry.save_checkpoint(job, oversized, JobProgress::default()),
Err(JobError::CheckpointTooLarge { .. })
));
}
#[test]
fn job_error_maps_onto_the_engine_error() {
assert!(matches!(
MongrelError::from(JobError::Cancelled),
MongrelError::Cancelled
));
assert!(matches!(
MongrelError::from(JobError::NotFound { job_id: 7 }),
MongrelError::NotFound(_)
));
assert!(matches!(
MongrelError::from(JobError::ConcurrencyLimit {
active: 3,
limit: 2
}),
MongrelError::ResourceLimitExceeded { .. }
));
assert!(matches!(
MongrelError::from(JobError::InjectedFault("x".to_string())),
MongrelError::Other(_)
));
}
#[test]
fn record_serde_uses_stable_text_encoding() {
let (_dir, registry) = open_temp();
let job = registry
.submit(JobKind::MaterializedViewRebuild, target())
.unwrap();
registry.admit(job).unwrap();
let record = registry.get(job).unwrap();
let json = serde_json::to_string(&record).unwrap();
assert!(json.contains("\"kind\":\"MaterializedViewRebuild\""));
assert!(json.contains("\"state\":\"Running\""));
let decoded: JobRecord = serde_json::from_str(&json).unwrap();
assert_eq!(decoded, record);
let unknown = json.replace("\"Running\"", "\"Napping\"");
assert!(serde_json::from_str::<JobRecord>(&unknown).is_err());
}
#[test]
fn checkpoint_codec_round_trip_and_bounds() {
let bytes = encode_build_checkpoint(3, b"impl-state").unwrap();
let decoded = decode_build_checkpoint(&bytes).unwrap();
assert_eq!(decoded.completed_phases, 3);
assert_eq!(decoded.state, b"impl-state");
let corrupt = encode_build_checkpoint(8, b"").unwrap();
assert!(matches!(
decode_build_checkpoint(&corrupt),
Err(JobError::Storage(_))
));
}
}