aion-store 0.27.1

//! Shared workloop persistence scenarios (workloop brief Leg 2).

use std::collections::BTreeMap;
use std::sync::Arc;
use std::time::Duration;

use aion_core::{
    ContentType, InvariantSpec, Payload, ToleranceSpec, WorkflowId, WorkloopArming, WorkloopSpec,
};
use chrono::{DateTime, TimeZone, Utc};

use super::{WorkloopRawWrite, WorkloopReopen};
use crate::StoreError;
use crate::workloop::{InvariantHealthState, InvariantStateRecord, WorkloopRecord, WorkloopStore};

pub(crate) async fn run(
    store: Arc<dyn WorkloopStore>,
    reopen: Option<WorkloopReopen>,
    write_raw: WorkloopRawWrite,
) -> Result<(), StoreError> {
    round_trip(&store).await?;
    due_sweep_is_bounded_by_next_check(&store).await?;
    removal_tombstones_the_loop(&store).await?;
    invariant_records_rotate_and_prune_within_window(&store).await?;
    current_record_survives_pruning_indefinitely(&store).await?;
    poisoned_row_is_visible_not_swept(&store, &write_raw).await?;
    drop(write_raw);
    restart_survival(store, reopen).await
}

fn instant(offset: i64) -> Result<DateTime<Utc>, StoreError> {
    Utc.with_ymd_and_hms(2026, 8, 25, 12, 0, 0)
        .single()
        .map(|base| base + chrono::Duration::seconds(offset))
        .ok_or_else(|| StoreError::Backend("conformance instant is invalid".to_owned()))
}

fn spec() -> Result<WorkloopSpec, StoreError> {
    let arming = WorkloopArming::every(Duration::from_secs(1500))
        .map_err(|error| StoreError::Serialization(error.to_string()))?;
    let tolerance = ToleranceSpec::both(3, Duration::from_secs(2700))
        .map_err(|error| StoreError::Serialization(error.to_string()))?;
    WorkloopSpec::new(
        arming,
        vec![InvariantSpec {
            name: String::from("serving"),
            record_type: String::from("ServeState"),
            tolerance,
            confirms: vec![String::from("sweep")],
        }],
        Duration::from_secs(14 * 86_400),
    )
    .map_err(|error| StoreError::Serialization(error.to_string()))
}

fn record(loop_id: WorkflowId, next_check_offset: i64) -> Result<WorkloopRecord, StoreError> {
    Ok(WorkloopRecord {
        loop_id,
        namespace: String::from("default"),
        spec: spec()?,
        window_seq: 2,
        next_window_at: Some(instant(next_check_offset)?),
        next_check_at: Some(instant(next_check_offset)?),
        last_iteration_closed_window: Some(2),
        invariant_health: BTreeMap::from([(
            String::from("serving"),
            InvariantHealthState {
                last_confirmed_at: Some(instant(-60)?),
                consecutive_unconfirmed: 0,
                last_evidence: None,
                alarmed: false,
            },
        )]),
        registered_at: instant(-3000)?,
        updated_at: instant(0)?,
    })
}

fn state_record(
    loop_id: WorkflowId,
    label: &str,
    offset: i64,
) -> Result<InvariantStateRecord, StoreError> {
    Ok(InvariantStateRecord {
        loop_id,
        invariant: String::from("serving"),
        payload: Payload::new(
            ContentType::Json,
            format!("{{\"label\":\"{label}\"}}").into_bytes(),
        ),
        record_type: String::from("ServeState"),
        window_seq: Some(2),
        recorded_at: instant(offset)?,
    })
}

async fn round_trip(store: &Arc<dyn WorkloopStore>) -> Result<(), StoreError> {
    let loop_id = WorkflowId::new_v4();
    let expected = record(loop_id.clone(), 300)?;
    store.put_workloop(expected.clone()).await?;

    let loaded = store.get_workloop(&loop_id).await?;
    if loaded.as_ref() != Some(&expected) {
        return Err(StoreError::Backend(format!(
            "workloop round trip mismatch: {loaded:?}"
        )));
    }

    let absent = store.get_workloop(&WorkflowId::new_v4()).await?;
    if absent.is_some() {
        return Err(StoreError::Backend(
            "absent workloop must answer None".to_owned(),
        ));
    }
    Ok(())
}

