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();