perf(agent): persist the vector index graph; incremental catch-up

open() marked the HNSW index dirty, so the first search of every session
rebuilt it from scratch — 36 s at 100K records with the (better, slower)
heuristic build. First query after open is now 1.7 / 15 / 159 ms at
1K / 10K / 100K; what remains is the one-off keyword index build.

- clawhdf5-ann: HnswIndex::graph_to_bytes / from_graph_bytes serialize the
  graph only (levels, tombstones, adjacency as u32, CRC32). The existing HDF5
  serializer embeds a full copy of every vector, which would double a store
  that already holds them. Loading validates everything — counts, levels vs
  layer count, connection limits, every neighbour id and the layer it must
  exist on — so a damaged graph, or a hostile one with a valid checksum, is an
  error rather than an out-of-bounds walk during search.
- clawhdf5-agent: each checkpoint writes the graph to <store>.h5.ann (synced,
  atomic, before the .h5) and records a fresh generation id in /meta. open()
  loads the sidecar only if its generation matches that checkpoint; missing,
  stale, damaged or mismatched sidecars are ignored and the index rebuilt.
  Records appended through WAL replay join the loaded index incrementally; a
  replayed Update or Tombstone invalidates it. snapshot() copies it. Only an
  index that exactly mirrors the cache is saved; otherwise a stale sidecar is
  removed.
- ensure_hnsw_fresh inserts records appended since the last sync instead of
  rebuilding, so save_batch no longer marks the whole index dirty.
- CheckpointMeta { wal_applied, ann_generation } with *_with_meta build/write/
  read functions; the *_with_mark ones delegate.
- Harness reports the one-off cold index build separately from the first query
  after a reopen.

