grafos_pipeline/

checkpoint.rs

//! Durable checkpoint support for pipeline stage state.

extern crate alloc;
use alloc::boxed::Box;
use alloc::format;
use alloc::vec::Vec;

use grafos_std::block::{BlockLease, BLOCK_SIZE};
use serde::{de::DeserializeOwned, Serialize};

use crate::error::EdgeError;

/// Magic bytes identifying a pipeline checkpoint record.
const MAGIC: [u8; 4] = *b"PCHK";
const VERSION: u32 = 1;
const HEADER_BLOCK: u64 = 0;
const DATA_START_BLOCK: u64 = 1;

// ---------------------------------------------------------------------------
// CheckpointFlushObserver — Phase 218.2 / slice 96
// ---------------------------------------------------------------------------

/// Sideband observer notified at checkpoint flush start and complete.
///
/// Slice 96 wires this so the scheduler-service `PreemptionManager`
/// can mark a lease as `NonPreemptibleReason::CheckpointInProgress`
/// for the duration of a flush. The trait is intentionally minimal
/// and observational — implementations MUST NOT block, return errors
/// the flush can act on, or otherwise change the flush's success
/// semantics. A failed observer notification (e.g. lease not tracked
/// by the preemption manager because it was released / expired /
/// fenced concurrently) is a sideband condition the observer logs
/// internally; it does NOT abort the checkpoint flush. The
/// checkpoint is the source of truth for stage durability; the
/// observer is a sideband signal to the scheduler.
///
/// **Why a trait callback rather than a direct cross-crate call?**
/// `grafos-pipeline` is `no_std + alloc` and intentionally has no
/// dependency on `grafos-scheduler-service` (which is `std`-only and
/// depends on the http server, audit chain, etc.). A direct call
/// would either force the pipeline crate into `std` or pull
/// scheduler-service into `no_std` builds. The trait keeps both
/// sides cycle-free: the producer (pipeline) calls a small trait;
/// the consumer (scheduler-service) implements it.
///
/// **Lifecycle contract:**
/// - `on_flush_start(lease_id)` is called exactly ONCE at the
///   beginning of a `checkpoint()` call when both the observer and
///   `lease_id` are wired.
/// - `on_flush_complete(lease_id)` is called exactly ONCE at the
///   end of `checkpoint()` — on BOTH the success path and every
///   failure path. The implementation is responsible for the RAII
///   guard via a `Drop`-based scope; producers don't need to
///   sprinkle clears at every error site.
pub trait CheckpointFlushObserver {
    /// Called at the START of a checkpoint flush. Implementations
    /// should mark the lease as transiently non-preemptible.
    fn on_flush_start(&self, lease_id: u128);

    /// Called at the END of a checkpoint flush (both success AND
    /// failure paths). Implementations should clear the
    /// non-preemptible marker.
    fn on_flush_complete(&self, lease_id: u128);
}

/// RAII guard that calls `on_flush_complete` on drop. Ensures the
/// marker is cleared even if the flush errors part-way through. Per
/// the slice 96 requirement: "if a flush fails partway, the marker
/// MUST be cleared so the lease becomes preemptible again."
struct FlushGuard<'a> {
    observer: &'a dyn CheckpointFlushObserver,
    lease_id: u128,
}

impl<'a> Drop for FlushGuard<'a> {
    fn drop(&mut self) {
        self.observer.on_flush_complete(self.lease_id);
    }
}

/// A stage state wrapper with durable checkpoint/restore support.
///
/// Serializes the state to a [`BlockLease`] on checkpoint and
/// deserializes on restore, using a format compatible with the
/// grafos-collections Durable pattern.
///
/// # Fail-closed
///
/// If a checkpoint write fails, the error is propagated immediately.
/// Callers should halt the stage rather than continue with an
/// uncheckpointed state.
///
/// # Preemption protection (Phase 218.2 / slice 96)
///
/// When constructed with [`CheckpointedStageState::with_observer`],
/// the wrapper notifies a [`CheckpointFlushObserver`] at flush start
/// and complete so the scheduler-service can mark the lease as
/// `NonPreemptibleReason::CheckpointInProgress`. The marker is
/// cleared on BOTH the success path and every failure path via an
/// RAII guard. Wrappers constructed via [`CheckpointedStageState::new`]
/// continue to flush without the sideband notification (back-compat).
pub struct CheckpointedStageState<T> {
    state: T,
    block_lease: BlockLease,
    /// Identifier of the lease backing this stage's durable state.
    /// `None` for back-compat callers; `Some(_)` when the producer
    /// has named the lease whose preemption window the flush
    /// occupies.
    lease_id: Option<u128>,
    /// Sideband observer notified at flush start / complete. `None`
    /// for back-compat callers; `Some(_)` when wired via
    /// [`Self::with_observer`].
    observer: Option<Box<dyn CheckpointFlushObserver + Send + Sync>>,
}

