hayson_engine/storage/

log_set.rs

use std::collections::BTreeMap;
use std::fs;
use std::io;
use std::path::{Path, PathBuf};

use super::log_file::{LogFile, Slot};

pub const LOG_SUFFIX: &str = ".jsonl";

enum Layout {
    /// One log file at exactly this path, held under the [`crate::schema::UNPARTITIONED`] name.
    Flat(PathBuf),
    /// One log file per partition value, all inside this directory.
    Partitioned(PathBuf),
}

/// Every log file of one table: a single file when the table declares no `partitionBy`, one file
/// per partition value when it does. Routing only — which file a line goes to and which file a slot
/// came from; it holds no registry and knows nothing about records, which is why it sits in
/// `storage/` next to [`LogFile`] rather than in `table/`.
pub struct LogSet {
    layout: Layout,
    logs: BTreeMap<String, LogFile>,
    /// Passed on to every [`LogFile`] this set opens, now and on a later [`LogSet::discover`], so a
    /// readonly table holds descriptors that cannot write rather than descriptors that merely never
    /// do — which is what keeps a readonly open from cutting a half-written tail off somebody
    /// else's file (see [`LogFile::discard_partial_tail`]).
    readonly: bool,
    /// Carried the same way as `readonly`, and for the same reason: a partition file first opened
    /// by [`LogSet::ensure`]/[`LogSet::discover`], long after the table was configured, has to end
    /// up with the durability its table was given rather than the [`LogFile`] default.
    fsync: bool,
}

/// The partition names `dir` already holds a log file for, sorted. A directory that does not exist
/// yet is simply a table nothing was written to.
pub fn partition_names_in(dir: &Path) -> Vec<String> {
    let Ok(entries) = fs::read_dir(dir) else {
        return Vec::new();
    };
    let mut names: Vec<String> = entries
        .filter_map(Result::ok)
        .filter_map(|entry| entry.file_name().into_string().ok())
        .filter_map(|name| name.strip_suffix(LOG_SUFFIX).map(str::to_string))
        .filter(|name| !name.is_empty())
        .collect();
    names.sort_unstable();
    names
}

impl LogSet {
    /// The layout of a table with no `partitionBy`: the single log file at `path`, created there
    /// and then when this set is writable and absent — a readonly set refuses a log that is not on
    /// disk, since it has nothing to create it with.
    pub fn flat(path: impl AsRef<Path>, readonly: bool) -> io::Result<Self> {
        let mut set = LogSet {
            layout: Layout::Flat(path.as_ref().to_path_buf()),
            logs: BTreeMap::new(),
            readonly,
            fsync: true,
        };
        set.load(crate::schema::UNPARTITIONED)?;
        Ok(set)
    }

    /// The layout of a table that declares `partitionBy`: every `<value>.jsonl` already in `dir`.
    /// A table whose partitions were never written yet opens with no log file at all — the first
    /// [`LogSet::append`] creates both the directory and the file it needs.
    pub fn partitioned(dir: impl AsRef<Path>, readonly: bool) -> io::Result<Self> {
        let dir = dir.as_ref().to_path_buf();
        let mut set = LogSet {
            layout: Layout::Partitioned(dir.clone()),
            logs: BTreeMap::new(),
            readonly,
            fsync: true,
        };
        for part in partition_names_in(&dir) {
            set.load(&part)?;
        }
        Ok(set)
    }

    /// Whether these descriptors can write at all — `false` for a set a readonly open built, which
    /// refuses an append rather than merely never making one.
    pub fn is_readonly(&self) -> bool {
        self.readonly
    }

    pub fn is_partitioned(&self) -> bool {
        matches!(self.layout, Layout::Partitioned(_))
    }

    /// Every open partition, in name order.
    pub fn parts(&self) -> Vec<String> {
        self.logs.keys().cloned().collect()
    }

    pub fn log_of(&self, part: &str) -> Option<&LogFile> {
        self.logs.get(part)
    }

    pub fn log_of_mut(&mut self, part: &str) -> Option<&mut LogFile> {
        self.logs.get_mut(part)
    }

