agentplane/runtime/

batch.rs

//! Driving a batch: reserve, run, record, repeat.
//!
//! The loop is deliberately boring, because the interesting decisions are all in
//! [`crate::batch`]'s types. What this file has to get right is the *ordering* —
//! an item's run id becomes durable before the run performs anything — and the
//! refusal to stop at the first failure.

use std::sync::Arc;

use crate::batch::{BatchCensus, BatchItem, BatchReport, BatchStatus, ItemOutcome, ItemSource};
use crate::core::{BatchId, PlanIR, RunId, RuntimeError, Spend};

use super::executor::{RunOutcome, RunStatus, Runtime};
use super::{Mode, metrics};

/// How a batch is to be run.
#[derive(Debug)]
pub struct BatchSpec {
    /// The plan every item runs. Frozen once, for the whole batch: a batch whose
    /// plan could change between items would be several acts wearing one name,
    /// and its audit record would say nothing about what actually happened to
    /// item 60,000.
    pub plan: PlanIR,
    pub source: Arc<dyn ItemSource>,
    /// Items fetched per call to the source.
    ///
    /// A paging hint, not a limit on the batch. See [`BatchSpec::page`].
    pub page: usize,
    /// Stop after this many items reach a terminal outcome, if set.
    ///
    /// For operating a very large batch in windows — "settle ten thousand and
    /// let me look". The batch stays `Running`; the next call resumes.
    pub max_items: Option<u64>,
}

impl BatchSpec {
    pub fn new(plan: PlanIR, source: Arc<dyn ItemSource>) -> Self {
        Self {
            plan,
            source,
            page: 256,
            max_items: None,
        }
    }

    /// How many items to fetch from the source at a time.
    #[must_use]
    pub const fn page(mut self, n: usize) -> Self {
        self.page = n;
        self
    }

    /// Process at most this many items before returning.
    #[must_use]
    pub const fn max_items(mut self, n: u64) -> Self {
        self.max_items = Some(n);
        self
    }
}

impl Runtime {
    /// Start a batch, or continue one.
    ///
    /// Passing an existing [`BatchId`] resumes it: processing picks up after the
    /// stored cursor, and any item that was reserved but never finished is
    /// replayed under its original run id rather than started again.
    ///
    /// Returns when the source is exhausted or `max_items` is reached. **A
    /// failing item does not stop the batch** — that is the failure isolation a
    /// batch exists for, and a caller who wants the opposite should be running
    /// one run, not many.
    ///
    /// What does stop it is an admission that never happened: a halt, a
    /// tenant ceiling, a store outage. Those hold for every item behind the
    /// one they hit, and they are returned as the error rather than written
    /// over the item, so the item is still pending when the next pass runs.
    pub async fn run_batch(
        &self,
        id: BatchId,
        spec: &BatchSpec,
    ) -> Result<BatchReport, RuntimeError> {
        let store = self.batches().ok_or_else(|| {
            RuntimeError::PlanContract(
                "batches need a batch store — build the runtime with `.batches(store)`".into(),
            )
        })?;

        crate::plan::validate(&spec.plan, &self.contract())
            .map_err(|e| RuntimeError::PlanContract(e.to_string()))?;
        store
            .open(id, &spec.plan.digest().to_hex())
            .await
            .map_err(RuntimeError::from_store)?;

        // Resume from where a previous pass got to. If this is lost the batch is
        // slower, not wrong: every item's reservation makes re-processing a
        // replay. See the module docs in `crate::batch`.
        let mut after = store.cursor(id).await.map_err(RuntimeError::from_store)?;
        let mut processed = 0_u64;

        loop {
            if spec.max_items.is_some_and(|m| processed >= m) {
                break;
            }
            let page = store_page(spec, after.as_deref()).await?;
            if page.is_empty() {
                // The source is done. Recorded durably: it is the difference
                // between a finished batch and one that stopped early with
                // every stored item terminal.
                store
                    .mark_exhausted(id)
                    .await
                    .map_err(RuntimeError::from_store)?;
                break;
            }

            for item in page {
                after = Some(item.key.clone());
                let outcome = self.run_item(id, spec, &item).await?;
                if outcome.is_terminal() {
                    processed += 1;
                }
                if spec.max_items.is_some_and(|m| processed >= m) {
                    break;
                }
            }
        }

        self.batch_report(id).await
    }

