krabka_raft/controller/

checkpoint.rs

//! The on-disk KIP-630 metadata checkpoints: the `FetchSnapshot` answer that
//! the engine gives a broker listener, the byte window of the newest
//! `<end_offset>-<epoch>.checkpoint`, the read of that file, and the
//! manual snapshot trigger. The checkpoint file layout is read straight from
//! disk here rather than through the engine, so it is kept apart from the rest
//! of the handle.

use {
    super::ControllerHandle,
    crate::{
        error::RaftError,
        kraft::{
            controller::checkpoint::{
                CheckpointChunk, latest_checkpoint_id, load_checkpoint_by_id, open_checkpoint,
            },
            snapshot_fetch::SnapshotId,
            transport::Inbound,
        },
    },
    std::path::Path,
    tokio::sync::oneshot,
};

/// A contiguous byte window of the latest metadata `.checkpoint`, returned by
/// [`ControllerHandle::read_snapshot_range`].
#[derive(Debug, PartialEq, Eq)]
pub struct SnapshotSlice {
    pub end_offset: i64,
    pub epoch: i32,
    pub total_size: i64,
    pub bytes: bytes::Bytes,
}

/// Outcome of [`ControllerHandle::read_snapshot_range`].
pub enum SnapshotRange {
    /// No `.checkpoint` exists yet.
    NoSnapshot,
    /// `position` is at or past the snapshot's end byte, which Kafka's
    /// `handleFetchSnapshotRequest` answers with `POSITION_OUT_OF_RANGE`.
    OutOfRange,
    /// The requested byte window.
    Slice(SnapshotSlice),
}

impl ControllerHandle {
    /// Read up to `max_bytes` of the latest metadata snapshot starting at
    /// `position`. Reads the engine's `.checkpoint` artifacts directly (the
    /// engine writes a bare KIP-630 checkpoint, no `.meta` sidecar).
    ///
    /// `position` and `max_bytes` stay the raw KIP-595 `FetchSnapshot` `int64`
    /// and `int32`: both are byte offsets into an on-disk checkpoint that the
    /// broker's handler forwards straight off the wire, so there is no domain
    /// layer between the decode and the slice index for a quantity to occupy.
    #[must_use]
    pub fn read_snapshot_range(&self, position: i64, max_bytes: i32) -> SnapshotRange {
        let Some(id) = latest_checkpoint_id(&self.data_dir) else {
            return SnapshotRange::NoSnapshot;
        };
        // Retention can delete the id between the scan and the open, which
        // reads as no checkpoint, as it did before the scan saw it.
        let Some(mut checkpoint) = open_checkpoint(&self.data_dir, id.0, id.1) else {
            return SnapshotRange::NoSnapshot;
        };
        match checkpoint.read_chunk(position, max_bytes) {
            Ok(CheckpointChunk::Bytes(bytes)) => SnapshotRange::Slice(SnapshotSlice {
                end_offset: id.0,
                epoch: id.1,
                total_size: i64::try_from(checkpoint.len()).unwrap_or(i64::MAX),
                bytes: bytes::Bytes::from(bytes),
            }),
            Ok(CheckpointChunk::OutOfRange) => SnapshotRange::OutOfRange,
            Err(_) => SnapshotRange::NoSnapshot,
        }
    }

    /// Answer one KIP-630 `FetchSnapshot` request through the engine, as the
    /// controller listener answers it.
    ///
    /// `body` is the request body at `version`, and the answer is the response
    /// body at that version. The engine follows Kafka's
    /// `KafkaRaftClient.handleFetchSnapshotRequest`: it serves exactly the
    /// snapshot the request names, and refuses another topic or partition, a
    /// stale or unknown `CurrentLeaderEpoch`, another cluster id, and a
    /// position outside the snapshot. It also fills `CurrentLeader` and the
    /// leader's endpoint. A broker listener that answered from the checkpoint
    /// directly would answer the same request differently.
    ///
    /// # Errors
    /// Returns [`RaftError::Shutdown`] if the engine task is gone, or if it
    /// dropped the request because the body does not decode.
    pub async fn fetch_snapshot(
        &self,
        version: i16,
        body: bytes::Bytes,
    ) -> Result<bytes::Bytes, RaftError> {
        let (reply, answer) = oneshot::channel();
        self.engine
            .deliver(Inbound::FetchSnapshot {
                req: body,
                version,
                reply,
            })
            .await?;
        answer.await.map_err(|_| RaftError::Shutdown)
    }

