qrusty 0.21.1

// src/payload_store.rs
// Implements: PER-0016, PER-0017, PER-0018

//! # Payload Store
//!
//! Manages message payloads in append-only memory-mapped files, keeping
//! them out of RocksDB values.  The OS page cache decides which pages
//! stay resident — hot payloads stay in RAM, cold ones get paged out.
//!
//! ## File Layout
//!
//! Each segment file (`payload_NNNN.dat`) is a sequence of length-prefixed
//! records:
//!
//! ```text
//! [u32 len][payload bytes][u32 len][payload bytes]...
//! ```
//!
//! ## Compaction
//!
//! A background task periodically rewrites live payloads to a new segment
//! and atomically swaps references.  This reclaims space from ack'd
//! messages without blocking normal operations.

use anyhow::Result;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::fs::{self, File, OpenOptions};
use std::io::Write;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex, RwLock};

/// Reference to a payload stored in a segment file.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct PayloadRef {
    /// Segment file number.
    pub file_id: u32,
    /// Byte offset within the segment (points to the length prefix).
    pub offset: u64,
    /// Payload length in bytes (excluding the 4-byte length prefix).
    pub length: u32,
}

/// Boundary produced by [`PayloadStore::seal_active`]. Every segment
/// with id `<= sealed_id` is now immutable and a candidate for
/// compaction; `compaction_id` is the id reserved for the rewritten
/// output. Payloads appended after the seal land in a segment with a
/// strictly greater id and are therefore never deleted by the
/// compaction this handle drives.
#[derive(Debug, Clone, Copy)]
pub struct SealHandle {
    pub sealed_id: u32,
    pub compaction_id: u32,
}

/// Maximum segment file size before rotating to a new file (default 256 MB).
const DEFAULT_SEGMENT_MAX_BYTES: u64 = 256 * 1024 * 1024;

/// A memory-mapped segment file for reading payloads.
struct MappedSegment {
    #[allow(dead_code)]
    mmap: memmap2::Mmap,
}

impl MappedSegment {
    fn open(path: &Path) -> Result<Self> {
        let file = File::open(path)?;
        // SAFETY: we treat the mmap as read-only and never modify it.
        let mmap = unsafe { memmap2::Mmap::map(&file)? };
        Ok(Self { mmap })
    }

    fn read(&self, offset: u64, length: u32) -> Option<&[u8]> {
        let start = offset as usize + 4; // skip length prefix
        let end = start + length as usize;
        if end <= self.mmap.len() {
            Some(&self.mmap[start..end])
        } else {
            None
        }
    }
}

/// Manages append-only payload segment files with mmap-backed reads.
pub struct PayloadStore {
    /// Directory containing segment files.
    dir: PathBuf,
    /// Currently active segment for writes.
    active_segment: Mutex<ActiveSegment>,
    /// Memory-mapped segments for reads.  Keyed by file_id.
    segments: RwLock<HashMap<u32, Arc<MappedSegment>>>,
    /// Maximum segment size before rotation.
    segment_max_bytes: u64,
}

struct ActiveSegment {
    file_id: u32,
    file: File,
    offset: u64,
}

