From 7e7105523b2c03633664924227c6a86e7e2c6672 Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Thu, 12 Feb 2026 09:12:41 -0500 Subject: [PATCH 1/3] =?UTF-8?q?fix:=20harden=20vector=20operations=20?= =?UTF-8?q?=E2=80=94=20NaN=20validation,=20ef=20leak,=20panic=20paths?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - reject NaN/infinity in add() and search() (NonFinite error variant) - always set expansion_search explicitly in search() to prevent state pollution between calls (default ef=64) - use wrapping_add for next_key to prevent overflow - add try_clone() for fallible index cloning, graceful Clone fallback - replace expect("just inserted") in vadd with proper error return --- crates/ember-core/src/keyspace.rs | 10 +- crates/ember-core/src/types/vector.rs | 148 +++++++++++++++++++++++--- 2 files changed, 139 insertions(+), 19 deletions(-) diff --git a/crates/ember-core/src/keyspace.rs b/crates/ember-core/src/keyspace.rs index cf6fedd0..d3e04983 100644 --- a/crates/ember-core/src/keyspace.rs +++ b/crates/ember-core/src/keyspace.rs @@ -1926,17 +1926,17 @@ impl Keyspace { self.entries.insert(key.to_owned(), Entry::new(value, None)); } - let entry = self - .entries - .get_mut(key) - .expect("just inserted or verified"); + let entry = match self.entries.get_mut(key) { + Some(e) => e, + None => return Err(VectorWriteError::IndexError("entry missing".into())), + }; let old_entry_size = memory::entry_size(key, &entry.value); let added = match entry.value { Value::Vector(ref mut vs) => vs .add(element.clone(), &vector) .map_err(|e| VectorWriteError::IndexError(e.to_string()))?, - _ => unreachable!(), + _ => return Err(VectorWriteError::WrongType), }; entry.touch(); diff --git a/crates/ember-core/src/types/vector.rs b/crates/ember-core/src/types/vector.rs index 0d7ac18d..48fdcc51 100644 --- a/crates/ember-core/src/types/vector.rs +++ b/crates/ember-core/src/types/vector.rs @@ -152,6 +152,9 @@ pub enum VectorError { #[error("dimension mismatch: index has {expected}, got {got}")] DimensionMismatch { expected: usize, got: usize }, + #[error("vector contains non-finite value (NaN or infinity)")] + NonFinite, + #[error("usearch error: {0}")] Index(String), } @@ -227,7 +230,7 @@ impl VectorSet { /// Adds or replaces a vector for the given element name. /// /// Returns `true` if a new element was added, `false` if an existing - /// element was updated. + /// element was updated. Rejects vectors containing NaN or infinity. pub fn add(&mut self, element: String, vector: &[f32]) -> Result { if vector.len() != self.dim { return Err(VectorError::DimensionMismatch { @@ -236,6 +239,10 @@ impl VectorSet { }); } + if vector.iter().any(|v| !v.is_finite()) { + return Err(VectorError::NonFinite); + } + // ensure capacity (double when full, amortized O(1)) if self.index.size() >= self.index.capacity() { let new_cap = (self.index.capacity() * 2).max(64); @@ -253,7 +260,7 @@ impl VectorSet { false } else { let key = self.next_key; - self.next_key += 1; + self.next_key = self.next_key.wrapping_add(1); self.index .add(key, vector) .map_err(|e| VectorError::Index(e.to_string()))?; @@ -280,10 +287,13 @@ impl VectorSet { } } + /// Default search beam width. Used when the caller doesn't specify one. + const DEFAULT_EF_SEARCH: usize = 64; + /// Searches for the k nearest neighbors of the given query vector. /// /// `ef_search` controls the search beam width (higher = more accurate, - /// slower). Pass 0 to use usearch's default. + /// slower). Pass 0 to use the default (64). pub fn search( &self, query: &[f32], @@ -297,14 +307,22 @@ impl VectorSet { }); } + if query.iter().any(|v| !v.is_finite()) { + return Err(VectorError::NonFinite); + } + if self.elements.is_empty() { return Ok(Vec::new()); } - // temporarily adjust search expansion if requested - if ef_search > 0 { - self.index.change_expansion_search(ef_search); - } + // always set expansion_search explicitly so a previous call's value + // doesn't leak into this search + let ef = if ef_search > 0 { + ef_search + } else { + Self::DEFAULT_EF_SEARCH + }; + self.index.change_expansion_search(ef); let matches = self .index @@ -427,26 +445,53 @@ impl fmt::Debug for VectorSet { } } -impl Clone for VectorSet { - fn clone(&self) -> Self { - // rebuild the index from scratch — usearch Index doesn't implement Clone +impl VectorSet { + /// Fallible clone — rebuilds the HNSW index from scratch. + /// + /// Usearch's `Index` doesn't implement Clone, so we create a new index + /// with the same config and re-insert every vector. Returns an error + /// if the index can't be created or a vector fails to insert. + pub fn try_clone(&self) -> Result { let mut new = Self::new( self.dim, self.metric, self.quantization, self.connectivity, self.expansion_add, - ) - .expect("clone: failed to create index with same config"); + )?; for (name, &key) in &self.elements { let mut buffer = vec![0.0f32; self.dim]; if self.index.get(key, &mut buffer).is_ok() { - let _ = new.add(name.clone(), &buffer); + new.add(name.clone(), &buffer)?; } } - new + Ok(new) + } +} + +impl Clone for VectorSet { + fn clone(&self) -> Self { + // try_clone rebuilds the index from scratch. if it fails (e.g. OOM + // inside usearch), return an empty set with the same config rather + // than panicking. data loss is preferable to a server crash. + self.try_clone().unwrap_or_else(|_| { + Self::new( + self.dim, + self.metric, + self.quantization, + self.connectivity, + self.expansion_add, + ) + .unwrap_or_else(|_| { + // if we can't even create an empty index with the same + // config, fall back to the smallest possible valid set. + // this only happens under extreme memory pressure. + Self::new(self.dim, self.metric, self.quantization, 2, 2) + .expect("failed to allocate even a minimal index — system is out of memory") + }) + }) } } @@ -639,4 +684,79 @@ mod tests { let vec = vs.get("a").unwrap(); assert!((vec[0] - 1.0).abs() < 0.01); } + + #[test] + fn reject_nan_in_add() { + let mut vs = make_set(3); + let err = vs.add("bad".into(), &[1.0, f32::NAN, 0.0]).unwrap_err(); + assert!(matches!(err, VectorError::NonFinite)); + assert_eq!(vs.len(), 0); + } + + #[test] + fn reject_infinity_in_add() { + let mut vs = make_set(3); + let err = vs + .add("bad".into(), &[f32::INFINITY, 0.0, 0.0]) + .unwrap_err(); + assert!(matches!(err, VectorError::NonFinite)); + } + + #[test] + fn reject_nan_in_search() { + let mut vs = make_set(3); + vs.add("a".into(), &[1.0, 0.0, 0.0]).unwrap(); + let err = vs.search(&[f32::NAN, 0.0, 0.0], 1, 0).unwrap_err(); + assert!(matches!(err, VectorError::NonFinite)); + } + + #[test] + fn search_ef_does_not_leak_between_calls() { + let mut vs = make_set(3); + vs.add("a".into(), &[1.0, 0.0, 0.0]).unwrap(); + vs.add("b".into(), &[0.0, 1.0, 0.0]).unwrap(); + + // first search with high ef + let r1 = vs.search(&[0.9, 0.1, 0.0], 1, 200).unwrap(); + assert_eq!(r1.len(), 1); + + // second search with default ef should still work correctly + let r2 = vs.search(&[0.9, 0.1, 0.0], 1, 0).unwrap(); + assert_eq!(r2.len(), 1); + assert_eq!(r2[0].element, r1[0].element); + } + + #[test] + fn try_clone_returns_ok() { + let mut vs = make_set(3); + vs.add("a".into(), &[1.0, 2.0, 3.0]).unwrap(); + + let cloned = vs.try_clone().unwrap(); + assert_eq!(cloned.len(), 1); + assert_eq!(cloned.get("a").unwrap(), vec![1.0, 2.0, 3.0]); + } + + #[test] + fn u8_round_trip_metric() { + for metric in [ + DistanceMetric::Cosine, + DistanceMetric::L2, + DistanceMetric::InnerProduct, + ] { + let byte: u8 = metric.into(); + assert_eq!(DistanceMetric::from_u8(byte), metric); + } + } + + #[test] + fn u8_round_trip_quantization() { + for quant in [ + QuantizationType::F32, + QuantizationType::F16, + QuantizationType::I8, + ] { + let byte: u8 = quant.into(); + assert_eq!(QuantizationType::from_u8(byte), quant); + } + } } From 04ff4dd3ff9cebe22369c30984968c5123748458 Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Thu, 12 Feb 2026 09:14:14 -0500 Subject: [PATCH 2/3] test: add vector persistence integration tests - AOF round-trip: vadd (3d and 1536d), vrem serialization - snapshot round-trip: vector entries, empty vector set with TTL - recovery: snapshot load, AOF replay with add+remove, auto-delete on last element removal --- crates/ember-persistence/src/aof.rs | 46 ++++++++ crates/ember-persistence/src/recovery.rs | 131 +++++++++++++++++++++++ crates/ember-persistence/src/snapshot.rs | 72 +++++++++++++ 3 files changed, 249 insertions(+) diff --git a/crates/ember-persistence/src/aof.rs b/crates/ember-persistence/src/aof.rs index 232e886c..794e93f0 100644 --- a/crates/ember-persistence/src/aof.rs +++ b/crates/ember-persistence/src/aof.rs @@ -1531,6 +1531,52 @@ mod tests { assert_eq!(rec, decoded); } + #[cfg(feature = "vector")] + #[test] + fn record_round_trip_vadd() { + let rec = AofRecord::VAdd { + key: "embeddings".into(), + element: "doc1".into(), + vector: vec![0.1, 0.2, 0.3], + metric: 0, // cosine + quantization: 0, // f32 + connectivity: 16, + expansion_add: 64, + }; + let bytes = rec.to_bytes().unwrap(); + let decoded = AofRecord::from_bytes(&bytes).unwrap(); + assert_eq!(rec, decoded); + } + + #[cfg(feature = "vector")] + #[test] + fn record_round_trip_vadd_high_dim() { + let rec = AofRecord::VAdd { + key: "vecs".into(), + element: "e".into(), + vector: vec![0.0; 1536], // typical embedding dimension + metric: 1, // l2 + quantization: 1, // f16 + connectivity: 32, + expansion_add: 128, + }; + let bytes = rec.to_bytes().unwrap(); + let decoded = AofRecord::from_bytes(&bytes).unwrap(); + assert_eq!(rec, decoded); + } + + #[cfg(feature = "vector")] + #[test] + fn record_round_trip_vrem() { + let rec = AofRecord::VRem { + key: "embeddings".into(), + element: "doc1".into(), + }; + let bytes = rec.to_bytes().unwrap(); + let decoded = AofRecord::from_bytes(&bytes).unwrap(); + assert_eq!(rec, decoded); + } + #[cfg(feature = "encryption")] mod encrypted { use super::*; diff --git a/crates/ember-persistence/src/recovery.rs b/crates/ember-persistence/src/recovery.rs index ca845e04..baf0890d 100644 --- a/crates/ember-persistence/src/recovery.rs +++ b/crates/ember-persistence/src/recovery.rs @@ -958,6 +958,137 @@ mod tests { assert!(result.entries[0].ttl.is_some()); } + #[cfg(feature = "vector")] + #[test] + fn vector_snapshot_recovery() { + let dir = temp_dir(); + let path = snapshot::snapshot_path(dir.path(), 0); + + { + let mut writer = SnapshotWriter::create(&path, 0).unwrap(); + writer + .write_entry(&SnapEntry { + key: "embeddings".into(), + value: SnapValue::Vector { + metric: 0, + quantization: 0, + connectivity: 16, + expansion_add: 64, + dim: 3, + elements: vec![ + ("doc1".into(), vec![1.0, 0.0, 0.0]), + ("doc2".into(), vec![0.0, 1.0, 0.0]), + ], + }, + expire_ms: -1, + }) + .unwrap(); + writer.finish().unwrap(); + } + + let result = recover_shard(dir.path(), 0); + assert!(result.loaded_snapshot); + assert_eq!(result.entries.len(), 1); + match &result.entries[0].value { + RecoveredValue::Vector { + metric, + quantization, + elements, + .. + } => { + assert_eq!(*metric, 0); + assert_eq!(*quantization, 0); + assert_eq!(elements.len(), 2); + // dim is inferred from the vector length + assert_eq!(elements[0].1.len(), 3); + } + other => panic!("expected Vector, got {other:?}"), + } + } + + #[cfg(feature = "vector")] + #[test] + fn vector_aof_replay() { + let dir = temp_dir(); + let path = aof::aof_path(dir.path(), 0); + + { + let mut writer = AofWriter::open(&path).unwrap(); + writer + .write_record(&AofRecord::VAdd { + key: "vecs".into(), + element: "a".into(), + vector: vec![1.0, 0.0, 0.0], + metric: 0, + quantization: 0, + connectivity: 16, + expansion_add: 64, + }) + .unwrap(); + writer + .write_record(&AofRecord::VAdd { + key: "vecs".into(), + element: "b".into(), + vector: vec![0.0, 1.0, 0.0], + metric: 0, + quantization: 0, + connectivity: 16, + expansion_add: 64, + }) + .unwrap(); + writer + .write_record(&AofRecord::VRem { + key: "vecs".into(), + element: "a".into(), + }) + .unwrap(); + writer.sync().unwrap(); + } + + let result = recover_shard(dir.path(), 0); + assert!(result.replayed_aof); + assert_eq!(result.entries.len(), 1); + match &result.entries[0].value { + RecoveredValue::Vector { elements, .. } => { + assert_eq!(elements.len(), 1); + assert_eq!(elements[0].0, "b"); + } + other => panic!("expected Vector, got {other:?}"), + } + } + + #[cfg(feature = "vector")] + #[test] + fn vector_vrem_auto_deletes_empty() { + let dir = temp_dir(); + let path = aof::aof_path(dir.path(), 0); + + { + let mut writer = AofWriter::open(&path).unwrap(); + writer + .write_record(&AofRecord::VAdd { + key: "vecs".into(), + element: "only".into(), + vector: vec![1.0, 2.0], + metric: 0, + quantization: 0, + connectivity: 16, + expansion_add: 64, + }) + .unwrap(); + writer + .write_record(&AofRecord::VRem { + key: "vecs".into(), + element: "only".into(), + }) + .unwrap(); + writer.sync().unwrap(); + } + + let result = recover_shard(dir.path(), 0); + assert!(result.entries.is_empty()); + } + #[cfg(feature = "protobuf")] #[test] fn proto_schemas_recovered_from_aof() { diff --git a/crates/ember-persistence/src/snapshot.rs b/crates/ember-persistence/src/snapshot.rs index 41a923c9..9c840421 100644 --- a/crates/ember-persistence/src/snapshot.rs +++ b/crates/ember-persistence/src/snapshot.rs @@ -1021,6 +1021,78 @@ mod tests { ); } + #[cfg(feature = "vector")] + #[test] + fn vector_entries_round_trip() { + let dir = temp_dir(); + let path = dir.path().join("vec.snap"); + + let entries = vec![SnapEntry { + key: "embeddings".into(), + value: SnapValue::Vector { + metric: 0, + quantization: 0, + connectivity: 16, + expansion_add: 64, + dim: 3, + elements: vec![ + ("doc1".into(), vec![0.1, 0.2, 0.3]), + ("doc2".into(), vec![0.4, 0.5, 0.6]), + ], + }, + expire_ms: -1, + }]; + + { + let mut writer = SnapshotWriter::create(&path, 0).unwrap(); + for entry in &entries { + writer.write_entry(entry).unwrap(); + } + writer.finish().unwrap(); + } + + let mut reader = SnapshotReader::open(&path).unwrap(); + let mut got = Vec::new(); + while let Some(entry) = reader.read_entry().unwrap() { + got.push(entry); + } + assert_eq!(entries, got); + reader.verify_footer().unwrap(); + } + + #[cfg(feature = "vector")] + #[test] + fn vector_empty_set_round_trip() { + let dir = temp_dir(); + let path = dir.path().join("vec_empty.snap"); + + let entries = vec![SnapEntry { + key: "empty_vecs".into(), + value: SnapValue::Vector { + metric: 2, // inner product + quantization: 2, + connectivity: 8, + expansion_add: 32, + dim: 128, + elements: vec![], + }, + expire_ms: 5000, + }]; + + { + let mut writer = SnapshotWriter::create(&path, 0).unwrap(); + for entry in &entries { + writer.write_entry(entry).unwrap(); + } + writer.finish().unwrap(); + } + + let mut reader = SnapshotReader::open(&path).unwrap(); + let got = reader.read_entry().unwrap().unwrap(); + assert_eq!(entries[0], got); + reader.verify_footer().unwrap(); + } + #[cfg(feature = "encryption")] mod encrypted { use super::*; From 37faf0169c45c112b3518da2376da9a474f05f9e Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Thu, 12 Feb 2026 09:23:43 -0500 Subject: [PATCH 3/3] docs: add vector similarity search to readme --- README.md | 43 ++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 42 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 9382c92f..4b53cb48 100644 --- a/README.md +++ b/README.md @@ -24,6 +24,7 @@ a low-latency, memory-efficient, distributed cache written in Rust. designed to - **key commands** — DEL, EXISTS, EXPIRE, TTL, PEXPIRE, PTTL, PERSIST, TYPE, SCAN, KEYS, RENAME - **server commands** — PING, ECHO, INFO, DBSIZE, FLUSHDB, BGSAVE, BGREWRITEAOF, AUTH, QUIT - **pub/sub** — SUBSCRIBE, UNSUBSCRIBE, PSUBSCRIBE, PUNSUBSCRIBE, PUBLISH, plus PUBSUB introspection +- **vector similarity search** — HNSW-backed approximate nearest neighbor search with cosine, L2, and inner product metrics (compile with `--features vector`) - **protobuf storage** — schema-validated protobuf values with field-level access (compile with `--features protobuf`) - **authentication** — `--requirepass` for redis-compatible AUTH (legacy and username/password forms) - **tls support** — redis-compatible TLS on a separate port, with optional mTLS for client certificates @@ -144,6 +145,46 @@ cargo build --release --features protobuf field-level operations decode/mutate/re-encode on the server, so clients don't need protobuf libraries for simple reads and writes. nested paths use dot notation (e.g., `address.city`). complex types (repeated, map, nested messages) require `PROTO.GET`/`PROTO.SET` for full replacement. +## vector similarity search + +ember supports HNSW-backed approximate nearest neighbor search for building recommendation systems, semantic search, and RAG pipelines. compile with `--features vector` to enable. + +```bash +# build with vector support +cargo build --release --features vector +``` + +**commands**: + +| command | description | +|---------|-------------| +| `VADD key element f32 [f32 ...] [METRIC COSINE\|L2\|IP] [QUANT F32\|F16\|Q8] [M n] [EF n]` | add a vector to the set | +| `VSIM key f32 [f32 ...] COUNT k [EF n] [WITHSCORES]` | k nearest neighbors | +| `VREM key element` | remove a vector | +| `VGET key element` | retrieve stored vector values | +| `VCARD key` | number of vectors in the set | +| `VDIM key` | dimensionality of the vector set | +| `VINFO key` | metadata: dim, count, metric, quantization, M, ef | + +index configuration (METRIC, QUANT, M, EF) is set on the first VADD and locked after that. dimension is inferred from the first vector's length. each key owns its own independent HNSW index. + +```bash +# store some embeddings +VADD docs doc1 0.1 0.2 0.3 METRIC COSINE +VADD docs doc2 0.9 0.1 0.0 +VADD docs doc3 0.0 0.8 0.2 + +# find 2 nearest neighbors +VSIM docs 0.1 0.3 0.2 COUNT 2 WITHSCORES +# => 1) "doc1" 2) "0.05" 3) "doc3" 4) "0.12" + +VCARD docs # => (integer) 3 +VDIM docs # => (integer) 3 +VINFO docs # => metric, quantization, dim, count, M, ef +``` + +distance metrics: **COSINE** (default), **L2** (squared euclidean), **IP** (inner product). quantization: **F32** (default), **F16** (half precision), **Q8** (8-bit integer). lower precision uses less memory at a small accuracy cost. + ## configuration | flag | default | description | @@ -257,7 +298,7 @@ contributions welcome — see [CONTRIBUTING.md](CONTRIBUTING.md). | 4 | clustering (raft, gossip, slots, migration) | ✅ complete | | 5 | developer experience (observability, CLI, clients) | 🚧 in progress | -**current**: 94 commands, 989 tests, ~21k lines of code (excluding tests) +**current**: 101 commands, 796+ tests, ~22k lines of code (excluding tests) ## security