aion_store/

workloop.rs

//! Durable workloop records and persistence contract (workloop brief Leg 2).
//!
//! Two record families:
//!
//! - **[`WorkloopRecord`]** — one per registered loop: the declared
//!   [`WorkloopSpec`], the engine's cadence bookkeeping (window sequence, next
//!   window, next check), and per-invariant health accounting. A sleeping loop
//!   IS this record plus its event history — no resident process, no scheduler
//!   presence (R13.3). The cadence sweeper drives itself off
//!   [`WorkloopStore::due_workloops`].
//! - **[`InvariantStateRecord`]** — the current-state record of one invariant
//!   (R7): a store document of its declared type, type-erased here as a
//!   [`Payload`] plus the declared type name. Exactly one CURRENT record per
//!   invariant survives indefinitely; prior generations are retention-bounded
//!   on the loop's declared window (R8.1) and live in the same single-key slot
//!   so every mutation is one atomic value swap on any backend.

use aion_core::{Payload, WorkflowId, WorkloopSpec};
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;

use crate::StoreError;

/// Per-invariant health accounting the engine persists between sweeps.
///
/// The values are COMPUTED by the engine's health machinery (pure, in
/// `aion-rs`); the store only carries them so tolerance accounting survives a
/// restart — an engine that forgot how many windows a loop had already missed
/// would silently re-grant the whole tolerance budget on every boot.
#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct InvariantHealthState {
    /// When the invariant was last confirmed, if ever.
    pub last_confirmed_at: Option<DateTime<Utc>>,
    /// Consecutive unconfirmed samples/missed windows since the last
    /// confirmation.
    pub consecutive_unconfirmed: u64,
    /// The kind of unconfirmed evidence last accrued (`sample-red` or
    /// `window-missed`), so a duration-form expiry can name the evidence it
    /// sits on; `None` when the invariant has accrued no unconfirmed evidence
    /// since its last confirmation — the `unconfirmed-unknown` case.
    pub last_evidence: Option<aion_core::AlarmCause>,
    /// Whether the tolerance-exceeded alarm has already been raised for the
    /// current unconfirmed run (edge-triggered: reset on confirmation, so a
    /// sustained breach alarms once, not once per sweep).
    pub alarmed: bool,
}

/// One durable workloop registration: the declared spec plus the engine's
/// cadence and health bookkeeping.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct WorkloopRecord {
    /// The loop's workflow identity (primary key).
    pub loop_id: WorkflowId,
    /// Namespace the loop runs in.
    pub namespace: String,
    /// The declared spec — re-validated on decode by [`WorkloopSpec`]'s own
    /// serde boundary, so a stored record cannot drift out of its declared
    /// invariants.
    pub spec: WorkloopSpec,
    /// Sequence number of the last cadence window that fired (0 = none yet;
    /// fired windows are one-based).
    pub window_seq: u64,
    /// When the next cadence window fires; `None` on a signal-only loop.
    pub next_window_at: Option<DateTime<Utc>>,
    /// When the sweeper must next evaluate this loop: the minimum of
    /// `next_window_at` and the earliest duration-form tolerance deadline.
    /// `None` = nothing left to check engine-side (a signal-only loop whose
    /// alarms are all latched); sleeping loops with a future value cost the
    /// sweeper nothing.
    pub next_check_at: Option<DateTime<Utc>>,
    /// The window sequence of the last recorded iteration close, if any —
    /// what the dead-man switch compares against the fired window (R4.3).
    pub last_iteration_closed_window: Option<u64>,
    /// Per-invariant health accounting, keyed by invariant name.
    pub invariant_health: BTreeMap<String, InvariantHealthState>,
    /// When the loop was registered.
    pub registered_at: DateTime<Utc>,
    /// Most recent mutation instant.
    pub updated_at: DateTime<Utc>,
}

impl WorkloopRecord {
    /// Encode the stable backend-neutral representation.
    ///
    /// # Errors
    ///
    /// Returns [`StoreError::Serialization`] when serialization fails.
    pub fn encode(&self) -> Result<Vec<u8>, StoreError> {
        serde_json::to_vec(self).map_err(|error| StoreError::Serialization(error.to_string()))
    }