impl<T: Serialize + DeserializeOwned> CheckpointedStageState<T> {
    /// Create a new checkpointed state wrapper.
    ///
    /// Does **not** write an initial checkpoint. Call [`checkpoint`](Self::checkpoint)
    /// to persist.
    ///
    /// Constructed without a [`CheckpointFlushObserver`] — checkpoint
    /// flushes do NOT emit preemption-protection sideband notifications.
    /// Use [`Self::with_observer`] to attach the marker producer for
    /// `NonPreemptibleReason::CheckpointInProgress` (Phase 218.2 /
    /// slice 96).
    pub fn new(state: T, block_lease: BlockLease) -> Self {
        Self {
            state,
            block_lease,
            lease_id: None,
            observer: None,
        }
    }

    /// Phase 218.2 / slice 96 — attach a [`CheckpointFlushObserver`]
    /// and the `lease_id` of the lease whose preemption window the
    /// flush occupies. Calls to [`Self::checkpoint`] then notify the
    /// observer at flush start AND on flush complete (both success
    /// AND failure paths, via an RAII guard).
    ///
    /// The producer is responsible for naming the correct
    /// `lease_id` — the lease backing the stage's durable
    /// checkpoint storage ([`BlockLease`]). The observer is
    /// typically the scheduler-service `PreemptionManager`'s
    /// checkpoint adapter (see
    /// `grafos_scheduler_service::preemption::PreemptionCheckpointObserver`).
    pub fn with_observer(
        state: T,
        block_lease: BlockLease,
        lease_id: u128,
        observer: Box<dyn CheckpointFlushObserver + Send + Sync>,
    ) -> Self {
        Self {
            state,
            block_lease,
            lease_id: Some(lease_id),
            observer: Some(observer),
        }
    }

    /// Serialize the current state to block storage.
    ///
    /// Uses a header block + data block layout:
    /// - Block 0: magic (4) + version (4) + data_len (8)
    /// - Blocks 1..N: postcard-serialized state bytes
    ///
    /// # Preemption protection (Phase 218.2 / slice 96)
    ///
    /// When this `CheckpointedStageState` was constructed via
    /// [`Self::with_observer`], the observer's `on_flush_start` is
    /// called at the top of this method and `on_flush_complete` is
    /// called when the flush returns — on BOTH the success path and
    /// every failure path. The notification is implemented via an
    /// RAII guard so partway-through errors don't leave the lease
    /// stuck in `CheckpointInProgress`.
    ///
    /// # Errors
    ///
    /// Returns [`EdgeError::CheckpointFailed`] on serialization or I/O failure.
    pub fn checkpoint(&mut self) -> Result<(), EdgeError> {
        // Phase 218.2 / slice 96 — sideband flush-protection
        // notification. The guard's Drop ensures `on_flush_complete`
        // runs on every exit path (Ok, Err, or panic-unwind).
        let _flush_guard = match (self.observer.as_deref(), self.lease_id) {
            (Some(observer), Some(lease_id)) => {
                observer.on_flush_start(lease_id);
                Some(FlushGuard { observer, lease_id })
            }
            _ => None,
        };

        let data = postcard::to_allocvec(&self.state)
            .map_err(|e| EdgeError::CheckpointFailed(format!("serialize: {e}")))?;

        let data_len = data.len() as u64;

        // Write header block
        let mut header = [0u8; BLOCK_SIZE];
        header[0..4].copy_from_slice(&MAGIC);
        header[4..8].copy_from_slice(&VERSION.to_le_bytes());
        header[8..16].copy_from_slice(&data_len.to_le_bytes());
        self.block_lease
            .block()
            .write_block(HEADER_BLOCK, &header)
            .map_err(|e| EdgeError::CheckpointFailed(format!("write header: {e}")))?;

        // Write data blocks
        let num_blocks = data.len().div_ceil(BLOCK_SIZE);
        let available = self.block_lease.block().num_blocks() as usize;
        if num_blocks + 1 > available {
            return Err(EdgeError::CheckpointFailed(
                "state too large for block lease".into(),
            ));
        }

        for i in 0..num_blocks {
            let start = i * BLOCK_SIZE;
            let end = core::cmp::min(start + BLOCK_SIZE, data.len());
            let mut block = [0u8; BLOCK_SIZE];
            block[..end - start].copy_from_slice(&data[start..end]);
            self.block_lease
                .block()
                .write_block(DATA_START_BLOCK + i as u64, &block)
                .map_err(|e| EdgeError::CheckpointFailed(format!("write data block {i}: {e}")))?;
        }

        Ok(())
    }