    /// Reserve, run, record — one item.
    ///
    /// The order is the effect protocol: the run id is durable before anything
    /// is performed, so an interrupted item is findable and replayable rather
    /// than a question nobody can answer.
    async fn run_item(
        &self,
        id: BatchId,
        spec: &BatchSpec,
        item: &BatchItem,
    ) -> Result<ItemOutcome, RuntimeError> {
        let store = self
            .batches()
            .ok_or_else(|| RuntimeError::PlanContract("no batch store".into()))?;

        let reserved = store
            .reserve(id, &item.key, RunId::generate())
            .await
            .map_err(RuntimeError::from_store)?;

        // Already terminal from an earlier pass: nothing to do, and re-running
        // would be the duplicate work the reservation exists to prevent.
        if let Some(done) = reserved.outcome.as_ref().filter(|o| o.is_terminal()) {
            return Ok(done.clone());
        }

        // A reservation with no terminal outcome is either a fresh item or one
        // interrupted mid-flight. Both are addressed by the same call: `replay`
        // resumes an existing journal, and a run that never wrote one is
        // admitted here for the first time.
        let outcome = if reserved.outcome.is_some() || self.has_journal(reserved.run).await? {
            self.replay(reserved.run, Mode::Resume).await
        } else {
            self.admit_plan_as(
                reserved.run,
                spec.plan.clone(),
                // Untrusted from `batch:<id>` unless the item says otherwise:
                // a source is somebody else's data far more often than it is
                // the operator's own.
                crate::core::Tainted::with_label(item.input.clone(), item.admission_label(id)),
                // A batch item's at-most-once identity is its reservation, not
                // an admission key: the run id is written to the batch store
                // before the run starts, so a retry replays that journal rather
                // than admitting again. Two mechanisms for one invariant would
                // be two places to disagree.
                super::executor::RunTerms::default(),
                // A batch pass is started from outside this process, so its
                // items stop when the plane stops admitting. The refusal is
                // not recorded, the reservation keeps no outcome, and the next
                // instance's pass re-admits exactly these items.
                super::executor::Entry::Outside,
            )
            .await
        };

        // Only a run that reached an outcome is an item outcome. An error
        // here is the plane refusing to *admit* — a halt, a ceiling, an
        // instance shutting down, a store that could not be reached — and it
        // holds for the next item as much
        // as this one, so the pass stops with it. Recording it would write a
        // terminal `Failed` over an item nothing ever ran: the halt lifts, the
        // ceiling frees, and the item stays failed forever. The reservation
        // is left without an outcome, which is exactly the state the next
        // pass reads as *reserved, never finished* and admits again.
        let out = outcome?;
        let (result, spend) = classify_item(out);
        store
            .record(id, &item.key, &result, spend)
            .await
            .map_err(RuntimeError::from_store)?;

        self.meter().count(metrics::BATCH_ITEMS, result.as_str());
        Ok(result)
    }

    /// Whether a run has already written anything.
    async fn has_journal(&self, run: RunId) -> Result<bool, RuntimeError> {
        match self.journal().head(run).await {
            Ok(head) => Ok(head.seq > 0),
            Err(e) => Err(RuntimeError::from_store(e)),
        }
    }

    /// The batch's current tally, without processing anything.
    pub async fn batch_report(&self, id: BatchId) -> Result<BatchReport, RuntimeError> {
        let store = self
            .batches()
            .ok_or_else(|| RuntimeError::PlanContract("batches need a batch store".into()))?;
        // Existence first, from the batch's own row. A census cannot answer
        // it — no such batch and a batch with no items yet both count zero
        // rows — and answering a mistyped id with an empty `Running` report
        // sends an operator watching a batch that will never exist.
        if store
            .plan_digest(id)
            .await
            .map_err(RuntimeError::from_store)?
            .is_none()
        {
            return Err(RuntimeError::from_store(crate::core::StoreError::NotFound(
                format!("batch {id}"),
            )));
        }
        let census = store.census(id).await.map_err(RuntimeError::from_store)?;
        let cursor = store.cursor(id).await.map_err(RuntimeError::from_store)?;
        let exhausted = store
            .is_exhausted(id)
            .await
            .map_err(RuntimeError::from_store)?;
        Ok(BatchReport {
            id,
            status: status_of(&census, exhausted),
            in_flight: census.in_flight + census.suspended,
            exhausted: census.exhausted,
            withheld: census.withheld,
            spend: census.spend,
            cursor,
        })
    }
}

/// A batch is complete only when nothing is reserved and nothing is waiting.
///
/// An exhausted item is waiting too — paused at a ceiling, resumable the
/// moment somebody raises it — and it holds the cursor for exactly that
/// reason. Reporting `Completed` over one would let `everything_settled()`
/// answer yes while an item's meters are unsettled, which is the "mostly
/// worked read as worked" this type exists to prevent.
///
/// Note there is no path to a status that says "succeeded": see `crate::batch`.
fn status_of(c: &BatchCensus, source_exhausted: bool) -> BatchStatus {
    if !source_exhausted || c.in_flight > 0 || c.suspended > 0 || c.exhausted > 0 || c.withheld > 0
    {
        return BatchStatus::Running;
    }
    BatchStatus::Completed {
        succeeded: c.succeeded,
        failed: c.failed,
        quarantined: c.quarantined,
    }
}