    pub fn path_of(&self, part: &str) -> PathBuf {
        match &self.layout {
            Layout::Flat(path) => path.clone(),
            Layout::Partitioned(dir) => dir.join(format!("{part}{LOG_SUFFIX}")),
        }
    }

    pub fn total_size(&self) -> u64 {
        self.logs.values().map(LogFile::size).sum()
    }

    /// The log for `part`, created on disk the first time a record names it.
    pub fn ensure(&mut self, part: &str) -> io::Result<&mut LogFile> {
        if !self.logs.contains_key(part) {
            if let Layout::Partitioned(dir) = &self.layout {
                fs::create_dir_all(dir)?;
            }
            self.load(part)?;
        }
        self.logs
            .get_mut(part)
            .ok_or_else(|| io::Error::other(format!("partition \"{part}\" is not open")))
    }

    pub fn append(&mut self, part: &str, lines: &[String]) -> io::Result<Vec<Slot>> {
        if lines.is_empty() {
            return Ok(Vec::new());
        }
        self.ensure(part)?.append(lines)
    }

    pub fn read_slice(&self, part: &str, slot: &Slot) -> io::Result<String> {
        self.logs
            .get(part)
            .ok_or_else(|| io::Error::other(format!("partition \"{part}\" is not open")))?
            .read_slice(slot)
    }

    /// Sets [`LogFile::set_fsync`] on every partition already open and on every one this set opens
    /// later — the whole table shares one durability setting, since the catalog declares `sync` per
    /// table and not per partition.
    pub fn set_fsync(&mut self, fsync: bool) {
        self.fsync = fsync;
        for log in self.logs.values_mut() {
            log.set_fsync(fsync);
        }
    }

    pub fn sync(&self) -> io::Result<()> {
        for log in self.logs.values() {
            log.sync()?;
        }
        Ok(())
    }

    /// Opens partition files that appeared in the directory since this handle last looked, and
    /// reports their names — an unpartitioned set never has any.
    pub fn discover(&mut self) -> io::Result<Vec<String>> {
        let Layout::Partitioned(dir) = &self.layout else {
            return Ok(Vec::new());
        };
        let mut found = Vec::new();
        for part in partition_names_in(&dir.clone()) {
            if self.logs.contains_key(&part) {
                continue;
            }
            self.load(&part)?;
            found.push(part);
        }
        Ok(found)
    }