    /// Restore a checkpointed state from block storage.
    ///
    /// Reads the header block to determine data length, then reads and
    /// deserializes the state from the data blocks.
    ///
    /// # Errors
    ///
    /// Returns [`EdgeError::CheckpointFailed`] on corrupt data, version
    /// mismatch, or deserialization failure.
    pub fn restore(block_lease: BlockLease) -> Result<Self, EdgeError> {
        let header = block_lease
            .block()
            .read_block(HEADER_BLOCK)
            .map_err(|e| EdgeError::CheckpointFailed(format!("read header: {e}")))?;

        if header[0..4] != MAGIC {
            return Err(EdgeError::CheckpointFailed("bad magic".into()));
        }
        let version = u32::from_le_bytes([header[4], header[5], header[6], header[7]]);
        if version != VERSION {
            return Err(EdgeError::CheckpointFailed(format!(
                "unsupported version: {version}"
            )));
        }
        let data_len = u64::from_le_bytes([
            header[8], header[9], header[10], header[11], header[12], header[13], header[14],
            header[15],
        ]) as usize;

        let num_blocks = data_len.div_ceil(BLOCK_SIZE);
        let mut data = Vec::with_capacity(data_len);
        for i in 0..num_blocks {
            let block = block_lease
                .block()
                .read_block(DATA_START_BLOCK + i as u64)
                .map_err(|e| EdgeError::CheckpointFailed(format!("read data block {i}: {e}")))?;
            let remaining = data_len - data.len();
            let take = core::cmp::min(BLOCK_SIZE, remaining);
            data.extend_from_slice(&block[..take]);
        }

        let state: T = postcard::from_bytes(&data)
            .map_err(|e| EdgeError::CheckpointFailed(format!("deserialize: {e}")))?;

        Ok(Self {
            state,
            block_lease,
            lease_id: None,
            observer: None,
        })
    }

    /// Borrow the current state.
    pub fn state(&self) -> &T {
        &self.state
    }

    /// Mutably borrow the current state.
    pub fn state_mut(&mut self) -> &mut T {
        &mut self.state
    }

