diff --git a/CHANGELOG.md b/CHANGELOG.md index 7cc517b..6a4eeb6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -19,6 +19,54 @@ - `clawhdf5-agent`: `benches/bench.rs` and `benches/memory_bench.rs` no longer compiled against the current `strategy`/`consolidation` APIs. +### Durability & Integrity +- `clawhdf5-agent`: a crash between writing a checkpoint and truncating the WAL + no longer **duplicates every pending entry** on the next open. Each + checkpoint records a `WalMark` (byte length + chained CRC of the WAL prefix it + folded in) in `/meta`; `open()` skips exactly that prefix when it is still + present. No WAL format change for this; older files behave as before. +- `clawhdf5-agent`: checkpoints and snapshots are durable as a unit — the temp + file is synced before the rename and the directory after it. Individual WAL + appends remain unsynced by design (documented in `CLAUDE.md`). +- `clawhdf5-agent`: `save_or_update` hits are logged as a new `Update` WAL + record, so replay updates in place instead of appending a duplicate. WAL + header version 3 → 4 (so older builds refuse the file rather than truncating + a record they can't parse); v3 files are read and upgraded in place. +- `clawhdf5-agent`: loading validates every per-record dataset length (a + truncated store is now `MemoryError::Schema`, not a later panic), fixes the + `n.len() == n.len()` tautology that trusted a norms dataset of any length, + and rejects `embedding_dim == 0` with records present. +- `clawhdf5-agent`: eight behavioural `MemoryConfig` fields are now persisted in + `/meta`. Previously they reset to defaults on every open — a compressed store + was rewritten uncompressed, `wal_enabled = false` flipped back to `true`. +- `clawhdf5-agent`: `compression = true` never worked in a default build (it + requested Zstd without enabling the feature, so every checkpoint failed with + `unsupported filter: 32015`). Default builds now use deflate; Zstd is the new + opt-in `zstd` feature. +- `clawhdf5-agent`: **single-writer lock** (`.h5.lock`, + `MemoryError::Locked`) — two handles on one store used to silently destroy + each other's data. New `HDF5Memory::open_read_only` gives a lock-free, + never-writing view; the CLI's read-only subcommands use it. +- `clawhdf5-agent`: an unreadable WAL (torn header / bad magic) is quarantined + (`HDF5Memory::quarantined_wal()`) instead of blocking `open()` of a healthy + store. A WAL from an unknown newer version still fails and is left intact. +- `clawhdf5-agent`: provenance records are renumbered on compaction (they + weren't, so every later `save_or_update` raised a false High integrity + alert); pending anomaly alerts and tracked sessions are bounded; + `snapshot()` includes entries still in the WAL. +- `clawhdf5-agent`: hybrid ranking is deterministic (index tie-breaks instead + of `HashMap` order); a set of identical positive scores — including a single + candidate — normalises to 1.0 rather than 0.0; the Hebbian boost no longer + reinforces zero-score filler results. +- `clawhdf5-format`: chunked/VDS/hyperslab reads size their buffers with + overflow-checked arithmetic and fallible allocation, so crafted dimensions + are `FormatError::Overflow` instead of a wrapped size or a process abort; + `parallel_read` bounds checks use `checked_add`. +- `clawhdf5`: a malformed filter-pipeline message is an error instead of being + treated as "no filters" (which returned compressed bytes as data); + `FileBuilder::write` is atomic and synced instead of truncating the + destination first. + ### CI / Testing - CI now lints every target (`cargo clippy --all-targets`) plus `clawhdf5-format`'s optional features, compiles all benches, and tests the diff --git a/CLAUDE.md b/CLAUDE.md index a787b4d..041cfa3 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -40,6 +40,25 @@ Cargo workspace with 16 crates under `crates/` (plus `libaec-sys`, an internal F format (v2) is still fully readable; the oldest no-CRC format (v1) is only reachable through the one-time migration path in `HDF5Memory::open`, not through the public `WalFile::read_entries`. + **What the WAL guarantees:** integrity, ordering, and recovery from a + *process* crash at any point — including between a checkpoint and the WAL + truncate (each checkpoint records a `WalMark` in `/meta`, and `open()` skips + the WAL prefix the `.h5` already contains, so entries are never applied + twice). Checkpoints and snapshots are made durable as a unit (temp file + synced, renamed, directory synced). **What it does not guarantee:** + individual WAL appends are *not* fsynced (a deliberate latency trade-off), so + saves made since the last checkpoint can be lost on power failure or kernel + panic. Current header version is 4 (adds the `Update` record used by + `save_or_update`); v3 files are read and upgraded in place. +- A store has a **single writer**: `HDF5Memory::create`/`open` hold an exclusive + advisory lock on `.h5.lock` and a second opener gets + `MemoryError::Locked`. Use `HDF5Memory::open_read_only` for a lock-free, + never-writing point-in-time view (the CLI's `recall`/`stats`/`agents-md`/ + `export` do). An unreadable WAL (torn header, bad magic) is quarantined to + `.h5.wal.corrupt-` rather than blocking `open()`; a WAL with an + unknown *newer* version still fails and is left untouched. +- `MemoryConfig::compression` uses deflate by default; enable the agent's + `zstd` feature to compress embeddings with Zstd instead (links libzstd). - `Dataset::verify_provenance()` (clawhdf5 facade, `provenance` feature, on by default) recomputes a dataset's SHA-256 and compares it against the `_provenance_sha256` attribute written automatically on save when diff --git a/crates/clawhdf5-agent/Cargo.toml b/crates/clawhdf5-agent/Cargo.toml index d7ba0ff..4c4caad 100644 --- a/crates/clawhdf5-agent/Cargo.toml +++ b/crates/clawhdf5-agent/Cargo.toml @@ -48,6 +48,9 @@ harness = false default = ["float16", "hnsw"] float16 = ["half"] parallel = ["rayon"] +# Compress embeddings with Zstd instead of deflate when +# `MemoryConfig::compression` is on. Off by default: it links libzstd (C). +zstd = ["clawhdf5/zstd"] # HNSW approximate-nearest-neighbour acceleration for the vector stage of # hybrid_search. On by default; the index is rebuilt from the cache on demand # and stays self-consistent with the persisted memory store. Disable with diff --git a/crates/clawhdf5-agent/src/anomaly.rs b/crates/clawhdf5-agent/src/anomaly.rs index 9564341..bc8975c 100644 --- a/crates/clawhdf5-agent/src/anomaly.rs +++ b/crates/clawhdf5-agent/src/anomaly.rs @@ -161,6 +161,9 @@ pub struct WriteEvent { // WriteAnomalyDetector // --------------------------------------------------------------------------- +/// Upper bound on distinct session ids the detector tracks at once. +const MAX_TRACKED_SESSIONS: usize = 4096; + /// Tracks write events and raises alerts for suspicious behaviour. #[derive(Debug)] pub struct WriteAnomalyDetector { @@ -189,6 +192,23 @@ impl WriteAnomalyDetector { if event.timestamp > self.last_timestamp { self.last_timestamp = event.timestamp; } + // Bound the per-session map: a long-lived process sees an unbounded + // number of distinct session ids. When it overflows, forget the + // sessions with the fewest writes (they are furthest from the limit + // this map exists to enforce); the current one is re-added below. + if self.session_counts.len() >= MAX_TRACKED_SESSIONS + && !self.session_counts.contains_key(&event.session_id) + { + let mut counts: Vec = self.session_counts.values().copied().collect(); + let keep_from = counts.len() / 2; + counts.select_nth_unstable(keep_from); + let threshold = counts[keep_from]; + self.session_counts.retain(|_, c| *c >= threshold); + if self.session_counts.len() >= MAX_TRACKED_SESSIONS { + // Every session had the same count: drop them all. + self.session_counts.clear(); + } + } *self .session_counts .entry(event.session_id.clone()) diff --git a/crates/clawhdf5-agent/src/async_memory.rs b/crates/clawhdf5-agent/src/async_memory.rs index e6200e0..1f665ad 100644 --- a/crates/clawhdf5-agent/src/async_memory.rs +++ b/crates/clawhdf5-agent/src/async_memory.rs @@ -408,6 +408,10 @@ impl AsyncHDF5Memory { let (tx, rx) = oneshot::channel(); let _ = self.write_tx.send(WriteCmd::Shutdown(tx)).await; let _ = rx.await; + // The writer task has stopped, so nothing can write through this + // handle any more: release the single-writer lock now rather than at + // drop, so the store can be reopened while `self` is still in scope. + self.inner.lock().await.release_store_lock(); Ok(()) } } diff --git a/crates/clawhdf5-agent/src/hybrid.rs b/crates/clawhdf5-agent/src/hybrid.rs index 0416a9b..20a9f2d 100644 --- a/crates/clawhdf5-agent/src/hybrid.rs +++ b/crates/clawhdf5-agent/src/hybrid.rs @@ -91,14 +91,22 @@ pub fn merge_vector_keyword( } let mut results: Vec<(usize, f32)> = merged.into_iter().collect(); - results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); + // Index tie-break: `merged` is a HashMap, so without it the ties that + // survive `truncate` differ from run to run. + results.sort_by(|a, b| { + b.1.partial_cmp(&a.1) + .unwrap_or(std::cmp::Ordering::Equal) + .then(a.0.cmp(&b.0)) + }); results.truncate(k); results } /// Normalize a set of scores to the [0, 1] range using min-max normalization. /// -/// If all scores are identical, returns 0.0 for each entry. +/// If all scores are identical there is no spread to normalise: each entry +/// gets 1.0 when that score is positive (all equally the best match) and 0.0 +/// otherwise (nothing matched). fn normalize_scores(scores: &[(usize, f32)]) -> Vec<(usize, f32)> { if scores.is_empty() { return Vec::new(); @@ -112,7 +120,13 @@ fn normalize_scores(scores: &[(usize, f32)]) -> Vec<(usize, f32)> { let range = max - min; if range == 0.0 { - return scores.iter().map(|(idx, _)| (*idx, 0.0)).collect(); + // All candidates scored the same (including the single-candidate + // case), so min-max has no spread to work with. They are all equally + // the best match if that score is positive, and all non-matches + // otherwise. This used to return 0.0 unconditionally, which erased a + // lone perfect match from the fused score. + let level = if max > 0.0 { 1.0 } else { 0.0 }; + return scores.iter().map(|(idx, _)| (*idx, level)).collect(); } scores @@ -324,10 +338,18 @@ mod tests { #[test] fn normalize_scores_single() { + // A lone positive score is the best match there is, not a non-match. let result = normalize_scores(&[(0, 5.0)]); assert_eq!(result.len(), 1); - // Single score normalizes to 0.0 (range is 0) - assert_eq!(result[0].1, 0.0); + assert_eq!(result[0].1, 1.0); + } + + #[test] + fn normalize_scores_all_equal() { + let matched = normalize_scores(&[(0, 0.4), (1, 0.4)]); + assert!(matched.iter().all(|(_, s)| *s == 1.0)); + let unmatched = normalize_scores(&[(0, 0.0), (1, 0.0)]); + assert!(unmatched.iter().all(|(_, s)| *s == 0.0)); } #[test] diff --git a/crates/clawhdf5-agent/src/lib.rs b/crates/clawhdf5-agent/src/lib.rs index d2be326..f07c755 100644 --- a/crates/clawhdf5-agent/src/lib.rs +++ b/crates/clawhdf5-agent/src/lib.rs @@ -37,6 +37,7 @@ pub mod schema; pub mod search; pub mod session; pub mod storage; +mod store_lock; pub mod temporal; pub mod wal; @@ -86,6 +87,8 @@ pub enum MemoryError { Hdf5(String), Schema(String), NotFound(String), + /// Another `HDF5Memory` (in this or another process) has the store open. + Locked(String), } impl std::fmt::Display for MemoryError { @@ -95,6 +98,7 @@ impl std::fmt::Display for MemoryError { MemoryError::Hdf5(e) => write!(f, "HDF5 error: {e}"), MemoryError::Schema(e) => write!(f, "schema error: {e}"), MemoryError::NotFound(e) => write!(f, "not found: {e}"), + MemoryError::Locked(e) => write!(f, "store is locked: {e}"), } } } @@ -201,6 +205,9 @@ pub trait AgentMemory { fn get_session_summary(&self, session_id: &str) -> Result>; } +/// Most anomaly alerts kept between `take_anomaly_alerts` calls. +const MAX_PENDING_ALERTS: usize = 1024; + // --- HDF5Memory --- pub struct HDF5Memory { @@ -240,6 +247,15 @@ pub struct HDF5Memory { /// via [`HDF5Memory::take_anomaly_alerts`]. Saves are never blocked on /// these — surfacing is opt-in for callers that want to act on them. anomaly_alerts: Vec, + /// Opened with [`HDF5Memory::open_read_only`]: nothing may reach the disk. + read_only: bool, + /// A WAL that `open()` could not read and moved aside; see + /// [`HDF5Memory::quarantined_wal`]. + quarantined_wal: Option, + /// Single-writer guard. Declared last so it is released only after the + /// WAL and everything else has been dropped. `None` once a wrapper that + /// has stopped all writes released it early (see `release_store_lock`). + _lock: Option, } impl std::fmt::Debug for HDF5Memory { @@ -251,6 +267,7 @@ impl std::fmt::Debug for HDF5Memory { impl HDF5Memory { /// Create a new HDF5 memory file with the given configuration. pub fn create(config: MemoryConfig) -> Result { + let lock = store_lock::StoreLock::acquire(&config.path)?; let cache = MemoryCache::new(config.embedding_dim); let sessions = SessionCache::new(); let knowledge = KnowledgeCache::new(); @@ -282,20 +299,110 @@ impl HDF5Memory { provenance: provenance::ProvenanceStore::new(), anomaly: anomaly::WriteAnomalyDetector::new(anomaly::AnomalyConfig::default()), anomaly_alerts: Vec::new(), + read_only: false, + quarantined_wal: None, + _lock: Some(lock), }) } /// Open an existing HDF5 memory file. + /// If the WAL at `wal_path` can't possibly be replayed — its header is + /// torn (crash while the file was being created) or isn't a WAL header at + /// all — move it aside so a healthy `.h5` still opens, and return where it + /// went. A well-formed header with an *unknown version* is left alone and + /// still fails `open()`: that WAL was most likely written by a newer + /// build, and discarding it would lose data this binary merely can't read. + fn quarantine_unreadable_wal(wal_path: &Path) -> Result> { + if !wal_path.exists() { + return Ok(None); + } + let reason = match wal::wal_header_status(wal_path)? { + wal::WalHeaderStatus::Readable | wal::WalHeaderStatus::UnknownVersion(_) => { + return Ok(None); + } + wal::WalHeaderStatus::Torn => "truncated header", + wal::WalHeaderStatus::BadMagic => "bad magic bytes", + }; + let ts = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_secs(); + let dest = wal_path.with_extension(format!("wal.corrupt-{ts}")); + std::fs::rename(wal_path, &dest)?; + eprintln!( + "clawhdf5-agent: WAL {} is unreadable ({reason}); moved to {} and continuing \ + from the last checkpoint", + wal_path.display(), + dest.display() + ); + Ok(Some(dest)) + } + + /// Give up the single-writer lock before this value is dropped. Only for + /// wrappers that have already stopped every write path but keep the handle + /// alive (`AsyncHDF5Memory::shutdown`), so the store can be reopened. + #[cfg_attr(not(feature = "async"), allow(dead_code))] + pub(crate) fn release_store_lock(&mut self) { + self._lock = None; + } + + /// Where `open()` moved an unreadable WAL, if it had to. Entries that were + /// only in that WAL are not in this store; the file is kept for forensics. + pub fn quarantined_wal(&self) -> Option<&Path> { + self.quarantined_wal.as_deref() + } + pub fn open(path: &Path) -> Result { - let (config, mut cache, sessions, knowledge) = storage::read_from_disk(path)?; + Self::open_impl(path, false) + } + + /// Open a store for reading only, without taking the single-writer lock — + /// so it works while another `HDF5Memory` (in this or another process) + /// has the store open for writing, e.g. to inspect what is on disk. + /// + /// It loads the last checkpoint plus whatever the WAL held at that + /// moment; it is a point-in-time view and does not follow later writes. + /// Nothing is written: the WAL file is not repaired, upgraded or moved, + /// and every operation that would persist state returns an error. + pub fn open_read_only(path: &Path) -> Result { + Self::open_impl(path, true) + } + + fn open_impl(path: &Path, read_only: bool) -> Result { + let lock = if read_only { + None + } else { + Some(store_lock::StoreLock::acquire(path)?) + }; + let ((config, mut cache, sessions, knowledge), wal_applied) = + storage::read_from_disk_with_mark(path)?; // Replay WAL if present let wal_path = path.with_extension("h5.wal"); - let wal = if wal_path.exists() { + let quarantined_wal = if read_only { + None + } else { + Self::quarantine_unreadable_wal(&wal_path)? + }; + let wal = if read_only { + // Replay in memory only. `WalFile::open` would truncate a torn + // tail and may rewrite the header — both belong to the writer. An + // unreadable WAL is simply skipped: the writer will deal with it. + if wal_path.exists() + && let Ok(entries) = + wal::WalFile::read_entries_for_migration(&wal_path, wal_applied) + { + wal::replay_into_cache(&entries, &mut cache); + } + None + } else if wal_path.exists() { // Uses the migration-only reader since this is the one legitimate // path that may need to read a legacy (pre-CRC) WAL file — see // WalFile::read_entries_for_migration. - let entries = wal::WalFile::read_entries_for_migration(&wal_path)?; + // `wal_applied` drops the prefix a checkpoint already folded in, + // in case the process died between writing the .h5 and + // truncating the WAL. + let entries = wal::WalFile::read_entries_for_migration(&wal_path, wal_applied)?; wal::replay_into_cache(&entries, &mut cache); Some(wal::WalFile::open(&wal_path)?) } else if config.wal_enabled { @@ -327,6 +434,9 @@ impl HDF5Memory { provenance: provenance::ProvenanceStore::new(), anomaly: anomaly::WriteAnomalyDetector::new(anomaly::AnomalyConfig::default()), anomaly_alerts: Vec::new(), + read_only, + quarantined_wal, + _lock: lock, }) } @@ -336,12 +446,22 @@ impl HDF5Memory { /// also clear the WAL, otherwise `open()` will replay stale entries /// on top of the already-persisted data, duplicating them. fn flush(&mut self) -> Result<()> { - storage::write_to_disk( + if self.read_only { + return Err(MemoryError::Io(std::io::Error::new( + std::io::ErrorKind::PermissionDenied, + "store was opened read-only", + ))); + } + // Record which WAL prefix this checkpoint contains, so a crash before + // the truncate below can't replay those entries a second time. + let wal_applied = self.wal.as_ref().map(|w| w.mark()); + storage::write_to_disk_with_mark( &self.config.path, &self.config, &self.cache, &self.sessions, &self.knowledge, + wal_applied, )?; if let Some(ref mut w) = self.wal { w.truncate()?; @@ -410,7 +530,7 @@ impl HDF5Memory { .into_iter() .flatten() { - self.anomaly_alerts.push(alert); + self.push_anomaly_alert(alert); } } @@ -431,7 +551,7 @@ impl HDF5Memory { .provenance .verify_integrity(record_id as u64, current_chunk) { - self.anomaly_alerts.push(anomaly::AnomalyAlert { + self.push_anomaly_alert(anomaly::AnomalyAlert { severity: anomaly::Severity::High, message: format!( "provenance integrity mismatch for record {record_id}: stored content no \ @@ -442,6 +562,18 @@ impl HDF5Memory { } } + /// Queue an alert, keeping only the most recent [`MAX_PENDING_ALERTS`]. + /// Alerts never block a save, so a caller that never drains them — or a + /// session stuck over its write limit, which alerts on every write — + /// must not be able to grow this without bound. + fn push_anomaly_alert(&mut self, alert: anomaly::AnomalyAlert) { + if self.anomaly_alerts.len() >= MAX_PENDING_ALERTS { + let excess = self.anomaly_alerts.len() + 1 - MAX_PENDING_ALERTS; + self.anomaly_alerts.drain(..excess); + } + self.anomaly_alerts.push(alert); + } + /// Alerts raised by anomaly detection / provenance checks since the last /// call, draining the internal queue. pub fn take_anomaly_alerts(&mut self) -> Vec { @@ -621,7 +753,7 @@ impl HDF5Memory { if let Some(existing_idx) = self.cache.find_by_tags(&entry.tags) { if let Some(ref mut w) = self.wal { let wal_entry = wal::WalEntry { - entry_type: wal::WalEntryType::Save, + entry_type: wal::WalEntryType::Update, timestamp: entry.timestamp, chunk: entry.chunk.clone(), embedding: entry.embedding.clone(), @@ -629,6 +761,7 @@ impl HDF5Memory { session_id: entry.session_id.clone(), tags: entry.tags.clone(), tombstone_index: None, + update_index: Some(existing_idx), }; w.append_save(&wal_entry)?; } @@ -660,9 +793,6 @@ impl HDF5Memory { .is_none_or(|w| w.pending_count() as usize > self.config.wal_max_entries); if needs_flush { self.flush()?; - if let Some(ref mut w) = self.wal { - w.truncate()?; - } } return Ok(existing_idx); } @@ -683,6 +813,7 @@ impl AgentMemory for HDF5Memory { session_id: entry.session_id.clone(), tags: entry.tags.clone(), tombstone_index: None, + update_index: None, }; w.append_save(&wal_entry)?; } @@ -708,9 +839,6 @@ impl AgentMemory for HDF5Memory { .is_none_or(|w| w.pending_count() as usize > self.config.wal_max_entries); if needs_flush { self.flush()?; - if let Some(ref mut w) = self.wal { - w.truncate()?; - } } Ok(idx) } @@ -761,8 +889,10 @@ impl AgentMemory for HDF5Memory { } fn compact(&mut self) -> Result { - let (removed, _index_map) = self.cache.compact(); + let (removed, index_map) = self.cache.compact(); if removed > 0 { + // Record ids are cache indices, which compaction just renumbered. + self.provenance.remap(&index_map); // Compaction renumbers cache indices; rebuild the index to match. self.hnsw_mark_dirty(); self.flush()?; @@ -779,7 +909,18 @@ impl AgentMemory for HDF5Memory { } fn snapshot(&self, dest: &Path) -> Result { - storage::snapshot_file(&self.config.path, dest) + let snapshot = storage::snapshot_file(&self.config.path, dest)?; + // Entries saved since the last checkpoint live only in the WAL. Copy + // it alongside (where `open()` looks for it) so the snapshot is the + // store as it is now, not as of the last checkpoint. The .h5 is + // copied first: if a checkpoint lands in between, the WAL copy is + // empty or its prefix is skipped via the checkpoint mark — never + // applied twice. + let wal_path = self.config.path.with_extension("h5.wal"); + if self.wal.as_ref().is_some_and(|w| !w.is_empty()) && wal_path.exists() { + storage::snapshot_file(&wal_path, &snapshot.with_extension("h5.wal"))?; + } + Ok(snapshot) } fn add_session( @@ -884,9 +1025,6 @@ impl HDF5Memory { *w *= d; } self.flush()?; - if let Some(ref mut w) = self.wal { - w.truncate()?; - } Ok(()) } @@ -898,9 +1036,6 @@ impl HDF5Memory { /// Explicit WAL merge: flush .h5, truncate WAL. pub fn flush_wal(&mut self) -> Result<()> { self.flush()?; - if let Some(ref mut w) = self.wal { - w.truncate()?; - } Ok(()) } } @@ -1401,6 +1536,149 @@ mod tests { assert_eq!(mem.count(), 1); } + #[test] + fn compaction_does_not_cause_false_provenance_alerts() { + let dir = TempDir::new().unwrap(); + let mut mem = HDF5Memory::create(make_config(&dir)).unwrap(); + for name in ["a", "b", "c"] { + let mut e = make_entry(name, &[1.0, 0.0, 0.0, 0.0]); + e.tags = format!("tag-{name}"); + mem.save(e).unwrap(); + } + // 1 of 3 tombstoned exceeds compact_threshold, so delete() compacts. + mem.delete(0).unwrap(); + assert_eq!(mem.cache.chunks, ["b", "c"]); + mem.take_anomaly_alerts(); + + // "c" moved from id 2 to id 1. Its recorded hash must have moved too, + // or this update is checked against "b"'s hash and flagged. + let mut update = make_entry("c2", &[0.0, 1.0, 0.0, 0.0]); + update.tags = "tag-c".into(); + assert_eq!(mem.save_or_update(update).unwrap(), 1); + let alerts = mem.take_anomaly_alerts(); + assert!( + !alerts.iter().any(|a| a.message.contains("provenance")), + "{alerts:?}" + ); + } + + #[test] + fn pending_alerts_are_bounded() { + let dir = TempDir::new().unwrap(); + let mut mem = HDF5Memory::create(make_config(&dir)).unwrap(); + for i in 0..(MAX_PENDING_ALERTS + 50) { + mem.push_anomaly_alert(anomaly::AnomalyAlert { + severity: anomaly::Severity::Low, + message: format!("alert {i}"), + timestamp: i as f64, + }); + } + let alerts = mem.take_anomaly_alerts(); + assert_eq!(alerts.len(), MAX_PENDING_ALERTS); + assert_eq!(alerts[0].message, "alert 50", "oldest are dropped first"); + } + + #[test] + fn snapshot_includes_entries_still_in_the_wal() { + let dir = TempDir::new().unwrap(); + let mut config = make_config(&dir); + config.wal_enabled = true; + let mut mem = HDF5Memory::create(config).unwrap(); + mem.save(make_entry("checkpointed", &[1.0, 0.0, 0.0, 0.0])) + .unwrap(); + mem.flush_wal().unwrap(); + mem.save(make_entry("wal-only", &[0.0, 1.0, 0.0, 0.0])) + .unwrap(); + + let snap = mem.snapshot(&dir.path().join("snap.h5")).unwrap(); + let restored = HDF5Memory::open(&snap).unwrap(); + assert_eq!(restored.cache.chunks, ["checkpointed", "wal-only"]); + } + + #[test] + fn store_has_a_single_writer() { + let dir = TempDir::new().unwrap(); + let config = make_config(&dir); + let path = config.path.clone(); + let mem = HDF5Memory::create(config).unwrap(); + assert!(matches!( + HDF5Memory::open(&path), + Err(MemoryError::Locked(_)) + )); + drop(mem); + HDF5Memory::open(&path).unwrap(); + } + + #[test] + fn read_only_open_coexists_with_a_writer_and_never_writes() { + let dir = TempDir::new().unwrap(); + let mut config = make_config(&dir); + config.wal_enabled = true; + let path = config.path.clone(); + let wal_path = path.with_extension("h5.wal"); + let mut writer = HDF5Memory::create(config).unwrap(); + writer + .save(make_entry("pending", &[1.0, 0.0, 0.0, 0.0])) + .unwrap(); + let wal_before = std::fs::read(&wal_path).unwrap(); + let h5_before = std::fs::read(&path).unwrap(); + + // Sees the checkpoint plus the writer's un-checkpointed WAL entry. + let mut reader = HDF5Memory::open_read_only(&path).unwrap(); + assert_eq!(reader.cache.chunks, ["pending"]); + assert!(reader.save(make_entry("nope", &[0.0; 4])).is_err()); + assert!(reader.flush_wal().is_err()); + drop(reader); + + assert_eq!(std::fs::read(&wal_path).unwrap(), wal_before); + assert_eq!(std::fs::read(&path).unwrap(), h5_before); + // The writer is unaffected. + writer + .save(make_entry("more", &[0.0, 1.0, 0.0, 0.0])) + .unwrap(); + } + + #[test] + fn unreadable_wal_is_quarantined_not_fatal() { + let dir = TempDir::new().unwrap(); + let config = make_config(&dir); + let path = config.path.clone(); + let wal_path = path.with_extension("h5.wal"); + { + let mut mem = HDF5Memory::create(config).unwrap(); + mem.save(make_entry("kept", &[1.0, 0.0, 0.0, 0.0])).unwrap(); + mem.flush_wal().unwrap(); + } + std::fs::write(&wal_path, b"not a wal at all").unwrap(); + + let mem = HDF5Memory::open(&path).unwrap(); + assert_eq!(mem.cache.chunks, ["kept"]); + let moved = mem.quarantined_wal().expect("WAL should be quarantined"); + assert_eq!(std::fs::read(moved).unwrap(), b"not a wal at all"); + // A fresh, valid WAL took its place. + assert!(wal::WalFile::read_entries(&wal_path).unwrap().is_empty()); + } + + #[test] + fn wal_from_a_newer_build_is_refused_not_discarded() { + let dir = TempDir::new().unwrap(); + let mut config = make_config(&dir); + config.wal_enabled = true; + let path = config.path.clone(); + let wal_path = path.with_extension("h5.wal"); + drop(HDF5Memory::create(config).unwrap()); + let mut bytes = std::fs::read(&wal_path).unwrap(); + bytes[4] = 200; // a version this build has never heard of + std::fs::write(&wal_path, &bytes).unwrap(); + + assert!(HDF5Memory::open(&path).is_err()); + assert_eq!( + std::fs::read(&wal_path).unwrap(), + bytes, + "WAL left untouched" + ); + } + #[test] fn empty_file_operations() { let dir = TempDir::new().unwrap(); @@ -1409,6 +1687,7 @@ mod tests { let mem = HDF5Memory::create(config).unwrap(); assert_eq!(mem.count(), 0); assert_eq!(mem.count_active(), 0); + drop(mem); // a store has a single writer; release it before reopening let mem2 = HDF5Memory::open(&path).unwrap(); assert_eq!(mem2.count(), 0); diff --git a/crates/clawhdf5-agent/src/provenance.rs b/crates/clawhdf5-agent/src/provenance.rs index 946940d..bbed6cc 100644 --- a/crates/clawhdf5-agent/src/provenance.rs +++ b/crates/clawhdf5-agent/src/provenance.rs @@ -105,6 +105,23 @@ impl ProvenanceStore { self.records.insert(provenance.record_id, provenance); } + /// Renumber records after the store was compacted. `index_map[old]` is + /// the record's new id, or `None` if it was removed. Without this, every + /// surviving record's hash ends up filed under some other record's id and + /// the next integrity check reports a bogus mismatch. + pub fn remap(&mut self, index_map: &[Option]) { + let old = std::mem::take(&mut self.records); + for (old_id, mut prov) in old { + let new_id = usize::try_from(old_id) + .ok() + .and_then(|i| index_map.get(i).copied().flatten()); + if let Some(new_id) = new_id { + prov.record_id = new_id as u64; + self.records.insert(new_id as u64, prov); + } + } + } + /// Retrieve by record ID. pub fn get(&self, record_id: u64) -> Option<&MemoryProvenance> { self.records.get(&record_id) diff --git a/crates/clawhdf5-agent/src/schema.rs b/crates/clawhdf5-agent/src/schema.rs index 6219bb7..5376d13 100644 --- a/crates/clawhdf5-agent/src/schema.rs +++ b/crates/clawhdf5-agent/src/schema.rs @@ -12,16 +12,36 @@ use crate::MemoryError; use crate::cache::MemoryCache; use crate::knowledge::KnowledgeCache; use crate::session::SessionCache; +use crate::wal::WalMark; pub const SCHEMA_VERSION: &str = "1.0"; pub const ZEROCLAW_VERSION: &str = "0.8.0"; +/// `/meta` attributes holding the [`WalMark`] of the WAL prefix already folded +/// into this file. Absent on files written before the mark existed, and when +/// the checkpoint was taken with an empty WAL. +const WAL_APPLIED_LEN_ATTR: &str = "wal_applied_len"; +const WAL_APPLIED_CRC_ATTR: &str = "wal_applied_crc"; + /// Build a complete HDF5 file from the in-memory state. pub fn build_hdf5_file( config: &MemoryConfig, cache: &MemoryCache, sessions: &SessionCache, knowledge: &KnowledgeCache, +) -> Result, MemoryError> { + build_hdf5_file_with_mark(config, cache, sessions, knowledge, None) +} + +/// [`build_hdf5_file`], recording which WAL prefix this state already +/// contains (see [`WalMark`]) so a crash before the WAL is truncated doesn't +/// replay those entries a second time. +pub fn build_hdf5_file_with_mark( + config: &MemoryConfig, + cache: &MemoryCache, + sessions: &SessionCache, + knowledge: &KnowledgeCache, + wal_applied: Option, ) -> Result, MemoryError> { let mut builder = clawhdf5::FileBuilder::new(); @@ -34,10 +54,35 @@ pub fn build_hdf5_file( meta.set_attr("embedding_dim", AttrValue::I64(config.embedding_dim as i64)); meta.set_attr("chunk_size", AttrValue::I64(config.chunk_size as i64)); meta.set_attr("overlap", AttrValue::I64(config.overlap as i64)); + // Behavioural settings. These used to live only in memory, so reopening a + // store silently reset them to defaults — e.g. a compressed store was + // rewritten uncompressed by the first checkpoint after a reopen. Loaders + // treat each one as optional so older files keep opening. + meta.set_attr("float16", AttrValue::I64(config.float16.into())); + meta.set_attr("compression", AttrValue::I64(config.compression.into())); + meta.set_attr( + "compression_level", + AttrValue::I64(config.compression_level.into()), + ); + meta.set_attr( + "compact_threshold", + AttrValue::F64(config.compact_threshold.into()), + ); + meta.set_attr("hebbian_boost", AttrValue::F64(config.hebbian_boost.into())); + meta.set_attr("decay_factor", AttrValue::F64(config.decay_factor.into())); + meta.set_attr("wal_enabled", AttrValue::I64(config.wal_enabled.into())); + meta.set_attr( + "wal_max_entries", + AttrValue::I64(config.wal_max_entries as i64), + ); meta.set_attr( "edgehdf5_version", AttrValue::String(ZEROCLAW_VERSION.into()), ); + if let Some(mark) = wal_applied.filter(|m| m.len > 0) { + meta.set_attr(WAL_APPLIED_LEN_ATTR, AttrValue::I64(mark.len as i64)); + meta.set_attr(WAL_APPLIED_CRC_ATTR, AttrValue::I64(i64::from(mark.crc))); + } // Need at least one dataset in the group for it to be a proper group meta.create_dataset("_marker").with_u8_data(&[1]).compact(); let finished_meta = meta.finish(); @@ -83,15 +128,33 @@ fn build_memory_group( let rows_per_chunk = (target_chunk_bytes / (d * 4)).max(1).min(n); ds.with_chunks(&[rows_per_chunk, d]); - // Compression: Zstd for embeddings — faster than deflate at same ratio. - // Shuffle is applied automatically (auto-shuffle pre-filter). + // Compression. Shuffle is applied automatically (auto-shuffle + // pre-filter). Zstd is faster than deflate at the same ratio but + // pulls in libzstd, so it is opt-in via the `zstd` feature; the + // default build uses deflate, which is always available. (This + // used to call `with_zstd` unconditionally, so without the + // feature every checkpoint of a compressed store failed with + // "unsupported filter: 32015".) Both are standard HDF5 filters; + // reading a zstd-compressed store needs a zstd-enabled build. if config.compression { - let level = if config.compression_level > 0 { - config.compression_level.min(22) - } else { - 3 // Zstd level 3: fast + good ratio for f32 embeddings - }; - ds.with_zstd(level); + #[cfg(feature = "zstd")] + { + let level = if config.compression_level > 0 { + config.compression_level.min(22) + } else { + 3 // fast + good ratio for f32 embeddings + }; + ds.with_zstd(level); + } + #[cfg(not(feature = "zstd"))] + { + let level = if config.compression_level > 0 { + config.compression_level.min(9) + } else { + 4 + }; + ds.with_deflate(level); + } } } @@ -309,6 +372,20 @@ fn write_string_dataset( } /// Validate an HDF5 file has the correct schema and load all data. +/// Read the checkpoint's [`WalMark`] from `/meta`, if it has one. +pub fn read_wal_mark(file: &clawhdf5::File) -> Option { + let attrs = file.group("meta").ok()?.attrs().ok()?; + let len = match attrs.get(WAL_APPLIED_LEN_ATTR)? { + AttrValue::I64(v) => u64::try_from(*v).ok()?, + _ => return None, + }; + let crc = match attrs.get(WAL_APPLIED_CRC_ATTR)? { + AttrValue::I64(v) => u32::try_from(*v).ok()?, + _ => return None, + }; + Some(WalMark { len, crc }) +} + pub fn validate_and_load( file: &clawhdf5::File, ) -> Result<(MemoryConfig, MemoryCache, SessionCache, KnowledgeCache), MemoryError> { @@ -344,15 +421,19 @@ pub fn validate_and_load( embedding_dim, chunk_size, overlap, - float16: false, - compression: false, - compression_level: 0, - compact_threshold: 0.3, - hebbian_boost: 0.15, - decay_factor: 0.98, + float16: optional_bool_attr(&attrs, "float16", false), + compression: optional_bool_attr(&attrs, "compression", false), + compression_level: optional_i64_attr(&attrs, "compression_level") + .and_then(|v| u32::try_from(v).ok()) + .unwrap_or(0), + compact_threshold: optional_f32_attr(&attrs, "compact_threshold", 0.3), + hebbian_boost: optional_f32_attr(&attrs, "hebbian_boost", 0.15), + decay_factor: optional_f32_attr(&attrs, "decay_factor", 0.98), created_at, - wal_enabled: true, - wal_max_entries: 500, + wal_enabled: optional_bool_attr(&attrs, "wal_enabled", true), + wal_max_entries: optional_i64_attr(&attrs, "wal_max_entries") + .and_then(|v| usize::try_from(v).ok()) + .unwrap_or(500), }; // Load /memory group @@ -391,20 +472,46 @@ fn load_memory_group( let tags = read_string_dataset_from_group(&group, "tags")?; let tombstones = read_u8_dataset(&group, "tombstones")?; - // Read norms if present, otherwise compute from embeddings - let norms = match read_f32_dataset(&group, "norms") { - Ok(n) if n.len() == n.len() => n, - _ => { - // Compute norms from flat embeddings - flat_embeddings - .chunks(embedding_dim) - .map(|chunk| { - let sq_sum: f32 = chunk.iter().map(|x| x * x).sum(); - sq_sum.sqrt() - }) - .collect() + // Every per-record dataset must describe exactly `n` records. Without + // this, a truncated or hand-edited file loads "successfully" and then + // panics on the first out-of-bounds index during search/delete. + if embedding_dim == 0 { + return Err(MemoryError::Schema(format!( + "/memory has {n} records but embedding_dim is 0" + ))); + } + let expected_flat = n.checked_mul(embedding_dim).ok_or_else(|| { + MemoryError::Schema(format!("/memory size overflow: {n} x {embedding_dim}")) + })?; + let check_len = |name: &str, actual: usize, expected: usize| { + if actual == expected { + Ok(()) + } else { + Err(MemoryError::Schema(format!( + "/memory/{name} has {actual} entries, expected {expected} \ + ({n} records)" + ))) } }; + check_len("embeddings", flat_embeddings.len(), expected_flat)?; + check_len("source_channel", source_channels.len(), n)?; + check_len("timestamps", timestamps.len(), n)?; + check_len("session_ids", session_ids.len(), n)?; + check_len("tags", tags.len(), n)?; + check_len("tombstones", tombstones.len(), n)?; + + // Norms are derived data: use the stored ones only if they are present + // and the right length, otherwise recompute from the embeddings. + let norms = match read_f32_dataset(&group, "norms") { + Ok(stored) if stored.len() == n => stored, + _ => flat_embeddings + .chunks(embedding_dim) + .map(|chunk| { + let sq_sum: f32 = chunk.iter().map(|x| x * x).sum(); + sq_sum.sqrt() + }) + .collect(), + }; // Unflatten embeddings let embeddings: Vec> = flat_embeddings @@ -531,6 +638,27 @@ fn extract_string_attr( } } +type MetaAttrs = std::collections::HashMap; + +fn optional_i64_attr(attrs: &MetaAttrs, name: &str) -> Option { + match attrs.get(name) { + Some(AttrValue::I64(v)) => Some(*v), + _ => None, + } +} + +fn optional_bool_attr(attrs: &MetaAttrs, name: &str, default: bool) -> bool { + optional_i64_attr(attrs, name).map_or(default, |v| v != 0) +} + +/// Finite values only: a NaN threshold/decay would poison every comparison. +fn optional_f32_attr(attrs: &MetaAttrs, name: &str, default: f32) -> f32 { + match attrs.get(name) { + Some(AttrValue::F64(v)) if v.is_finite() => *v as f32, + _ => default, + } +} + fn extract_i64_attr( attrs: &std::collections::HashMap, name: &str, @@ -616,3 +744,108 @@ fn read_u8_dataset(group: &clawhdf5::Group<'_>, name: &str) -> Result, M .map_err(|e| MemoryError::Hdf5(format!("cannot read u8 from {name}: {e}")))?; Ok(data.into_iter().map(|v| v as u8).collect()) } + +#[cfg(test)] +mod tests { + use super::*; + + fn config() -> MemoryConfig { + MemoryConfig::new(std::path::PathBuf::from("unused.h5"), "agent", 4) + } + + fn cache_with(n: usize) -> MemoryCache { + let mut cache = MemoryCache::new(4); + for i in 0..n { + cache.push( + format!("chunk {i}"), + vec![i as f32 + 1.0, 0.0, 0.0, 0.0], + "user".into(), + i as f64, + "s".into(), + "t".into(), + ); + } + cache + } + + fn roundtrip(cache: &MemoryCache) -> Result { + let bytes = build_hdf5_file( + &config(), + cache, + &SessionCache::new(), + &KnowledgeCache::new(), + )?; + let file = + clawhdf5::File::from_bytes(bytes).map_err(|e| MemoryError::Hdf5(e.to_string()))?; + validate_and_load(&file).map(|(_, cache, _, _)| cache) + } + + #[test] + fn behavioural_config_survives_a_reopen() { + let mut cfg = config(); + cfg.compression = true; + cfg.compression_level = 7; + cfg.compact_threshold = 0.5; + cfg.hebbian_boost = 0.25; + cfg.decay_factor = 0.9; + cfg.wal_enabled = false; + cfg.wal_max_entries = 42; + let bytes = build_hdf5_file( + &cfg, + &cache_with(2), + &SessionCache::new(), + &KnowledgeCache::new(), + ) + .unwrap(); + let file = clawhdf5::File::from_bytes(bytes).unwrap(); + let (loaded, loaded_cache, ..) = validate_and_load(&file).unwrap(); + // The compressed embeddings must also read back intact. + assert_eq!(loaded_cache.embeddings, cache_with(2).embeddings); + assert!(loaded.compression); + assert_eq!(loaded.compression_level, 7); + assert_eq!(loaded.compact_threshold, 0.5); + assert_eq!(loaded.hebbian_boost, 0.25); + assert_eq!(loaded.decay_factor, 0.9); + assert!(!loaded.wal_enabled); + assert_eq!(loaded.wal_max_entries, 42); + } + + #[test] + fn consistent_store_loads() { + let loaded = roundtrip(&cache_with(3)).unwrap(); + assert_eq!(loaded.chunks.len(), 3); + assert_eq!(loaded.norms, vec![1.0, 2.0, 3.0]); + } + + #[test] + fn wrong_length_norms_are_recomputed_not_trusted() { + // Regression: the guard used to be `n.len() == n.len()`, so a norms + // dataset of any length was accepted and corrupted every cosine score. + let mut cache = cache_with(3); + cache.norms = vec![99.0]; + let loaded = roundtrip(&cache).unwrap(); + assert_eq!(loaded.norms, vec![1.0, 2.0, 3.0]); + } + + #[test] + fn mismatched_per_record_datasets_are_schema_errors() { + type Corrupt = fn(&mut MemoryCache); + let cases: [(&str, Corrupt); 5] = [ + ("tombstones", |c| c.tombstones.truncate(1)), + ("timestamps", |c| c.timestamps.truncate(1)), + ("tags", |c| c.tags.truncate(1)), + ("session_ids", |c| c.session_ids.truncate(1)), + ("source_channel", |c| c.source_channels.truncate(1)), + ]; + for (name, corrupt) in cases { + let mut cache = cache_with(3); + corrupt(&mut cache); + match roundtrip(&cache) { + Err(MemoryError::Schema(msg)) => { + assert!(msg.contains(name), "{name}: unexpected message {msg}") + } + other => panic!("{name}: expected Schema error, got {:?}", other.map(|_| ())), + } + } + } +} diff --git a/crates/clawhdf5-agent/src/search.rs b/crates/clawhdf5-agent/src/search.rs index d958f85..b66f2f2 100644 --- a/crates/clawhdf5-agent/src/search.rs +++ b/crates/clawhdf5-agent/src/search.rs @@ -113,13 +113,24 @@ impl HDF5Memory { } }) .collect(); + // Ties broken by index so results (and therefore which records get + // boosted) don't depend on HashMap iteration order upstream. results.sort_by(|a, b| { b.score .partial_cmp(&a.score) .unwrap_or(std::cmp::Ordering::Equal) + .then(a.index.cmp(&b.index)) }); - let hit_indices: Vec = results.iter().map(|r| r.index).collect(); + // Only reinforce records that actually matched. When fewer than `k` + // records are relevant, the rest of the list is zero-score filler; + // boosting it would teach the store that arbitrary records are + // important just because they were nearby in iteration order. + let hit_indices: Vec = results + .iter() + .filter(|r| r.score > 0.0) + .map(|r| r.index) + .collect(); self.apply_hebbian_boost(&hit_indices); self.flush().ok(); diff --git a/crates/clawhdf5-agent/src/storage.rs b/crates/clawhdf5-agent/src/storage.rs index 8d8c1b0..c9000a1 100644 --- a/crates/clawhdf5-agent/src/storage.rs +++ b/crates/clawhdf5-agent/src/storage.rs @@ -11,6 +11,7 @@ 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( @@ -20,7 +21,20 @@ pub fn write_to_disk( sessions: &SessionCache, knowledge: &KnowledgeCache, ) -> Result<(), MemoryError> { - let bytes = schema::build_hdf5_file(config, cache, sessions, knowledge)?; + 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, +) -> Result<(), MemoryError> { + let bytes = schema::build_hdf5_file_with_mark(config, cache, sessions, knowledge, wal_applied)?; if bytes.is_empty() { return Err(MemoryError::Hdf5("build_hdf5_file produced 0 bytes".into())); @@ -28,9 +42,41 @@ pub fn write_to_disk( // Write to a temp file first, then rename for atomicity let tmp_path = path.with_extension("h5.tmp"); - std::fs::write(&tmp_path, &bytes).map_err(MemoryError::Io)?; - std::fs::rename(&tmp_path, path).map_err(MemoryError::Io)?; + write_synced(&tmp_path, &bytes)?; + rename_synced(&tmp_path, path) +} +/// Write `bytes` to `path` and flush them to stable storage. +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. +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(()) } @@ -42,6 +88,15 @@ pub fn write_to_disk( 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), MemoryError> { let mmap = clawhdf5_io::MmapReader::open(path).map_err(MemoryError::Io)?; // Advise the OS we'll need the whole file for parsing @@ -53,8 +108,9 @@ pub fn read_from_disk( 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)) + Ok(((config, cache, sessions, knowledge), wal_applied)) } /// Copy an HDF5 file atomically to a destination. @@ -78,7 +134,10 @@ pub fn snapshot_file(src: &Path, dest: &Path) -> Result.h5.lock`, held for the lifetime of +/// the owning `HDF5Memory` and released when it is dropped (or when the +/// process dies — the OS drops the lock with the file descriptor, so a crash +/// never leaves a stale lock behind; the empty lock file itself is harmless). +#[derive(Debug)] +pub(crate) struct StoreLock { + _file: File, +} + +impl StoreLock { + pub(crate) fn lock_path(store: &Path) -> PathBuf { + store.with_extension("h5.lock") + } + + pub(crate) fn acquire(store: &Path) -> Result { + let path = Self::lock_path(store); + let file = OpenOptions::new() + .create(true) + .truncate(false) + .write(true) + .open(&path)?; + // A previous owner may be mid-teardown (e.g. an `AsyncHDF5Memory` + // dropped without `shutdown()`: its background task releases the + // store a moment later), so give the lock a short, bounded grace + // period before reporting a genuine second writer. + let mut attempts_left = LOCK_RETRIES; + loop { + match file.try_lock() { + Ok(()) => return Ok(Self { _file: file }), + Err(TryLockError::WouldBlock) if attempts_left > 0 => { + attempts_left -= 1; + std::thread::sleep(LOCK_RETRY_DELAY); + } + Err(TryLockError::WouldBlock) => { + return Err(MemoryError::Locked(format!( + "{} is already open in this or another process (lock file {})", + store.display(), + path.display() + ))); + } + Err(TryLockError::Error(e)) => return Err(MemoryError::Io(e)), + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn second_acquire_fails_until_first_is_dropped() { + let dir = tempfile::TempDir::new().unwrap(); + let store = dir.path().join("s.h5"); + let first = StoreLock::acquire(&store).unwrap(); + assert!(matches!( + StoreLock::acquire(&store), + Err(MemoryError::Locked(_)) + )); + drop(first); + StoreLock::acquire(&store).unwrap(); + } +} diff --git a/crates/clawhdf5-agent/src/wal.rs b/crates/clawhdf5-agent/src/wal.rs index 5f1c78b..43bcd2f 100644 --- a/crates/clawhdf5-agent/src/wal.rs +++ b/crates/clawhdf5-agent/src/wal.rs @@ -28,7 +28,18 @@ const WAL_HEADER_LEN: u64 = WAL_MAGIC.len() as u64 + 1 + 4; /// its target Save) — the moved/inserted entry's stored CRC was computed /// against a different predecessor than the one now in front of it on disk, /// so the chain breaks at that point and replay stops there. -const WAL_VERSION: u8 = 3; +const WAL_VERSION: u8 = 4; + +/// The chained-CRC format before [`WalEntryType::Update`] records existed. +/// Byte-for-byte the same framing as [`WAL_VERSION`], so it is read by the +/// same code, and `WalFile::open` upgrades it in place by rewriting the +/// header's version byte (the header is not covered by the CRC chain). +/// +/// The bump exists for *older binaries*: they don't know record type 0x04, +/// would treat it as a torn tail, and would truncate it — and everything +/// after it — away. An unknown header version makes them refuse the file +/// with a clear error instead. +const WAL_VERSION_CHAINED_NO_UPDATE: u8 = 3; /// The previous WAL format version: still a CRC32 per entry (so a bit-flip /// within one entry is caught), but not chained to the previous entry's CRC @@ -67,6 +78,10 @@ pub enum WalEntryType { Save = 0x01, Tombstone = 0x02, ActivationUpdate = 0x03, + /// Replace the record at `update_index` in place (`save_or_update` hit). + /// Logged as a plain `Save` before this existed, so replay appended a + /// duplicate instead of updating. + Update = 0x04, } impl WalEntryType { @@ -75,6 +90,7 @@ impl WalEntryType { 0x01 => Some(Self::Save), 0x02 => Some(Self::Tombstone), 0x03 => Some(Self::ActivationUpdate), + 0x04 => Some(Self::Update), _ => None, } } @@ -91,6 +107,8 @@ pub struct WalEntry { pub tags: String, /// For tombstone entries: the index of the entry to delete. pub tombstone_index: Option, + /// For update entries: the index of the record to replace. + pub update_index: Option, } /// How many entries to accumulate before updating the header entry_count. @@ -112,6 +130,62 @@ pub struct WalFile { /// Reset to 0 by `truncate()`/`create_fresh_wal_file`, and re-derived by /// scanning existing entries when `open()` attaches to a non-empty file. running_crc: u32, + /// Bytes of verified entries after the header (the length of the chain + /// `running_crc` covers). Together they form the [`WalMark`]. + chain_len: u64, +} + +/// What a WAL file's 9-byte header looks like, without reading any entries. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum WalHeaderStatus { + /// A version this build can read (current or legacy). + Readable, + /// Shorter than a header — e.g. a crash while the file was being created. + /// It cannot contain entries. + Torn, + /// Not a WAL file at all. + BadMagic, + /// Well-formed header from a version this build doesn't know — most + /// likely written by a *newer* build. Never discard this: the entries are + /// probably fine, this binary just can't read them. + UnknownVersion(u8), +} + +/// Classify the header of the WAL at `path`. +pub fn wal_header_status(path: &Path) -> std::io::Result { + let mut header = [0u8; WAL_HEADER_LEN as usize]; + let mut f = File::open(path)?; + let mut filled = 0; + while filled < header.len() { + match f.read(&mut header[filled..])? { + 0 => return Ok(WalHeaderStatus::Torn), + n => filled += n, + } + } + if header[0..4] != WAL_MAGIC { + return Ok(WalHeaderStatus::BadMagic); + } + Ok(match header[4] { + WAL_VERSION + | WAL_VERSION_CHAINED_NO_UPDATE + | WAL_VERSION_CRC_UNCHAINED + | WAL_VERSION_LEGACY_NO_CRC => WalHeaderStatus::Readable, + v => WalHeaderStatus::UnknownVersion(v), + }) +} + +/// A position in a WAL's CRC chain: `len` bytes of entries after the header, +/// whose chained CRC is `crc`. +/// +/// A checkpoint stores the mark of the WAL prefix it folded into the `.h5` +/// file. If the process dies after the new `.h5` is in place but before the +/// WAL is truncated, the next `open()` finds that exact prefix still in the +/// WAL and skips it instead of replaying it on top of data that already +/// contains it (which used to duplicate every pending entry). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct WalMark { + pub len: u64, + pub crc: u32, } impl WalFile { @@ -138,7 +212,15 @@ impl WalFile { let mut ver = [0u8; 1]; f.read_exact(&mut ver)?; match ver[0] { - WAL_VERSION => { + WAL_VERSION | WAL_VERSION_CHAINED_NO_UPDATE => { + if ver[0] == WAL_VERSION_CHAINED_NO_UPDATE { + // Same framing; stamp the current version so an older + // binary refuses this file rather than truncating an + // Update record it can't parse. See the constant. + f.seek(SeekFrom::Start(4))?; + f.write_all(&[WAL_VERSION])?; + f.seek(SeekFrom::Start(5))?; + } let mut count_buf = [0u8; 4]; f.read_exact(&mut count_buf)?; let header_count = u32::from_le_bytes(count_buf); @@ -147,7 +229,8 @@ impl WalFile { // be stale from deferred group-commit sync, same // tolerance `read_entries` already has, so the scanned // count is also the more accurate of the two). - let (entries, running_crc, verified_bytes) = read_chained_entries(&mut f, 0); + let (entries, running_crc, verified_bytes) = + read_chained_entries(&mut f, 0, None); let entry_count = if entries.is_empty() { header_count } else { @@ -190,6 +273,7 @@ impl WalFile { entry_count, pending_header_sync: 0, running_crc, + chain_len: verified_bytes, }) } WAL_VERSION_CRC_UNCHAINED | WAL_VERSION_LEGACY_NO_CRC => { @@ -201,6 +285,7 @@ impl WalFile { entry_count: 0, pending_header_sync: 0, running_crc: 0, + chain_len: 0, }) } v => Err(MemoryError::Schema(format!("unsupported WAL version {v}"))), @@ -213,6 +298,7 @@ impl WalFile { entry_count: 0, pending_header_sync: 0, running_crc: 0, + chain_len: 0, }) } } @@ -236,8 +322,20 @@ impl WalFile { 4 + entry.session_id.len() + 4 + entry.tags.len(), ); - buf.push(WalEntryType::Save as u8); - buf.extend_from_slice(&entry.timestamp.to_le_bytes()); + match entry.update_index { + Some(index) => { + let index = u32::try_from(index).map_err(|_| { + MemoryError::Schema(format!("WAL update index {index} exceeds u32")) + })?; + buf.push(WalEntryType::Update as u8); + buf.extend_from_slice(&entry.timestamp.to_le_bytes()); + buf.extend_from_slice(&index.to_le_bytes()); + } + None => { + buf.push(WalEntryType::Save as u8); + buf.extend_from_slice(&entry.timestamp.to_le_bytes()); + } + } serialize_str(&mut buf, &entry.chunk); buf.extend_from_slice(&(emb_len as u32).to_le_bytes()); for &val in &entry.embedding { @@ -258,6 +356,7 @@ impl WalFile { .as_mut() .ok_or_else(|| MemoryError::Io(std::io::Error::other("WAL file not open")))?; f.write_all(&buf)?; + self.chain_len += buf.len() as u64; self.running_crc = crc; self.entry_count += 1; @@ -282,6 +381,7 @@ impl WalFile { .as_mut() .ok_or_else(|| MemoryError::Io(std::io::Error::other("WAL file not open")))?; f.write_all(&buf)?; + self.chain_len += buf.len() as u64; self.running_crc = crc; self.entry_count += 1; @@ -311,7 +411,7 @@ impl WalFile { /// legacy-no-CRC file returns a typed error instead of silently /// downgrading to the unverified parser. pub fn read_entries(path: &Path) -> Result, MemoryError> { - Self::read_entries_impl(path, false) + Self::read_entries_impl(path, false, None) } /// Like [`WalFile::read_entries`], but also accepts @@ -320,13 +420,23 @@ impl WalFile { /// legitimate caller is `HDF5Memory::open`'s one-time migration of a /// pre-CRC WAL file, which immediately recreates it in the current /// format afterward. Do not use this for anything else. - pub(crate) fn read_entries_for_migration(path: &Path) -> Result, MemoryError> { - Self::read_entries_impl(path, true) + /// + /// `applied` is the checkpoint mark read from the `.h5` file, if any: if + /// the WAL's chain passes through it (same byte length, same chained + /// CRC), everything up to that point is already in the `.h5` and is + /// dropped. If it never does — the normal case, because the WAL was + /// truncated after the checkpoint — every entry is returned. + pub(crate) fn read_entries_for_migration( + path: &Path, + applied: Option, + ) -> Result, MemoryError> { + Self::read_entries_impl(path, true, applied) } fn read_entries_impl( path: &Path, allow_legacy_no_crc: bool, + applied: Option, ) -> Result, MemoryError> { if !path.exists() { return Ok(Vec::new()); @@ -342,8 +452,9 @@ impl WalFile { let entry_count_hint = u32::from_le_bytes([header[5], header[6], header[7], header[8]]); match header[4] { - WAL_VERSION => { - let (entries, _final_crc, _verified_bytes) = read_chained_entries(&mut f, 0); + WAL_VERSION | WAL_VERSION_CHAINED_NO_UPDATE => { + let (entries, _final_crc, _verified_bytes) = + read_chained_entries(&mut f, 0, applied); Ok(entries) } WAL_VERSION_CRC_UNCHAINED => { @@ -406,9 +517,19 @@ impl WalFile { self.entry_count = 0; self.pending_header_sync = 0; self.running_crc = 0; + self.chain_len = 0; Ok(()) } + /// The mark covering every entry currently in this WAL. Store it with a + /// checkpoint taken from the state those entries produced. + pub fn mark(&self) -> WalMark { + WalMark { + len: self.chain_len, + crc: self.running_crc, + } + } + /// Number of pending entries. pub fn pending_count(&self) -> u32 { self.entry_count @@ -448,6 +569,28 @@ pub fn replay_into_cache(entries: &[WalEntry], cache: &mut crate::cache::MemoryC entry.tags.clone(), ); } + WalEntryType::Update => match entry.update_index { + // The index was valid when the record was written; if the + // store no longer has it, keep the data rather than drop it. + Some(idx) if idx < cache.len() => cache.update( + idx, + entry.chunk.clone(), + entry.embedding.clone(), + entry.source_channel.clone(), + entry.timestamp, + entry.session_id.clone(), + ), + _ => { + cache.push( + entry.chunk.clone(), + entry.embedding.clone(), + entry.source_channel.clone(), + entry.timestamp, + entry.session_id.clone(), + entry.tags.clone(), + ); + } + }, WalEntryType::Tombstone => { if let Some(idx) = entry.tombstone_index { cache.mark_deleted(idx); @@ -524,7 +667,16 @@ fn chained_crc(entry_bytes: &[u8], prev_crc: u32) -> u32 { /// The byte count is what lets `open()` position an append at the end of the /// VERIFIED prefix rather than at end-of-file. Appending past a torn tail /// writes entries that replay can never reach — see `open`. -fn read_chained_entries(f: &mut R, start_crc: u32) -> (Vec, u32, u64) { +/// +/// `applied`, when given, is a checkpoint mark: once the chain reaches exactly +/// that position, the entries collected so far are discarded (they are +/// already in the `.h5` file). A zero-length mark matches nothing. +fn read_chained_entries( + f: &mut R, + start_crc: u32, + applied: Option, +) -> (Vec, u32, u64) { + let applied = applied.filter(|m| m.len > 0); let mut entries = Vec::new(); let mut running_crc = start_crc; let mut verified_bytes: u64 = 0; @@ -554,6 +706,14 @@ fn read_chained_entries(f: &mut R, start_crc: u32) -> (Vec, u if let Some(entry) = entry_opt { entries.push(entry); } + if applied + == Some(WalMark { + len: verified_bytes, + crc: running_crc, + }) + { + entries.clear(); + } } (entries, running_crc, verified_bytes) } @@ -615,7 +775,14 @@ fn read_one_entry(r: &mut R) -> Result, ()> { let timestamp = f64::from_le_bytes(ts_buf); match entry_type { - WalEntryType::Save => { + WalEntryType::Save | WalEntryType::Update => { + let update_index = if entry_type == WalEntryType::Update { + let mut idx_buf = [0u8; 4]; + r.read_exact(&mut idx_buf).map_err(|_| ())?; + Some(u32::from_le_bytes(idx_buf) as usize) + } else { + None + }; let chunk = read_len_prefixed_str(r).map_err(|_| ())?; let embedding = read_embedding(r).map_err(|_| ())?; let source_channel = read_len_prefixed_str(r).map_err(|_| ())?; @@ -630,6 +797,7 @@ fn read_one_entry(r: &mut R) -> Result, ()> { session_id, tags, tombstone_index: None, + update_index, })) } WalEntryType::Tombstone => { @@ -645,6 +813,7 @@ fn read_one_entry(r: &mut R) -> Result, ()> { session_id: String::new(), tags: String::new(), tombstone_index: Some(idx), + update_index: None, })) } WalEntryType::ActivationUpdate => Ok(None), @@ -668,6 +837,7 @@ mod tests { session_id: "sess-001".to_string(), tags: "tag1,tag2".to_string(), tombstone_index: None, + update_index: None, } } @@ -798,6 +968,7 @@ mod tests { session_id: "sess-öö-123".to_string(), tags: "α,β,γ".to_string(), tombstone_index: None, + update_index: None, }; wal.append_save(&entry).unwrap(); } @@ -932,6 +1103,148 @@ mod tests { assert!(entries.is_empty()); } + /// Reopen `path` and return the stored chunks in order. + fn reopen_chunks(path: &std::path::Path) -> Vec { + let mem = HDF5Memory::open(path).unwrap(); + mem.cache.chunks.clone() + } + + #[test] + fn crash_between_checkpoint_and_wal_truncate_does_not_duplicate() { + // flush() writes the new .h5 and only then truncates the WAL. Dying in + // between leaves BOTH a .h5 that contains the pending entries and a + // WAL that still lists them; replaying blindly used to double them. + let dir = TempDir::new().unwrap(); + let config = make_config(&dir); + let h5_path = config.path.clone(); + let wal_path = h5_path.with_extension("h5.wal"); + let stale_wal = dir.path().join("stale.wal"); + + { + let mut mem = HDF5Memory::create(config).unwrap(); + for name in ["a", "b", "c"] { + mem.save(make_entry(name, &[1.0, 0.0, 0.0, 0.0])).unwrap(); + } + assert_eq!(mem.wal_pending_count(), 3); + std::fs::copy(&wal_path, &stale_wal).unwrap(); + mem.flush_wal().unwrap(); + } + // Undo the truncate: this is the on-disk state right after the crash. + std::fs::copy(&stale_wal, &wal_path).unwrap(); + assert_eq!(WalFile::read_entries(&wal_path).unwrap().len(), 3); + + assert_eq!(reopen_chunks(&h5_path), ["a", "b", "c"]); + + // Entries appended to that same WAL after recovery are still replayed. + { + let mut mem = HDF5Memory::open(&h5_path).unwrap(); + mem.save(make_entry("d", &[0.0, 1.0, 0.0, 0.0])).unwrap(); + } + assert_eq!(reopen_chunks(&h5_path), ["a", "b", "c", "d"]); + } + + #[test] + fn entries_written_after_a_completed_checkpoint_are_all_replayed() { + // Normal case: the checkpoint's mark refers to a WAL that has since + // been truncated, so it must not suppress anything in the new one — + // including when the new WAL grows past the old mark's length. + let dir = TempDir::new().unwrap(); + let config = make_config(&dir); + let h5_path = config.path.clone(); + { + let mut mem = HDF5Memory::create(config).unwrap(); + mem.save(make_entry("a", &[1.0, 0.0, 0.0, 0.0])).unwrap(); + mem.flush_wal().unwrap(); + for name in ["b", "c", "d"] { + mem.save(make_entry(name, &[1.0, 0.0, 0.0, 0.0])).unwrap(); + } + } + assert_eq!(reopen_chunks(&h5_path), ["a", "b", "c", "d"]); + } + + #[test] + fn save_or_update_replays_as_update_not_duplicate() { + let dir = TempDir::new().unwrap(); + let config = make_config(&dir); + let h5_path = config.path.clone(); + { + let mut mem = HDF5Memory::create(config).unwrap(); + let mut first = make_entry("v1", &[1.0, 0.0, 0.0, 0.0]); + first.tags = "key".into(); + let mut second = make_entry("v2", &[0.0, 1.0, 0.0, 0.0]); + second.tags = "key".into(); + let a = mem.save_or_update(first).unwrap(); + mem.save(make_entry("other", &[0.0, 0.0, 1.0, 0.0])) + .unwrap(); + let b = mem.save_or_update(second).unwrap(); + assert_eq!(a, b); + assert_eq!(mem.cache.chunks, ["v2", "other"]); + // Dropped without a checkpoint: all three records live in the WAL. + } + let mem = HDF5Memory::open(&h5_path).unwrap(); + assert_eq!(mem.cache.chunks, ["v2", "other"]); + assert_eq!(mem.cache.embeddings[0], [0.0, 1.0, 0.0, 0.0]); + } + + #[test] + fn v3_wal_is_read_and_upgraded_in_place() { + let dir = TempDir::new().unwrap(); + let wal_path = dir.path().join("old.wal"); + { + let mut wal = WalFile::open(&wal_path).unwrap(); + wal.append_save(&make_wal_entry("kept", &[1.0])).unwrap(); + } + // Rewrite the header as the pre-Update chained format. + let mut bytes = std::fs::read(&wal_path).unwrap(); + bytes[4] = WAL_VERSION_CHAINED_NO_UPDATE; + std::fs::write(&wal_path, &bytes).unwrap(); + + assert_eq!(WalFile::read_entries(&wal_path).unwrap().len(), 1); + { + let mut wal = WalFile::open(&wal_path).unwrap(); + assert_eq!(wal.pending_count(), 1); + wal.append_save(&make_wal_entry("new", &[2.0])).unwrap(); + } + assert_eq!(std::fs::read(&wal_path).unwrap()[4], WAL_VERSION); + let chunks: Vec<_> = WalFile::read_entries(&wal_path) + .unwrap() + .into_iter() + .map(|e| e.chunk) + .collect(); + assert_eq!(chunks, ["kept", "new"]); + } + + #[test] + fn mark_matching_is_exact() { + let dir = TempDir::new().unwrap(); + let wal_path = dir.path().join("m.wal"); + let mut wal = WalFile::open(&wal_path).unwrap(); + wal.append_save(&make_wal_entry("one", &[1.0])).unwrap(); + let after_one = wal.mark(); + wal.append_save(&make_wal_entry("two", &[2.0])).unwrap(); + let after_two = wal.mark(); + drop(wal); + + let read = |m| { + WalFile::read_entries_for_migration(&wal_path, m) + .unwrap() + .into_iter() + .map(|e| e.chunk) + .collect::>() + }; + assert_eq!(read(None), ["one", "two"]); + assert_eq!(read(Some(after_one)), ["two"]); + assert!(read(Some(after_two)).is_empty()); + // Right length, wrong CRC (a different WAL generation): skip nothing. + let foreign = WalMark { + crc: after_one.crc ^ 1, + ..after_one + }; + assert_eq!(read(Some(foreign)), ["one", "two"]); + // Reopening resumes the same mark. + assert_eq!(WalFile::open(&wal_path).unwrap().mark(), after_two); + } + #[test] fn test_wal_replay_on_open() { // Test WAL replay using read_entries + replay_into_cache directly, @@ -1269,7 +1582,7 @@ mod tests { std::fs::write(&wal_path, build_legacy_v1_wal_bytes()).unwrap(); // Only the migration-only reader may read a legacy no-CRC file. - let entries = WalFile::read_entries_for_migration(&wal_path).unwrap(); + let entries = WalFile::read_entries_for_migration(&wal_path, None).unwrap(); assert_eq!(entries.len(), 1); assert_eq!(entries[0].chunk, "legacy-chunk"); assert_eq!(entries[0].embedding, vec![1.0, 2.0]); diff --git a/crates/clawhdf5-agent/tests/e2e_tests.rs b/crates/clawhdf5-agent/tests/e2e_tests.rs index d151ece..3321d8b 100644 --- a/crates/clawhdf5-agent/tests/e2e_tests.rs +++ b/crates/clawhdf5-agent/tests/e2e_tests.rs @@ -196,7 +196,7 @@ fn test_migration_round_trip() { mem.add_relation(e1, e2, "discusses", 0.8).unwrap(); // Verify all data transferred by reopening - let reopened = HDF5Memory::open(&path).unwrap(); + let reopened = HDF5Memory::open_read_only(&path).unwrap(); assert_eq!(reopened.count(), 500); // Verify sessions @@ -266,7 +266,7 @@ fn test_knowledge_graph_workflow() { assert_eq!(entity.entity_type, "library"); // Persistence - let reopened = HDF5Memory::open(&path).unwrap(); + let reopened = HDF5Memory::open_read_only(&path).unwrap(); assert_eq!(reopened.knowledge().entities.len(), 4); assert_eq!(reopened.knowledge().relations.len(), 4); @@ -316,7 +316,7 @@ fn test_multi_session_workflow() { assert_eq!(mem.count(), 100); // 5 sessions * 20 entries // Reopen and verify sessions - let reopened = HDF5Memory::open(&path).unwrap(); + let reopened = HDF5Memory::open_read_only(&path).unwrap(); for sess in 0..5 { let summary = reopened .get_session_summary(&format!("sess_{sess}")) @@ -460,7 +460,7 @@ fn test_snapshot_and_continue() { assert_eq!(snap_mem.count(), 50); // Original should have 100 - let orig_mem = HDF5Memory::open(&path).unwrap(); + let orig_mem = HDF5Memory::open_read_only(&path).unwrap(); assert_eq!(orig_mem.count(), 100); } @@ -483,7 +483,7 @@ fn test_config_persistence_across_ops() { mem.add_session("s1", 0, 0, "ch", "summary").unwrap(); mem.add_entity("Entity", "type", -1).unwrap(); - let reopened = HDF5Memory::open(&path).unwrap(); + let reopened = HDF5Memory::open_read_only(&path).unwrap(); assert_eq!(reopened.config().embedding_dim, 128); assert_eq!(reopened.config().embedder, "custom:my-embedder-v2"); assert_eq!(reopened.config().chunk_size, 2048); @@ -695,7 +695,7 @@ fn test_large_text_chunks() { mem.save_batch(entries).unwrap(); // Reopen and verify - let reopened = HDF5Memory::open(&path).unwrap(); + let reopened = HDF5Memory::open_read_only(&path).unwrap(); assert_eq!(reopened.count(), 10); let (_, cache, _, _) = read_cache(&path); @@ -752,7 +752,7 @@ fn test_interleaved_sessions_entries() { mem.flush_wal().unwrap(); // Verify - let reopened = HDF5Memory::open(&path).unwrap(); + let reopened = HDF5Memory::open_read_only(&path).unwrap(); assert_eq!(reopened.count(), 6); assert_eq!( reopened.get_session_summary("s1").unwrap().as_deref(), @@ -806,7 +806,7 @@ fn test_knowledge_graph_with_embeddings() { mem.add_relation(e_python, e_hdf5, "reads", 0.9).unwrap(); // Verify entity-embedding linkage persists - let reopened = HDF5Memory::open(&path).unwrap(); + let reopened = HDF5Memory::open_read_only(&path).unwrap(); let rust_entity = reopened.knowledge().get_entity(e_rust).unwrap(); assert_eq!(rust_entity.embedding_idx, idx0 as i64); diff --git a/crates/clawhdf5-agent/tests/stress_tests.rs b/crates/clawhdf5-agent/tests/stress_tests.rs index eb09d84..572da6f 100644 --- a/crates/clawhdf5-agent/tests/stress_tests.rs +++ b/crates/clawhdf5-agent/tests/stress_tests.rs @@ -105,7 +105,7 @@ fn test_heavy_tombstoning() { assert_eq!(mem.count_active(), 5000); // Verify persistence - let reopened = HDF5Memory::open(&path).unwrap(); + let reopened = HDF5Memory::open_read_only(&path).unwrap(); assert_eq!(reopened.count(), 5000); } @@ -163,7 +163,7 @@ fn test_large_embeddings_1536() { assert_eq!(mem.count(), 10_000); // Verify persistence - let reopened = HDF5Memory::open(&path).unwrap(); + let reopened = HDF5Memory::open_read_only(&path).unwrap(); assert_eq!(reopened.count(), 10_000); // Verify search works on large dims @@ -545,7 +545,7 @@ fn test_delete_all_entries() { assert_eq!(mem.count(), 0); // Verify persistence - let reopened = HDF5Memory::open(&path).unwrap(); + let reopened = HDF5Memory::open_read_only(&path).unwrap(); assert_eq!(reopened.count(), 0); } @@ -639,7 +639,7 @@ fn test_unicode_content() { ]; mem.save_batch(entries).unwrap(); - let reopened = HDF5Memory::open(&path).unwrap(); + let reopened = HDF5Memory::open_read_only(&path).unwrap(); assert_eq!(reopened.count(), 3); let (_, cache, _, _) = clawhdf5_agent::storage::read_from_disk(&path).unwrap(); @@ -685,6 +685,6 @@ fn test_rapid_save_delete_cycles() { assert_eq!(removed, 250); assert_eq!(mem.count(), 250); - let reopened = HDF5Memory::open(&path).unwrap(); + let reopened = HDF5Memory::open_read_only(&path).unwrap(); assert_eq!(reopened.count(), 250); } diff --git a/crates/clawhdf5-cli/src/main.rs b/crates/clawhdf5-cli/src/main.rs index e6053b9..03076a2 100644 --- a/crates/clawhdf5-cli/src/main.rs +++ b/crates/clawhdf5-cli/src/main.rs @@ -146,7 +146,7 @@ fn run(cli: Cli) -> Result<(), Box> { } Commands::Recall { index } => { - let mem = HDF5Memory::open(&cli.path)?; + let mem = HDF5Memory::open_read_only(&cli.path)?; match mem.get_chunk(index) { Some(content) => { let j = serde_json::json!({ "index": index, "chunk": content }); @@ -160,7 +160,7 @@ fn run(cli: Cli) -> Result<(), Box> { } Commands::Stats => { - let mem = HDF5Memory::open(&cli.path)?; + let mem = HDF5Memory::open_read_only(&cli.path)?; let cfg = mem.config(); let j = serde_json::json!({ "path": cli.path.display().to_string(), @@ -187,7 +187,7 @@ fn run(cli: Cli) -> Result<(), Box> { } Commands::AgentsMd { output } => { - let mem = HDF5Memory::open(&cli.path)?; + let mem = HDF5Memory::open_read_only(&cli.path)?; let md = mem.generate_agents_md(); match output { Some(p) => { @@ -199,7 +199,7 @@ fn run(cli: Cli) -> Result<(), Box> { } Commands::Export => { - let mem = HDF5Memory::open(&cli.path)?; + let mem = HDF5Memory::open_read_only(&cli.path)?; for i in 0..mem.count() { if let Some(chunk) = mem.get_chunk(i) { let j = serde_json::json!({ "index": i, "chunk": chunk }); diff --git a/crates/clawhdf5-format/src/chunked_read.rs b/crates/clawhdf5-format/src/chunked_read.rs index 43184a3..edb1705 100644 --- a/crates/clawhdf5-format/src/chunked_read.rs +++ b/crates/clawhdf5-format/src/chunked_read.rs @@ -132,6 +132,47 @@ fn ensure_len(data: &[u8], offset: usize, needed: usize) -> Result<(), FormatErr Ok(()) } +/// `elements * elem_size` for sizes that come from the file. Dataspace and +/// chunk dimensions are untrusted 64-bit fields, so a crafted file can make +/// the plain product wrap to a small number (or to something enormous). +pub(crate) fn checked_byte_len(elements: u64, elem_size: usize) -> Result { + usize::try_from(elements) + .ok() + .and_then(|n| n.checked_mul(elem_size)) + .ok_or_else(|| { + FormatError::Overflow(format!( + "{elements} elements of {elem_size} bytes exceeds the addressable size" + )) + }) +} + +/// Product of chunk dimensions times the element size, overflow-checked. +pub(crate) fn checked_chunk_byte_len( + chunk_dims: &[usize], + elem_size: usize, +) -> Result { + chunk_dims + .iter() + .try_fold(elem_size, |acc, &d| acc.checked_mul(d)) + .ok_or_else(|| { + FormatError::Overflow(format!( + "chunk dimensions {chunk_dims:?} x {elem_size} bytes exceeds the addressable size" + )) + }) +} + +/// A zero-filled output buffer of `len` bytes. `vec![0; len]` aborts the +/// process when the allocation fails; a size taken from the file must surface +/// as an error instead. +pub(crate) fn alloc_output(len: usize) -> Result, FormatError> { + let mut out = Vec::new(); + out.try_reserve_exact(len).map_err(|_| { + FormatError::Overflow(format!("cannot allocate {len} bytes for dataset output")) + })?; + out.resize(len, 0); + Ok(out) +} + fn read_offset(data: &[u8], pos: usize, size: u8) -> Result { let s = size as usize; if pos.checked_add(s).is_none_or(|end| end > data.len()) { @@ -393,7 +434,7 @@ pub fn read_chunked_data( } (4, Some(1)) => { // Single chunk — one chunk covering the entire dataset - let chunk_byte_size: usize = chunk_dims.iter().product::() * elem_size; + let chunk_byte_size = checked_chunk_byte_len(&chunk_dims, elem_size)?; let (csize, fmask) = if let Some(fs) = single_filtered_size { (fs as u32, single_filter_mask.unwrap_or(0)) } else { @@ -454,9 +495,13 @@ pub fn read_chunked_data( }; // Assemble output - let total_elements = dataspace.num_elements() as usize; - let total_bytes = total_elements * elem_size; - let mut output = vec![0u8; total_bytes]; + let total_bytes = checked_byte_len(dataspace.checked_num_elements()?, elem_size)?; + if total_bytes == 0 { + // Also keeps the stride products below in range: with a zero-sized + // dimension the total is 0 even if other dimensions are huge. + return Ok(Vec::new()); + } + let mut output = alloc_output(total_bytes)?; let mut ds_strides = vec![1usize; rank]; for i in (0..rank.saturating_sub(1)).rev() { @@ -468,8 +513,7 @@ pub fn read_chunked_data( chunk_strides[i] = chunk_strides[i + 1] * chunk_dims[i + 1]; } - let chunk_total_elements: usize = chunk_dims.iter().product(); - let chunk_total_bytes = chunk_total_elements * elem_size; + let chunk_total_bytes = checked_chunk_byte_len(&chunk_dims, elem_size)?; // Fast path: no filters — copy directly from file_data without intermediate alloc if pipeline.is_none() { @@ -623,7 +667,7 @@ pub fn read_chunked_data_cached( let chunks = match (version, chunk_index_type) { (3, _) => collect_chunk_info(file_data, addr, ndims, offset_size, length_size)?, (4, Some(1)) => { - let chunk_byte_size: usize = chunk_dims.iter().product::() * elem_size; + let chunk_byte_size = checked_chunk_byte_len(&chunk_dims, elem_size)?; let (csize, fmask) = if let Some(fs) = single_filtered_size { (fs as u32, single_filter_mask.unwrap_or(0)) } else { @@ -689,9 +733,13 @@ pub fn read_chunked_data_cached( let chunks = cache.all_indexed_chunks().unwrap_or_default(); // Assemble output - let total_elements = dataspace.num_elements() as usize; - let total_bytes = total_elements * elem_size; - let mut output = vec![0u8; total_bytes]; + let total_bytes = checked_byte_len(dataspace.checked_num_elements()?, elem_size)?; + if total_bytes == 0 { + // Also keeps the stride products below in range: with a zero-sized + // dimension the total is 0 even if other dimensions are huge. + return Ok(Vec::new()); + } + let mut output = alloc_output(total_bytes)?; let mut ds_strides = vec![1usize; rank]; for i in (0..rank.saturating_sub(1)).rev() { @@ -703,8 +751,7 @@ pub fn read_chunked_data_cached( chunk_strides[i] = chunk_strides[i + 1] * chunk_dims[i + 1]; } - let chunk_total_elements: usize = chunk_dims.iter().product(); - let chunk_total_bytes = chunk_total_elements * elem_size; + let chunk_total_bytes = checked_chunk_byte_len(&chunk_dims, elem_size)?; for chunk_info in &chunks { let coord: Vec = chunk_info.offsets.iter().take(rank).copied().collect(); @@ -976,7 +1023,7 @@ pub fn read_chunked_data_sweep( let chunks = match (version, chunk_index_type) { (3, _) => collect_chunk_info(file_data, addr, ndims, offset_size, length_size)?, (4, Some(1)) => { - let chunk_byte_size: usize = chunk_dims.iter().product::() * elem_size; + let chunk_byte_size = checked_chunk_byte_len(&chunk_dims, elem_size)?; let (csize, fmask) = if let Some(fs) = single_filtered_size { (fs as u32, single_filter_mask.unwrap_or(0)) } else { @@ -1042,9 +1089,13 @@ pub fn read_chunked_data_sweep( let chunks = cache.all_indexed_chunks().unwrap_or_default(); // Assemble output - let total_elements = dataspace.num_elements() as usize; - let total_bytes = total_elements * elem_size; - let mut output = vec![0u8; total_bytes]; + let total_bytes = checked_byte_len(dataspace.checked_num_elements()?, elem_size)?; + if total_bytes == 0 { + // Also keeps the stride products below in range: with a zero-sized + // dimension the total is 0 even if other dimensions are huge. + return Ok(Vec::new()); + } + let mut output = alloc_output(total_bytes)?; let mut ds_strides = vec![1usize; rank]; for i in (0..rank.saturating_sub(1)).rev() { @@ -1056,8 +1107,7 @@ pub fn read_chunked_data_sweep( chunk_strides[i] = chunk_strides[i + 1] * chunk_dims[i + 1]; } - let chunk_total_elements: usize = chunk_dims.iter().product(); - let chunk_total_bytes = chunk_total_elements * elem_size; + let chunk_total_bytes = checked_chunk_byte_len(&chunk_dims, elem_size)?; for chunk_info in &chunks { let coord: Vec = chunk_info.offsets.iter().take(rank).copied().collect(); @@ -1199,7 +1249,7 @@ pub fn read_chunked_data_indexed( let chunks = match (version, chunk_index_type) { (3, _) => collect_chunk_info(file_data, addr, ndims, offset_size, length_size)?, (4, Some(1)) => { - let chunk_byte_size: usize = chunk_dims.iter().product::() * elem_size; + let chunk_byte_size = checked_chunk_byte_len(&chunk_dims, elem_size)?; let (csize, fmask) = if let Some(fs) = single_filtered_size { (fs as u32, single_filter_mask.unwrap_or(0)) } else { @@ -1463,6 +1513,64 @@ fn copy_chunk_to_output( mod tests { use super::*; + fn simple_space(dimensions: Vec) -> Dataspace { + Dataspace { + space_type: crate::dataspace::DataspaceType::Simple, + rank: dimensions.len() as u8, + dimensions, + max_dimensions: None, + } + } + + #[test] + fn crafted_dimensions_are_errors_not_wraparound() { + // 2^63 * 2 wraps to 0 with a plain product; 2^40 * 2^40 wraps too. + for dims in [ + vec![1u64 << 63, 2], + vec![1 << 40, 1 << 40], + vec![u64::MAX, u64::MAX], + ] { + let space = simple_space(dims.clone()); + assert!( + matches!(space.checked_num_elements(), Err(FormatError::Overflow(_))), + "{dims:?}" + ); + // The infallible accessor saturates instead of wrapping. + assert_eq!(space.num_elements(), u64::MAX, "{dims:?}"); + } + assert_eq!(simple_space(vec![3, 4]).checked_num_elements().unwrap(), 12); + // A zero-sized dimension makes the whole product 0, not an overflow. + assert_eq!( + simple_space(vec![0, 1 << 40, 1 << 40]) + .checked_num_elements() + .unwrap(), + 0 + ); + } + + #[test] + fn byte_length_helpers_check_overflow() { + assert_eq!(checked_byte_len(10, 8).unwrap(), 80); + assert!(matches!( + checked_byte_len(u64::MAX, 8), + Err(FormatError::Overflow(_)) + )); + assert_eq!(checked_chunk_byte_len(&[10, 10], 4).unwrap(), 400); + assert!(matches!( + checked_chunk_byte_len(&[usize::MAX, 2], 4), + Err(FormatError::Overflow(_)) + )); + } + + #[test] + fn unallocatable_output_is_an_error_not_an_abort() { + assert_eq!(alloc_output(16).unwrap(), vec![0u8; 16]); + assert!(matches!( + alloc_output(usize::MAX / 2), + Err(FormatError::Overflow(_)) + )); + } + fn write_offset(buf: &mut Vec, val: u64, size: u8) { match size { 4 => buf.extend_from_slice(&(val as u32).to_le_bytes()), diff --git a/crates/clawhdf5-format/src/data_read.rs b/crates/clawhdf5-format/src/data_read.rs index a469084..9f33175 100644 --- a/crates/clawhdf5-format/src/data_read.rs +++ b/crates/clawhdf5-format/src/data_read.rs @@ -475,8 +475,10 @@ fn read_virtual_data( use crate::selection::Selection; let elem_size = datatype.type_size() as usize; - let total_elems = dataspace.num_elements() as usize; - let mut out = vec![0u8; total_elems.saturating_mul(elem_size)]; + let mut out = crate::chunked_read::alloc_output(crate::chunked_read::checked_byte_len( + dataspace.checked_num_elements()?, + elem_size, + )?)?; let virtual_dims = &dataspace.dimensions; @@ -616,12 +618,14 @@ fn extract_selection_from_buffer( block, } => { let rank = dims.len(); - let output_elements: usize = count + let output_elements = count .iter() .zip(block.iter()) - .map(|(&c, &b)| (c * b) as usize) - .product(); - let mut output = vec![0u8; output_elements * elem_size]; + .try_fold(1u64, |acc, (&c, &b)| acc.checked_mul(c.checked_mul(b)?)) + .ok_or_else(|| FormatError::Overflow("hyperslab count x block overflows".into()))?; + let mut output = crate::chunked_read::alloc_output( + crate::chunked_read::checked_byte_len(output_elements, elem_size)?, + )?; // Compute dataset strides (row-major) let mut ds_strides = vec![1usize; rank]; diff --git a/crates/clawhdf5-format/src/dataspace.rs b/crates/clawhdf5-format/src/dataspace.rs index efece6d..a08ca9d 100644 --- a/crates/clawhdf5-format/src/dataspace.rs +++ b/crates/clawhdf5-format/src/dataspace.rs @@ -1,5 +1,7 @@ //! HDF5 Dataspace message parsing (message type 0x0001). +#[cfg(not(feature = "std"))] +use alloc::format; #[cfg(not(feature = "std"))] use alloc::vec::Vec; @@ -167,6 +169,27 @@ impl Dataspace { } } + /// [`Dataspace::num_elements`] with the product overflow-checked. The + /// dimensions are untrusted 64-bit fields; read paths that size a buffer + /// from them must use this one. + pub fn checked_num_elements(&self) -> Result { + match self.space_type { + DataspaceType::Null => Ok(0), + DataspaceType::Scalar => Ok(1), + DataspaceType::Simple if self.dimensions.is_empty() => Ok(0), + DataspaceType::Simple => self + .dimensions + .iter() + .try_fold(1u64, |acc, &d| acc.checked_mul(d)) + .ok_or_else(|| { + FormatError::Overflow(format!( + "dataspace dimensions {:?} overflow the element count", + self.dimensions + )) + }), + } + } + /// Total number of elements. Scalar = 1, Null = 0. pub fn num_elements(&self) -> u64 { match self.space_type { @@ -176,7 +199,12 @@ impl Dataspace { if self.dimensions.is_empty() { 0 } else { - self.dimensions.iter().product() + // Saturate rather than wrap: a wrapped product could + // under-size a buffer. Size-critical callers use + // `checked_num_elements`. + self.dimensions + .iter() + .fold(1u64, |acc, &d| acc.saturating_mul(d)) } } } diff --git a/crates/clawhdf5-format/src/parallel_read.rs b/crates/clawhdf5-format/src/parallel_read.rs index c294203..14f5613 100644 --- a/crates/clawhdf5-format/src/parallel_read.rs +++ b/crates/clawhdf5-format/src/parallel_read.rs @@ -73,9 +73,12 @@ pub fn decompress_chunks_lane_partitioned( let c_addr = chunk_info.address as usize; let size = chunk_info.chunk_size as usize; - if c_addr + size > file_data.len() { + if c_addr + .checked_add(size) + .is_none_or(|end| end > file_data.len()) + { return Err(FormatError::UnexpectedEof { - expected: c_addr + size, + expected: c_addr.saturating_add(size), available: file_data.len(), }); } @@ -144,9 +147,12 @@ pub fn decompress_chunks_parallel( .map(|(index, chunk_info)| { let c_addr = chunk_info.address as usize; let size = chunk_info.chunk_size as usize; - if c_addr + size > file_data.len() { + if c_addr + .checked_add(size) + .is_none_or(|end| end > file_data.len()) + { return Err(FormatError::UnexpectedEof { - expected: c_addr + size, + expected: c_addr.saturating_add(size), available: file_data.len(), }); } @@ -182,9 +188,12 @@ pub fn decompress_chunks_sequential( for chunk_info in chunks { let c_addr = chunk_info.address as usize; let size = chunk_info.chunk_size as usize; - if c_addr + size > file_data.len() { + if c_addr + .checked_add(size) + .is_none_or(|end| end > file_data.len()) + { return Err(FormatError::UnexpectedEof { - expected: c_addr + size, + expected: c_addr.saturating_add(size), available: file_data.len(), }); } diff --git a/crates/clawhdf5/src/lazy.rs b/crates/clawhdf5/src/lazy.rs index e561deb..192a05c 100644 --- a/crates/clawhdf5/src/lazy.rs +++ b/crates/clawhdf5/src/lazy.rs @@ -436,19 +436,24 @@ impl<'f, R: HDF5Read> LazyDataset<'f, R> { )?) } - fn filter_pipeline(&self) -> Option { + /// `Ok(None)` means the dataset has no filter pipeline. A pipeline message + /// that is present but unparseable is an error: treating it as "no + /// filters" would hand the caller the still-compressed bytes as if they + /// were the data. + fn filter_pipeline(&self) -> Result, Error> { self.header .messages .iter() .find(|m| m.msg_type == MessageType::FilterPipeline) - .and_then(|msg| FilterPipeline::parse(&msg.data).ok()) + .map(|msg| FilterPipeline::parse(&msg.data).map_err(Error::Format)) + .transpose() } fn read_raw(&self) -> Result, Error> { let dt = self.datatype()?; let ds = self.dataspace()?; let dl = self.data_layout()?; - let pipeline = self.filter_pipeline(); + let pipeline = self.filter_pipeline()?; let data = self.file.reader.as_bytes(); Ok(data_read::read_raw_data_full( data, diff --git a/crates/clawhdf5/src/mmap_file.rs b/crates/clawhdf5/src/mmap_file.rs index 972501d..1573780 100644 --- a/crates/clawhdf5/src/mmap_file.rs +++ b/crates/clawhdf5/src/mmap_file.rs @@ -377,19 +377,24 @@ impl<'f> MmapDataset<'f> { )?) } - fn filter_pipeline(&self) -> Option { + /// `Ok(None)` means the dataset has no filter pipeline. A pipeline message + /// that is present but unparseable is an error: treating it as "no + /// filters" would hand the caller the still-compressed bytes as if they + /// were the data. + fn filter_pipeline(&self) -> Result, Error> { self.header .messages .iter() .find(|m| m.msg_type == MessageType::FilterPipeline) - .and_then(|msg| FilterPipeline::parse(&msg.data).ok()) + .map(|msg| FilterPipeline::parse(&msg.data).map_err(Error::Format)) + .transpose() } fn read_raw(&self) -> Result, Error> { let dt = self.datatype()?; let ds = self.dataspace()?; let dl = self.data_layout()?; - let pipeline = self.filter_pipeline(); + let pipeline = self.filter_pipeline()?; Ok(data_read::read_raw_data_full( self.file.reader.as_bytes(), &dl, diff --git a/crates/clawhdf5/src/reader.rs b/crates/clawhdf5/src/reader.rs index 223664d..fee9594 100644 --- a/crates/clawhdf5/src/reader.rs +++ b/crates/clawhdf5/src/reader.rs @@ -447,7 +447,7 @@ impl<'f> Dataset<'f> { let dt = self.datatype()?; let ds = self.dataspace()?; let dl = self.data_layout()?; - let pipeline = self.filter_pipeline(); + let pipeline = self.filter_pipeline()?; Ok(data_read::read_raw_data_selection( self.file.data.as_bytes(), &dl, @@ -743,19 +743,24 @@ impl<'f> Dataset<'f> { )?) } - fn filter_pipeline(&self) -> Option { + /// `Ok(None)` means the dataset has no filter pipeline. A pipeline message + /// that is present but unparseable is an error: treating it as "no + /// filters" would hand the caller the still-compressed bytes as if they + /// were the data. + fn filter_pipeline(&self) -> Result, Error> { self.header .messages .iter() .find(|m| m.msg_type == MessageType::FilterPipeline) - .and_then(|msg| FilterPipeline::parse(&msg.data).ok()) + .map(|msg| FilterPipeline::parse(&msg.data).map_err(Error::Format)) + .transpose() } fn read_raw(&self) -> Result, Error> { let dt = self.datatype()?; let ds = self.dataspace()?; let dl = self.data_layout()?; - let pipeline = self.filter_pipeline(); + let pipeline = self.filter_pipeline()?; // Virtual datasets are assembled from source datasets; the per-file // chunk cache does not apply. Route them through the resolver path so diff --git a/crates/clawhdf5/src/writer.rs b/crates/clawhdf5/src/writer.rs index a9aa64d..59904ba 100644 --- a/crates/clawhdf5/src/writer.rs +++ b/crates/clawhdf5/src/writer.rs @@ -72,7 +72,7 @@ impl FileBuilder { /// Serialize and write the file to the given path. pub fn write>(self, path: P) -> Result<(), Error> { let bytes = self.finish()?; - std::fs::write(path, bytes).map_err(Error::Io) + write_file_atomically(path.as_ref(), &bytes).map_err(Error::Io) } } @@ -186,3 +186,86 @@ pub fn create_datasets_parallel(specs: Vec) -> Result, Erro let bytes = clawhdf5_format::file_writer::finalize_parallel(blocks)?; Ok(bytes) } + +/// Write `bytes` to `path` so that a crash or power loss leaves either the old +/// file or the complete new one — never a truncated mix. `std::fs::write` +/// truncates the destination first, so dying mid-write used to destroy the +/// existing file. +fn write_file_atomically(path: &std::path::Path, bytes: &[u8]) -> std::io::Result<()> { + use std::io::Write; + + // Same directory as the target, so the rename stays on one filesystem. + let mut tmp_name = path + .file_name() + .ok_or_else(|| { + std::io::Error::new(std::io::ErrorKind::InvalidInput, "path has no file name") + })? + .to_os_string(); + tmp_name.push(format!(".tmp-{}", std::process::id())); + let tmp_path = path.with_file_name(tmp_name); + + let result = (|| { + let mut f = std::fs::File::create(&tmp_path)?; + f.write_all(bytes)?; + f.sync_all()?; + std::fs::rename(&tmp_path, path) + })(); + if result.is_err() { + let _ = std::fs::remove_file(&tmp_path); + return result; + } + // Make the rename itself durable. Best-effort: not every filesystem + // supports syncing a directory, and the new file is already in place. + #[cfg(unix)] + if let Some(dir) = path.parent() { + let dir = if dir.as_os_str().is_empty() { + std::path::Path::new(".") + } else { + dir + }; + if let Ok(d) = std::fs::File::open(dir) { + let _ = d.sync_all(); + } + } + Ok(()) +} + +#[cfg(test)] +mod atomic_write_tests { + use super::write_file_atomically; + + fn entries(dir: &std::path::Path) -> Vec { + let mut names: Vec = std::fs::read_dir(dir) + .unwrap() + .map(|e| e.unwrap().file_name().to_string_lossy().into_owned()) + .collect(); + names.sort(); + names + } + + #[test] + fn replaces_existing_file_and_leaves_no_temp_behind() { + let dir = tempfile::TempDir::new().unwrap(); + let path = dir.path().join("out.h5"); + std::fs::write(&path, b"old contents").unwrap(); + + write_file_atomically(&path, b"new").unwrap(); + + assert_eq!(std::fs::read(&path).unwrap(), b"new"); + assert_eq!(entries(dir.path()), ["out.h5"]); + } + + #[test] + fn failure_leaves_the_existing_file_untouched() { + let dir = tempfile::TempDir::new().unwrap(); + // The target is a directory, so the final rename cannot succeed. + let path = dir.path().join("taken"); + std::fs::create_dir(&path).unwrap(); + std::fs::write(path.join("keep"), b"x").unwrap(); + + assert!(write_file_atomically(&path, b"new").is_err()); + + assert!(path.is_dir()); + assert_eq!(entries(dir.path()), ["taken"], "temp file cleaned up"); + } +}