    fn load(&mut self, part: &str) -> io::Result<()> {
        let mut log = LogFile::open(self.path_of(part), self.readonly)?;
        log.set_fsync(self.fsync);
        self.logs.insert(part.to_string(), log);
        Ok(())
    }
}

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

    fn line(key: &str) -> String {
        format!("{{\"_k\":\"{key}\",\"_op\":\"put\",\"_seq\":1,\"_ts\":1,\"id\":\"{key}\"}}")
    }

    #[test]
    fn a_flat_set_holds_exactly_the_file_it_was_given() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("notes.jsonl");
        let mut set = LogSet::flat(&path, false).unwrap();

        assert!(!set.is_partitioned());
        assert_eq!(set.parts(), vec![String::new()]);
        assert_eq!(set.path_of(""), path);
        set.append("", &[line("n1")]).unwrap();
        assert_eq!(set.total_size(), fs::metadata(&path).unwrap().len());
    }

    #[test]
    fn a_new_partition_value_creates_its_own_file_on_first_append() {
        let dir = tempfile::tempdir().unwrap();
        let mut set = LogSet::partitioned(dir.path().join("notes"), false).unwrap();
        assert!(set.parts().is_empty());

        set.append("note", &[line("n1")]).unwrap();
        set.append("adr", &[line("a1")]).unwrap();

        assert_eq!(set.parts(), vec!["adr".to_string(), "note".to_string()]);
        assert!(dir.path().join("notes/note.jsonl").is_file());
        assert!(dir.path().join("notes/adr.jsonl").is_file());
    }

    #[test]
    fn appending_to_a_known_value_reuses_its_file_instead_of_making_another() {
        let dir = tempfile::tempdir().unwrap();
        let mut set = LogSet::partitioned(dir.path().join("notes"), false).unwrap();
        set.append("note", &[line("n1")]).unwrap();
        set.append("note", &[line("n2")]).unwrap();

        assert_eq!(set.parts(), vec!["note".to_string()]);
        let content = fs::read_to_string(dir.path().join("notes/note.jsonl")).unwrap();
        assert_eq!(content.trim().lines().count(), 2);
    }

    #[test]
    fn reopening_finds_every_partition_file_already_on_disk() {
        let dir = tempfile::tempdir().unwrap();
        {
            let mut set = LogSet::partitioned(dir.path().join("notes"), false).unwrap();
            set.append("note", &[line("n1")]).unwrap();
            set.append("rule", &[line("r1")]).unwrap();
            set.sync().unwrap();
        }

        let set = LogSet::partitioned(dir.path().join("notes"), false).unwrap();
        assert_eq!(set.parts(), vec!["note".to_string(), "rule".to_string()]);
        assert!(set.log_of("note").unwrap().size() > 0);
    }

    #[test]
    fn discover_opens_a_partition_file_another_writer_created_since_the_last_look() {
        let dir = tempfile::tempdir().unwrap();
        let mut reader = LogSet::partitioned(dir.path().join("notes"), false).unwrap();
        {
            let mut writer = LogSet::partitioned(dir.path().join("notes"), false).unwrap();
            writer.append("adr", &[line("a1")]).unwrap();
            writer.sync().unwrap();
        }

        assert!(reader.log_of("adr").is_none());
        assert_eq!(reader.discover().unwrap(), vec!["adr".to_string()]);
        assert!(reader.log_of("adr").unwrap().size() > 0);
        assert!(reader.discover().unwrap().is_empty());
    }

    #[test]
    fn reads_a_slice_back_from_the_partition_it_was_written_to() {
        let dir = tempfile::tempdir().unwrap();
        let mut set = LogSet::partitioned(dir.path().join("notes"), false).unwrap();
        let slots = set.append("adr", &[line("a1")]).unwrap();

        assert_eq!(set.read_slice("adr", &slots[0]).unwrap(), line("a1"));
        assert!(set.read_slice("note", &slots[0]).is_err());
    }

    #[test]
    fn turning_fsync_off_reaches_partitions_opened_before_and_after_the_call() {
        let dir = tempfile::tempdir().unwrap();
        let mut set = LogSet::partitioned(dir.path().join("notes"), false).unwrap();
        set.append("adr", &[line("a1")]).unwrap();

        set.set_fsync(false);
        set.append("note", &[line("n1")]).unwrap();
        set.sync().unwrap();

        assert_eq!(
            set.log_of("adr").unwrap().fsyncs(),
            0,
            "a partition open before the call"
        );
        assert_eq!(
            set.log_of("note").unwrap().fsyncs(),
            0,
            "a partition created after the call"
        );
    }

    #[test]
    fn a_set_left_alone_keeps_fsyncing_every_partition_it_holds() {
        let dir = tempfile::tempdir().unwrap();
        let mut set = LogSet::partitioned(dir.path().join("notes"), false).unwrap();
        set.append("adr", &[line("a1")]).unwrap();
        set.append("note", &[line("n1")]).unwrap();
        set.sync().unwrap();

        assert_eq!(set.log_of("adr").unwrap().fsyncs(), 1);
        assert_eq!(set.log_of("note").unwrap().fsyncs(), 1);
    }

    #[test]
    fn a_file_that_is_not_a_log_never_becomes_a_partition() {
        let dir = tempfile::tempdir().unwrap();
        fs::create_dir_all(dir.path().join("notes")).unwrap();
        fs::write(dir.path().join("notes/note.jsonl"), "").unwrap();
        fs::write(dir.path().join("notes/README.md"), "not a log").unwrap();
        fs::write(dir.path().join("notes/.jsonl"), "").unwrap();

        let set = LogSet::partitioned(dir.path().join("notes"), false).unwrap();
        assert_eq!(set.parts(), vec!["note".to_string()]);
    }
}