async fn due_sweep_is_bounded_by_next_check(
    store: &Arc<dyn WorkloopStore>,
) -> Result<(), StoreError> {
    let due_id = WorkflowId::new_v4();
    let boundary_id = WorkflowId::new_v4();
    let sleeping_id = WorkflowId::new_v4();
    store.put_workloop(record(due_id.clone(), -10)?).await?;
    store.put_workloop(record(boundary_id.clone(), 0)?).await?;
    store
        .put_workloop(record(sleeping_id.clone(), 3600)?)
        .await?;

    let due = store.due_workloops(instant(0)?).await?;
    let due_ids: Vec<_> = due.iter().map(|record| record.loop_id.clone()).collect();
    if !due_ids.contains(&due_id) || !due_ids.contains(&boundary_id) {
        return Err(StoreError::Backend(
            "past-due and at-boundary loops must both be due (due is inclusive)".to_owned(),
        ));
    }
    if due_ids.contains(&sleeping_id) {
        return Err(StoreError::Backend(
            "a sleeping loop with a future next_check_at must not appear in the due sweep"
                .to_owned(),
        ));
    }
    Ok(())
}

async fn removal_tombstones_the_loop(store: &Arc<dyn WorkloopStore>) -> Result<(), StoreError> {
    let loop_id = WorkflowId::new_v4();
    store.put_workloop(record(loop_id.clone(), -10)?).await?;

    if !store.remove_workloop(&loop_id).await? {
        return Err(StoreError::Backend(
            "removing an existing workloop must report existence".to_owned(),
        ));
    }
    if store.remove_workloop(&loop_id).await? {
        return Err(StoreError::Backend(
            "removing an absent workloop must report absence".to_owned(),
        ));
    }
    if store.get_workloop(&loop_id).await?.is_some() {
        return Err(StoreError::Backend(
            "a removed workloop must not answer get".to_owned(),
        ));
    }
    let due = store.due_workloops(instant(0)?).await?;
    if due.iter().any(|record| record.loop_id == loop_id) {
        return Err(StoreError::Backend(
            "a removed workloop must not appear in the due sweep".to_owned(),
        ));
    }
    Ok(())
}

async fn invariant_records_rotate_and_prune_within_window(
    store: &Arc<dyn WorkloopStore>,
) -> Result<(), StoreError> {
    let loop_id = WorkflowId::new_v4();
    let first = state_record(loop_id.clone(), "first", -100)?;
    let second = state_record(loop_id.clone(), "second", -50)?;
    let third = state_record(loop_id.clone(), "third", -5)?;
    let fourth = state_record(loop_id.clone(), "fourth", -1)?;

    // A retention window older than every record: these three writes rotate
    // and prune NOTHING, so the rotation order below is measured on its own.
    let unbounded = instant(-1_000_000)?;
    for record in [&first, &second, &third] {
        let removed = store
            .put_invariant_record(record.clone(), unbounded)
            .await?;
        if removed != 0 {
            return Err(StoreError::Backend(format!(
                "a write whose retention window predates every generation must prune \
                 nothing, got {removed}"
            )));
        }
    }

    let current = store.current_invariant_record(&loop_id, "serving").await?;
    if current.as_ref() != Some(&third) {
        return Err(StoreError::Backend(format!(
            "current invariant record mismatch: {current:?}"
        )));
    }
    let generations = store
        .invariant_record_generations(&loop_id, "serving")
        .await?;
    if generations != vec![first, second.clone()] {
        return Err(StoreError::Backend(format!(
            "generation order mismatch (expected oldest first): {generations:?}"
        )));
    }

    // 🔴 THE INSTALL AND ITS RETENTION ARE ONE WRITE. This put rotates `third`
    // into the generation list AND applies a window that keeps only what is
    // newer than -60 — so the `first` generation must be gone when this single
    // call returns, with no separate prune anywhere.
    let removed = store
        .put_invariant_record(fourth.clone(), instant(-60)?)
        .await?;
    if removed != 1 {
        return Err(StoreError::Backend(format!(
            "the write must report exactly the generation count its retention window \
             removed, got {removed}"
        )));
    }
    let current = store.current_invariant_record(&loop_id, "serving").await?;
    if current.as_ref() != Some(&fourth) {
        return Err(StoreError::Backend(format!(
            "the installed record must be current after a pruning write: {current:?}"
        )));
    }
    let generations = store
        .invariant_record_generations(&loop_id, "serving")
        .await?;
    if generations != vec![second, third] {
        return Err(StoreError::Backend(format!(
            "post-prune generations mismatch: {generations:?}"
        )));
    }
    Ok(())
}