    /// Manually trigger a metadata snapshot (KIP-630 checkpoint) on this node.
    ///
    /// # Errors
    /// Returns [`RaftError`] if serialization or the file write fails.
    pub async fn trigger_snapshot(&self) -> Result<(), RaftError> {
        self.engine.trigger_snapshot().await
    }
}

/// The highest `(end_offset, epoch)` `.checkpoint` in `dir` plus its raw
/// bytes. Matches the bare-checkpoint format the engine writes (no `.meta`
/// sidecar).
pub(super) fn load_latest_checkpoint(dir: &Path) -> Option<(SnapshotId, Vec<u8>)> {
    let id = latest_checkpoint_id(dir)?;
    let bytes = load_checkpoint_by_id(dir, id.0, id.1)?;
    Some((id, bytes))
}

#[cfg(test)]
mod tests {
    use {
        super::*,
        crate::{controller::test_support::joining_controller, metadata_partition_dir},
        std::fs,
        tempfile::TempDir,
    };

    #[test]
    fn load_latest_checkpoint_picks_highest_offset_then_epoch() {
        let dir = TempDir::new().unwrap();
        fs::write(
            dir.path()
                .join("00000000000000000001-0000000009.checkpoint"),
            b"old-offset",
        )
        .unwrap();
        fs::write(
            dir.path()
                .join("00000000000000000002-0000000001.checkpoint"),
            b"old-epoch",
        )
        .unwrap();
        fs::write(
            dir.path()
                .join("00000000000000000002-0000000003.checkpoint"),
            b"best",
        )
        .unwrap();
        fs::write(dir.path().join("9-9.txt"), b"ignored suffix").unwrap();
        fs::write(
            dir.path().join("not-a-checkpoint.checkpoint"),
            b"ignored name",
        )
        .unwrap();

        let latest = load_latest_checkpoint(dir.path()).expect("checkpoint");
        assert2::assert!(latest == ((2, 3), b"best".to_vec()));
    }

    #[yare::parameterized(
        zero = { 0, Some(b"abc".as_slice()) },
        last_byte = { 2, Some(b"c".as_slice()) },
        exact_end = { 3, None },
        past_end = { 4, None },
        negative = { -1, None },
    )]
    #[test_macro(tokio::test)]
    async fn read_snapshot_range_rejects_the_end_and_everything_past_it(
        position: i64,
        expected: Option<&'static [u8]>,
    ) {
        let (dir, ctrl) = joining_controller("join start").await;
        let checkpoint_dir = metadata_partition_dir(dir.path());
        fs::create_dir_all(&checkpoint_dir).unwrap();
        fs::write(
            checkpoint_dir.join("00000000000000000010-0000000004.checkpoint"),
            b"abc",
        )
        .unwrap();

        // Kafka's `handleFetchSnapshotRequest` refuses `position < 0` and
        // `position >= size`, so the end of the snapshot is not a valid position.
        let read = |position| match ctrl.read_snapshot_range(position, 10) {
            SnapshotRange::Slice(slice) => Some(slice),
            SnapshotRange::OutOfRange => None,
            SnapshotRange::NoSnapshot => panic!("the checkpoint was written"),
        };
        let window = |bytes: &'static [u8]| SnapshotSlice {
            end_offset: 10,
            epoch: 4,
            total_size: 3,
            bytes: bytes::Bytes::from_static(bytes),
        };
        assert2::assert!(
            read(position) == expected.map(window),
            "position {position}"
        );
        ctrl.shutdown().await;
    }
}