/// Turn a run's ending into an item's.
///
/// A store or contract error is *the item's* failure, not the batch's: one
/// malformed item must not take down a settlement run for the other 99,999. The
/// detail goes on the item record, and the item's own journal holds the rest.
fn classify_item(out: RunOutcome) -> (ItemOutcome, Spend) {
    let spend = out.spend();
    let item = match out.status {
        RunStatus::Succeeded => ItemOutcome::Succeeded,
        RunStatus::Suspended(r) => ItemOutcome::Suspended(r.to_string()),
        RunStatus::Quarantined(why) => ItemOutcome::Quarantined(why),
        // Also `Quarantined`, because this enum is deliberately coarse about
        // *what a human must do next* and the answer here is the same one:
        // hands off. An abandoned run may hold a mutation nobody could account
        // for, so filing it under `Failed` — the outcome this batch's readers
        // re-run — is the one mapping that could take money twice. The reason
        // names who gave up and why, which is what tells the two apart.
        RunStatus::Abandoned { actor, reason } => {
            ItemOutcome::Quarantined(format!("abandoned by '{actor}': {reason}"))
        }
        // `Replanning` is never observed by a caller — the executor resolves
        // it internally — but if one ever escaped it is an item that did not
        // settle, which is what `Failed` means here.
        RunStatus::Failed(why) | RunStatus::Replanning(why) => ItemOutcome::Failed(why),
        // A pause, not a fault — mirrored from the run's own semantics: the
        // work stands, and raising the ceiling resumes it. Filing it under
        // `Failed` taught operators to re-run items whose work was intact.
        RunStatus::Exhausted(limit) => ItemOutcome::Exhausted(limit.to_string()),
        // The same shape, from the other cause: the item's work stands and the
        // remedy is lifting the halt on the subject, not raising a budget — so
        // it is its own outcome rather than a reason string under `Exhausted`.
        RunStatus::Withheld { subject, reason } => {
            ItemOutcome::Withheld(format!("authority '{subject}' withdrawn: {reason}"))
        }
        // Unreachable by construction — a batch item is admitted as an ordinary
        // run and none of these is ever admitted — so reaching it means the
        // reservation points at a run this batch did not start. `Quarantined`
        // for the same reason abandonment is: the answer to *what does a person
        // do next* is hands off, and filing it under `Failed` would hand it to
        // the readers who re-run.
        RunStatus::Swept
        | RunStatus::BrokeGlass { .. }
        | RunStatus::HaltLifted { .. }
        | RunStatus::HoldReleased { .. }
        | RunStatus::Observed => ItemOutcome::Quarantined(
            "this item's run id resolves to one of the plane's own records, or to \
                 a session it only observed, rather than to the item's run — the \
                 reservation and the journal disagree"
                .to_owned(),
        ),
        // An item somebody stopped did not settle, and the batch must not
        // report otherwise — but the reason names the person, so a partial
        // batch can be told apart from one that hit a wall.
        RunStatus::Cancelled { actor, reason } => {
            ItemOutcome::Failed(format!("cancelled by '{actor}': {reason}"))
        }
    };
    (item, spend)
}

/// Fetch the next page, mapping a source failure onto the runtime's errors.
async fn store_page(spec: &BatchSpec, after: Option<&str>) -> Result<Vec<BatchItem>, RuntimeError> {
    let page = spec
        .source
        .next(after, spec.page)
        .await
        // Surfaced, never swallowed into an empty page: a source failure that
        // read as "no more items" would truncate a settlement run silently,
        // which is the exact failure this crate exists to make loud.
        .map_err(|e| RuntimeError::PlanContract(e.to_string()))?;
    in_cursor_order(after, &page)?;
    Ok(page)
}

/// Refuse a page whose keys do not strictly increase in byte order past
/// `after`.
///
/// The resume cursor is the greatest key of the finished prefix, compared as
/// bytes. A source that pages in any other order — a locale collation, a
/// numeric sort of `"9"` before `"10"`, a repeated key — resumes from a cursor
/// that skips items it never ran or re-offers ones it did, and the batch
/// reports complete either way.
fn in_cursor_order(after: Option<&str>, page: &[BatchItem]) -> Result<(), RuntimeError> {
    let mut previous = after;
    for item in page {
        if let Some(p) = previous
            && item.key.as_bytes() <= p.as_bytes()
        {
            return Err(RuntimeError::PlanContract(format!(
                "batch source returned key '{}' after '{p}': keys must strictly \
                 increase in byte order, which is the order the resume cursor \
                 compares — a source paging in any other order skips or repeats \
                 items on resume",
                item.key
            )));
        }
        previous = Some(item.key.as_str());
    }
    Ok(())
}