qrusty 0.21.0

// 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,
}

/// 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
    }

    /// Compacts the store: rewrites only the live payloads referenced by
    /// the given set of PayloadRefs into a fresh segment, then removes
    /// old segments.
    ///
    /// Returns a mapping from old PayloadRef to new PayloadRef so callers
    /// can update RocksDB references atomically.
    pub fn compact(&self, live_refs: &[PayloadRef]) -> Result<HashMap<(u32, u64), PayloadRef>> {
        if live_refs.is_empty() {
            // PER-0021: Still clean up old segments even with zero live refs.
            self.cleanup_all_segments()?;
            return Ok(HashMap::new());
        }

        // Determine new segment ID.
        let new_id = {
            let active = self.active_segment.lock().unwrap();
            active.file_id + 1
        };

        let new_path = self.dir.join(format!("payload_{:04}.dat", new_id));
        let mut new_file = OpenOptions::new()
            .create(true)
            .truncate(true)
            .write(true)
            .open(&new_path)?;
        let mut new_offset: u64 = 0;
        let mut ref_map: HashMap<(u32, u64), PayloadRef> = HashMap::new();

        for pref in live_refs {
            if let Some(data) = self.read(pref) {
                let length = data.len() as u32;
                new_file.write_all(&length.to_le_bytes())?;
                new_file.write_all(&data)?;

                ref_map.insert(
                    (pref.file_id, pref.offset),
                    PayloadRef {
                        file_id: new_id,
                        offset: new_offset,
                        length,
                    },
                );
                new_offset += 4 + data.len() as u64;
            }
        }
        new_file.flush()?;

        // Mmap the new segment.
        if let Ok(seg) = MappedSegment::open(&new_path) {
            self.segments.write().unwrap().insert(new_id, Arc::new(seg));
        }

        // Swap active segment to the new one.
        {
            let mut active = self.active_segment.lock().unwrap();
            let old_ids: Vec<u32> = (0..=active.file_id).collect();

            active.file_id = new_id;
            active.file = OpenOptions::new().append(true).open(&new_path)?;
            active.offset = new_offset;

            // Remove old segment mmaps and files.
            let mut segments = self.segments.write().unwrap();
            for old_id in old_ids {
                segments.remove(&old_id);
                let old_path = self.dir.join(format!("payload_{:04}.dat", old_id));
                if old_path.exists() && old_id != new_id {
                    let _ = fs::remove_file(&old_path);
                }
            }
        }

        Ok(ref_map)
    }

    /// Removes all segment files and their mmap regions, resetting to a
    /// clean state.  Called when compaction finds zero live references.
    // Implements: PER-0021
    fn cleanup_all_segments(&self) -> Result<()> {
        let mut active = self.active_segment.lock().unwrap();
        let mut segments = self.segments.write().unwrap();

        // Collect all segment file IDs to delete.
        let old_ids: Vec<u32> = segments.keys().copied().collect();
        let active_id = active.file_id;

        // Remove all mmap'd segments.
        for id in &old_ids {
            segments.remove(id);
        }

        // Delete all segment files from disk.
        if let Ok(entries) = fs::read_dir(&self.dir) {
            for entry in entries.flatten() {
                let name = entry.file_name();
                let name_str = name.to_string_lossy();
                if name_str.starts_with("payload_") && name_str.ends_with(".dat") {
                    let _ = fs::remove_file(entry.path());
                }
            }
        }

        // Reset active segment to a fresh, empty file (id 0).
        let fresh_path = self.dir.join("payload_0000.dat");
        active.file_id = 0;
        active.offset = 0;
        active.file = OpenOptions::new()
            .create(true)
            .truncate(true)
            .write(true)
            .open(&fresh_path)?;

        tracing::info!(
            "Payload store cleanup: removed {} mmap segment(s) + active segment {}",
            old_ids.len(),
            active_id
        );
        Ok(())
    }

    /// 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 ref_map = store.compact(&[p1.clone(), p3.clone()]).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");
    }

    #[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
            );
        }
    }
}