hayson_engine/compaction/

compact.rs

use std::collections::HashMap;
use std::io;

use crate::atomic::write_atomic;
use crate::registry::table_registry::PrimaryEntry;
use crate::storage::{LogFile, LogOp, Slot};

pub struct CompactionResult {
    pub(crate) slots: HashMap<String, PrimaryEntry>,
    pub records: usize,
    pub bytes_before: u64,
    pub bytes_after: u64,
    pub reclaimed_bytes: u64,
}

#[cfg(unix)]
const PERMISSION_BITS: u32 = 0o777;

#[cfg(unix)]
fn permissions_of(path: &std::path::Path) -> io::Result<Option<u32>> {
    use std::os::unix::fs::PermissionsExt;
    match std::fs::metadata(path) {
        Ok(metadata) => Ok(Some(metadata.permissions().mode() & PERMISSION_BITS)),
        Err(error) if error.kind() == io::ErrorKind::NotFound => Ok(None),
        Err(error) => Err(error),
    }
}

// NOTE: wasi has no permission bits at all; the host sandbox bounds fs access instead.
#[cfg(not(unix))]
fn permissions_of(_path: &std::path::Path) -> io::Result<Option<u32>> {
    Ok(None)
}

fn plan(log: &LogFile) -> io::Result<(Vec<u8>, HashMap<String, PrimaryEntry>, usize)> {
    let scan = log.scan(0)?;

    let mut order: Vec<String> = Vec::new();
    let mut position_of: HashMap<String, usize> = HashMap::new();
    let mut kept: Vec<Option<Slot>> = Vec::new();
    let mut seq_of: Vec<u64> = Vec::new();

    for scanned in &scan.entries {
        let key = &scanned.entry.meta.key;
        match scanned.entry.meta.op {
            LogOp::Del => {
                if let Some(pos) = position_of.remove(key) {
                    kept[pos] = None;
                }
            }
            LogOp::Put => {
                if let Some(&pos) = position_of.get(key) {
                    kept[pos] = Some(scanned.slot);
                    seq_of[pos] = scanned.entry.meta.seq;
                } else {
                    let pos = order.len();
                    order.push(key.clone());
                    kept.push(Some(scanned.slot));
                    seq_of.push(scanned.entry.meta.seq);
                    position_of.insert(key.clone(), pos);
                }
            }
        }
    }

    let mut buffer = Vec::new();
    let mut slots = HashMap::with_capacity(order.len());
    for (pos, key) in order.into_iter().enumerate() {
        let Some(slot) = kept[pos] else { continue };
        let line = log.read_slice(&slot)?;
        let new_offset = buffer.len() as u64;
        buffer.extend_from_slice(line.as_bytes());
        buffer.push(b'\n');
        slots.insert(
            key,
            PrimaryEntry {
                offset: new_offset,
                length: slot.length,
                seq: seq_of[pos],
            },
        );
    }

    let records = slots.len();
    Ok((buffer, slots, records))
}