    /// Decode and re-validate a stored representation (the spec's serde
    /// boundary re-runs every declaration check).
    ///
    /// # Errors
    ///
    /// Returns [`StoreError::Serialization`] for malformed bytes or a spec
    /// that no longer satisfies its declaration invariants.
    pub fn decode(bytes: &[u8]) -> Result<Self, StoreError> {
        serde_json::from_slice(bytes).map_err(|error| StoreError::Serialization(error.to_string()))
    }
}

/// The current-state record of one invariant (R7): a store document of the
/// invariant's declared type, type-erased as payload bytes + type name.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct InvariantStateRecord {
    /// Loop the invariant belongs to.
    pub loop_id: WorkflowId,
    /// Invariant name.
    pub invariant: String,
    /// The record value, typed at the surface by `record_type` and carried
    /// type-erased here (the load-bearing type-erasure invariant).
    pub payload: Payload,
    /// Declared AWL type name of the value — provenance, never a schema.
    pub record_type: String,
    /// Cadence window the value was produced in; `None` on signal-only loops.
    pub window_seq: Option<u64>,
    /// When the value was recorded.
    pub recorded_at: DateTime<Utc>,
}

/// The single-key slot holding one invariant's CURRENT record plus its
/// retention-bounded prior generations (oldest first).
///
/// One value per (loop, invariant) key keeps every mutation — rotation on a
/// new current, pruning to the declared window — a single atomic value swap,
/// which every backend (including the distributed CAS path) supports without
/// multi-key transactions. Retention is what keeps the slot bounded: the
/// declared window is mandatory, so unbounded growth is unrepresentable.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct InvariantRecordSlot {
    /// The current record — survives indefinitely, NEVER pruned (R8.1).
    pub current: InvariantStateRecord,
    /// Prior generations within the retention window, oldest first.
    pub previous: Vec<InvariantStateRecord>,
}

impl InvariantRecordSlot {
    /// Encode the stable backend-neutral representation.
    ///
    /// # Errors
    ///
    /// Returns [`StoreError::Serialization`] when serialization fails.
    pub fn encode(&self) -> Result<Vec<u8>, StoreError> {
        serde_json::to_vec(self).map_err(|error| StoreError::Serialization(error.to_string()))
    }

    /// Decode a stored representation.
    ///
    /// # Errors
    ///
    /// Returns [`StoreError::Serialization`] for malformed bytes.
    pub fn decode(bytes: &[u8]) -> Result<Self, StoreError> {
        serde_json::from_slice(bytes).map_err(|error| StoreError::Serialization(error.to_string()))
    }

    /// Rotate `record` in as the new current, pushing the old current onto the
    /// generation list.
    #[must_use]
    pub fn rotated(mut self, record: InvariantStateRecord) -> Self {
        self.previous.push(self.current);
        self.current = record;
        self
    }

    /// Drop prior generations recorded before `older_than`, returning how many
    /// were removed. The current record is never touched (R8.1).
    pub fn prune(&mut self, older_than: DateTime<Utc>) -> u64 {
        let before = self.previous.len();
        self.previous
            .retain(|record| record.recorded_at >= older_than);
        u64::try_from(before.saturating_sub(self.previous.len())).unwrap_or(u64::MAX)
    }
}

/// The one field the due sweep filters on, decodable WITHOUT validating the
/// whole record: a thousand sleeping loops must cost the sweeper a comparison
/// each, not a full spec re-validation each (the R13.3 measurement is what
/// caught the difference — ~20µs per full decode versus nanoseconds per
/// comparison).
#[derive(Clone, Copy, Debug, Deserialize)]
pub struct WorkloopDueProbe {
    /// Mirror of [`WorkloopRecord::next_check_at`].
    pub next_check_at: Option<DateTime<Utc>>,
}

impl WorkloopDueProbe {
    /// Probe-decode just the due instant from stored workloop bytes.
    ///
    /// # Errors
    ///
    /// Returns [`StoreError::Serialization`] for bytes that are not a JSON
    /// object carrying the field — a poisoned row, which the sweep skips and
    /// the listing surfaces.
    pub fn decode(bytes: &[u8]) -> Result<Self, StoreError> {
        serde_json::from_slice(bytes).map_err(|error| StoreError::Serialization(error.to_string()))
    }
}

