diff --git a/src/common/di.rs b/src/common/di.rs index fc046c71..da5615ef 100644 --- a/src/common/di.rs +++ b/src/common/di.rs @@ -436,6 +436,15 @@ impl AppServiceFactory { // have landed. let job_registry = Arc::new(JobRegistry::new()); + // Recoverable-run engine's PG provider (`jobs.recoverable_runs`). + // Runs on the maintenance pool so a long-running scan's cursor + // updates never contend with the request-path pool. Boot-time + // crash-recovery sweep fires in `build_app_state` after the + // provider is placed on `AppState`. + let job_store_provider = Arc::new( + crate::infrastructure::scheduler::PgJobStoreProvider::new(maintenance_pool.clone()), + ); + Ok(CoreServices { path_service, file_content_cache, @@ -449,6 +458,7 @@ impl AppServiceFactory { zip_service: None, // Placeholder - replaced after app services init config: self.config.clone(), job_registry, + job_store_provider, }) } @@ -2126,6 +2136,36 @@ impl AppServiceFactory { } } + // Recoverable-run engine crash recovery. Any row still marked + // Running or CancelRequested when the previous process died gets + // flipped to Paused with `error_message = 'server restart mid-run'`. + // We do NOT auto-resume — operators explicitly re-trigger via + // `POST /api/admin/jobs/{name}/trigger`, which resumes from the + // persisted cursor. Runs BEFORE the scheduler starts so a + // periodic-triggered recoverable job's first tick sees a clean + // slate. See `docs/plan/job-registry.md` Part 2. + use crate::infrastructure::scheduler::JobStoreProvider as _; + match app_state.core.job_store_provider.boot_recovery_sweep().await { + Ok(0) => tracing::debug!( + target: "oxicloud::scheduler", + event = "recoverable.boot_recovery", + flipped = 0, + "no orphaned recoverable runs found at boot" + ), + Ok(n) => tracing::warn!( + target: "oxicloud::scheduler", + event = "recoverable.boot_recovery", + flipped = n, + "flipped {n} orphaned recoverable run(s) Running/CancelRequested → Paused (previous process died mid-run)" + ), + Err(e) => tracing::error!( + target: "oxicloud::scheduler", + event = "recoverable.boot_recovery.failed", + error = %e, + "boot recovery sweep failed — orphaned runs may remain in Running state" + ), + } + // Start the periodic-job scheduler AFTER every native service has // finished registering its jobs on `core.job_registry`. Starting // it earlier would race the first tick against late registrations. @@ -2166,6 +2206,14 @@ pub struct CoreServices { /// themselves here during their creation; `SchedulerEngine::start` /// spins up the supervisor loop at the end of `build_app_state`. pub job_registry: Arc, + /// PG-backed provider for `jobs.recoverable_runs`. Recoverable + /// tenants (storage migration, reextract, consistency checks — + /// Part 2 of `docs/plan/job-registry.md`) plug into this via + /// `svc.register_recoverable_job(®istry, &job_store_provider).await`. + /// Boot-time crash-recovery sweep is run in `build_app_state` right + /// after this provider is created. + pub job_store_provider: + Arc, } /// Container for repository services diff --git a/src/infrastructure/scheduler/mod.rs b/src/infrastructure/scheduler/mod.rs index 6eb449cc..2b5cd20d 100644 --- a/src/infrastructure/scheduler/mod.rs +++ b/src/infrastructure/scheduler/mod.rs @@ -24,10 +24,17 @@ mod engine; mod handler; +mod pg_job_store; +mod recoverable; mod registry; mod types; pub use engine::SchedulerEngine; pub use handler::JobHandler; +pub use pg_job_store::{PgJobStore, PgJobStoreProvider}; +pub use recoverable::{ + JobStore, JobStoreProvider, OpenedRun, RecoverableAdapter, RecoverableJobHandler, RunOutcome, + RunStatus, run_or_resume, +}; pub use registry::{JobEntry, JobRegistry, JobSummary, RegisterError}; pub use types::{ErrCause, JobOutcome, JobRunArgs}; diff --git a/src/infrastructure/scheduler/pg_job_store.rs b/src/infrastructure/scheduler/pg_job_store.rs new file mode 100644 index 00000000..5b798e9a --- /dev/null +++ b/src/infrastructure/scheduler/pg_job_store.rs @@ -0,0 +1,375 @@ +//! PostgreSQL adapter for the recoverable-run engine +//! ([`super::recoverable`]). Concrete impls of [`JobStore`] and +//! [`JobStoreProvider`] backed by `jobs.recoverable_runs`. +//! +//! Both types are cheap to construct (just an `Arc` plus, for +//! `PgJobStore`, the bound run's id + started_at). One `PgJobStoreProvider` +//! lives on `AppState.core.job_store_provider`; per-run `PgJobStore` +//! instances are built by `open_or_start`. + +use std::sync::Arc; + +use async_trait::async_trait; +use chrono::{DateTime, Utc}; +use sqlx::PgPool; +use uuid::Uuid; + +use crate::common::errors::DomainError; + +use super::recoverable::{JobStore, JobStoreProvider, OpenedRun, RunStatus}; + +// ─── PgJobStore — bound to one run ────────────────────────────────────────── + +/// A `JobStore` bound to a specific `jobs.recoverable_runs.id`. Every +/// method issues one small UPDATE / SELECT against that row. +pub struct PgJobStore { + pool: Arc, + run_id: Uuid, + started_at: DateTime, +} + +impl PgJobStore { + /// Called only from [`PgJobStoreProvider::open_or_start`] and its + /// test helpers — implementors never construct one directly. + pub(super) fn new(pool: Arc, run_id: Uuid, started_at: DateTime) -> Self { + Self { + pool, + run_id, + started_at, + } + } +} + +fn map_sqlx_err(op: &'static str, e: sqlx::Error) -> DomainError { + DomainError::internal_error("JobStore", format!("{op}: {e}")) +} + +#[async_trait] +impl JobStore for PgJobStore { + fn run_id(&self) -> Uuid { + self.run_id + } + + fn started_at(&self) -> DateTime { + self.started_at + } + + async fn status(&self) -> Result { + let row: Option<(String,)> = + sqlx::query_as("SELECT status FROM jobs.recoverable_runs WHERE id = $1") + .bind(self.run_id) + .fetch_optional(self.pool.as_ref()) + .await + .map_err(|e| map_sqlx_err("status", e))?; + let raw = row + .ok_or_else(|| { + DomainError::internal_error("JobStore", format!("run vanished: {}", self.run_id)) + })? + .0; + RunStatus::parse(&raw).ok_or_else(|| { + DomainError::internal_error("JobStore", format!("unknown status value: {raw}")) + }) + } + + async fn checkpoint(&self, cursor: Vec, delta_count: u64) -> Result<(), DomainError> { + // stats.scanned_count += delta_count. jsonb_set expects the new + // value serialised as jsonb; the cast chain from bigint → text + // → jsonb is the standard way to bump a numeric counter without + // pulling the whole JSONB into Rust. + let delta = delta_count as i64; + sqlx::query( + r#" + UPDATE jobs.recoverable_runs + SET cursor = $2, + stats = jsonb_set( + stats, + '{scanned_count}', + ((COALESCE(stats->>'scanned_count', '0')::bigint + $3)::text)::jsonb + ), + last_progress_at = NOW() + WHERE id = $1 + "#, + ) + .bind(self.run_id) + .bind(&cursor[..]) + .bind(delta) + .execute(self.pool.as_ref()) + .await + .map_err(|e| map_sqlx_err("checkpoint", e))?; + Ok(()) + } + + async fn mark_completed(&self) -> Result<(), DomainError> { + sqlx::query( + r#" + UPDATE jobs.recoverable_runs + SET status = 'Completed', + completed_at = NOW(), + last_progress_at = NOW() + WHERE id = $1 + "#, + ) + .bind(self.run_id) + .execute(self.pool.as_ref()) + .await + .map_err(|e| map_sqlx_err("mark_completed", e))?; + Ok(()) + } + + async fn mark_paused(&self, cursor: Option>) -> Result<(), DomainError> { + // Two-query variant would be simpler but this preserves the + // final cursor value in one statement whether or not the + // handler advanced it since the last checkpoint. + if let Some(c) = cursor { + sqlx::query( + r#" + UPDATE jobs.recoverable_runs + SET status = 'Paused', + cursor = $2, + last_progress_at = NOW() + WHERE id = $1 + "#, + ) + .bind(self.run_id) + .bind(&c[..]) + .execute(self.pool.as_ref()) + .await + .map_err(|e| map_sqlx_err("mark_paused", e))?; + } else { + sqlx::query( + r#" + UPDATE jobs.recoverable_runs + SET status = 'Paused', + last_progress_at = NOW() + WHERE id = $1 + "#, + ) + .bind(self.run_id) + .execute(self.pool.as_ref()) + .await + .map_err(|e| map_sqlx_err("mark_paused", e))?; + } + Ok(()) + } + + async fn mark_failed(&self, message: &str) -> Result<(), DomainError> { + sqlx::query( + r#" + UPDATE jobs.recoverable_runs + SET status = 'Failed', + completed_at = NOW(), + last_progress_at = NOW(), + error_message = $2 + WHERE id = $1 + "#, + ) + .bind(self.run_id) + .bind(message) + .execute(self.pool.as_ref()) + .await + .map_err(|e| map_sqlx_err("mark_failed", e))?; + Ok(()) + } +} + +// ─── PgJobStoreProvider — registry-level ops ──────────────────────────────── + +/// The `JobStoreProvider` PG-backed implementation. Constructs +/// `PgJobStore` handles via `open_or_start`, and drives the boot-time +/// crash-recovery sweep. +pub struct PgJobStoreProvider { + pool: Arc, +} + +impl PgJobStoreProvider { + pub fn new(pool: Arc) -> Self { + Self { pool } + } +} + +#[async_trait] +impl JobStoreProvider for PgJobStoreProvider { + async fn open_or_start(&self, job_name: &str) -> Result { + // Two-shot: look up the latest non-terminal row; if found, + // dispatch on its status; if not, INSERT a fresh Running row. + // + // Concurrent-insert race is caught by the partial unique index + // `one_active_run_per_job` — the losing INSERT falls back to + // re-querying and dispatching on whatever the winner wrote. + // We retry once because after the first losing insert, the + // winning row is guaranteed to exist and no third caller can + // race in ahead of us (they'd hit the same unique index). + for attempt in 0..2 { + match self.try_open_or_start(job_name).await { + Ok(opened) => return Ok(opened), + Err(OpenErr::Retry) => { + tracing::debug!( + target: "oxicloud::scheduler", + event = "recoverable.open_or_start.race", + job = job_name, + attempt = attempt, + "open_or_start lost to a concurrent INSERT; retrying" + ); + continue; + } + Err(OpenErr::Fatal(e)) => return Err(e), + } + } + Err(DomainError::internal_error( + "JobStore", + format!("open_or_start({job_name}): retry budget exhausted"), + )) + } + + async fn boot_recovery_sweep(&self) -> Result { + // Every row abandoned in Running / CancelRequested by the + // previous process flips to Paused with a synthetic + // error_message. We DO NOT auto-resume — operators trigger + // the resume explicitly per the trait doc. + let result = sqlx::query( + r#" + UPDATE jobs.recoverable_runs + SET status = 'Paused', + error_message = COALESCE(error_message, 'server restart mid-run') + WHERE status IN ('Running', 'CancelRequested') + "#, + ) + .execute(self.pool.as_ref()) + .await + .map_err(|e| map_sqlx_err("boot_recovery_sweep", e))?; + Ok(result.rows_affected()) + } +} + +/// Row shape returned by `open_or_start`'s SELECT — factored out +/// so clippy's `type_complexity` lint doesn't yell at the query. +type ExistingRun = (Uuid, String, DateTime, Option>); + +/// Internal error surface for the two-shot open_or_start retry loop. +enum OpenErr { + /// Lost to a concurrent INSERT — caller retries. + Retry, + /// Any other DB error — surfaces to caller unchanged. + Fatal(DomainError), +} + +impl PgJobStoreProvider { + /// One attempt of open_or_start. Returns `Err(Retry)` on the + /// unique-index-conflict path so the outer loop re-queries. + async fn try_open_or_start(&self, job_name: &str) -> Result { + // Latest non-terminal row for this job_name, if any. + let existing: Option = sqlx::query_as( + r#" + SELECT id, status, started_at, cursor + FROM jobs.recoverable_runs + WHERE job_name = $1 + AND status IN ('Running', 'Paused', 'CancelRequested') + ORDER BY started_at DESC + LIMIT 1 + "#, + ) + .bind(job_name) + .fetch_optional(self.pool.as_ref()) + .await + .map_err(|e| OpenErr::Fatal(map_sqlx_err("open_or_start.select", e)))?; + + match existing { + Some((id, raw_status, _started_at, _cursor)) => { + let status = RunStatus::parse(&raw_status).ok_or_else(|| { + OpenErr::Fatal(DomainError::internal_error( + "JobStore", + format!("unknown status: {raw_status}"), + )) + })?; + match status { + RunStatus::Running | RunStatus::CancelRequested => { + Ok(OpenedRun::AlreadyActive { + run_id: id, + status, + }) + } + RunStatus::Paused => { + // Flip to Running and hand back the cursor. + // Race note: another concurrent caller could + // race the same UPDATE. Both would succeed + // (Paused → Running is idempotent), but only + // one caller's dispatch would then race the + // partial unique index on subsequent + // operations. Acceptable — the loser's + // handler will observe `status = Running` + // (via `store.status()`) and can early-exit. + // In practice this is a rare edge case that + // ONLY hits if two admin triggers land in + // the same microsecond. + let row: Option<(DateTime, Option>)> = sqlx::query_as( + r#" + UPDATE jobs.recoverable_runs + SET status = 'Running', + last_progress_at = NOW() + WHERE id = $1 + RETURNING started_at, cursor + "#, + ) + .bind(id) + .fetch_optional(self.pool.as_ref()) + .await + .map_err(|e| { + OpenErr::Fatal(map_sqlx_err("open_or_start.resume", e)) + })?; + let (started_at, cursor_bytes) = row.ok_or_else(|| { + OpenErr::Fatal(DomainError::internal_error( + "JobStore", + format!("run vanished during resume: {id}"), + )) + })?; + let store: Arc = Arc::new(PgJobStore::new( + self.pool.clone(), + id, + started_at, + )); + Ok(OpenedRun::Resumed { + store, + cursor: cursor_bytes.unwrap_or_default(), + }) + } + // Terminal states shouldn't appear here (WHERE + // clause filters them). Defensive branch. + _ => Err(OpenErr::Fatal(DomainError::internal_error( + "JobStore", + format!("terminal status leaked into open_or_start: {status:?}"), + ))), + } + } + None => { + // No non-terminal row → INSERT a fresh one. The + // partial unique index protects against a concurrent + // second INSERT; on conflict we retry. + let run_id = Uuid::new_v4(); + let now = Utc::now(); + let result = sqlx::query( + r#" + INSERT INTO jobs.recoverable_runs + (id, job_name, status, started_at, last_progress_at) + VALUES ($1, $2, 'Running', $3, $3) + ON CONFLICT ON CONSTRAINT one_active_run_per_job DO NOTHING + "#, + ) + .bind(run_id) + .bind(job_name) + .bind(now) + .execute(self.pool.as_ref()) + .await + .map_err(|e| OpenErr::Fatal(map_sqlx_err("open_or_start.insert", e)))?; + + if result.rows_affected() == 1 { + let store: Arc = + Arc::new(PgJobStore::new(self.pool.clone(), run_id, now)); + Ok(OpenedRun::Fresh { store }) + } else { + // Someone raced us. Retry to pick up their row. + Err(OpenErr::Retry) + } + } + } + } +} diff --git a/src/infrastructure/scheduler/recoverable.rs b/src/infrastructure/scheduler/recoverable.rs new file mode 100644 index 00000000..204b2cf3 --- /dev/null +++ b/src/infrastructure/scheduler/recoverable.rs @@ -0,0 +1,816 @@ +//! Part 2 of `docs/plan/job-registry.md` — the recoverable-run engine. +//! +//! Sibling to Part 1's [`JobHandler`](super::handler::JobHandler): where +//! `JobHandler` covers one-shot periodic jobs whose outcome is a +//! `JobOutcome`, this module covers long-running iteration that must +//! survive process restarts. State lives in `jobs.recoverable_runs` +//! and is threaded to the handler via a [`JobStore`]. +//! +//! # Layering +//! +//! A [`RecoverableJobHandler`] is wrapped by [`RecoverableAdapter`] +//! to expose a `JobHandler` face; the wrapper is what registers with +//! the existing [`JobRegistry`](super::registry::JobRegistry). Part 1 +//! knows nothing about cursors — every recoverable job appears to the +//! supervisor as a normal `JobHandler` whose `run()` calls +//! [`run_or_resume`] under the hood. +//! +//! # Persistence contract +//! +//! - [`JobStoreProvider::open_or_start`] is the sole entry into +//! `jobs.recoverable_runs`. It enforces the "one non-terminal run +//! per `job_name`" invariant via the DB's partial unique index. +//! - [`JobStoreProvider::boot_recovery_sweep`] runs once at server +//! startup to flip `Running`/`CancelRequested` rows abandoned by a +//! previous process to `Paused`, so an operator can resume them +//! explicitly. +//! +//! # For future implementors +//! +//! - Implement [`RecoverableJobHandler`] on your service. Write a +//! cursor-based scan loop that polls [`JobStore::status`] between +//! batches for cooperative cancellation and calls +//! [`JobStore::checkpoint`] every ~30 s or ~1 000 rows. +//! - Register via `svc.register_recoverable_job(®istry, &provider).await` +//! (see the ergonomic helper on the service — same shape as +//! Part 1's `register_job`). +//! - `docs/architecture/jobs.md` will cover this in operator-facing +//! detail once Slice 2 (admin endpoints) lands. + +use std::sync::Arc; +use std::time::Duration; + +use async_trait::async_trait; +use chrono::{DateTime, Utc}; +use serde::Serialize; +use uuid::Uuid; + +use crate::common::errors::DomainError; + +use super::handler::JobHandler; +use super::types::{JobOutcome, JobRunArgs}; + +// ─── Run status ───────────────────────────────────────────────────────────── + +/// Mirror of the `TEXT` values allowed in `jobs.recoverable_runs.status`. +/// +/// Terminal set = `{Completed, Failed}`. Non-terminal set (the one the +/// exclusivity partial unique index scopes) = +/// `{Running, Paused, CancelRequested}`. +/// +/// `CancelRequested` IS non-terminal — the run is still shutting down. +/// A second trigger arriving during cancel MUST NOT spawn a parallel +/// run; the trigger endpoint returns the surviving row instead. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +pub enum RunStatus { + Running, + Paused, + CancelRequested, + Completed, + Failed, +} + +impl RunStatus { + /// Stable label matching the SQL storage form. + pub fn as_str(self) -> &'static str { + match self { + RunStatus::Running => "Running", + RunStatus::Paused => "Paused", + RunStatus::CancelRequested => "CancelRequested", + RunStatus::Completed => "Completed", + RunStatus::Failed => "Failed", + } + } + + /// Parse from the SQL `status` column value; returns `None` for + /// unknown strings (schema drift signal). + pub fn parse(s: &str) -> Option { + match s { + "Running" => Some(RunStatus::Running), + "Paused" => Some(RunStatus::Paused), + "CancelRequested" => Some(RunStatus::CancelRequested), + "Completed" => Some(RunStatus::Completed), + "Failed" => Some(RunStatus::Failed), + _ => None, + } + } + + /// The set the exclusivity partial index scopes. Read: "a run in + /// this state blocks a fresh dispatch." + pub fn is_non_terminal(self) -> bool { + matches!( + self, + RunStatus::Running | RunStatus::Paused | RunStatus::CancelRequested + ) + } +} + +// ─── Run outcome (handler → engine) ───────────────────────────────────────── + +/// What a [`RecoverableJobHandler`] returns from `run_resumable`. +/// Translated by [`run_or_resume`] into a [`JobOutcome`] for uniform +/// supervisor logging + last-outcome storage. +/// +/// - `Completed` — walked the whole space; engine writes `status = Completed`. +/// - `Paused` — cooperative pause (cancel poll or graceful shutdown); +/// engine persists cursor + writes `status = Paused` so a future +/// resume picks up from here. +/// - `Failed` — irrecoverable error; cursor NOT advanced; engine +/// writes `status = Failed` with the message. +#[derive(Debug, Clone)] +pub enum RunOutcome { + Completed, + Paused { cursor: Vec }, + Failed { message: String }, +} + +// ─── Traits — implementor + port ──────────────────────────────────────────── + +/// The implementor-facing contract for a long-running, restart-tolerant +/// job. Sibling of [`JobHandler`]; NOT a subtrait — a stateless job +/// that only implements `JobHandler` never needs to know Part 2 exists. +/// +/// # Contract +/// +/// - **`name()` must be stable.** Appears in `jobs.recoverable_runs.job_name`, +/// log lines, and admin URLs (`POST /api/admin/jobs/{name}/trigger`). +/// Renaming after release is a breaking change. +/// - **Poll `store.status()` between batches** — the operator-cancel +/// path sets `status = CancelRequested`, and the handler MUST +/// observe that and return `RunOutcome::Paused { cursor }` at the +/// next safe boundary. Failing to poll means cancel doesn't work. +/// - **Checkpoint periodically.** Every ~30 s OR ~1 000 rows, +/// whichever comes first. Cheaper thresholds waste DB traffic; +/// coarser thresholds leak more work on crash. +/// - **Do NOT catch panics inside `run_resumable`.** The Part 1 +/// supervisor's `tokio::spawn` + `catch_unwind` boundary covers +/// panics uniformly — masking one loses the `cause=panicked` +/// diagnostic. +/// - **Do NOT accept a wall-clock timeout.** Part 1's `timeout` +/// knob is applied by the supervisor only for `JobHandler` +/// dispatches. A `tokio::time::timeout` fired mid-scan aborts the +/// task without letting the handler persist the cursor — the +/// cooperative `status()` poll is the ONLY safe cancel path for +/// recoverable jobs. +/// - **Do NOT call the terminal-write methods** (`mark_completed`, +/// `mark_paused`, `mark_failed`) on the store — [`run_or_resume`] +/// owns those, driven by your `RunOutcome` return value. Calling +/// them yourself risks leaving the row in a state that disagrees +/// with what you return. +#[async_trait] +pub trait RecoverableJobHandler: Send + Sync { + /// Stable snake_case identifier. Must match the eventual admin + /// URL fragment: `POST /api/admin/jobs/{name}/trigger`. + fn name(&self) -> &str; + + /// Long-running scan. See trait-level doc for the contract. + /// + /// `store` — bound to THIS run (a single row in + /// `jobs.recoverable_runs`). Use it for cancel polling + + /// checkpointing + finding recording. + /// `args` — per-dispatch parameters forwarded from the trigger + /// endpoint (`?force=true` maps to `args.force`). + /// `resume_cursor` — the cursor persisted by a prior Paused run, + /// or `None` for a fresh run. Decode into your own key type + /// (blob hash, file_id UUID, ltree path, …). + async fn run_resumable( + &self, + store: &dyn JobStore, + args: &JobRunArgs, + resume_cursor: Option>, + ) -> RunOutcome; +} + +/// Bound-to-a-run handle. The handler polls status + writes +/// checkpoints; [`run_or_resume`] alone drives the terminal +/// transitions (marked in the trait doc as engine-only). +/// +/// Terminal writes are ON this trait (not a separate one) to keep +/// the concrete impl monolithic — but handler code must not call +/// them. See the `RecoverableJobHandler` trait doc. +#[async_trait] +pub trait JobStore: Send + Sync { + /// UUID identifying this specific run (`jobs.recoverable_runs.id`). + fn run_id(&self) -> Uuid; + + /// Fixed at run start. Long-running consistency scans use this + /// as their grace-window reference — NOT `chrono::Utc::now()`, + /// which would drift across a multi-hour scan. + fn started_at(&self) -> DateTime; + + /// Current status of the run's row. Between batches the handler + /// polls this; if it returns [`RunStatus::CancelRequested`], the + /// handler MUST return [`RunOutcome::Paused`] at the next safe + /// boundary. + async fn status(&self) -> Result; + + /// Advance cursor + accumulate `delta_count` into + /// `stats.scanned_count`, bump `last_progress_at`. Called between + /// batches — the run's heartbeat. + async fn checkpoint(&self, cursor: Vec, delta_count: u64) -> Result<(), DomainError>; + + // ─── Terminal writes — engine-only. Do not call from handler code. + + /// Engine-only. Called by [`run_or_resume`] on + /// [`RunOutcome::Completed`]. Handler code MUST NOT call this. + async fn mark_completed(&self) -> Result<(), DomainError>; + + /// Engine-only. Called by [`run_or_resume`] on + /// [`RunOutcome::Paused`]. `cursor` = the resume key the handler + /// returned. Handler code MUST NOT call this. + async fn mark_paused(&self, cursor: Option>) -> Result<(), DomainError>; + + /// Engine-only. Called by [`run_or_resume`] on + /// [`RunOutcome::Failed`]. Handler code MUST NOT call this. + async fn mark_failed(&self, message: &str) -> Result<(), DomainError>; +} + +/// Registry-level operations on `jobs.recoverable_runs` — NOT bound +/// to a specific run. Provides the entry point [`run_or_resume`] uses +/// to look up / create a run, and the boot-time crash-recovery sweep. +#[async_trait] +pub trait JobStoreProvider: Send + Sync { + /// Called by [`run_or_resume`]. Behaviour: + /// + /// - No non-terminal row for `job_name`: INSERT a fresh Running + /// row (`cursor = NULL`, `started_at = NOW()`), return + /// [`OpenedRun::Fresh`]. + /// - Latest non-terminal row is `Paused`: UPDATE to Running, + /// return [`OpenedRun::Resumed`] with the persisted cursor. + /// - Latest non-terminal row is `Running` or `CancelRequested`: + /// return [`OpenedRun::AlreadyActive`] — caller MUST NOT + /// dispatch a parallel run. + /// + /// A concurrent INSERT race is handled internally via the DB's + /// partial unique index — the losing INSERT falls back to reading + /// the winning row. + async fn open_or_start(&self, job_name: &str) -> Result; + + /// Boot-time crash recovery. Any row abandoned in `Running` or + /// `CancelRequested` when the previous process died gets flipped + /// to `Paused` with `error_message = 'server restart mid-run'`. + /// Returns the number of rows updated. + /// + /// Does NOT auto-resume — the bug that killed the previous run + /// may still be present. Operators trigger the resume explicitly + /// via `POST /api/admin/jobs/{name}/trigger`, which calls + /// `open_or_start` and picks up the Paused cursor. + async fn boot_recovery_sweep(&self) -> Result; +} + +/// Result of [`JobStoreProvider::open_or_start`]. +pub enum OpenedRun { + /// Fresh run — new row inserted, cursor is None (start from scratch). + Fresh { store: Arc }, + /// Existing Paused run resumed. `cursor` is the last-persisted + /// resume key; the handler decodes it into its own type. + Resumed { + store: Arc, + cursor: Vec, + }, + /// A non-terminal run is already active; the caller must NOT + /// spawn a parallel dispatch. Returned to admin/trigger callers + /// as `Ok { count: 0, extra: {"skipped": "already_running", …} }`. + AlreadyActive { run_id: Uuid, status: RunStatus }, +} + +// ─── Engine glue ──────────────────────────────────────────────────────────── + +/// The single entry point for running a `RecoverableJobHandler` +/// outside test code. Coordinates row lookup/creation, dispatches +/// the handler, translates `RunOutcome` → `JobOutcome`, writes the +/// terminal status. +/// +/// Called by [`RecoverableAdapter::run`] (the Part 1 JobHandler face) +/// so recoverable jobs slot into the existing scheduler unchanged. +pub async fn run_or_resume( + job: Arc, + provider: Arc, + args: &JobRunArgs, +) -> JobOutcome { + let opened = match provider.open_or_start(job.name()).await { + Ok(o) => o, + Err(e) => return JobOutcome::err(format!("open_or_start failed: {e}")), + }; + let (store, resume_cursor) = match opened { + OpenedRun::AlreadyActive { run_id, status } => { + return JobOutcome::ok_with( + 0, + serde_json::json!({ + "skipped": "already_running", + "run_id": run_id.to_string(), + "status": status.as_str(), + }), + ); + } + OpenedRun::Fresh { store } => (store, None), + OpenedRun::Resumed { store, cursor } => (store, Some(cursor)), + }; + let run_id = store.run_id(); + + // Dispatch. Terminal writes to `jobs.recoverable_runs` happen + // here (NOT in the handler) so the row always ends in a state + // that matches what the handler returned. + match job.run_resumable(&*store, args, resume_cursor).await { + RunOutcome::Completed => { + log_terminal_write_err("mark_completed", run_id, store.mark_completed().await); + JobOutcome::ok_with( + 0, + serde_json::json!({ + "completed": true, + "run_id": run_id.to_string(), + }), + ) + } + RunOutcome::Paused { cursor } => { + let cursor_hex = hex::encode(&cursor); + log_terminal_write_err( + "mark_paused", + run_id, + store.mark_paused(Some(cursor)).await, + ); + JobOutcome::ok_with( + 0, + serde_json::json!({ + "paused": true, + "run_id": run_id.to_string(), + "cursor_hex": cursor_hex, + }), + ) + } + RunOutcome::Failed { message } => { + log_terminal_write_err( + "mark_failed", + run_id, + store.mark_failed(&message).await, + ); + JobOutcome::err(format!("{message} (run_id={run_id})")) + } + } +} + +fn log_terminal_write_err(op: &str, run_id: Uuid, res: Result<(), DomainError>) { + if let Err(e) = res { + tracing::warn!( + target: "oxicloud::scheduler", + event = "recoverable.terminal_write_failed", + op = op, + run_id = %run_id, + error = %e, + "failed to write terminal status for recoverable run" + ); + } +} + +// ─── Adapter — bridge to Part 1's JobHandler ──────────────────────────────── + +/// Wraps a `RecoverableJobHandler` behind a `JobHandler` face so it +/// registers with the existing `JobRegistry` unchanged. The Part 1 +/// supervisor's dispatch loop calls the adapter's `run()`, which +/// delegates to `run_or_resume(inner, provider, args)`. +/// +/// Constructed by `service.register_recoverable_job(®istry, +/// &provider)` — see the ergonomic helper on each recoverable +/// service. +pub struct RecoverableAdapter { + inner: Arc, + provider: Arc, + name: String, +} + +impl RecoverableAdapter { + pub fn new(inner: Arc, provider: Arc) -> Self { + let name = inner.name().to_string(); + Self { + inner, + provider, + name, + } + } +} + +#[async_trait] +impl JobHandler for RecoverableAdapter { + fn name(&self) -> &str { + &self.name + } + async fn run(&self, args: &JobRunArgs) -> JobOutcome { + run_or_resume(self.inner.clone(), self.provider.clone(), args).await + } +} + +// ─── Ergonomics: JobRegistry extension for recoverable jobs ───────────────── + +impl super::registry::JobRegistry { + /// Register a recoverable job. Wraps the handler in a + /// [`RecoverableAdapter`] and delegates to the standard + /// [`register`](super::registry::JobRegistry::register) — so a + /// recoverable job appears to the supervisor as a normal + /// `JobHandler` at `name`. + /// + /// `interval` follows the same semantic as periodic jobs: + /// - `Some(dur)` — supervisor fires it periodically (and admin + /// triggers land on the same `run_or_resume` dispatch). + /// - `None` — admin-triggered only. Typical for long-running + /// tenants (storage migration, reextract, consistency checks). + /// + /// Timeout is force-None — recoverable jobs use cooperative + /// cancellation via `store.status()` polling, NOT wall-clock + /// timeouts. See `RecoverableJobHandler` trait doc. + pub async fn register_recoverable_job( + &self, + handler: Arc, + provider: Arc, + interval: Option, + ) { + let adapter = Arc::new(RecoverableAdapter::new(handler, provider)); + self.register(adapter, interval, None).await; + } +} + +// ─── Tests — in-memory JobStore mock + run_or_resume paths ────────────────── + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Mutex; + + // ─── In-memory JobStore ──────────────────────────────────────────────── + + struct MemStore { + run_id: Uuid, + started_at: DateTime, + state: Mutex, + } + + struct MemStoreState { + status: RunStatus, + cursor: Option>, + scanned_count: u64, + error_message: Option, + } + + #[async_trait] + impl JobStore for MemStore { + fn run_id(&self) -> Uuid { + self.run_id + } + fn started_at(&self) -> DateTime { + self.started_at + } + async fn status(&self) -> Result { + Ok(self.state.lock().unwrap().status) + } + async fn checkpoint( + &self, + cursor: Vec, + delta_count: u64, + ) -> Result<(), DomainError> { + let mut s = self.state.lock().unwrap(); + s.cursor = Some(cursor); + s.scanned_count += delta_count; + Ok(()) + } + async fn mark_completed(&self) -> Result<(), DomainError> { + self.state.lock().unwrap().status = RunStatus::Completed; + Ok(()) + } + async fn mark_paused(&self, cursor: Option>) -> Result<(), DomainError> { + let mut s = self.state.lock().unwrap(); + s.status = RunStatus::Paused; + if let Some(c) = cursor { + s.cursor = Some(c); + } + Ok(()) + } + async fn mark_failed(&self, message: &str) -> Result<(), DomainError> { + let mut s = self.state.lock().unwrap(); + s.status = RunStatus::Failed; + s.error_message = Some(message.to_string()); + Ok(()) + } + } + + // ─── In-memory JobStoreProvider ──────────────────────────────────────── + // + // Simplified: one job_name at a time, no cross-job isolation. Enough + // to exercise the run_or_resume control flow. + + struct MemProvider { + stores: Mutex>>, + } + + impl MemProvider { + fn new() -> Self { + Self { + stores: Mutex::new(Vec::new()), + } + } + + /// Test-only helper — seed a Running row without going through + /// `open_or_start`. Lets tests set up the "concurrent trigger + /// hits already-active" scenario without racing. + fn seed_running(&self) -> Uuid { + let store = Arc::new(MemStore { + run_id: Uuid::new_v4(), + started_at: Utc::now(), + state: Mutex::new(MemStoreState { + status: RunStatus::Running, + cursor: None, + scanned_count: 0, + error_message: None, + }), + }); + let id = store.run_id; + self.stores.lock().unwrap().push(store); + id + } + + /// Test-only read — last-created run's status, for post-hoc + /// assertions. + fn last_status(&self) -> Option { + let stores = self.stores.lock().unwrap(); + stores + .last() + .map(|s| s.state.lock().unwrap().status) + } + + /// Test-only read — last-created run's cursor. + fn last_cursor(&self) -> Option> { + let stores = self.stores.lock().unwrap(); + stores.last().and_then(|s| s.state.lock().unwrap().cursor.clone()) + } + } + + #[async_trait] + impl JobStoreProvider for MemProvider { + async fn open_or_start(&self, _job_name: &str) -> Result { + let mut stores = self.stores.lock().unwrap(); + if let Some(store) = stores.last() { + let state = store.state.lock().unwrap(); + if state.status.is_non_terminal() { + return match state.status { + RunStatus::Paused => { + let cursor = state.cursor.clone().unwrap_or_default(); + drop(state); + store.state.lock().unwrap().status = RunStatus::Running; + Ok(OpenedRun::Resumed { + store: store.clone(), + cursor, + }) + } + _ => Ok(OpenedRun::AlreadyActive { + run_id: store.run_id, + status: state.status, + }), + }; + } + } + let store = Arc::new(MemStore { + run_id: Uuid::new_v4(), + started_at: Utc::now(), + state: Mutex::new(MemStoreState { + status: RunStatus::Running, + cursor: None, + scanned_count: 0, + error_message: None, + }), + }); + stores.push(store.clone()); + Ok(OpenedRun::Fresh { store }) + } + + async fn boot_recovery_sweep(&self) -> Result { + let stores = self.stores.lock().unwrap(); + let mut n = 0u64; + for s in stores.iter() { + let mut state = s.state.lock().unwrap(); + if matches!( + state.status, + RunStatus::Running | RunStatus::CancelRequested + ) { + state.status = RunStatus::Paused; + state.error_message = Some("server restart mid-run".into()); + n += 1; + } + } + Ok(n) + } + } + + // ─── Handlers ────────────────────────────────────────────────────────── + + struct CompletingHandler; + #[async_trait] + impl RecoverableJobHandler for CompletingHandler { + fn name(&self) -> &str { + "completer" + } + async fn run_resumable( + &self, + store: &dyn JobStore, + _args: &JobRunArgs, + _resume_cursor: Option>, + ) -> RunOutcome { + store.checkpoint(vec![1, 2, 3], 5).await.unwrap(); + RunOutcome::Completed + } + } + + struct PausingHandler; + #[async_trait] + impl RecoverableJobHandler for PausingHandler { + fn name(&self) -> &str { + "pauser" + } + async fn run_resumable( + &self, + _store: &dyn JobStore, + _args: &JobRunArgs, + _resume_cursor: Option>, + ) -> RunOutcome { + RunOutcome::Paused { + cursor: b"halfway".to_vec(), + } + } + } + + struct FailingHandler; + #[async_trait] + impl RecoverableJobHandler for FailingHandler { + fn name(&self) -> &str { + "failer" + } + async fn run_resumable( + &self, + _store: &dyn JobStore, + _args: &JobRunArgs, + _resume_cursor: Option>, + ) -> RunOutcome { + RunOutcome::Failed { + message: "boom".into(), + } + } + } + + struct ResumeInspectHandler { + saw_cursor: Arc>>>, + } + #[async_trait] + impl RecoverableJobHandler for ResumeInspectHandler { + fn name(&self) -> &str { + "resumer" + } + async fn run_resumable( + &self, + _store: &dyn JobStore, + _args: &JobRunArgs, + resume_cursor: Option>, + ) -> RunOutcome { + *self.saw_cursor.lock().unwrap() = resume_cursor; + RunOutcome::Completed + } + } + + // ─── Tests ───────────────────────────────────────────────────────────── + + #[tokio::test] + async fn fresh_run_completes_and_marks_status_completed() { + let provider = Arc::new(MemProvider::new()); + let provider_trait: Arc = provider.clone(); + + let outcome = run_or_resume( + Arc::new(CompletingHandler), + provider_trait, + &JobRunArgs::default(), + ) + .await; + + assert!(outcome.is_ok(), "expected Ok, got {outcome:?}"); + if let JobOutcome::Ok { extra, .. } = outcome { + assert_eq!(extra["completed"], true); + assert!(extra["run_id"].is_string()); + } + assert_eq!(provider.last_status(), Some(RunStatus::Completed)); + } + + #[tokio::test] + async fn paused_run_persists_cursor_and_marks_status_paused() { + let provider = Arc::new(MemProvider::new()); + let provider_trait: Arc = provider.clone(); + + let outcome = run_or_resume( + Arc::new(PausingHandler), + provider_trait, + &JobRunArgs::default(), + ) + .await; + + assert!(outcome.is_ok()); + if let JobOutcome::Ok { extra, .. } = outcome { + assert_eq!(extra["paused"], true); + assert_eq!(extra["cursor_hex"], hex::encode(b"halfway")); + } + assert_eq!(provider.last_status(), Some(RunStatus::Paused)); + assert_eq!(provider.last_cursor(), Some(b"halfway".to_vec())); + } + + #[tokio::test] + async fn failed_run_marks_status_failed_and_returns_err() { + let provider = Arc::new(MemProvider::new()); + let provider_trait: Arc = provider.clone(); + + let outcome = run_or_resume( + Arc::new(FailingHandler), + provider_trait, + &JobRunArgs::default(), + ) + .await; + + assert!(!outcome.is_ok(), "expected Err, got {outcome:?}"); + if let JobOutcome::Err { message } = outcome { + assert!(message.starts_with("boom (run_id=")); + } + assert_eq!(provider.last_status(), Some(RunStatus::Failed)); + } + + #[tokio::test] + async fn resume_hands_cursor_back_to_handler() { + let provider = Arc::new(MemProvider::new()); + let provider_trait: Arc = provider.clone(); + + // Run 1 pauses with cursor. + run_or_resume( + Arc::new(PausingHandler), + provider_trait.clone(), + &JobRunArgs::default(), + ) + .await; + assert_eq!(provider.last_status(), Some(RunStatus::Paused)); + + // Run 2 must see resume_cursor = the paused cursor. + let seen = Arc::new(Mutex::new(None)); + run_or_resume( + Arc::new(ResumeInspectHandler { + saw_cursor: seen.clone(), + }), + provider_trait, + &JobRunArgs::default(), + ) + .await; + assert_eq!(*seen.lock().unwrap(), Some(b"halfway".to_vec())); + } + + #[tokio::test] + async fn concurrent_trigger_hits_already_active() { + let provider = Arc::new(MemProvider::new()); + let provider_trait: Arc = provider.clone(); + + // Seed a Running row (simulates an in-flight prior dispatch). + let seeded_run_id = provider.seed_running(); + + let outcome = run_or_resume( + Arc::new(CompletingHandler), + provider_trait, + &JobRunArgs::default(), + ) + .await; + + // Must be Ok with skipped=already_running, NOT a fresh dispatch. + assert!(outcome.is_ok()); + if let JobOutcome::Ok { extra, .. } = &outcome { + assert_eq!(extra["skipped"], "already_running"); + assert_eq!(extra["run_id"], seeded_run_id.to_string()); + assert_eq!(extra["status"], "Running"); + } + // Seeded run's status untouched (no parallel dispatch happened). + assert_eq!(provider.last_status(), Some(RunStatus::Running)); + } + + #[tokio::test] + async fn boot_recovery_sweep_flips_running_to_paused() { + let provider = Arc::new(MemProvider::new()); + let provider_trait: Arc = provider.clone(); + + provider.seed_running(); + provider.seed_running(); + + let flipped = provider_trait.boot_recovery_sweep().await.unwrap(); + assert_eq!(flipped, 2); + assert_eq!(provider.last_status(), Some(RunStatus::Paused)); + } + + #[tokio::test] + async fn runstatus_parse_is_symmetric() { + for s in [ + RunStatus::Running, + RunStatus::Paused, + RunStatus::CancelRequested, + RunStatus::Completed, + RunStatus::Failed, + ] { + assert_eq!(RunStatus::parse(s.as_str()), Some(s)); + } + assert!(RunStatus::parse("garbage").is_none()); + } +}