clawhdf5/crates/clawhdf5-agent/src/storage.rs

//! Disk I/O operations for HDF5 memory files.
//!
//! Uses memory-mapped I/O via `clawhdf5_io::MmapReader` for efficient
//! file reading with OS-managed paging.

use std::path::Path;

use crate::MemoryConfig;
use crate::MemoryError;
use crate::cache::MemoryCache;
use crate::knowledge::KnowledgeCache;
use crate::schema;
use crate::session::SessionCache;
use crate::wal::WalMark;

/// Write all in-memory state to an HDF5 file on disk.
pub fn write_to_disk(
    path: &Path,
    config: &MemoryConfig,
    cache: &MemoryCache,
    sessions: &SessionCache,
    knowledge: &KnowledgeCache,
) -> Result<(), MemoryError> {
    write_to_disk_with_mark(path, config, cache, sessions, knowledge, None)
}

/// [`write_to_disk`] for a checkpoint: `wal_applied` is the mark of the WAL
/// prefix whose entries `cache` already contains.
pub fn write_to_disk_with_mark(
    path: &Path,
    config: &MemoryConfig,
    cache: &MemoryCache,
    sessions: &SessionCache,
    knowledge: &KnowledgeCache,
    wal_applied: Option<WalMark>,
) -> Result<(), MemoryError> {
    let meta = schema::CheckpointMeta {
        wal_applied,
        ann_generation: None,
    };
    write_to_disk_with_meta(path, config, cache, sessions, knowledge, &meta)
}

/// [`write_to_disk`] with full checkpoint bookkeeping.
pub fn write_to_disk_with_meta(
    path: &Path,
    config: &MemoryConfig,
    cache: &MemoryCache,
    sessions: &SessionCache,
    knowledge: &KnowledgeCache,
    checkpoint: &schema::CheckpointMeta,
) -> Result<(), MemoryError> {
    let bytes = schema::build_hdf5_file_with_meta(config, cache, sessions, knowledge, checkpoint)?;

    if bytes.is_empty() {
        return Err(MemoryError::Hdf5("build_hdf5_file produced 0 bytes".into()));
    }

    // Write to a temp file first, then rename for atomicity
    let tmp_path = path.with_extension("h5.tmp");
    write_synced(&tmp_path, &bytes)?;
    rename_synced(&tmp_path, path)
}

/// Write `bytes` to `path` and flush them to stable storage.
pub(crate) fn write_synced(path: &Path, bytes: &[u8]) -> Result<(), MemoryError> {
    use std::io::Write;
    let mut f = std::fs::File::create(path).map_err(MemoryError::Io)?;
    f.write_all(bytes).map_err(MemoryError::Io)?;
    f.sync_all().map_err(MemoryError::Io)
}

/// Rename `from` over `to`, then sync the parent directory so the rename
/// itself survives a power loss. `from` must already be synced: without that,
/// the rename can reach disk before the data and leave an empty or partial
/// file under the final name.
///
/// This is per-checkpoint/snapshot cost only (each is already a full file
/// write). Individual WAL appends are deliberately not synced — see the
/// durability notes in the crate docs.
pub(crate) fn rename_synced(from: &Path, to: &Path) -> Result<(), MemoryError> {
    std::fs::rename(from, to).map_err(MemoryError::Io)?;
    #[cfg(unix)]
    if let Some(dir) = to.parent() {
        let dir = if dir.as_os_str().is_empty() {
            Path::new(".")
        } else {
            dir
        };
        // Directory fsync is best-effort: some filesystems refuse it, and the
        // rename has already happened.
        if let Ok(d) = std::fs::File::open(dir) {
            let _ = d.sync_all();
        }
    }
    Ok(())
}

/// Read an HDF5 file and return all state.
///
/// Uses memory-mapped I/O via `clawhdf5_io::MmapReader` for efficient
/// file access. The OS pages in data on demand rather than reading the
/// entire file into a contiguous buffer upfront.
pub fn read_from_disk(
    path: &Path,
) -> Result<(MemoryConfig, MemoryCache, SessionCache, KnowledgeCache), MemoryError> {
    read_from_disk_with_mark(path).map(|(state, _mark)| state)
}

/// Everything [`read_from_disk`] returns.
pub type StoreState = (MemoryConfig, MemoryCache, SessionCache, KnowledgeCache);

/// [`read_from_disk`], plus the checkpoint's [`WalMark`] (if any) so the
/// caller can skip WAL entries this file already contains.
pub fn read_from_disk_with_mark(path: &Path) -> Result<(StoreState, Option<WalMark>), MemoryError> {
    // `File::open` memory-maps the file itself (the facade's `mmap` feature is
    // on by default). Mapping it here and handing over `as_bytes().to_vec()`
    // did the same work and then copied the whole store — a second full copy
    // of the file, live for the whole parse, on top of the mapping.
    let file = clawhdf5::File::open(path)
        .map_err(|e| MemoryError::Hdf5(format!("cannot open {}: {e}", path.display())))?;

    let (mut config, cache, sessions, knowledge) = schema::validate_and_load(&file)?;
    config.path = path.to_path_buf();
    let wal_applied = schema::read_wal_mark(&file);

    Ok(((config, cache, sessions, knowledge), wal_applied))
}

/// [`read_from_disk`], plus all checkpoint bookkeeping.
pub fn read_from_disk_with_meta(
    path: &Path,
) -> Result<(StoreState, schema::CheckpointMeta), MemoryError> {
    let file = clawhdf5::File::open(path)
        .map_err(|e| MemoryError::Hdf5(format!("cannot open {}: {e}", path.display())))?;
    let (mut config, cache, sessions, knowledge) = schema::validate_and_load(&file)?;
    config.path = path.to_path_buf();
    let meta = schema::read_checkpoint_meta(&file);
    Ok(((config, cache, sessions, knowledge), meta))
}

/// Copy an HDF5 file atomically to a destination.
pub fn snapshot_file(src: &Path, dest: &Path) -> Result<std::path::PathBuf, MemoryError> {
    let dest_file = if dest.is_dir() {
        let filename = src.file_name().ok_or_else(|| {
            MemoryError::Io(std::io::Error::new(
                std::io::ErrorKind::InvalidInput,
                "source has no filename",
            ))
        })?;
        let ts = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs();
        dest.join(format!("snapshot_{ts}_{}", filename.to_string_lossy()))
    } else {
        dest.to_path_buf()
    };

    // Atomic copy: write to temp, then rename
    let tmp_path = dest_file.with_extension("h5.tmp");
    std::fs::copy(src, &tmp_path).map_err(MemoryError::Io)?;
    std::fs::File::open(&tmp_path)
        .and_then(|f| f.sync_all())
        .map_err(MemoryError::Io)?;
    rename_synced(&tmp_path, &dest_file)?;

    Ok(dest_file)
}