/// A workloop row that was present but could not be decoded.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct UndecodableWorkloop {
    /// Key under which the poisoned row is stored (the loop id's text form).
    pub loop_id: String,
    /// Decode failure rendered for operator diagnosis.
    pub error: String,
}

/// Complete workloop listing, including poisoned-row visibility.
#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct WorkloopListing {
    /// Successfully decoded records, ordered by loop id text.
    pub workloops: Vec<WorkloopRecord>,
    /// Present rows that could not be decoded, ordered by loop id text.
    pub undecodable: Vec<UndecodableWorkloop>,
}

/// Durable workloop persistence contract.
#[async_trait]
pub trait WorkloopStore: Send + Sync + 'static {
    /// Create or replace a loop's registration record.
    async fn put_workloop(&self, record: WorkloopRecord) -> Result<(), StoreError>;

    /// Look up one loop by id.
    async fn get_workloop(
        &self,
        loop_id: &WorkflowId,
    ) -> Result<Option<WorkloopRecord>, StoreError>;

    /// List decodable loops and report every undecodable row, both ordered by
    /// loop id text.
    async fn list_workloops(&self) -> Result<WorkloopListing, StoreError>;

    /// Loops whose `next_check_at` is `Some` and at or before `as_of` — the
    /// sweeper's work set. A sleeping loop with a future (or absent) check
    /// instant never appears, which is what makes a thousand sleeping loops
    /// cost the sweeper nothing. Undecodable rows are excluded here (they
    /// surface via [`WorkloopStore::list_workloops`], never silently in the
    /// hot path).
    async fn due_workloops(&self, as_of: DateTime<Utc>) -> Result<Vec<WorkloopRecord>, StoreError>;

    /// Remove a loop's registration (retirement). Returns whether a row
    /// existed. Invariant current-state records are NOT removed: the current
    /// record survives indefinitely by declaration (R8.1).
    async fn remove_workloop(&self, loop_id: &WorkflowId) -> Result<bool, StoreError>;

    /// Install `record` as the invariant's current record — rotating any prior
    /// current into the generation list — and prune prior generations recorded
    /// before `prune_before`, in ONE read-modify-write. Returns how many prior
    /// generations the prune removed. The current record is never pruned; it
    /// survives indefinitely by declaration (R8.1).
    ///
    /// # 🔴 THE INSTALL AND THE PRUNE ARE ONE COMMIT, DELIBERATELY
    ///
    /// Both halves address the SAME key — one invariant's slot — and the close
    /// path performs them back to back on every park. Two separate
    /// read-modify-write commits made a park cost `2 + 2N` spine-linear
    /// commits for `N` declared invariants; folding them makes it `2 + N`. The
    /// second commit read back bytes the first had just written, paid a second
    /// durable round trip for them, and could interleave with nothing useful:
    /// a park between them would observe a slot whose retention window had not
    /// yet been applied. One commit removes both the cost and that window.
    ///
    /// Retention is therefore not a separate verb that a caller could forget:
    /// installing a record IS what prunes the ones it aged out, so declared
    /// retention is retention that happens.
    async fn put_invariant_record(
        &self,
        record: InvariantStateRecord,
        prune_before: DateTime<Utc>,
    ) -> Result<u64, StoreError>;

    /// The invariant's current record, if one was ever produced.
    async fn current_invariant_record(
        &self,
        loop_id: &WorkflowId,
        invariant: &str,
    ) -> Result<Option<InvariantStateRecord>, StoreError>;

    /// The invariant's prior generations within retention, oldest first.
    async fn invariant_record_generations(
        &self,
        loop_id: &WorkflowId,
        invariant: &str,
    ) -> Result<Vec<InvariantStateRecord>, StoreError>;
}

#[cfg(test)]
mod tests {
    use std::time::Duration;

    use aion_core::{ContentType, InvariantSpec, ToleranceSpec, WorkloopArming};
    use chrono::TimeZone;

    use super::*;