async fn current_record_survives_pruning_indefinitely(
    store: &Arc<dyn WorkloopStore>,
) -> Result<(), StoreError> {
    let loop_id = WorkflowId::new_v4();
    let only = state_record(loop_id.clone(), "only", -1_000_000)?;
    let successor = state_record(loop_id.clone(), "successor", -1)?;

    // An absent slot answers emptily rather than erroring, and a retention
    // window far newer than the record being installed prunes nothing: the
    // record a write installs is never aged out by its own window.
    let removed = store
        .put_invariant_record(only.clone(), instant(0)?)
        .await?;
    if removed != 0 {
        return Err(StoreError::Backend(
            "the first write to a slot has no prior generations to remove".to_owned(),
        ));
    }
    let current = store.current_invariant_record(&loop_id, "serving").await?;
    if current.as_ref() != Some(&only) {
        return Err(StoreError::Backend(
            "the current invariant record must survive any retention window".to_owned(),
        ));
    }

    // 🔴 THE DISCRIMINATING HALF. The same cutoff that left `only` alone as the
    // CURRENT record removes it the moment it is rotated into a prior
    // generation. Without this, the assertion above is satisfied by a store
    // that never prunes anything at all.
    let removed = store
        .put_invariant_record(successor.clone(), instant(0)?)
        .await?;
    if removed != 1 {
        return Err(StoreError::Backend(format!(
            "a rotated-out generation older than the window must be removed by the write \
             that rotated it, got {removed}"
        )));
    }
    let current = store.current_invariant_record(&loop_id, "serving").await?;
    if current.as_ref() != Some(&successor) {
        return Err(StoreError::Backend(
            "the newly installed record must be current".to_owned(),
        ));
    }
    let generations = store
        .invariant_record_generations(&loop_id, "serving")
        .await?;
    if !generations.is_empty() {
        return Err(StoreError::Backend(format!(
            "every prior generation older than the window must be gone: {generations:?}"
        )));
    }
    Ok(())
}

async fn poisoned_row_is_visible_not_swept(
    store: &Arc<dyn WorkloopStore>,
    write_raw: &WorkloopRawWrite,
) -> Result<(), StoreError> {
    let poisoned_id = WorkflowId::new_v4();
    write_raw(
        poisoned_id.to_string(),
        b"{\"not\":\"a workloop\"}".to_vec(),
    )
    .await?;

    let listing = store.list_workloops().await?;
    if !listing
        .undecodable
        .iter()
        .any(|row| row.loop_id == poisoned_id.to_string())
    {
        return Err(StoreError::Backend(
            "a poisoned workloop row must be visible in the listing".to_owned(),
        ));
    }
    // The due sweep must not fail — and must not include — a poisoned row.
    let due = store.due_workloops(instant(0)?).await?;
    if due
        .iter()
        .any(|record| record.loop_id.to_string() == poisoned_id.to_string())
    {
        return Err(StoreError::Backend(
            "a poisoned workloop row must not be swept as due".to_owned(),
        ));
    }
    // Removal works without requiring the row to decode.
    if !store.remove_workloop(&poisoned_id).await? {
        return Err(StoreError::Backend(
            "a poisoned workloop row must be removable".to_owned(),
        ));
    }
    Ok(())
}

async fn restart_survival(
    store: Arc<dyn WorkloopStore>,
    reopen: Option<WorkloopReopen>,
) -> Result<(), StoreError> {
    let Some(reopen) = reopen else {
        // The in-memory backend has no reopenable storage and is explicitly
        // exempt from the restart scenario.
        return Ok(());
    };
    let loop_id = WorkflowId::new_v4();
    let expected = record(loop_id.clone(), 600)?;
    store.put_workloop(expected.clone()).await?;
    let state = state_record(loop_id.clone(), "durable", -5)?;
    store
        .put_invariant_record(state.clone(), instant(-1_000_000)?)
        .await?;
    drop(store);

    let reopened = reopen().await?;
    let loaded = reopened.get_workloop(&loop_id).await?;
    if loaded.as_ref() != Some(&expected) {
        return Err(StoreError::Backend(format!(
            "workloop record must survive a store reopen: {loaded:?}"
        )));
    }
    let current = reopened
        .current_invariant_record(&loop_id, "serving")
        .await?;
    if current.as_ref() != Some(&state) {
        return Err(StoreError::Backend(format!(
            "invariant current record must survive a store reopen: {current:?}"
        )));
    }
    Ok(())
}