    /// Consume the wrapper and return the underlying block lease.
    ///
    /// This allows transferring the lease to a subsequent
    /// [`restore`](Self::restore) call so that both sides access the
    /// same underlying block storage.
    pub fn into_block_lease(self) -> BlockLease {
        self.block_lease
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::EdgeCheckpoint;
    use grafos_std::block::BlockBuilder;
    use grafos_std::host;

    fn setup_block(num_blocks: u64) -> BlockLease {
        host::reset_mock();
        host::mock_set_fbbu_num_blocks(num_blocks);
        BlockBuilder::new().acquire().expect("acquire")
    }

    #[test]
    fn checkpoint_and_restore_roundtrip() {
        let block_lease = setup_block(64);

        let state = EdgeCheckpoint {
            input_offset: 42,
            output_commit: 99,
        };

        let mut ckpt = CheckpointedStageState::new(state, block_lease);
        ckpt.checkpoint().expect("checkpoint");

        // Transfer the block lease so restore reads the same blocks.
        let lease = ckpt.into_block_lease();
        let restored = CheckpointedStageState::<EdgeCheckpoint>::restore(lease).expect("restore");
        assert_eq!(restored.state().input_offset, 42);
        assert_eq!(restored.state().output_commit, 99);
    }

    #[test]
    fn restore_fails_on_corrupt_data() {
        let block_lease = setup_block(64);
        // Write garbage to block 0
        let mut garbage = [0u8; BLOCK_SIZE];
        garbage[0..4].copy_from_slice(b"BAAD");
        block_lease.block().write_block(0, &garbage).expect("write");

        let result = CheckpointedStageState::<EdgeCheckpoint>::restore(block_lease);
        match result {
            Err(EdgeError::CheckpointFailed(msg)) => assert_eq!(msg, "bad magic"),
            Err(other) => panic!("unexpected error: {other}"),
            Ok(_) => panic!("expected error but got Ok"),
        }
    }

    #[test]
    fn state_mut_allows_modification() {
        let block_lease = setup_block(64);
        let state = EdgeCheckpoint {
            input_offset: 0,
            output_commit: 0,
        };

        let mut ckpt = CheckpointedStageState::new(state, block_lease);
        ckpt.state_mut().input_offset = 100;
        ckpt.state_mut().output_commit = 200;

        assert_eq!(ckpt.state().input_offset, 100);
        assert_eq!(ckpt.state().output_commit, 200);
    }

    #[test]
    fn checkpoint_preserves_complex_state() {
        let block_lease = setup_block(64);

        #[derive(Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize)]
        struct StageState {
            name: alloc::string::String,
            items_processed: u64,
            checkpoint: EdgeCheckpoint,
        }

        let state = StageState {
            name: "transform-stage".into(),
            items_processed: 1024,
            checkpoint: EdgeCheckpoint {
                input_offset: 500,
                output_commit: 480,
            },
        };

        let mut ckpt = CheckpointedStageState::new(state.clone(), block_lease);
        ckpt.checkpoint().expect("checkpoint");

        // Transfer the block lease so restore reads the same blocks.
        let lease = ckpt.into_block_lease();
        let restored = CheckpointedStageState::<StageState>::restore(lease).expect("restore");
        assert_eq!(restored.state(), &state);
    }

    // -----------------------------------------------------------------
    // Phase 218.2 / slice 96 — CheckpointFlushObserver pin tests
    // -----------------------------------------------------------------
    //
    // The slice-96 contract: when a `CheckpointedStageState` is
    // constructed via `with_observer`, the observer's
    // `on_flush_start(lease_id)` runs at the top of `checkpoint()` and
    // `on_flush_complete(lease_id)` runs at the end — on BOTH the
    // success path and every failure path. Three pin tests cover:
    //   - mark on flush start
    //   - clear on success
    //   - clear on failure (the most load-bearing — without RAII
    //     discipline, a partway-through flush would leave the lease
    //     stuck in `CheckpointInProgress` forever)
    //
    // The observer is a sideband signal to the scheduler; it does NOT
    // change the flush's success semantics. A failed observer
    // notification (lease not tracked, etc.) is the implementation's
    // problem to handle internally.

    use core::sync::atomic::{AtomicU64, Ordering};

    /// Test helper. Records the order and lease_id of every
    /// `on_flush_start` / `on_flush_complete` call so tests can pin
    /// the lifecycle precisely.
    #[derive(Default)]
    struct RecordingObserver {
        start_count: AtomicU64,
        complete_count: AtomicU64,
        last_start_lease: AtomicU64,
        last_complete_lease: AtomicU64,
    }

    impl CheckpointFlushObserver for alloc::sync::Arc<RecordingObserver> {
        fn on_flush_start(&self, lease_id: u128) {
            self.start_count.fetch_add(1, Ordering::SeqCst);
            // Test leases all fit in u64; truncating is safe for the
            // recording assertion.
            self.last_start_lease
                .store(lease_id as u64, Ordering::SeqCst);
        }

        fn on_flush_complete(&self, lease_id: u128) {
            self.complete_count.fetch_add(1, Ordering::SeqCst);
            self.last_complete_lease
                .store(lease_id as u64, Ordering::SeqCst);
        }
    }

    #[test]
    fn checkpoint_flush_marks_lease_non_preemptible() {
        // Verify `on_flush_start` is called with the correct lease_id
        // at the BEGINNING of the flush. Pinned together with the
        // success-path clear test below — the start count is 1 and
        // the recorded lease_id matches the one the producer wired.
        let block_lease = setup_block(64);
        let observer = alloc::sync::Arc::new(RecordingObserver::default());
        let lease_id: u128 = 0xdead_beef_cafe_babe;

        let state = EdgeCheckpoint {
            input_offset: 7,
            output_commit: 7,
        };

        let mut ckpt = CheckpointedStageState::with_observer(
            state,
            block_lease,
            lease_id,
            alloc::boxed::Box::new(observer.clone()),
        );
        ckpt.checkpoint().expect("checkpoint");

        assert_eq!(
            observer.start_count.load(Ordering::SeqCst),
            1,
            "on_flush_start must be called exactly once at flush start"
        );
        assert_eq!(
            observer.last_start_lease.load(Ordering::SeqCst),
            lease_id as u64,
            "on_flush_start must receive the producer-named lease_id"
        );
    }

    #[test]
    fn checkpoint_flush_clears_marker_on_success() {
        // Verify `on_flush_complete` is called once with the same
        // lease_id at the END of a successful flush. The full
        // lifecycle: start once, complete once, both keyed on the
        // producer-named lease_id.
        let block_lease = setup_block(64);
        let observer = alloc::sync::Arc::new(RecordingObserver::default());
        let lease_id: u128 = 0x1234_5678_9abc_def0;

        let state = EdgeCheckpoint {
            input_offset: 11,
            output_commit: 22,
        };

        let mut ckpt = CheckpointedStageState::with_observer(
            state,
            block_lease,
            lease_id,
            alloc::boxed::Box::new(observer.clone()),
        );
        ckpt.checkpoint().expect("checkpoint");

        assert_eq!(observer.start_count.load(Ordering::SeqCst), 1);
        assert_eq!(
            observer.complete_count.load(Ordering::SeqCst),
            1,
            "on_flush_complete must be called exactly once on success"
        );
        assert_eq!(
            observer.last_complete_lease.load(Ordering::SeqCst),
            lease_id as u64,
            "on_flush_complete must receive the producer-named lease_id"
        );
    }

    #[test]
    fn checkpoint_flush_clears_marker_on_failure() {
        // The load-bearing test. Trigger a failure partway through
        // the flush by giving a block lease too small to hold the
        // serialized data — the "state too large for block lease"
        // path. Without the RAII guard discipline, a partway-through
        // failure would leave the lease stuck in
        // `CheckpointInProgress` forever; this test pins that
        // `on_flush_complete` runs even when the flush returns an
        // error.
        //
        // Setup: 1 block available (header only) so any non-empty
        // serialized state triggers the size check failure.
        let block_lease = setup_block(1);
        let observer = alloc::sync::Arc::new(RecordingObserver::default());
        let lease_id: u128 = 0xabcd_ef01_2345_6789;

        // A state that won't fit in 0 data blocks (postcard for
        // EdgeCheckpoint serializes to several bytes — needs at
        // least 1 data block, but we gave only 1 total block).
        let state = EdgeCheckpoint {
            input_offset: 99,
            output_commit: 100,
        };

        let mut ckpt = CheckpointedStageState::with_observer(
            state,
            block_lease,
            lease_id,
            alloc::boxed::Box::new(observer.clone()),
        );
        let result = ckpt.checkpoint();

        // The flush failed — the size check should have rejected it.
        assert!(
            matches!(result, Err(EdgeError::CheckpointFailed(_))),
            "expected CheckpointFailed, got {result:?}"
        );

        // The marker MUST still have been cleared, even though the
        // flush failed. RAII via FlushGuard's Drop is what makes
        // this work.
        assert_eq!(
            observer.start_count.load(Ordering::SeqCst),
            1,
            "on_flush_start runs before the size check"
        );
        assert_eq!(
            observer.complete_count.load(Ordering::SeqCst),
            1,
            "on_flush_complete MUST run on the failure path \
             (RAII guard via Drop) so the lease isn't stuck \
             non-preemptible after a failed flush"
        );
        assert_eq!(
            observer.last_complete_lease.load(Ordering::SeqCst),
            lease_id as u64,
            "on_flush_complete must clear the producer-named lease_id"
        );
    }

    #[test]
    fn checkpoint_without_observer_remains_silent() {
        // Back-compat pin: the original `new()` constructor produces
        // a `CheckpointedStageState` with no observer. Flushes do
        // NOT panic, do NOT call into a missing observer, and do NOT
        // change behavior relative to pre-slice-96 behavior.
        let block_lease = setup_block(64);
        let state = EdgeCheckpoint {
            input_offset: 5,
            output_commit: 5,
        };

        // No observer attached.
        let mut ckpt = CheckpointedStageState::new(state, block_lease);
        ckpt.checkpoint().expect("checkpoint");
        // If we got here, the no-observer path didn't dereference a
        // null observer or otherwise misbehave. Nothing else to
        // assert — the absence of a panic or missed-method dispatch
        // is the pin.
    }
}