    fn instant(offset: i64) -> Result<DateTime<Utc>, &'static str> {
        Utc.with_ymd_and_hms(2026, 8, 25, 1, 0, 0)
            .single()
            .map(|base| base + chrono::Duration::seconds(offset))
            .ok_or("test instant must be valid")
    }

    fn spec() -> Result<WorkloopSpec, Box<dyn std::error::Error>> {
        Ok(WorkloopSpec::new(
            WorkloopArming::every(Duration::from_secs(1500))?,
            vec![InvariantSpec {
                name: String::from("serving"),
                record_type: String::from("ServeState"),
                tolerance: ToleranceSpec::both(3, Duration::from_secs(2700))?,
                confirms: vec![String::from("sweep")],
            }],
            Duration::from_secs(14 * 86_400),
        )?)
    }

    fn record() -> Result<WorkloopRecord, Box<dyn std::error::Error>> {
        let registered_at = instant(0)?;
        Ok(WorkloopRecord {
            loop_id: WorkflowId::new(uuid::Uuid::from_u128(9)),
            namespace: String::from("default"),
            spec: spec()?,
            window_seq: 4,
            next_window_at: Some(instant(1500)?),
            next_check_at: Some(instant(1500)?),
            last_iteration_closed_window: Some(4),
            invariant_health: BTreeMap::from([(
                String::from("serving"),
                InvariantHealthState {
                    last_confirmed_at: Some(instant(60)?),
                    consecutive_unconfirmed: 1,
                    last_evidence: Some(aion_core::AlarmCause::WindowMissed),
                    alarmed: false,
                },
            )]),
            registered_at,
            updated_at: instant(90)?,
        })
    }

    fn state_record(offset: i64) -> Result<InvariantStateRecord, Box<dyn std::error::Error>> {
        Ok(InvariantStateRecord {
            loop_id: WorkflowId::new(uuid::Uuid::from_u128(9)),
            invariant: String::from("serving"),
            payload: Payload::new(ContentType::Json, b"{\"connected\":2}".to_vec()),
            record_type: String::from("ServeState"),
            window_seq: Some(4),
            recorded_at: instant(offset)?,
        })
    }

    #[test]
    fn workloop_record_round_trips() -> Result<(), Box<dyn std::error::Error>> {
        let expected = record()?;
        assert_eq!(WorkloopRecord::decode(&expected.encode()?)?, expected);
        Ok(())
    }

    #[test]
    fn a_stored_record_with_an_invalid_spec_refuses_decode()
    -> Result<(), Box<dyn std::error::Error>> {
        let encoded = record()?.encode()?;
        let mut value: serde_json::Value = serde_json::from_slice(&encoded)?;
        // Strip the tolerance declaration in the stored bytes: decode must
        // refuse rather than resurrect an undeclared-tolerance loop.
        value["spec"]["invariants"][0]["tolerance"] = serde_json::json!({
            "consecutive_windows": null,
            "unconfirmed_for": null,
        });
        let error = WorkloopRecord::decode(&serde_json::to_vec(&value)?).err();
        assert!(matches!(error, Some(StoreError::Serialization(_))));
        Ok(())
    }

    #[test]
    fn slot_rotation_keeps_current_and_orders_generations() -> Result<(), Box<dyn std::error::Error>>
    {
        let first = state_record(0)?;
        let second = state_record(10)?;
        let third = state_record(20)?;
        let slot = InvariantRecordSlot {
            current: first.clone(),
            previous: Vec::new(),
        }
        .rotated(second.clone())
        .rotated(third.clone());

        assert_eq!(slot.current, third);
        assert_eq!(slot.previous, vec![first, second]);
        Ok(())
    }

    #[test]
    fn prune_removes_only_out_of_window_generations_and_never_current()
    -> Result<(), Box<dyn std::error::Error>> {
        let stale_current = state_record(-100)?;
        let mut slot = InvariantRecordSlot {
            current: stale_current.clone(),
            previous: vec![state_record(-50)?, state_record(-10)?, state_record(5)?],
        };

        let removed = slot.prune(instant(0)?);

        assert_eq!(removed, 2);
        assert_eq!(slot.previous, vec![state_record(5)?]);
        // The current record is older than the cutoff and still survives:
        // exactly one current record per invariant survives indefinitely.
        assert_eq!(slot.current, stale_current);
        Ok(())
    }

    #[test]
    fn slot_round_trips() -> Result<(), Box<dyn std::error::Error>> {
        let slot = InvariantRecordSlot {
            current: state_record(0)?,
            previous: vec![state_record(-10)?],
        };
        assert_eq!(InvariantRecordSlot::decode(&slot.encode()?)?, slot);
        Ok(())
    }
}