impl PayloadStore {
    /// Opens or creates a payload store in the given directory.
    pub fn open(dir: &Path) -> Result<Self> {
        fs::create_dir_all(dir)?;

        // Find existing segments.
        let mut max_id: u32 = 0;
        let mut segment_files: Vec<u32> = Vec::new();
        for entry in fs::read_dir(dir)? {
            let entry = entry?;
            let name = entry.file_name();
            let name_str = name.to_string_lossy();
            if let Some(rest) = name_str.strip_prefix("payload_") {
                if let Some(num_str) = rest.strip_suffix(".dat") {
                    if let Ok(id) = num_str.parse::<u32>() {
                        segment_files.push(id);
                        if id > max_id {
                            max_id = id;
                        }
                    }
                }
            }
        }

        // Open (or create) the active segment for appending.
        let active_id = if segment_files.is_empty() { 0 } else { max_id };

        // Memory-map existing segments for reads — but NOT the active
        // segment, because its mmap would be frozen at the startup size
        // and wouldn't see data appended after open.  The active segment
        // uses the fs::read fallback path instead.
        let mut segments = HashMap::new();
        for id in &segment_files {
            if *id == active_id {
                continue; // active segment is read via fs::read fallback
            }
            let path = dir.join(format!("payload_{:04}.dat", id));
            match MappedSegment::open(&path) {
                Ok(seg) => {
                    segments.insert(*id, Arc::new(seg));
                }
                Err(e) => {
                    tracing::warn!("Failed to mmap segment {}: {}", id, e);
                }
            }
        }
        let active_path = dir.join(format!("payload_{:04}.dat", active_id));
        let file = OpenOptions::new()
            .create(true)
            .append(true)
            .open(&active_path)?;
        let offset = file.metadata()?.len();

        Ok(Self {
            dir: dir.to_path_buf(),
            active_segment: Mutex::new(ActiveSegment {
                file_id: active_id,
                file,
                offset,
            }),
            segments: RwLock::new(segments),
            segment_max_bytes: std::env::var("QRUSTY_SEGMENT_MAX_MB")
                .ok()
                .and_then(|v| v.parse::<u64>().ok())
                .map(|mb| mb * 1024 * 1024)
                .unwrap_or(DEFAULT_SEGMENT_MAX_BYTES),
        })
    }

    /// Appends a payload and returns a reference to it.
    pub fn append(&self, payload: &[u8]) -> Result<PayloadRef> {
        let mut active = self.active_segment.lock().unwrap();

        // Rotate if the active segment is too large.
        if active.offset > self.segment_max_bytes {
            let new_id = active.file_id + 1;
            let new_path = self.dir.join(format!("payload_{:04}.dat", new_id));
            let new_file = OpenOptions::new()
                .create(true)
                .append(true)
                .open(&new_path)?;

            // Re-mmap the old segment (which may have grown since initial open).
            let old_path = self.dir.join(format!("payload_{:04}.dat", active.file_id));
            if let Ok(seg) = MappedSegment::open(&old_path) {
                self.segments
                    .write()
                    .unwrap()
                    .insert(active.file_id, Arc::new(seg));
            }

            active.file_id = new_id;
            active.file = new_file;
            active.offset = 0;
        }

        let file_id = active.file_id;
        let offset = active.offset;
        let length = payload.len() as u32;

        // Write length-prefixed record.
        active.file.write_all(&length.to_le_bytes())?;
        active.file.write_all(payload)?;
        active.offset += 4 + payload.len() as u64;

        Ok(PayloadRef {
            file_id,
            offset,
            length,
        })
    }

    /// Reads a payload by reference.  Returns `None` if the segment or
    /// offset is invalid.
    pub fn read(&self, pref: &PayloadRef) -> Option<Vec<u8>> {
        // Try mapped segments first (all segments except possibly the active one).
        {
            let segments = self.segments.read().unwrap();
            if let Some(seg) = segments.get(&pref.file_id) {
                return seg.read(pref.offset, pref.length).map(|s| s.to_vec());
            }
        }

        // Fall back to reading the active segment's file directly (it may
        // not be mmap'd yet if it's still being written to).
        let active = self.active_segment.lock().unwrap();
        if active.file_id == pref.file_id {
            // Read from the OS file — the data was already flushed by write_all.
            let path = self.dir.join(format!("payload_{:04}.dat", pref.file_id));
            if let Ok(data) = fs::read(&path) {
                let start = pref.offset as usize + 4;
                let end = start + pref.length as usize;
                if end <= data.len() {
                    return Some(data[start..end].to_vec());
                }
            }
        }

        None
    }