Co-Authored-By: Claude Fable 5.1 <[email protected]>
This commit is contained in:
osobh
2026-09-19 08:00:57 -07:00
co-authored by Claude Fable 5.1
parent 2bfbb7fb4b
commit 0ee698accd
7 changed files with 641 additions and 20 deletions
+291 -11
View File
@@ -390,8 +390,13 @@ impl HDF5Memory {
} else {
Some(store_lock::StoreLock::acquire(path)?)
};
let ((config, mut cache, sessions, knowledge), wal_applied) =
storage::read_from_disk_with_mark(path)?;
let ((config, mut cache, sessions, knowledge), checkpoint) =
storage::read_from_disk_with_meta(path)?;
let wal_applied = checkpoint.wal_applied;
let n_checkpoint = cache.len();
// Set if WAL replay did anything other than append records; the saved
// vector index then no longer describes the first `n_checkpoint` ones.
let mut replay_only_appended = true;
// Replay WAL if present
let wal_path = path.with_extension("h5.wal");
@@ -408,6 +413,9 @@ impl HDF5Memory {
&& let Ok(entries) =
wal::WalFile::read_entries_for_migration(&wal_path, wal_applied)
{
replay_only_appended &= entries
.iter()
.all(|e| e.entry_type == wal::WalEntryType::Save);
wal::replay_into_cache(&entries, &mut cache);
}
None
@@ -419,6 +427,9 @@ impl HDF5Memory {
// 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)?;
replay_only_appended &= entries
.iter()
.all(|e| e.entry_type == wal::WalEntryType::Save);
wal::replay_into_cache(&entries, &mut cache);
Some(wal::WalFile::open(&wal_path)?)
} else if config.wal_enabled {
@@ -427,6 +438,25 @@ impl HDF5Memory {
None
};
#[cfg(feature = "hnsw")]
let loaded_index = if replay_only_appended {
Self::load_vector_index(path, checkpoint.ann_generation, &cache, n_checkpoint)
} else {
None
};
#[cfg(not(feature = "hnsw"))]
let _ = (
n_checkpoint,
replay_only_appended,
checkpoint.ann_generation,
);
#[cfg(feature = "hnsw")]
let synced_len = if loaded_index.is_some() {
cache.len()
} else {
0
};
Ok(Self {
config,
cache,
@@ -435,14 +465,14 @@ impl HDF5Memory {
wal,
strategy: None,
ephemeral: None,
// Existing data is loaded from disk + WAL replay; mark the index
// dirty so it is (re)built from the cache on the first search.
// Reuse the vector index saved with the checkpoint when there is
// one; otherwise mark it dirty so the first search builds it.
#[cfg(feature = "hnsw")]
hnsw: None,
hnsw_dirty: loaded_index.is_none(),
#[cfg(feature = "hnsw")]
hnsw_dirty: true,
hnsw_synced_len: synced_len,
#[cfg(feature = "hnsw")]
hnsw_synced_len: 0,
hnsw: loaded_index,
// No on-disk provenance ledger exists yet (see CLAUDE.md), so
// there's no historical hash to verify loaded records against —
// the store starts empty and is populated as records are
@@ -458,6 +488,99 @@ impl HDF5Memory {
})
}
/// Where the vector index graph is kept between sessions.
#[cfg_attr(not(feature = "hnsw"), allow(dead_code))]
fn vector_index_path(store: &Path) -> PathBuf {
store.with_extension("h5.ann")
}
/// Save the vector index graph next to the store, returning the generation
/// id the checkpoint must record for it. Only an index that exactly mirrors
/// the cache is saved; otherwise any stale sidecar is removed and `None`
/// returned, and the next session rebuilds. Failures are not errors — the
/// sidecar is a cache of derived data.
#[cfg(feature = "hnsw")]
fn persist_vector_index(&self) -> Option<u64> {
let path = Self::vector_index_path(&self.config.path);
let index = match self.hnsw.as_ref() {
Some(index)
if !self.hnsw_dirty
&& self.hnsw_synced_len == self.cache.embeddings.len()
&& index.len() == self.cache.embeddings.len() =>
{
index
}
_ => {
let _ = std::fs::remove_file(&path);
return None;
}
};
static COUNTER: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
let nanos = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map_or(0, |d| d.as_nanos() as u64);
let generation = nanos
^ (u64::from(std::process::id()) << 32)
^ COUNTER.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
let mut bytes = generation.to_le_bytes().to_vec();
bytes.extend_from_slice(&index.graph_to_bytes());
let tmp = path.with_extension("ann.tmp");
let written =
storage::write_synced(&tmp, &bytes).and_then(|()| storage::rename_synced(&tmp, &path));
match written {
Ok(()) => Some(generation),
Err(_) => {
let _ = std::fs::remove_file(&tmp);
let _ = std::fs::remove_file(&path);
None
}
}
}
#[cfg(not(feature = "hnsw"))]
fn persist_vector_index(&self) -> Option<u64> {
None
}
/// Load the vector index saved with the checkpoint identified by
/// `generation`, covering the first `n_checkpoint` records of `cache`.
/// Anything unexpected — no sidecar, another generation, a damaged or
/// mismatched graph — yields `None` and the index is rebuilt on demand.
#[cfg(feature = "hnsw")]
fn load_vector_index(
store: &Path,
generation: Option<u64>,
cache: &MemoryCache,
n_checkpoint: usize,
) -> Option<HnswIndex> {
let generation = generation?;
let bytes = std::fs::read(Self::vector_index_path(store)).ok()?;
let (stamp, graph) = bytes.split_at_checked(8)?;
if u64::from_le_bytes(stamp.try_into().ok()?) != generation {
return None;
}
let vectors = cache.embeddings.get(..n_checkpoint)?.to_vec();
let mut index = HnswIndex::from_graph_bytes(graph, vectors).ok()?;
if index.dimension() != cache.embedding_dim {
return None;
}
// Records appended since (replayed from the WAL) join incrementally.
for id in n_checkpoint..cache.embeddings.len() {
if cache.embeddings[id].len() != index.dimension()
|| index.insert(cache.embeddings[id].clone()) != id
{
return None;
}
}
for (id, &t) in cache.tombstones.iter().enumerate() {
if t != 0 {
index.mark_deleted(id);
}
}
Some(index)
}
/// Bring the keyword index in line with the cache and return it.
///
/// Appends need no hook: records the index hasn't seen yet (whatever path
@@ -517,13 +640,19 @@ impl HDF5Memory {
// 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(
// Written before the .h5 so a crash in between leaves a sidecar whose
// generation matches no checkpoint (ignored), never the reverse.
let ann_generation = self.persist_vector_index();
storage::write_to_disk_with_meta(
&self.config.path,
&self.config,
&self.cache,
&self.sessions,
&self.knowledge,
wal_applied,
&schema::CheckpointMeta {
wal_applied,
ann_generation,
},
)?;
if let Some(ref mut w) = self.wal {
w.truncate()?;
@@ -686,6 +815,30 @@ impl HDF5Memory {
#[cfg(feature = "hnsw")]
fn ensure_hnsw_fresh(&mut self) {
let n = self.cache.embeddings.len();
// Records appended since the index was last in sync (a batch save, or
// any path that pushes to the cache without a hook) are inserted
// incrementally rather than triggering a rebuild of the whole graph.
if !self.hnsw_dirty
&& self.hnsw_synced_len < n
&& let Some(index) = self.hnsw.as_mut()
&& index.len() == self.hnsw_synced_len
{
let dim = index.dimension();
let appended = (self.hnsw_synced_len..n).all(|id| {
self.cache.embeddings[id].len() == dim
&& index.insert(self.cache.embeddings[id].clone()) == id
});
if appended {
for id in self.hnsw_synced_len..n {
if self.cache.tombstones[id] != 0 {
index.mark_deleted(id);
}
}
self.hnsw_synced_len = n;
} else {
self.hnsw_dirty = true;
}
}
if self.hnsw_dirty || self.hnsw_synced_len != n {
self.hnsw = self.build_hnsw_from_cache();
self.hnsw_synced_len = n;
@@ -928,8 +1081,8 @@ impl AgentMemory for HDF5Memory {
);
indices.push(idx);
}
// Batch inserts rebuild the index once rather than node-by-node.
self.hnsw_mark_dirty();
// The vector and keyword indexes pick the new records up
// incrementally the next time they are needed.
self.flush()?;
Ok(indices)
}
@@ -987,6 +1140,13 @@ impl AgentMemory for HDF5Memory {
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"))?;
}
// The saved vector index belongs to the checkpoint just copied (its
// generation id is in that .h5), so it is valid for the snapshot too.
// Best effort: without it the snapshot simply rebuilds on first search.
let ann_path = Self::vector_index_path(&self.config.path);
if ann_path.exists() {
let _ = storage::snapshot_file(&ann_path, &Self::vector_index_path(&snapshot));
}
Ok(snapshot)
}
@@ -1674,6 +1834,126 @@ mod tests {
assert_eq!(restored.cache.chunks, ["checkpointed", "wal-only"]);
}
/// A store with `n` records spread over a few directions, WAL on.
#[cfg(feature = "hnsw")]
fn indexed_store(dir: &TempDir, n: usize) -> (HDF5Memory, PathBuf) {
let mut config = make_config(dir);
config.wal_enabled = true;
config.wal_max_entries = 10_000;
config.compact_threshold = 0.0;
let path = config.path.clone();
let mut mem = HDF5Memory::create(config).unwrap();
for i in 0..n {
let a = i as f32 * 0.37;
mem.save(make_entry(
&format!("rec{i}"),
&[a.cos(), a.sin(), (a * 0.5).cos(), 0.1],
))
.unwrap();
}
(mem, path)
}
#[cfg(feature = "hnsw")]
fn top_ids(mem: &mut HDF5Memory, q: &[f32]) -> Vec<usize> {
mem.hybrid_search(q, "", 1.0, 0.0, 5)
.into_iter()
.map(|r| r.index)
.collect()
}
#[cfg(feature = "hnsw")]
#[test]
fn vector_index_is_reloaded_not_rebuilt() {
let dir = TempDir::new().unwrap();
let (mut mem, path) = indexed_store(&dir, 60);
let q = [0.3f32.cos(), 0.3f32.sin(), 0.9, 0.1];
let expected = top_ids(&mut mem, &q); // builds the index
mem.flush_wal().unwrap(); // checkpoint + sidecar
drop(mem);
assert!(HDF5Memory::vector_index_path(&path).exists());
let mut reopened = HDF5Memory::open(&path).unwrap();
assert!(!reopened.hnsw_dirty, "index should come from the sidecar");
assert_eq!(reopened.hnsw.as_ref().unwrap().len(), 60);
assert_eq!(top_ids(&mut reopened, &q), expected);
}
#[cfg(feature = "hnsw")]
#[test]
fn records_appended_after_the_checkpoint_join_the_loaded_index() {
let dir = TempDir::new().unwrap();
let (mut mem, path) = indexed_store(&dir, 40);
top_ids(&mut mem, &[1.0, 0.0, 0.0, 0.0]);
mem.flush_wal().unwrap();
// Only in the WAL when the process "dies".
mem.save(make_entry("late", &[0.0, 0.0, 0.0, 1.0])).unwrap();
drop(mem);
let mut reopened = HDF5Memory::open(&path).unwrap();
assert!(!reopened.hnsw_dirty);
assert_eq!(reopened.hnsw.as_ref().unwrap().len(), 41);
assert_eq!(top_ids(&mut reopened, &[0.0, 0.0, 0.0, 1.0])[0], 40);
}
#[cfg(feature = "hnsw")]
#[test]
fn replayed_update_invalidates_the_saved_index() {
let dir = TempDir::new().unwrap();
let (mut mem, path) = indexed_store(&dir, 40);
top_ids(&mut mem, &[1.0, 0.0, 0.0, 0.0]);
mem.flush_wal().unwrap();
// An in-place update after the checkpoint changes record 0's vector;
// the saved graph was built over the old one.
let mut moved = make_entry("rec0 moved", &[0.0, 0.0, 0.0, 1.0]);
moved.tags = mem.cache.tags[0].clone();
mem.save_or_update(moved).unwrap();
let expected = top_ids(&mut mem, &[0.0, 0.0, 0.0, 1.0]);
std::mem::forget(mem); // die without the drop-time checkpoint
let mut reopened = HDF5Memory::open_read_only(&path).unwrap();
assert!(reopened.hnsw_dirty, "saved index must not be reused");
assert_eq!(top_ids(&mut reopened, &[0.0, 0.0, 0.0, 1.0]), expected);
}
#[cfg(feature = "hnsw")]
#[test]
fn stale_or_damaged_index_sidecar_is_ignored() {
let dir = TempDir::new().unwrap();
let (mut mem, path) = indexed_store(&dir, 40);
let q = [1.0, 0.0, 0.0, 0.0];
top_ids(&mut mem, &q); // builds the index
mem.flush_wal().unwrap();
let ann = HDF5Memory::vector_index_path(&path);
let first_sidecar = std::fs::read(&ann).unwrap();
// A second checkpoint gets a new generation.
mem.save(make_entry("more", &[0.5, 0.5, 0.0, 0.0])).unwrap();
top_ids(&mut mem, &q);
mem.flush_wal().unwrap();
let expected_after = top_ids(&mut mem, &q);
drop(mem);
// Sidecar from the earlier checkpoint: wrong generation.
std::fs::write(&ann, &first_sidecar).unwrap();
let mut reopened = HDF5Memory::open_read_only(&path).unwrap();
assert!(reopened.hnsw_dirty);
assert_eq!(top_ids(&mut reopened, &q), expected_after);
drop(reopened);
// Right generation, damaged graph.
let mut mem = HDF5Memory::open(&path).unwrap();
top_ids(&mut mem, &q);
mem.flush_wal().unwrap();
drop(mem);
let mut bytes = std::fs::read(&ann).unwrap();
let mid = bytes.len() / 2;
bytes[mid] ^= 0xFF;
std::fs::write(&ann, &bytes).unwrap();
let mut reopened = HDF5Memory::open_read_only(&path).unwrap();
assert!(reopened.hnsw_dirty);
assert_eq!(top_ids(&mut reopened, &q), expected_after);
}
#[test]
fn keyword_index_stays_in_sync_through_every_mutation() {
let dir = TempDir::new().unwrap();
+50
View File
@@ -22,6 +22,7 @@ pub const ZEROCLAW_VERSION: &str = "0.8.0";
/// 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";
const ANN_GENERATION_ATTR: &str = "ann_generation";
/// Build a complete HDF5 file from the in-memory state.
pub fn build_hdf5_file(
@@ -43,6 +44,34 @@ pub fn build_hdf5_file_with_mark(
knowledge: &KnowledgeCache,
wal_applied: Option<WalMark>,
) -> Result<Vec<u8>, MemoryError> {
let meta = CheckpointMeta {
wal_applied,
ann_generation: None,
};
build_hdf5_file_with_meta(config, cache, sessions, knowledge, &meta)
}
/// Bookkeeping a checkpoint records in `/meta` beside the store's contents.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct CheckpointMeta {
/// The WAL prefix this checkpoint already contains; see [`WalMark`].
pub wal_applied: Option<WalMark>,
/// Identifies the vector-index sidecar (`<store>.h5.ann`) written with this
/// checkpoint. A sidecar is loaded only if it carries the same value, so
/// one left over from another checkpoint can never be attached to records
/// it wasn't built from.
pub ann_generation: Option<u64>,
}
/// [`build_hdf5_file`] with checkpoint bookkeeping.
pub fn build_hdf5_file_with_meta(
config: &MemoryConfig,
cache: &MemoryCache,
sessions: &SessionCache,
knowledge: &KnowledgeCache,
checkpoint: &CheckpointMeta,
) -> Result<Vec<u8>, MemoryError> {
let wal_applied = checkpoint.wal_applied;
let mut builder = clawhdf5::FileBuilder::new();
// /meta group with schema attributes
@@ -83,6 +112,11 @@ pub fn build_hdf5_file_with_mark(
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)));
}
if let Some(generation) = checkpoint.ann_generation {
// Stored as the i64 with the same bits; attributes have no u64 scalar
// round trip through every reader.
meta.set_attr(ANN_GENERATION_ATTR, AttrValue::I64(generation as i64));
}
// 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();
@@ -386,6 +420,22 @@ pub fn read_wal_mark(file: &clawhdf5::File) -> Option<WalMark> {
Some(WalMark { len, crc })
}
/// Read the checkpoint bookkeeping from `/meta`.
pub fn read_checkpoint_meta(file: &clawhdf5::File) -> CheckpointMeta {
let ann_generation = file
.group("meta")
.ok()
.and_then(|g| g.attrs().ok())
.and_then(|attrs| match attrs.get(ANN_GENERATION_ATTR) {
Some(AttrValue::I64(v)) => Some(*v as u64),
_ => None,
});
CheckpointMeta {
wal_applied: read_wal_mark(file),
ann_generation,
}
}
pub fn validate_and_load(
file: &clawhdf5::File,
) -> Result<(MemoryConfig, MemoryCache, SessionCache, KnowledgeCache), MemoryError> {
+33 -3
View File
@@ -34,7 +34,23 @@ pub fn write_to_disk_with_mark(
knowledge: &KnowledgeCache,
wal_applied: Option<WalMark>,
) -> Result<(), MemoryError> {
let bytes = schema::build_hdf5_file_with_mark(config, cache, sessions, knowledge, wal_applied)?;
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()));
@@ -47,7 +63,7 @@ pub fn write_to_disk_with_mark(
}
/// Write `bytes` to `path` and flush them to stable storage.
fn write_synced(path: &Path, bytes: &[u8]) -> Result<(), MemoryError> {
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)?;
@@ -62,7 +78,7 @@ fn write_synced(path: &Path, bytes: &[u8]) -> Result<(), MemoryError> {
/// 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> {
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() {
@@ -113,6 +129,20 @@ pub fn read_from_disk_with_mark(path: &Path) -> Result<(StoreState, Option<WalMa
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 mmap = clawhdf5_io::MmapReader::open(path).map_err(MemoryError::Io)?;
mmap.advise_willneed(0, mmap.len());
let file = clawhdf5::File::from_bytes(mmap.as_bytes().to_vec())
.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() {