gitks/actor/snapshot.rs

//! Raft log snapshot mechanism for log compaction.
//!
//! When the Raft log grows beyond the size threshold, a snapshot is created
//! that captures the current state of all repositories. Old log entries before
//! the snapshot are then discarded.
//!
//! Snapshot format:
//! - `raft-snapshot.dat`: Contains the serialized state at a given log index
//!
//! The snapshot includes:
//! - All repository entries (path, role, last_commit)
//! - The log index at which the snapshot was taken
//! - The term at which the snapshot was taken

use std::collections::HashMap;
use std::io::{BufReader, BufWriter, Read, Write};
use std::path::{Path, PathBuf};

use crate::actor::handler::RepoEntry;
use crate::error::{GitError, GitResult};

/// Snapshot metadata and data.
#[derive(Debug, Clone)]
pub struct RaftSnapshot {
    /// The log index at which this snapshot was taken.
    pub last_included_index: u64,
    /// The term at which this snapshot was taken.
    pub last_included_term: u64,
    /// All repository entries at the time of the snapshot.
    pub repos: HashMap<String, RepoEntry>,
}

impl RaftSnapshot {
    /// Create a new snapshot from the current state.
    pub fn new(
        last_included_index: u64,
        last_included_term: u64,
        repos: HashMap<String, RepoEntry>,
    ) -> Self {
        Self {
            last_included_index,
            last_included_term,
            repos,
        }
    }

    /// Serialize the snapshot to bytes.
    pub fn encode(&self) -> Vec<u8> {
        let mut buf = Vec::new();
        // Header
        buf.extend(self.last_included_index.to_be_bytes());
        buf.extend(self.last_included_term.to_be_bytes());
        // Repository count
        buf.extend((self.repos.len() as u32).to_be_bytes());
        // Each repository entry
        for (path, entry) in &self.repos {
            encode_string(&mut buf, path);
            encode_string(&mut buf, &entry.role);
            encode_string(&mut buf, &entry.last_commit);
            buf.push(if entry.read_only { 1 } else { 0 });
        }
        buf
    }

    /// Deserialize a snapshot from bytes.
    pub fn decode(data: &[u8]) -> Option<Self> {
        if data.len() < 20 {
            return None;
        }
        let mut offset = 0;
        let last_included_index = read_u64(data, &mut offset)?;
        let last_included_term = read_u64(data, &mut offset)?;
        let repo_count = read_u32(data, &mut offset)? as usize;

        let mut repos = HashMap::with_capacity(repo_count);
        for _ in 0..repo_count {
            let path = read_string(data, &mut offset)?;
            let role = read_string(data, &mut offset)?;
            let last_commit = read_string(data, &mut offset)?;
            let read_only = data.get(offset).copied().unwrap_or(0) == 1;
            offset += 1;
            repos.insert(
                path,
                RepoEntry {
                    role,
                    last_commit,
                    read_only,
                },
            );
        }

        Some(Self {
            last_included_index,
            last_included_term,
            repos,
        })
    }
}

/// Snapshot storage manager.
pub struct SnapshotStorage {
    snapshot_path: PathBuf,
}

impl SnapshotStorage {
    pub fn new(data_dir: &Path) -> Self {
        Self {
            snapshot_path: data_dir.join("raft-snapshot.dat"),
        }
    }

    /// Save a snapshot to disk.
    pub fn save(&self, snapshot: &RaftSnapshot) -> GitResult<()> {
        let data = snapshot.encode();
        let file = std::fs::File::create(&self.snapshot_path).map_err(GitError::Io)?;
        let mut writer = BufWriter::new(file);
        writer.write_all(&data).map_err(GitError::Io)?;
        writer.flush().map_err(GitError::Io)?;

        tracing::info!(
            index = snapshot.last_included_index,
            term = snapshot.last_included_term,
            repos = snapshot.repos.len(),
            "raft snapshot saved"
        );
        Ok(())
    }

    /// Load a snapshot from disk.
    pub fn load(&self) -> GitResult<Option<RaftSnapshot>> {
        if !self.snapshot_path.exists() {
            return Ok(None);
        }

        let file = std::fs::File::open(&self.snapshot_path).map_err(GitError::Io)?;
        let mut reader = BufReader::new(file);
        let mut data = Vec::new();
        reader.read_to_end(&mut data).map_err(GitError::Io)?;

        match RaftSnapshot::decode(&data) {
            Some(snapshot) => {
                tracing::info!(
                    index = snapshot.last_included_index,
                    term = snapshot.last_included_term,
                    repos = snapshot.repos.len(),
                    "raft snapshot loaded"
                );
                Ok(Some(snapshot))
            }
            None => {
                tracing::warn!("failed to decode raft snapshot, ignoring");
                Ok(None)
            }
        }
    }

    /// Check if a snapshot exists.
    pub fn exists(&self) -> bool {
        self.snapshot_path.exists()
    }

    /// Delete the snapshot file.
    pub fn delete(&self) -> GitResult<()> {
        if self.snapshot_path.exists() {
            std::fs::remove_file(&self.snapshot_path).map_err(GitError::Io)?;
        }
        Ok(())
    }
}

// ── Helper functions ─────────────────────────────────────────

fn encode_string(buf: &mut Vec<u8>, s: &str) {
    let bytes = s.as_bytes();
    buf.extend((bytes.len() as u32).to_be_bytes());
    buf.extend(bytes);
}

fn read_u32(data: &[u8], offset: &mut usize) -> Option<u32> {
    if *offset + 4 > data.len() {
        return None;
    }
    let val = u32::from_be_bytes(data[*offset..*offset + 4].try_into().ok()?);
    *offset += 4;
    Some(val)
}

fn read_u64(data: &[u8], offset: &mut usize) -> Option<u64> {
    if *offset + 8 > data.len() {
        return None;
    }
    let val = u64::from_be_bytes(data[*offset..*offset + 8].try_into().ok()?);
    *offset += 8;
    Some(val)
}

fn read_string(data: &[u8], offset: &mut usize) -> Option<String> {
    let len = read_u32(data, offset)? as usize;
    if *offset + len > data.len() {
        return None;
    }
    let s = String::from_utf8_lossy(&data[*offset..*offset + len]).into_owned();
    *offset += len;
    Some(s)
}