    /// Seals the active segment and rotates writes to a fresh one, so any
    /// payload appended *after* this call lands in a segment that the
    /// ensuing compaction will never delete.
    ///
    /// Compaction MUST call this BEFORE it snapshots the live reference
    /// set (from RocksDB). Otherwise a payload appended into the active
    /// segment between the snapshot and the segment deletion would have
    /// its backing file removed out from under it — the
    /// "segment file missing or corrupted" data loss this guards against.
    pub fn seal_active(&self) -> Result<SealHandle> {
        let mut active = self.active_segment.lock().unwrap();
        let sealed_id = active.file_id;

        // mmap the now-sealed segment so compaction reads it via the
        // segment map rather than the active fs fallback, which we are
        // about to repoint at the fresh segment.
        let sealed_path = self.dir.join(format!("payload_{:04}.dat", sealed_id));
        if let Ok(seg) = MappedSegment::open(&sealed_path) {
            self.segments
                .write()
                .unwrap()
                .insert(sealed_id, Arc::new(seg));
        }

        // Reserve sealed_id+1 for the compacted output and start the new
        // write head at sealed_id+2 so neither collides with the other,
        // nor with future append rotations (which only ever increment).
        let compaction_id = sealed_id + 1;
        let fresh_id = sealed_id + 2;
        let fresh_path = self.dir.join(format!("payload_{:04}.dat", fresh_id));
        let fresh_file = OpenOptions::new()
            .create(true)
            .append(true)
            .open(&fresh_path)?;
        active.file_id = fresh_id;
        active.file = fresh_file;
        active.offset = 0;

        Ok(SealHandle {
            sealed_id,
            compaction_id,
        })
    }

    /// Rewrites the live payloads that live in sealed segments
    /// (`file_id <= seal.sealed_id`) into the reserved compacted segment,
    /// then deletes ONLY those sealed source segments. Refs that point at
    /// the fresh active segment (appended after the seal) are left in
    /// place and are not returned in the map.
    ///
    /// If any sealed live payload cannot be read, the whole compaction is
    /// ABORTED and nothing is deleted, so a transient read failure can
    /// never strand a reference against a deleted segment. Returns a map
    /// from each rewritten ref's `(file_id, offset)` to its new ref.
    pub fn compact_sealed(
        &self,
        live_refs: &[PayloadRef],
        seal: SealHandle,
    ) -> Result<HashMap<(u32, u64), PayloadRef>> {
        let comp_path = self
            .dir
            .join(format!("payload_{:04}.dat", seal.compaction_id));
        // Created lazily on the first rewritten payload so an empty
        // compaction leaves no stray zero-byte segment behind.
        let mut comp_file: Option<File> = None;
        let mut new_offset: u64 = 0;
        let mut ref_map: HashMap<(u32, u64), PayloadRef> = HashMap::new();

        for pref in live_refs {
            // Refs in the fresh active segment (or beyond) were appended
            // after the seal; leave them exactly where they are.
            if pref.file_id > seal.sealed_id {
                continue;
            }
            let data = match self.read(pref) {
                Some(d) => d,
                None => {
                    // Abort without deleting anything: the source segment
                    // stays, so the still-valid reference keeps resolving.
                    if let Some(f) = comp_file.take() {
                        drop(f);
                        let _ = fs::remove_file(&comp_path);
                    }
                    anyhow::bail!(
                        "payload compaction aborted: unreadable live payload \
                         file_id={} offset={} length={} — leaving sealed segments intact",
                        pref.file_id,
                        pref.offset,
                        pref.length
                    );
                }
            };
            if comp_file.is_none() {
                comp_file = Some(
                    OpenOptions::new()
                        .create(true)
                        .truncate(true)
                        .write(true)
                        .open(&comp_path)?,
                );
            }
            let file = comp_file.as_mut().unwrap();
            let length = data.len() as u32;
            file.write_all(&length.to_le_bytes())?;
            file.write_all(&data)?;
            ref_map.insert(
                (pref.file_id, pref.offset),
                PayloadRef {
                    file_id: seal.compaction_id,
                    offset: new_offset,
                    length,
                },
            );
            new_offset += 4 + data.len() as u64;
        }

        if let Some(mut f) = comp_file {
            f.flush()?;
            drop(f);
            if let Ok(seg) = MappedSegment::open(&comp_path) {
                self.segments
                    .write()
                    .unwrap()
                    .insert(seal.compaction_id, Arc::new(seg));
            }
        }

        // Delete ONLY the sealed source segments. The fresh active
        // segment and the compacted output are never removed here.
        {
            let mut segments = self.segments.write().unwrap();
            for old_id in 0..=seal.sealed_id {
                segments.remove(&old_id);
                let old_path = self.dir.join(format!("payload_{:04}.dat", old_id));
                if old_path.exists() {
                    let _ = fs::remove_file(&old_path);
                }
            }
        }

        Ok(ref_map)
    }