/// Rewrites `log` down to its live records, dropping overwritten versions and tombstoned keys.
/// The rewrite lands through [`write_atomic`] (temp file, `fsync`, rename), so a crash before the
/// rename leaves the original log exactly as it was. `log` is reopened at its (unchanged) path
/// afterwards, so the handle keeps working for further scans and appends.
pub fn compact_log(log: &mut LogFile) -> io::Result<CompactionResult> {
    let bytes_before = log.size();
    let (buffer, slots, records) = plan(log)?;
    let bytes_after = buffer.len() as u64;

    let mode = permissions_of(log.path())?;
    write_atomic(log.path(), &buffer, mode)?;
    log.reopen()?;

    Ok(CompactionResult {
        slots,
        records,
        bytes_before,
        bytes_after,
        reclaimed_bytes: bytes_before.saturating_sub(bytes_after),
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::atomic::open_temporary;
    use crate::registry::table_registry::TableRegistry;
    use crate::schema::{parse_schema, HaysonRecord};
    use crate::storage::{serialize_entry, LogMeta};
    use serde_json::json;
    use std::fs;
    use std::io::Write;

    fn put(key: &str, seq: u64, name: &str) -> String {
        let mut data = HaysonRecord::new();
        data.insert("name".into(), json!(name));
        serialize_entry(
            &LogMeta {
                key: key.into(),
                op: LogOp::Put,
                seq,
                ts: 1,
            },
            Some(&data),
        )
    }

    fn del(key: &str, seq: u64) -> String {
        serialize_entry(
            &LogMeta {
                key: key.into(),
                op: LogOp::Del,
                seq,
                ts: 1,
            },
            None,
        )
    }

    #[test]
    fn keeps_only_the_last_version_of_an_overwritten_key() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("users.jsonl");
        let mut log = LogFile::open(&path, false).unwrap();
        log.append(&[put("a", 1, "first"), put("a", 2, "second")]).unwrap();

        let result = compact_log(&mut log).unwrap();

        assert_eq!(result.records, 1);
        let scan = log.scan(0).unwrap();
        assert_eq!(scan.entries.len(), 1);
        assert_eq!(
            scan.entries[0].entry.data.as_ref().unwrap().get("name"),
            Some(&json!("second"))
        );
    }

    #[test]
    fn drops_a_key_that_was_put_then_deleted() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("users.jsonl");
        let mut log = LogFile::open(&path, false).unwrap();
        log.append(&[put("a", 1, "first"), put("b", 2, "kept"), del("a", 3)])
            .unwrap();

        let result = compact_log(&mut log).unwrap();

        assert_eq!(result.records, 1);
        assert!(!result.slots.contains_key("a"));
        let scan = log.scan(0).unwrap();
        assert_eq!(scan.entries.len(), 1);
        assert_eq!(scan.entries[0].entry.meta.key, "b");
    }

    #[test]
    fn compacts_an_all_tombstones_log_down_to_empty() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("users.jsonl");
        let mut log = LogFile::open(&path, false).unwrap();
        log.append(&[put("a", 1, "first"), del("a", 2)]).unwrap();

        let result = compact_log(&mut log).unwrap();

        assert_eq!(result.records, 0);
        assert_eq!(log.size(), 0);
        assert_eq!(result.bytes_after, 0);
    }

    #[test]
    fn reclaims_the_space_freed_by_dead_entries_and_reports_it() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("users.jsonl");
        let mut log = LogFile::open(&path, false).unwrap();
        log.append(&[put("a", 1, "first"), put("a", 2, "second"), put("b", 3, "kept")])
            .unwrap();
        let bytes_before = log.size();

        let result = compact_log(&mut log).unwrap();

        assert_eq!(result.bytes_before, bytes_before);
        assert!(result.reclaimed_bytes > 0);
        assert_eq!(result.bytes_after, log.size());
    }

    #[test]
    fn resulting_file_is_a_usable_log_file_that_still_accepts_appends() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("users.jsonl");
        let mut log = LogFile::open(&path, false).unwrap();
        log.append(&[put("a", 1, "first"), put("a", 2, "second")]).unwrap();

        compact_log(&mut log).unwrap();
        log.append(&[put("c", 3, "third")]).unwrap();

        let reopened = LogFile::open(&path, true).unwrap();
        let scan = reopened.scan(0).unwrap();
        assert_eq!(scan.entries.len(), 2);
        assert_eq!(scan.entries[0].entry.meta.key, "a");
        assert_eq!(scan.entries[1].entry.meta.key, "c");
    }

    #[test]
    fn original_log_is_untouched_if_the_process_dies_before_the_rename() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("users.jsonl");
        let mut log = LogFile::open(&path, false).unwrap();
        log.append(&[put("a", 1, "first"), put("a", 2, "second")]).unwrap();
        let before = fs::read(&path).unwrap();

        let (buffer, _, _) = plan(&log).unwrap();
        let temp_path = path.with_extension("jsonl.tmp");
        let mut temp = open_temporary(&temp_path, None).unwrap();
        temp.write_all(&buffer).unwrap();
        temp.sync_all().unwrap();
        drop(temp);

        assert_eq!(
            fs::read(&path).unwrap(),
            before,
            "crash before rename must leave the original log intact"
        );

        fs::rename(&temp_path, &path).unwrap();
        assert_eq!(fs::read(&path).unwrap(), buffer);
    }

    fn users_schema() -> crate::schema::TableSchema {
        parse_schema(&json!({
            "schemaVersion": 1,
            "table": "users",
            "primaryKey": "id",
            "fields": { "id": { "type": "string" }, "name": { "type": "string" } },
        }))
        .unwrap()
    }

    #[test]
    fn registry_remap_after_compaction_still_answers_queries_correctly() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("users.jsonl");
        let mut log = LogFile::open(&path, false).unwrap();
        log.append(&[
            put("a", 1, "first"),
            put("a", 2, "second"),
            put("b", 3, "kept"),
            del("b", 4),
            put("c", 5, "also-kept"),
        ])
        .unwrap();

        let mut registry = TableRegistry::build_from_log(users_schema(), &log).unwrap();
        assert_eq!(registry.records(), 2);

        let result = compact_log(&mut log).unwrap();
        registry.remap(&result.slots, "");

        assert_eq!(registry.records(), 2);
        assert_eq!(registry.stats().dead_bytes, 0);

        for key in ["a", "c"] {
            let entry = registry.slot_of(key).expect("key survives compaction");
            let slot = Slot {
                offset: entry.offset,
                length: entry.length,
            };
            let line = log.read_slice(&slot).unwrap();
            assert!(line.contains(&format!("\"_k\":\"{key}\"")));
        }
        assert!(registry.slot_of("b").is_none());
    }
}