    /// Returns the total size of all segment files on disk.
    pub fn disk_usage_bytes(&self) -> u64 {
        let mut total = 0u64;
        if let Ok(entries) = fs::read_dir(&self.dir) {
            for entry in entries.flatten() {
                if entry.file_name().to_string_lossy().starts_with("payload_") {
                    if let Ok(meta) = entry.metadata() {
                        total += meta.len();
                    }
                }
            }
        }
        total
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::TempDir;

    #[test]
    fn test_append_and_read() {
        let dir = TempDir::new().unwrap();
        let store = PayloadStore::open(dir.path()).unwrap();

        let pref = store.append(b"hello world").unwrap();
        assert_eq!(pref.file_id, 0);
        assert_eq!(pref.offset, 0);
        assert_eq!(pref.length, 11);

        let data = store.read(&pref).unwrap();
        assert_eq!(data, b"hello world");
    }

    #[test]
    fn test_multiple_appends() {
        let dir = TempDir::new().unwrap();
        let store = PayloadStore::open(dir.path()).unwrap();

        let p1 = store.append(b"aaa").unwrap();
        let p2 = store.append(b"bbb").unwrap();
        let p3 = store.append(b"ccc").unwrap();

        assert_eq!(store.read(&p1).unwrap(), b"aaa");
        assert_eq!(store.read(&p2).unwrap(), b"bbb");
        assert_eq!(store.read(&p3).unwrap(), b"ccc");

        // Verify sequential offsets: 4+3=7 per record
        assert_eq!(p1.offset, 0);
        assert_eq!(p2.offset, 7);
        assert_eq!(p3.offset, 14);
    }

    #[test]
    fn test_compact_rewrites_live_refs() {
        let dir = TempDir::new().unwrap();
        let store = PayloadStore::open(dir.path()).unwrap();

        let p1 = store.append(b"keep-me").unwrap();
        let _p2 = store.append(b"delete-me").unwrap();
        let p3 = store.append(b"keep-me-too").unwrap();

        // Compact with only p1 and p3 as live.
        let seal = store.seal_active().unwrap();
        let ref_map = store
            .compact_sealed(&[p1.clone(), p3.clone()], seal)
            .unwrap();

        // Old refs should map to new refs.
        let new_p1 = &ref_map[&(p1.file_id, p1.offset)];
        let new_p3 = &ref_map[&(p3.file_id, p3.offset)];

        assert_eq!(store.read(new_p1).unwrap(), b"keep-me");
        assert_eq!(store.read(new_p3).unwrap(), b"keep-me-too");
    }

    /// Regression for the compaction data-loss race: a payload appended
    /// AFTER the seal but BEFORE the sealed segments are deleted must
    /// survive (it lands in the fresh segment, which compaction never
    /// touches). Previously compact deleted segments 0..=active, taking
    /// the raced append's backing file with it -> "segment missing".
    #[test]
    fn test_append_after_seal_survives_compaction() {
        let dir = TempDir::new().unwrap();
        let store = PayloadStore::open(dir.path()).unwrap();

        let keep = store.append(b"live-before-seal").unwrap();

        // Seal first, exactly as Storage::compact_payloads now does.
        let seal = store.seal_active().unwrap();

        // A publish races in after the seal.
        let raced = store.append(b"appended-after-seal").unwrap();
        assert!(
            raced.file_id > seal.sealed_id,
            "raced append must land in the fresh, undeletable segment",
        );

        let ref_map = store.compact_sealed(std::slice::from_ref(&keep), seal).unwrap();

        // The pre-seal live payload was rewritten and is readable...
        let new_keep = &ref_map[&(keep.file_id, keep.offset)];
        assert_eq!(store.read(new_keep).unwrap(), b"live-before-seal");
        // ...and the raced append survived (its segment was not deleted).
        assert_eq!(store.read(&raced).unwrap(), b"appended-after-seal");
    }

    /// An unreadable live ref must abort the whole compaction without
    /// deleting anything, so the intact payloads keep resolving rather
    /// than being stranded against a deleted segment.
    #[test]
    fn test_compact_aborts_on_unreadable_ref_without_deleting() {
        let dir = TempDir::new().unwrap();
        let store = PayloadStore::open(dir.path()).unwrap();
        let good = store.append(b"intact").unwrap();

        let seal = store.seal_active().unwrap();
        // A bogus ref into the sealed segment that cannot be read.
        let bogus = PayloadRef {
            file_id: seal.sealed_id,
            offset: 999_999,
            length: 10,
        };
        let result = store.compact_sealed(&[good.clone(), bogus], seal);
        assert!(
            result.is_err(),
            "compaction must abort on an unreadable live ref",
        );

        // Nothing deleted: the good payload still resolves at its ref.
        assert_eq!(store.read(&good).unwrap(), b"intact");
    }

    #[test]
    fn test_reopen_reads_existing_data() {
        let dir = TempDir::new().unwrap();
        let pref;

        {
            let store = PayloadStore::open(dir.path()).unwrap();
            pref = store.append(b"persistent").unwrap();
        }

        // Reopen and read.
        let store = PayloadStore::open(dir.path()).unwrap();
        assert_eq!(store.read(&pref).unwrap(), b"persistent");
    }

    #[test]
    fn test_disk_usage() {
        let dir = TempDir::new().unwrap();
        let store = PayloadStore::open(dir.path()).unwrap();

        store.append(b"data").unwrap();
        // 4 bytes length prefix + 4 bytes data = 8 bytes
        assert_eq!(store.disk_usage_bytes(), 8);
    }

    /// Regression test for the stale-mmap bug: appending to the active
    /// segment after open must be readable immediately, not return None
    /// because the mmap snapshot is frozen at the startup file size.
    #[test]
    fn test_append_after_reopen_is_readable() {
        let dir = TempDir::new().unwrap();

        // Phase 1: write some data and close.
        let pref_old;
        {
            let store = PayloadStore::open(dir.path()).unwrap();
            pref_old = store.append(b"old-data").unwrap();
        }

        // Phase 2: reopen (active segment file exists on disk), append
        // new data, and read it back immediately.
        let store = PayloadStore::open(dir.path()).unwrap();

        // Old data should still be readable.
        assert_eq!(store.read(&pref_old).unwrap(), b"old-data");

        // New data appended after reopen must be readable.
        let pref_new = store.append(b"new-data-after-reopen").unwrap();
        let read_back = store.read(&pref_new);
        assert!(
            read_back.is_some(),
            "payload appended after reopen must be readable (stale mmap bug)"
        );
        assert_eq!(read_back.unwrap(), b"new-data-after-reopen");
    }

    /// Verifies that many appends to the same active segment are all
    /// readable — not just the first one within the initial mmap range.
    #[test]
    fn test_many_appends_all_readable() {
        let dir = TempDir::new().unwrap();

        // Seed with one record so the file exists on disk.
        {
            let store = PayloadStore::open(dir.path()).unwrap();
            store.append(b"seed").unwrap();
        }

        // Reopen and append many more.
        let store = PayloadStore::open(dir.path()).unwrap();
        let mut refs = Vec::new();
        for i in 0..100 {
            let data = format!("payload-{:04}", i);
            refs.push((store.append(data.as_bytes()).unwrap(), data));
        }

        // Every single one must be readable.
        for (pref, expected) in &refs {
            let actual = store
                .read(pref)
                .unwrap_or_else(|| panic!("failed to read payload at offset {}", pref.offset));
            assert_eq!(
                String::from_utf8_lossy(&actual),
                *expected,
                "payload mismatch at offset {}",
                pref.offset
            );
        }
    }
}