From 542b8c012abffba1eea7acf67b47ebfac98f4b8e Mon Sep 17 00:00:00 2001 From: happy-v587 Date: Sun, 19 Jul 2026 22:35:29 +0800 Subject: [PATCH 01/32] feat(storage): add vector data column family --- src/conf/src/raft_type.rs | 2 ++ src/raft/src/lib.rs | 14 ++++++++------ src/storage/src/batch.rs | 1 + src/storage/src/format_base_value.rs | 19 ++++++++++++++++--- src/storage/src/logindex/types.rs | 4 +++- src/storage/src/redis.rs | 8 +++++++- src/storage/src/storage.rs | 1 + src/storage/tests/redis_basic_test.rs | 21 +++++++++++++-------- 8 files changed, 51 insertions(+), 19 deletions(-) diff --git a/src/conf/src/raft_type.rs b/src/conf/src/raft_type.rs index 67934b10..c8fab979 100644 --- a/src/conf/src/raft_type.rs +++ b/src/conf/src/raft_type.rs @@ -111,6 +111,7 @@ pub enum ColumnFamilyIndex { ListsDataCF = 3, ZsetsDataCF = 4, ZsetsScoreCF = 5, + VectorDataCF = 6, } impl ColumnFamilyIndex { @@ -122,6 +123,7 @@ impl ColumnFamilyIndex { 3 => Some(Self::ListsDataCF), 4 => Some(Self::ZsetsDataCF), 5 => Some(Self::ZsetsScoreCF), + 6 => Some(Self::VectorDataCF), _ => None, } } diff --git a/src/raft/src/lib.rs b/src/raft/src/lib.rs index 556a57b6..dd334c9f 100644 --- a/src/raft/src/lib.rs +++ b/src/raft/src/lib.rs @@ -53,12 +53,13 @@ pub const COLUMN_FAMILY_COUNT: usize = storage::ColumnFamilyIndex::COUNT; /// List of CF names, in the same order as storage::ColumnFamilyIndex pub const CF_NAMES: [&str; COLUMN_FAMILY_COUNT] = [ - "default", // MetaCF = 0 - "hash_data_cf", // HashesDataCF = 1 - "set_data_cf", // SetsDataCF = 2 - "list_data_cf", // ListsDataCF = 3 - "zset_data_cf", // ZsetsDataCF = 4 - "zset_score_cf", // ZsetsScoreCF = 5 + "default", // MetaCF = 0 + "hash_data_cf", // HashesDataCF = 1 + "set_data_cf", // SetsDataCF = 2 + "list_data_cf", // ListsDataCF = 3 + "zset_data_cf", // ZsetsDataCF = 4 + "zset_score_cf", // ZsetsScoreCF = 5 + "vector_data_cf", // VectorDataCF = 6 ]; const _: () = assert!( @@ -80,6 +81,7 @@ mod tests { ColumnFamilyIndex::ListsDataCF, ColumnFamilyIndex::ZsetsDataCF, ColumnFamilyIndex::ZsetsScoreCF, + ColumnFamilyIndex::VectorDataCF, ]; for (i, cf_index) in variants.iter().enumerate() { assert_eq!( diff --git a/src/storage/src/batch.rs b/src/storage/src/batch.rs index 984d93f3..0bec4d9b 100644 --- a/src/storage/src/batch.rs +++ b/src/storage/src/batch.rs @@ -156,6 +156,7 @@ fn cf_index_to_usize(cf_idx: ColumnFamilyIndex) -> usize { ColumnFamilyIndex::ListsDataCF => 3, ColumnFamilyIndex::ZsetsDataCF => 4, ColumnFamilyIndex::ZsetsScoreCF => 5, + ColumnFamilyIndex::VectorDataCF => 6, } } diff --git a/src/storage/src/format_base_value.rs b/src/storage/src/format_base_value.rs index b27ff077..4ba80b4d 100644 --- a/src/storage/src/format_base_value.rs +++ b/src/storage/src/format_base_value.rs @@ -36,13 +36,16 @@ pub enum DataType { ZSet = 4, None = 5, All = 6, + VectorSet = 7, } impl DataType { pub fn min_meta_raw_len(self) -> Result { match self { DataType::String => Ok(TYPE_LENGTH + STRING_VALUE_SUFFIXLENGTH), - DataType::Hash | DataType::Set | DataType::ZSet => Ok(BASE_META_VALUE_LENGTH), + DataType::Hash | DataType::Set | DataType::ZSet | DataType::VectorSet => { + Ok(BASE_META_VALUE_LENGTH) + } DataType::List => Ok(LISTS_META_VALUE_LENGTH), _ => InvalidFormatSnafu { message: format!("data type: {self:?} should not be used as meta value"), @@ -64,6 +67,7 @@ impl TryFrom for DataType { 4 => Ok(DataType::ZSet), 5 => Ok(DataType::None), 6 => Ok(DataType::All), + 7 => Ok(DataType::VectorSet), _ => InvalidFormatSnafu { message: format!("Invalid data type byte: {value}"), } @@ -72,8 +76,17 @@ impl TryFrom for DataType { } } -pub const DATA_TYPE_STRINGS: [&str; 7] = ["string", "hash", "set", "list", "zset", "none", "all"]; -pub const DATA_TYPE_TAG: [char; 7] = ['k', 'h', 's', 'l', 'z', 'n', 'a']; +pub const DATA_TYPE_STRINGS: [&str; 8] = [ + "string", + "hash", + "set", + "list", + "zset", + "none", + "all", + "vectorset", +]; +pub const DATA_TYPE_TAG: [char; 8] = ['k', 'h', 's', 'l', 'z', 'n', 'a', 'v']; pub fn data_type_to_string(data_type: DataType) -> &'static str { DATA_TYPE_STRINGS[data_type as usize] diff --git a/src/storage/src/logindex/types.rs b/src/storage/src/logindex/types.rs index e6da3817..5542375b 100644 --- a/src/storage/src/logindex/types.rs +++ b/src/storage/src/logindex/types.rs @@ -51,7 +51,7 @@ pub type Result = std::result::Result; /// inconsistencies across db_access.rs, event_listener.rs, and cf_tracker.rs. pub mod cf_metadata { /// Number of column families - pub const COLUMN_FAMILY_COUNT: usize = 6; + pub const COLUMN_FAMILY_COUNT: usize = 7; /// CF names as byte slices (for comparison with rocksdb CF handles) /// Note: Using &[u8] instead of &[u8; N] to allow variable-length names @@ -62,6 +62,7 @@ pub mod cf_metadata { b"list_data_cf", b"zset_data_cf", b"zset_score_cf", + b"vector_data_cf", ]; /// CF names as &str (for convenience in some contexts) @@ -72,6 +73,7 @@ pub mod cf_metadata { "list_data_cf", "zset_data_cf", "zset_score_cf", + "vector_data_cf", ]; // Compile-time guards: ensure arrays match COLUMN_FAMILY_COUNT. diff --git a/src/storage/src/redis.rs b/src/storage/src/redis.rs index 43c345af..1ab28f7a 100644 --- a/src/storage/src/redis.rs +++ b/src/storage/src/redis.rs @@ -62,13 +62,14 @@ pub enum ColumnFamilyIndex { ListsDataCF = 3, // list data ZsetsDataCF = 4, // zset data ZsetsScoreCF = 5, // zset score + VectorDataCF = 6, // vector set data } impl ColumnFamilyIndex { /// Total number of column families. /// Update this constant when adding new column families. /// This constant is used by batch.rs for validation. - pub const COUNT: usize = 6; + pub const COUNT: usize = 7; pub fn name(&self) -> &'static str { match self { @@ -78,6 +79,7 @@ impl ColumnFamilyIndex { ColumnFamilyIndex::ListsDataCF => "list_data_cf", ColumnFamilyIndex::ZsetsDataCF => "zset_data_cf", ColumnFamilyIndex::ZsetsScoreCF => "zset_score_cf", + ColumnFamilyIndex::VectorDataCF => "vector_data_cf", } } @@ -89,6 +91,7 @@ impl ColumnFamilyIndex { ColumnFamilyIndex::ZsetsDataCF | ColumnFamilyIndex::ZsetsScoreCF => { Some(DataType::ZSet) } + ColumnFamilyIndex::VectorDataCF => Some(DataType::VectorSet), ColumnFamilyIndex::MetaCF => None, } } @@ -397,6 +400,7 @@ impl Redis { ("list_data_cf", true, None), // list: bloom filter ("zset_data_cf", false, Some(16 * 1024)), // zset data: 16KB block size ("zset_score_cf", false, Some(16 * 1024)), // zset score: 16KB block size + ("vector_data_cf", true, None), // vector set: bloom filter ]; let column_families: Vec = CF_CONFIGS .iter() @@ -497,6 +501,7 @@ impl Redis { ColumnFamilyIndex::ListsDataCF, ColumnFamilyIndex::ZsetsDataCF, ColumnFamilyIndex::ZsetsScoreCF, + ColumnFamilyIndex::VectorDataCF, ] .iter() .find(|cf| cf.name() == cf_name) @@ -620,6 +625,7 @@ impl Redis { self.get_cf_handle(ColumnFamilyIndex::ListsDataCF), self.get_cf_handle(ColumnFamilyIndex::ZsetsDataCF), self.get_cf_handle(ColumnFamilyIndex::ZsetsScoreCF), + self.get_cf_handle(ColumnFamilyIndex::VectorDataCF), ]; Ok(Box::new(crate::batch::RocksBatch::new( diff --git a/src/storage/src/storage.rs b/src/storage/src/storage.rs index 465e1256..7cf93675 100644 --- a/src/storage/src/storage.rs +++ b/src/storage/src/storage.rs @@ -532,6 +532,7 @@ impl Storage { 3 => ColumnFamilyIndex::ListsDataCF, 4 => ColumnFamilyIndex::ZsetsDataCF, 5 => ColumnFamilyIndex::ZsetsScoreCF, + 6 => ColumnFamilyIndex::VectorDataCF, _ => { return Err(crate::error::Error::RedisErr { message: format!("Invalid column family index: {}", entry.cf_idx), diff --git a/src/storage/tests/redis_basic_test.rs b/src/storage/tests/redis_basic_test.rs index d3294c0a..5526ab4c 100644 --- a/src/storage/tests/redis_basic_test.rs +++ b/src/storage/tests/redis_basic_test.rs @@ -146,9 +146,9 @@ mod redis_basic_test { assert!(!redis.is_starting.load(Ordering::SeqCst)); assert!(redis.db().is_some()); - assert_eq!(redis.handles.len(), 6); + assert_eq!(redis.handles.len(), 7); - for cf_index in 0..6 { + for cf_index in 0..7 { let cf_enum = match cf_index { 0 => ColumnFamilyIndex::MetaCF, 1 => ColumnFamilyIndex::HashesDataCF, @@ -156,6 +156,7 @@ mod redis_basic_test { 3 => ColumnFamilyIndex::ListsDataCF, 4 => ColumnFamilyIndex::ZsetsDataCF, 5 => ColumnFamilyIndex::ZsetsScoreCF, + 6 => ColumnFamilyIndex::VectorDataCF, _ => panic!("Invalid CF index"), }; @@ -168,12 +169,13 @@ mod redis_basic_test { } let expected_cf_names = [ - "default", // MetaCF - "hash_data_cf", // HashesDataCF - "set_data_cf", // SetsDataCF - "list_data_cf", // ListsDataCF - "zset_data_cf", // ZsetsDataCF - "zset_score_cf", // ZsetsScoreCF + "default", // MetaCF + "hash_data_cf", // HashesDataCF + "set_data_cf", // SetsDataCF + "list_data_cf", // ListsDataCF + "zset_data_cf", // ZsetsDataCF + "zset_score_cf", // ZsetsScoreCF + "vector_data_cf", // VectorDataCF ]; for (i, expected_name) in expected_cf_names.iter().enumerate() { @@ -377,6 +379,8 @@ mod redis_basic_test { assert_eq!(ColumnFamilyIndex::ListsDataCF as usize, 3); assert_eq!(ColumnFamilyIndex::ZsetsDataCF as usize, 4); assert_eq!(ColumnFamilyIndex::ZsetsScoreCF as usize, 5); + assert_eq!(ColumnFamilyIndex::VectorDataCF as usize, 6); + assert_eq!(ColumnFamilyIndex::COUNT, 7); assert_eq!(ColumnFamilyIndex::MetaCF.name(), "default"); assert_eq!(ColumnFamilyIndex::HashesDataCF.name(), "hash_data_cf"); @@ -384,6 +388,7 @@ mod redis_basic_test { assert_eq!(ColumnFamilyIndex::ListsDataCF.name(), "list_data_cf"); assert_eq!(ColumnFamilyIndex::ZsetsDataCF.name(), "zset_data_cf"); assert_eq!(ColumnFamilyIndex::ZsetsScoreCF.name(), "zset_score_cf"); + assert_eq!(ColumnFamilyIndex::VectorDataCF.name(), "vector_data_cf"); } #[test] From 3b309f675a7452d35eea573d57e0c4f85fcb530b Mon Sep 17 00:00:00 2001 From: happy-v587 Date: Sun, 19 Jul 2026 22:54:50 +0800 Subject: [PATCH 02/32] feat(storage): add vector codecs --- src/storage/src/format_vector.rs | 352 +++++++++++++++++++++++++++++++ src/storage/src/lib.rs | 3 + src/storage/src/vector.rs | 205 ++++++++++++++++++ 3 files changed, 560 insertions(+) create mode 100644 src/storage/src/format_vector.rs create mode 100644 src/storage/src/vector.rs diff --git a/src/storage/src/format_vector.rs b/src/storage/src/format_vector.rs new file mode 100644 index 00000000..3812d19b --- /dev/null +++ b/src/storage/src/format_vector.rs @@ -0,0 +1,352 @@ +// Copyright (c) 2024-present, arana-db Community. All rights reserved. +// +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use bytes::{Buf, BufMut, BytesMut}; +use chrono::Utc; +use snafu::ensure; + +use crate::{ + DataType, + error::{InvalidFormatSnafu, Result}, + storage_define::BASE_META_VALUE_LENGTH, + vector::CanonicalVector, +}; + +pub const VECTOR_META_FORMAT: u8 = 1; +pub const VECTOR_ENCODING_FP32_LE: u8 = 1; +pub const VECTOR_METRIC_COSINE: u8 = 1; +pub const VECTOR_VALUE_MAGIC: u8 = 0x56; +pub const VECTOR_VALUE_FORMAT: u8 = 1; + +const VECTOR_META_ZERO_RESERVE_LENGTH: usize = 8; +const VECTOR_VALUE_HEADER_LENGTH: usize = 10; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct VectorMeta { + count: u64, + pub(crate) version: u64, + dimension: u32, + ctime: u64, + etime: u64, +} + +impl VectorMeta { + pub(crate) fn new(count: u64, dimension: u32) -> Self { + let now = Utc::now().timestamp_micros() as u64; + Self { + count, + version: now, + dimension, + ctime: now, + etime: 0, + } + } + + pub(crate) fn encode(&self) -> BytesMut { + let mut output = BytesMut::with_capacity(BASE_META_VALUE_LENGTH); + output.put_u8(DataType::VectorSet as u8); + output.put_u64_le(self.count); + output.put_u64_le(self.version); + output.put_u8(VECTOR_META_FORMAT); + output.put_u8(VECTOR_ENCODING_FP32_LE); + output.put_u8(VECTOR_METRIC_COSINE); + output.put_u8(0); + output.put_u32_le(self.dimension); + output.put_bytes(0, VECTOR_META_ZERO_RESERVE_LENGTH); + output.put_u64_le(self.ctime); + output.put_u64_le(self.etime); + output + } + + pub(crate) fn decode(value: &[u8]) -> Result { + ensure!( + value.len() == BASE_META_VALUE_LENGTH, + InvalidFormatSnafu { + message: format!( + "invalid vector meta length: {} != {}", + value.len(), + BASE_META_VALUE_LENGTH + ) + } + ); + + let mut reader = value; + let data_type = DataType::try_from(reader.get_u8())?; + ensure!( + data_type == DataType::VectorSet, + InvalidFormatSnafu { + message: format!("invalid vector meta data type: {data_type:?}") + } + ); + + let count = reader.get_u64_le(); + let version = reader.get_u64_le(); + let format = reader.get_u8(); + let encoding = reader.get_u8(); + let metric = reader.get_u8(); + let flags = reader.get_u8(); + let dimension = reader.get_u32_le(); + let zero_reserve = &reader[..VECTOR_META_ZERO_RESERVE_LENGTH]; + reader.advance(VECTOR_META_ZERO_RESERVE_LENGTH); + let ctime = reader.get_u64_le(); + let etime = reader.get_u64_le(); + + ensure!( + format == VECTOR_META_FORMAT, + InvalidFormatSnafu { + message: format!("unsupported vector meta format: {format}") + } + ); + ensure!( + encoding == VECTOR_ENCODING_FP32_LE, + InvalidFormatSnafu { + message: format!("unsupported vector encoding: {encoding}") + } + ); + ensure!( + metric == VECTOR_METRIC_COSINE, + InvalidFormatSnafu { + message: format!("unsupported vector metric: {metric}") + } + ); + ensure!( + flags == 0 && zero_reserve.iter().all(|byte| *byte == 0), + InvalidFormatSnafu { + message: "invalid non-zero vector meta reserve".to_string() + } + ); + ensure!( + dimension != 0, + InvalidFormatSnafu { + message: "vector dimension must not be zero".to_string() + } + ); + + Ok(Self { + count, + version, + dimension, + ctime, + etime, + }) + } + + pub(crate) fn count(&self) -> u64 { + self.count + } + + pub(crate) fn set_count(&mut self, count: u64) { + self.count = count; + } + + pub(crate) fn version(&self) -> u64 { + self.version + } + + pub(crate) fn dimension(&self) -> u32 { + self.dimension + } + + pub(crate) fn etime(&self) -> u64 { + self.etime + } + + pub(crate) fn is_stale(&self) -> bool { + self.etime != 0 && self.etime < Utc::now().timestamp_micros() as u64 + } +} + +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct VectorDataValue { + canonical: CanonicalVector, +} + +impl VectorDataValue { + pub(crate) fn from_canonical(canonical: &CanonicalVector) -> Self { + Self { + canonical: canonical.clone(), + } + } + + pub(crate) fn encode(&self) -> BytesMut { + let mut output = BytesMut::with_capacity( + VECTOR_VALUE_HEADER_LENGTH + self.canonical.normalized().len() * size_of::(), + ); + output.put_u8(VECTOR_VALUE_MAGIC); + output.put_u8(VECTOR_VALUE_FORMAT); + output.put_u32_le(self.canonical.dimension()); + output.put_f32_le(self.canonical.original_l2()); + for component in self.canonical.normalized() { + output.put_f32_le(*component); + } + output + } + + pub(crate) fn decode(value: &[u8]) -> Result { + ensure!( + value.len() >= VECTOR_VALUE_HEADER_LENGTH, + InvalidFormatSnafu { + message: format!( + "invalid vector value length: {} < {}", + value.len(), + VECTOR_VALUE_HEADER_LENGTH + ) + } + ); + + let mut reader = value; + let magic = reader.get_u8(); + let format = reader.get_u8(); + let dimension = reader.get_u32_le(); + let original_l2 = reader.get_f32_le(); + + ensure!( + magic == VECTOR_VALUE_MAGIC, + InvalidFormatSnafu { + message: format!("invalid vector value magic: {magic:#04x}") + } + ); + ensure!( + format == VECTOR_VALUE_FORMAT, + InvalidFormatSnafu { + message: format!("unsupported vector value format: {format}") + } + ); + ensure!( + dimension != 0, + InvalidFormatSnafu { + message: "vector value dimension must not be zero".to_string() + } + ); + ensure!( + original_l2.is_finite() && original_l2 > 0.0, + InvalidFormatSnafu { + message: "vector value L2 norm must be finite and positive".to_string() + } + ); + + let payload_length = (dimension as usize).checked_mul(size_of::()); + ensure!( + payload_length.is_some_and(|length| reader.len() == length), + InvalidFormatSnafu { + message: format!( + "invalid vector payload length: {} for dimension {}", + reader.len(), + dimension + ) + } + ); + + let normalized = reader + .chunks_exact(size_of::()) + .map(|chunk| f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]])) + .collect::>(); + ensure!( + normalized.iter().all(|component| component.is_finite()), + InvalidFormatSnafu { + message: "vector payload components must be finite".to_string() + } + ); + + Ok(Self { + canonical: CanonicalVector::from_normalized_parts(dimension, original_l2, normalized), + }) + } + + pub(crate) fn dimension(&self) -> u32 { + self.canonical.dimension() + } + + pub(crate) fn original_l2(&self) -> f32 { + self.canonical.original_l2() + } + + pub(crate) fn canonical(&self) -> &CanonicalVector { + &self.canonical + } +} + +#[cfg(test)] +mod tests { + use crate::{ + format_member_data_key::{MemberDataKey, ParsedMemberDataKey}, + vector::CanonicalVector, + }; + + use super::*; + + #[test] + fn member_data_key_round_trips_empty_binary_element() { + let encoded = MemberDataKey::new(b"vectors\0key", 42, b"") + .encode() + .expect("encode member key"); + let decoded = ParsedMemberDataKey::new(&encoded).expect("decode member key"); + + assert_eq!(decoded.key(), b"vectors\0key"); + assert_eq!(decoded.version(), 42); + assert_eq!(decoded.data(), b""); + } + + #[test] + fn vector_data_value_round_trips() { + let canonical = CanonicalVector::from_values(&[3.0, 4.0]).expect("valid vector"); + let encoded = VectorDataValue::from_canonical(&canonical).encode(); + let decoded = VectorDataValue::decode(&encoded).expect("decode vector value"); + + assert_eq!(decoded.dimension(), 2); + assert!((decoded.original_l2() - 5.0).abs() < 1e-6); + assert_eq!(decoded.canonical(), &canonical); + } + + #[test] + fn vector_meta_round_trips() { + let mut meta = VectorMeta::new(2, 2); + meta.version = 42; + let encoded = meta.encode(); + let decoded = VectorMeta::decode(&encoded).expect("decode vector meta"); + + assert_eq!(decoded.count(), 2); + assert_eq!(decoded.version(), 42); + assert_eq!(decoded.dimension(), 2); + assert!(!decoded.is_stale()); + } + + #[test] + fn vector_codecs_reject_malformed_bytes() { + let canonical = CanonicalVector::from_values(&[3.0, 4.0]).expect("valid vector"); + let encoded_value = VectorDataValue::from_canonical(&canonical).encode(); + + let mut bad_magic = encoded_value.clone(); + bad_magic[0] = 0; + assert!(VectorDataValue::decode(&bad_magic).is_err()); + + let mut zero_dimension = encoded_value.clone(); + zero_dimension[2..6].copy_from_slice(&0_u32.to_le_bytes()); + assert!(VectorDataValue::decode(&zero_dimension).is_err()); + + let mut non_finite_payload = encoded_value; + non_finite_payload[10..14].copy_from_slice(&f32::NAN.to_le_bytes()); + assert!(VectorDataValue::decode(&non_finite_payload).is_err()); + + let encoded_meta = VectorMeta::new(2, 2).encode(); + assert!(VectorMeta::decode(&encoded_meta[..encoded_meta.len() - 1]).is_err()); + + let mut bad_meta_format = encoded_meta; + bad_meta_format[17] = 0; + assert!(VectorMeta::decode(&bad_meta_format).is_err()); + } +} diff --git a/src/storage/src/lib.rs b/src/storage/src/lib.rs index 0b99a5d4..34a946e3 100644 --- a/src/storage/src/lib.rs +++ b/src/storage/src/lib.rs @@ -19,6 +19,8 @@ extern crate core; mod format_base_data_value; mod format_member_data_key; +mod format_vector; +pub mod vector; mod data_compaction_filter; mod durable_fs; @@ -80,3 +82,4 @@ pub use statistics::KeyStatistics; pub use storage::{BgTask, BgTaskHandler}; pub use storage_impl::BeforeOrAfter; pub use util::{safe_cleanup_test_db, unique_test_db_path}; +pub use vector::{CanonicalVector, VectorHit, VectorQuery, VectorSearchMode, VectorSearchOptions}; diff --git a/src/storage/src/vector.rs b/src/storage/src/vector.rs new file mode 100644 index 00000000..11dfb5a0 --- /dev/null +++ b/src/storage/src/vector.rs @@ -0,0 +1,205 @@ +// Copyright (c) 2024-present, arana-db Community. All rights reserved. +// +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use snafu::ensure; + +use crate::error::{InvalidArgumentSnafu, Result}; + +#[derive(Debug, Clone, PartialEq)] +pub struct CanonicalVector { + dimension: u32, + original_l2: f32, + normalized: Vec, +} + +impl CanonicalVector { + pub fn from_fp32_le(raw: &[u8]) -> Result { + ensure!( + !raw.is_empty() && raw.len().is_multiple_of(size_of::()), + InvalidArgumentSnafu { + message: "vector blob must contain one or more little-endian FP32 values" + .to_string() + } + ); + + let values = raw + .chunks_exact(size_of::()) + .map(|chunk| f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]])) + .collect::>(); + Self::from_values(&values) + } + + pub fn from_values(values: &[f32]) -> Result { + ensure!( + !values.is_empty(), + InvalidArgumentSnafu { + message: "vector must not be empty".to_string() + } + ); + ensure!( + values.len() <= u32::MAX as usize, + InvalidArgumentSnafu { + message: "vector dimension exceeds u32::MAX".to_string() + } + ); + ensure!( + values.iter().all(|value| value.is_finite()), + InvalidArgumentSnafu { + message: "vector components must be finite".to_string() + } + ); + + let norm_squared = values + .iter() + .map(|value| f64::from(*value) * f64::from(*value)) + .sum::(); + ensure!( + norm_squared.is_finite() && norm_squared > 0.0, + InvalidArgumentSnafu { + message: "vector L2 norm must be finite and greater than zero".to_string() + } + ); + + let norm = norm_squared.sqrt(); + let original_l2 = norm as f32; + ensure!( + original_l2.is_finite() && original_l2 > 0.0, + InvalidArgumentSnafu { + message: "vector L2 norm cannot be represented as FP32".to_string() + } + ); + + let normalized = values + .iter() + .map(|value| (f64::from(*value) / norm) as f32) + .collect(); + + Ok(Self { + dimension: values.len() as u32, + original_l2, + normalized, + }) + } + + pub fn dimension(&self) -> u32 { + self.dimension + } + + pub fn original_l2(&self) -> f32 { + self.original_l2 + } + + pub(crate) fn normalized(&self) -> &[f32] { + &self.normalized + } + + pub(crate) fn from_normalized_parts( + dimension: u32, + original_l2: f32, + normalized: Vec, + ) -> Self { + Self { + dimension, + original_l2, + normalized, + } + } + + pub fn score(&self, other: &Self) -> Result { + ensure!( + self.dimension == other.dimension, + InvalidArgumentSnafu { + message: format!( + "vector dimension mismatch: expected {}, got {}", + self.dimension, other.dimension + ) + } + ); + + let dot = self + .normalized + .iter() + .zip(&other.normalized) + .map(|(left, right)| f64::from(*left) * f64::from(*right)) + .sum::() + .clamp(-1.0, 1.0); + Ok(((dot + 1.0) / 2.0).clamp(0.0, 1.0)) + } + + pub fn restore(&self) -> Vec { + self.normalized + .iter() + .map(|value| f64::from(*value) * f64::from(self.original_l2)) + .collect() + } +} + +#[derive(Debug, Clone, PartialEq)] +pub enum VectorQuery { + Element(Vec), + Vector(CanonicalVector), +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum VectorSearchMode { + Approximate, + Truth, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct VectorSearchOptions { + pub count: usize, + pub mode: VectorSearchMode, +} + +#[derive(Debug, Clone, PartialEq)] +pub struct VectorHit { + pub element: Vec, + pub score: f64, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn canonical_vector_normalizes_and_restores_values() { + let vector = CanonicalVector::from_values(&[3.0, 4.0]).expect("valid vector"); + assert_eq!(vector.dimension(), 2); + assert!((vector.original_l2() - 5.0).abs() < 1e-6); + let restored = vector.restore(); + assert!((restored[0] - 3.0).abs() < 1e-6); + assert!((restored[1] - 4.0).abs() < 1e-6); + } + + #[test] + fn canonical_vector_rejects_invalid_inputs() { + assert!(CanonicalVector::from_values(&[]).is_err()); + assert!(CanonicalVector::from_values(&[0.0, 0.0]).is_err()); + assert!(CanonicalVector::from_values(&[f32::NAN]).is_err()); + assert!(CanonicalVector::from_fp32_le(&[0, 1, 2]).is_err()); + } + + #[test] + fn cosine_score_maps_to_redis_range() { + let x = CanonicalVector::from_values(&[1.0, 0.0]).expect("valid x"); + let same = CanonicalVector::from_values(&[1.0, 0.0]).expect("valid same"); + let opposite = CanonicalVector::from_values(&[-1.0, 0.0]).expect("valid opposite"); + assert!((x.score(&same).expect("score") - 1.0).abs() < 1e-12); + assert!(x.score(&opposite).expect("score").abs() < 1e-12); + } +} From e822cab2f3ecdcc6ac0ca80fc8d983b3a18c8c83 Mon Sep 17 00:00:00 2001 From: happy-v587 Date: Sun, 19 Jul 2026 23:19:21 +0800 Subject: [PATCH 03/32] feat(storage): persist vector set members --- src/storage/src/format_vector.rs | 12 +- src/storage/src/lib.rs | 1 + src/storage/src/redis_vectors.rs | 301 +++++++++++++++++++++++++ src/storage/tests/redis_vector_test.rs | 216 ++++++++++++++++++ 4 files changed, 520 insertions(+), 10 deletions(-) create mode 100644 src/storage/src/redis_vectors.rs create mode 100644 src/storage/tests/redis_vector_test.rs diff --git a/src/storage/src/format_vector.rs b/src/storage/src/format_vector.rs index 3812d19b..d59e3e89 100644 --- a/src/storage/src/format_vector.rs +++ b/src/storage/src/format_vector.rs @@ -161,10 +161,6 @@ impl VectorMeta { self.dimension } - pub(crate) fn etime(&self) -> u64 { - self.etime - } - pub(crate) fn is_stale(&self) -> bool { self.etime != 0 && self.etime < Utc::now().timestamp_micros() as u64 } @@ -184,7 +180,7 @@ impl VectorDataValue { pub(crate) fn encode(&self) -> BytesMut { let mut output = BytesMut::with_capacity( - VECTOR_VALUE_HEADER_LENGTH + self.canonical.normalized().len() * size_of::(), + VECTOR_VALUE_HEADER_LENGTH + size_of_val(self.canonical.normalized()), ); output.put_u8(VECTOR_VALUE_MAGIC); output.put_u8(VECTOR_VALUE_FORMAT); @@ -271,10 +267,6 @@ impl VectorDataValue { self.canonical.dimension() } - pub(crate) fn original_l2(&self) -> f32 { - self.canonical.original_l2() - } - pub(crate) fn canonical(&self) -> &CanonicalVector { &self.canonical } @@ -308,7 +300,7 @@ mod tests { let decoded = VectorDataValue::decode(&encoded).expect("decode vector value"); assert_eq!(decoded.dimension(), 2); - assert!((decoded.original_l2() - 5.0).abs() < 1e-6); + assert!((decoded.canonical().original_l2() - 5.0).abs() < 1e-6); assert_eq!(decoded.canonical(), &canonical); } diff --git a/src/storage/src/lib.rs b/src/storage/src/lib.rs index 34a946e3..4ca63142 100644 --- a/src/storage/src/lib.rs +++ b/src/storage/src/lib.rs @@ -54,6 +54,7 @@ mod redis_hashes; mod redis_lists; mod redis_sets; mod redis_strings; +mod redis_vectors; pub mod error; mod format_zset_score_key; diff --git a/src/storage/src/redis_vectors.rs b/src/storage/src/redis_vectors.rs new file mode 100644 index 00000000..f25e3f78 --- /dev/null +++ b/src/storage/src/redis_vectors.rs @@ -0,0 +1,301 @@ +// Copyright (c) 2024-present, arana-db Community. All rights reserved. +// +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use kstd::lock_mgr::ScopeRecordLock; +use rocksdb::ReadOptions; +use snafu::{OptionExt, ResultExt}; + +use crate::{ + CanonicalVector, ColumnFamilyIndex, DataType, Redis, Result, TypeCheckState, + error::{ + InvalidArgumentSnafu, InvalidFormatSnafu, KeyNotFoundSnafu, OptionNoneSnafu, RedisErrSnafu, + RocksSnafu, + }, + format_base_key::BaseMetaKey, + format_member_data_key::MemberDataKey, + format_vector::{VectorDataValue, VectorMeta}, +}; + +impl Redis { + pub fn is_cluster_mode(&self) -> bool { + self.append_log_fn.get().is_some() + } + + fn ensure_vector_standalone(&self) -> Result<()> { + if self.is_cluster_mode() { + return RedisErrSnafu { + message: "ERR Vector Set is not supported in cluster mode".to_string(), + } + .fail(); + } + Ok(()) + } + + fn parse_vector_meta(&self, value: &[u8]) -> Result> { + if value.is_empty() { + return Ok(None); + } + + if value[0] == DataType::VectorSet as u8 { + let meta = VectorMeta::decode(value)?; + return Ok((!meta.is_stale() && meta.count() != 0).then_some(meta)); + } + + match self.check_type_state(value, DataType::VectorSet)? { + TypeCheckState::Missing | TypeCheckState::Stale => Ok(None), + TypeCheckState::Match => VectorMeta::decode(value).map(Some), + } + } + + pub fn vadd(&self, key: &[u8], element: &[u8], vector: &CanonicalVector) -> Result { + self.ensure_vector_standalone()?; + let db = self.db.as_ref().context(OptionNoneSnafu { + message: "db is not initialized".to_string(), + })?; + let meta_cf = self + .get_cf_handle(ColumnFamilyIndex::MetaCF) + .context(OptionNoneSnafu { + message: "MetaCF is not initialized".to_string(), + })?; + let vector_cf = self + .get_cf_handle(ColumnFamilyIndex::VectorDataCF) + .context(OptionNoneSnafu { + message: "VectorDataCF is not initialized".to_string(), + })?; + + let lock_key = String::from_utf8_lossy(key); + let _lock = ScopeRecordLock::new(self.lock_mgr.as_ref(), &lock_key); + let meta_key = BaseMetaKey::new(key).encode()?; + let stored_meta = db.get_cf(&meta_cf, &meta_key).context(RocksSnafu)?; + let live_meta = stored_meta + .as_deref() + .map(|value| self.parse_vector_meta(value)) + .transpose()? + .flatten(); + + let is_new_set = live_meta.is_none(); + let mut meta = match live_meta { + Some(meta) => { + if meta.dimension() != vector.dimension() { + return InvalidArgumentSnafu { + message: format!( + "vector dimension mismatch: expected {}, got {}", + meta.dimension(), + vector.dimension() + ), + } + .fail(); + } + meta + } + None => VectorMeta::new(1, vector.dimension()), + }; + + let member_key = MemberDataKey::new(key, meta.version(), element).encode()?; + let inserted = if is_new_set { + true + } else { + db.get_cf(&vector_cf, &member_key) + .context(RocksSnafu)? + .is_none() + }; + if inserted && !is_new_set { + let count = meta.count().checked_add(1); + let Some(count) = count else { + return InvalidArgumentSnafu { + message: "vector set size overflow".to_string(), + } + .fail(); + }; + meta.set_count(count); + } + + let member_value = VectorDataValue::from_canonical(vector).encode(); + let meta_value = meta.encode(); + let mut batch = self.create_batch()?; + batch.put(ColumnFamilyIndex::VectorDataCF, &member_key, &member_value)?; + batch.put(ColumnFamilyIndex::MetaCF, &meta_key, &meta_value)?; + batch.commit()?; + Ok(inserted) + } + + pub fn vrem(&self, key: &[u8], element: &[u8]) -> Result { + self.ensure_vector_standalone()?; + let db = self.db.as_ref().context(OptionNoneSnafu { + message: "db is not initialized".to_string(), + })?; + let meta_cf = self + .get_cf_handle(ColumnFamilyIndex::MetaCF) + .context(OptionNoneSnafu { + message: "MetaCF is not initialized".to_string(), + })?; + let vector_cf = self + .get_cf_handle(ColumnFamilyIndex::VectorDataCF) + .context(OptionNoneSnafu { + message: "VectorDataCF is not initialized".to_string(), + })?; + + let lock_key = String::from_utf8_lossy(key); + let _lock = ScopeRecordLock::new(self.lock_mgr.as_ref(), &lock_key); + let meta_key = BaseMetaKey::new(key).encode()?; + let Some(meta_raw) = db.get_cf(&meta_cf, &meta_key).context(RocksSnafu)? else { + return Ok(false); + }; + let Some(mut meta) = self.parse_vector_meta(&meta_raw)? else { + return Ok(false); + }; + + let member_key = MemberDataKey::new(key, meta.version(), element).encode()?; + if db + .get_cf(&vector_cf, &member_key) + .context(RocksSnafu)? + .is_none() + { + return Ok(false); + } + + let mut batch = self.create_batch()?; + batch.delete(ColumnFamilyIndex::VectorDataCF, &member_key)?; + if meta.count() > 1 { + meta.set_count(meta.count() - 1); + let meta_value = meta.encode(); + batch.put(ColumnFamilyIndex::MetaCF, &meta_key, &meta_value)?; + } else { + batch.delete(ColumnFamilyIndex::MetaCF, &meta_key)?; + } + batch.commit()?; + Ok(true) + } + + pub fn vcard(&self, key: &[u8]) -> Result { + self.ensure_vector_standalone()?; + Ok(self.read_vector_meta(key)?.map_or(0, |meta| meta.count())) + } + + pub fn vdim(&self, key: &[u8]) -> Result { + self.ensure_vector_standalone()?; + match self.read_vector_meta(key)? { + Some(meta) => Ok(meta.dimension()), + None => KeyNotFoundSnafu { + key: String::from_utf8_lossy(key).to_string(), + } + .fail(), + } + } + + pub fn vemb(&self, key: &[u8], element: &[u8]) -> Result>> { + self.ensure_vector_standalone()?; + let db = self.db.as_ref().context(OptionNoneSnafu { + message: "db is not initialized".to_string(), + })?; + let meta_cf = self + .get_cf_handle(ColumnFamilyIndex::MetaCF) + .context(OptionNoneSnafu { + message: "MetaCF is not initialized".to_string(), + })?; + let vector_cf = self + .get_cf_handle(ColumnFamilyIndex::VectorDataCF) + .context(OptionNoneSnafu { + message: "VectorDataCF is not initialized".to_string(), + })?; + let snapshot = db.snapshot(); + let mut read_options = ReadOptions::default(); + read_options.set_snapshot(&snapshot); + + let meta_key = BaseMetaKey::new(key).encode()?; + let Some(meta_raw) = db + .get_cf_opt(&meta_cf, &meta_key, &read_options) + .context(RocksSnafu)? + else { + return Ok(None); + }; + let Some(meta) = self.parse_vector_meta(&meta_raw)? else { + return Ok(None); + }; + let member_key = MemberDataKey::new(key, meta.version(), element).encode()?; + let Some(value_raw) = db + .get_cf_opt(&vector_cf, &member_key, &read_options) + .context(RocksSnafu)? + else { + return Ok(None); + }; + let value = VectorDataValue::decode(&value_raw)?; + if value.dimension() != meta.dimension() { + return InvalidFormatSnafu { + message: format!( + "vector member dimension {} does not match meta dimension {}", + value.dimension(), + meta.dimension() + ), + } + .fail(); + } + Ok(Some(value.canonical().restore())) + } + + pub fn vismember(&self, key: &[u8], element: &[u8]) -> Result { + self.ensure_vector_standalone()?; + let db = self.db.as_ref().context(OptionNoneSnafu { + message: "db is not initialized".to_string(), + })?; + let meta_cf = self + .get_cf_handle(ColumnFamilyIndex::MetaCF) + .context(OptionNoneSnafu { + message: "MetaCF is not initialized".to_string(), + })?; + let vector_cf = self + .get_cf_handle(ColumnFamilyIndex::VectorDataCF) + .context(OptionNoneSnafu { + message: "VectorDataCF is not initialized".to_string(), + })?; + let snapshot = db.snapshot(); + let mut read_options = ReadOptions::default(); + read_options.set_snapshot(&snapshot); + + let meta_key = BaseMetaKey::new(key).encode()?; + let Some(meta_raw) = db + .get_cf_opt(&meta_cf, &meta_key, &read_options) + .context(RocksSnafu)? + else { + return Ok(false); + }; + let Some(meta) = self.parse_vector_meta(&meta_raw)? else { + return Ok(false); + }; + let member_key = MemberDataKey::new(key, meta.version(), element).encode()?; + Ok(db + .get_cf_opt(&vector_cf, &member_key, &read_options) + .context(RocksSnafu)? + .is_some()) + } + + fn read_vector_meta(&self, key: &[u8]) -> Result> { + let db = self.db.as_ref().context(OptionNoneSnafu { + message: "db is not initialized".to_string(), + })?; + let meta_cf = self + .get_cf_handle(ColumnFamilyIndex::MetaCF) + .context(OptionNoneSnafu { + message: "MetaCF is not initialized".to_string(), + })?; + let meta_key = BaseMetaKey::new(key).encode()?; + match db.get_cf(&meta_cf, &meta_key).context(RocksSnafu)? { + Some(value) => self.parse_vector_meta(&value), + None => Ok(None), + } + } +} diff --git a/src/storage/tests/redis_vector_test.rs b/src/storage/tests/redis_vector_test.rs new file mode 100644 index 00000000..30630ed1 --- /dev/null +++ b/src/storage/tests/redis_vector_test.rs @@ -0,0 +1,216 @@ +// Copyright (c) 2024-present, arana-db Community. All rights reserved. +// +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#![allow(clippy::unwrap_used)] + +use std::{path::PathBuf, sync::Arc}; + +use kstd::lock_mgr::LockMgr; +use rocksdb::{IteratorMode, ReadOptions}; +use storage::{ + BaseMetaKey, BgTaskHandler, CanonicalVector, ColumnFamilyIndex, Redis, StorageOptions, + safe_cleanup_test_db, unique_test_db_path, +}; + +fn open_redis(path: &PathBuf) -> Redis { + let storage_options = Arc::new(StorageOptions::default()); + let (bg_task_handler, _) = BgTaskHandler::new(); + let lock_mgr = Arc::new(LockMgr::new(1000)); + let mut redis = Redis::new(storage_options, 1, Arc::new(bg_task_handler), lock_mgr); + redis + .open(path.to_str().expect("test path must be UTF-8")) + .expect("open test db"); + redis +} + +fn with_redis(test: impl FnOnce(&Redis)) { + let path = unique_test_db_path(); + safe_cleanup_test_db(&path); + let redis = open_redis(&path); + + test(&redis); + + redis.set_need_close(true); + drop(redis); + safe_cleanup_test_db(&path); +} + +#[test] +fn test_vadd_create_update_and_dimension_guard() { + with_redis(|redis| { + let a = CanonicalVector::from_values(&[1.0, 0.0]).expect("valid vector"); + let b = CanonicalVector::from_values(&[0.0, 1.0]).expect("valid vector"); + let wrong_dimension = CanonicalVector::from_values(&[1.0, 0.0, 0.0]).expect("valid vector"); + + assert!(redis.vadd(b"vectors", b"a", &a).expect("insert a")); + assert!(redis.vadd(b"vectors", b"b", &b).expect("insert b")); + assert!(!redis.vadd(b"vectors", b"a", &b).expect("update a")); + assert_eq!(redis.vcard(b"vectors").expect("card"), 2); + assert_eq!(redis.vdim(b"vectors").expect("dim"), 2); + assert_eq!( + redis.vemb(b"vectors", b"a").expect("emb"), + Some(vec![0.0, 1.0]) + ); + + assert!(redis.vadd(b"vectors", b"a", &wrong_dimension).is_err()); + assert_eq!(redis.vcard(b"vectors").expect("card after error"), 2); + assert_eq!( + redis.vemb(b"vectors", b"a").expect("emb after error"), + Some(vec![0.0, 1.0]) + ); + }); +} + +#[test] +fn test_vadd_is_binary_safe_and_accepts_empty_element() { + with_redis(|redis| { + let vector = CanonicalVector::from_values(&[1.0, 0.0]).expect("valid vector"); + + assert!( + redis + .vadd(b"vectors\0key", b"\0binary", &vector) + .expect("binary member") + ); + assert!( + redis + .vismember(b"vectors\0key", b"\0binary") + .expect("binary membership") + ); + assert!( + redis + .vadd(b"empty-element", b"", &vector) + .expect("empty member") + ); + assert!( + redis + .vismember(b"empty-element", b"") + .expect("empty membership") + ); + }); +} + +#[test] +fn test_vcard_vdim_vemb_and_vismember_missing_semantics() { + with_redis(|redis| { + assert_eq!(redis.vcard(b"missing").expect("missing card"), 0); + assert!(redis.vdim(b"missing").is_err()); + assert_eq!( + redis.vemb(b"missing", b"member").expect("missing emb"), + None + ); + assert!( + !redis + .vismember(b"missing", b"member") + .expect("missing membership") + ); + + let vector = CanonicalVector::from_values(&[3.0, 4.0]).expect("valid vector"); + assert!(redis.vadd(b"vectors", b"member", &vector).expect("insert")); + assert_eq!(redis.vcard(b"vectors").expect("card"), 1); + assert_eq!(redis.vdim(b"vectors").expect("dimension"), 2); + assert!(redis.vismember(b"vectors", b"member").expect("membership")); + let restored = redis + .vemb(b"vectors", b"member") + .expect("embedding") + .expect("member exists"); + assert!((restored[0] - 3.0).abs() < 1e-6); + assert!((restored[1] - 4.0).abs() < 1e-6); + assert_eq!(redis.vemb(b"vectors", b"absent").expect("absent emb"), None); + }); +} + +#[test] +fn test_vrem_deletes_last_member_and_meta() { + with_redis(|redis| { + let vector = CanonicalVector::from_values(&[1.0, 0.0]).expect("valid vector"); + assert!(redis.vadd(b"vectors", b"a", &vector).expect("insert a")); + assert!(redis.vadd(b"vectors", b"b", &vector).expect("insert b")); + + assert!(!redis.vrem(b"vectors", b"absent").expect("remove absent")); + assert!(redis.vrem(b"vectors", b"a").expect("remove a")); + assert_eq!(redis.vcard(b"vectors").expect("card"), 1); + assert!(redis.vrem(b"vectors", b"b").expect("remove b")); + assert_eq!(redis.vcard(b"vectors").expect("missing card"), 0); + assert!(redis.get_key_type(b"vectors").is_err()); + }); +} + +#[test] +fn test_vector_commands_return_wrongtype_for_string_key() { + with_redis(|redis| { + let key = b"string-key"; + let vector = CanonicalVector::from_values(&[1.0, 0.0]).expect("valid vector"); + redis.set(key, b"value").expect("set string"); + + let errors = [ + redis.vadd(key, b"member", &vector).unwrap_err(), + redis.vrem(key, b"member").unwrap_err(), + redis.vcard(key).unwrap_err(), + redis.vdim(key).unwrap_err(), + redis.vemb(key, b"member").unwrap_err(), + redis.vismember(key, b"member").unwrap_err(), + ]; + assert!( + errors + .iter() + .all(|error| error.to_string().contains("WRONGTYPE")) + ); + }); +} + +#[test] +fn test_vector_meta_and_member_are_committed_together() { + let path = unique_test_db_path(); + safe_cleanup_test_db(&path); + let redis = open_redis(&path); + let vector = CanonicalVector::from_values(&[3.0, 4.0]).expect("valid vector"); + assert!(redis.vadd(b"vectors", b"member", &vector).expect("insert")); + + { + let db = redis.db.as_ref().expect("db is initialized"); + let meta_cf = redis + .get_cf_handle(ColumnFamilyIndex::MetaCF) + .expect("MetaCF exists"); + let vector_cf = redis + .get_cf_handle(ColumnFamilyIndex::VectorDataCF) + .expect("VectorDataCF exists"); + let snapshot = db.snapshot(); + let mut meta_options = ReadOptions::default(); + meta_options.set_snapshot(&snapshot); + let mut vector_options = ReadOptions::default(); + vector_options.set_snapshot(&snapshot); + let meta_key = BaseMetaKey::new(b"vectors").encode().expect("meta key"); + + assert!( + db.get_cf_opt(&meta_cf, &meta_key, &meta_options) + .expect("read meta") + .is_some() + ); + assert_eq!( + db.iterator_cf_opt(&vector_cf, vector_options, IteratorMode::Start) + .map(|entry| entry.expect("read vector entry")) + .count(), + 1 + ); + } + assert_eq!(redis.vcard(b"vectors").expect("card"), 1); + assert!(redis.vismember(b"vectors", b"member").expect("membership")); + + redis.set_need_close(true); + drop(redis); + safe_cleanup_test_db(&path); +} From 6837573c43dd3e0b25fe05ed4be88e6327b2edb6 Mon Sep 17 00:00:00 2001 From: happy-v587 Date: Sun, 19 Jul 2026 23:28:15 +0800 Subject: [PATCH 04/32] feat(storage): add flat vector similarity search --- src/storage/src/redis_vectors.rs | 170 ++++++++++++++++++++++++- src/storage/tests/redis_vector_test.rs | 153 +++++++++++++++++++++- 2 files changed, 320 insertions(+), 3 deletions(-) diff --git a/src/storage/src/redis_vectors.rs b/src/storage/src/redis_vectors.rs index f25e3f78..67441a65 100644 --- a/src/storage/src/redis_vectors.rs +++ b/src/storage/src/redis_vectors.rs @@ -15,12 +15,15 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::{cmp::Ordering, collections::BinaryHeap}; + use kstd::lock_mgr::ScopeRecordLock; -use rocksdb::ReadOptions; +use rocksdb::{Direction, IteratorMode, ReadOptions}; use snafu::{OptionExt, ResultExt}; use crate::{ - CanonicalVector, ColumnFamilyIndex, DataType, Redis, Result, TypeCheckState, + CanonicalVector, ColumnFamilyIndex, DataType, Redis, Result, TypeCheckState, VectorHit, + VectorQuery, VectorSearchOptions, error::{ InvalidArgumentSnafu, InvalidFormatSnafu, KeyNotFoundSnafu, OptionNoneSnafu, RedisErrSnafu, RocksSnafu, @@ -28,8 +31,38 @@ use crate::{ format_base_key::BaseMetaKey, format_member_data_key::MemberDataKey, format_vector::{VectorDataValue, VectorMeta}, + storage_define::SUFFIX_RESERVE_LENGTH, }; +#[derive(Debug)] +struct HeapHit { + element: Vec, + score: f64, +} + +impl PartialEq for HeapHit { + fn eq(&self, other: &Self) -> bool { + self.score.total_cmp(&other.score) == Ordering::Equal && self.element == other.element + } +} + +impl Eq for HeapHit {} + +impl PartialOrd for HeapHit { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for HeapHit { + fn cmp(&self, other: &Self) -> Ordering { + other + .score + .total_cmp(&self.score) + .then_with(|| self.element.cmp(&other.element)) + } +} + impl Redis { pub fn is_cluster_mode(&self) -> bool { self.append_log_fn.get().is_some() @@ -283,6 +316,139 @@ impl Redis { .is_some()) } + pub fn vsim( + &self, + key: &[u8], + query: VectorQuery, + options: VectorSearchOptions, + ) -> Result> { + self.ensure_vector_standalone()?; + if options.count == 0 { + return InvalidArgumentSnafu { + message: "vector search count must be greater than zero".to_string(), + } + .fail(); + } + + let db = self.db.as_ref().context(OptionNoneSnafu { + message: "db is not initialized".to_string(), + })?; + let meta_cf = self + .get_cf_handle(ColumnFamilyIndex::MetaCF) + .context(OptionNoneSnafu { + message: "MetaCF is not initialized".to_string(), + })?; + let vector_cf = self + .get_cf_handle(ColumnFamilyIndex::VectorDataCF) + .context(OptionNoneSnafu { + message: "VectorDataCF is not initialized".to_string(), + })?; + let snapshot = db.snapshot(); + let mut point_read_options = ReadOptions::default(); + point_read_options.set_snapshot(&snapshot); + + let meta_key = BaseMetaKey::new(key).encode()?; + let Some(meta_raw) = db + .get_cf_opt(&meta_cf, &meta_key, &point_read_options) + .context(RocksSnafu)? + else { + return Ok(Vec::new()); + }; + let Some(meta) = self.parse_vector_meta(&meta_raw)? else { + return Ok(Vec::new()); + }; + + let query_vector = match query { + VectorQuery::Element(element) => { + let query_key = MemberDataKey::new(key, meta.version(), &element).encode()?; + let Some(query_raw) = db + .get_cf_opt(&vector_cf, &query_key, &point_read_options) + .context(RocksSnafu)? + else { + return KeyNotFoundSnafu { + key: String::from_utf8_lossy(&element).to_string(), + } + .fail(); + }; + VectorDataValue::decode(&query_raw)?.canonical().clone() + } + VectorQuery::Vector(vector) => vector, + }; + if query_vector.dimension() != meta.dimension() { + return InvalidArgumentSnafu { + message: format!( + "vector dimension mismatch: expected {}, got {}", + meta.dimension(), + query_vector.dimension() + ), + } + .fail(); + } + + let prefix = MemberDataKey::new(key, meta.version(), b"").encode_seek_key()?; + let mut scan_options = ReadOptions::default(); + scan_options.set_snapshot(&snapshot); + let iterator = db.iterator_cf_opt( + &vector_cf, + scan_options, + IteratorMode::From(&prefix, Direction::Forward), + ); + let mut heap = BinaryHeap::new(); + + for entry in iterator { + let (encoded_key, encoded_value) = entry.context(RocksSnafu)?; + if !encoded_key.starts_with(&prefix) { + break; + } + if encoded_key.len() < prefix.len() + SUFFIX_RESERVE_LENGTH { + return InvalidFormatSnafu { + message: "vector member key is shorter than its generation prefix".to_string(), + } + .fail(); + } + + let element_end = encoded_key.len() - SUFFIX_RESERVE_LENGTH; + let element = encoded_key[prefix.len()..element_end].to_vec(); + let value = VectorDataValue::decode(&encoded_value)?; + if value.dimension() != meta.dimension() { + return InvalidFormatSnafu { + message: format!( + "vector member dimension {} does not match meta dimension {}", + value.dimension(), + meta.dimension() + ), + } + .fail(); + } + let hit = HeapHit { + element, + score: query_vector.score(value.canonical())?, + }; + + if heap.len() < options.count { + heap.push(hit); + } else if heap.peek().is_some_and(|worst| hit < *worst) { + heap.pop(); + heap.push(hit); + } + } + + let mut hits = heap + .into_iter() + .map(|hit| VectorHit { + element: hit.element, + score: hit.score, + }) + .collect::>(); + hits.sort_by(|left, right| { + right + .score + .total_cmp(&left.score) + .then_with(|| left.element.cmp(&right.element)) + }); + Ok(hits) + } + fn read_vector_meta(&self, key: &[u8]) -> Result> { let db = self.db.as_ref().context(OptionNoneSnafu { message: "db is not initialized".to_string(), diff --git a/src/storage/tests/redis_vector_test.rs b/src/storage/tests/redis_vector_test.rs index 30630ed1..eb693477 100644 --- a/src/storage/tests/redis_vector_test.rs +++ b/src/storage/tests/redis_vector_test.rs @@ -23,7 +23,7 @@ use kstd::lock_mgr::LockMgr; use rocksdb::{IteratorMode, ReadOptions}; use storage::{ BaseMetaKey, BgTaskHandler, CanonicalVector, ColumnFamilyIndex, Redis, StorageOptions, - safe_cleanup_test_db, unique_test_db_path, + VectorQuery, VectorSearchMode, VectorSearchOptions, safe_cleanup_test_db, unique_test_db_path, }; fn open_redis(path: &PathBuf) -> Redis { @@ -49,6 +49,21 @@ fn with_redis(test: impl FnOnce(&Redis)) { safe_cleanup_test_db(&path); } +fn populate_search_vectors(redis: &Redis) -> CanonicalVector { + let x = CanonicalVector::from_values(&[1.0, 0.0]).expect("x"); + let y = CanonicalVector::from_values(&[0.0, 1.0]).expect("y"); + let neg_x = CanonicalVector::from_values(&[-1.0, 0.0]).expect("negative x"); + redis.vadd(b"search", b"b", &y).expect("insert b"); + redis.vadd(b"search", b"a", &y).expect("insert a"); + redis.vadd(b"search", b"x", &x).expect("insert x"); + redis.vadd(b"search", b"neg", &neg_x).expect("insert neg"); + x +} + +fn search_options(count: usize, mode: VectorSearchMode) -> VectorSearchOptions { + VectorSearchOptions { count, mode } +} + #[test] fn test_vadd_create_update_and_dimension_guard() { with_redis(|redis| { @@ -214,3 +229,139 @@ fn test_vector_meta_and_member_are_committed_together() { drop(redis); safe_cleanup_test_db(&path); } + +#[test] +fn test_vsim_direct_vector_returns_exact_top_k() { + with_redis(|redis| { + let x = populate_search_vectors(redis); + let hits = redis + .vsim( + b"search", + VectorQuery::Vector(x), + search_options(3, VectorSearchMode::Approximate), + ) + .expect("search"); + + assert_eq!( + hits.iter() + .map(|hit| hit.element.as_slice()) + .collect::>(), + vec![b"x".as_slice(), b"a".as_slice(), b"b".as_slice()] + ); + assert!((hits[0].score - 1.0).abs() < 1e-12); + assert!((hits[1].score - 0.5).abs() < 1e-12); + assert!((hits[2].score - 0.5).abs() < 1e-12); + }); +} + +#[test] +fn test_vsim_ele_uses_stored_member_as_query() { + with_redis(|redis| { + populate_search_vectors(redis); + let hits = redis + .vsim( + b"search", + VectorQuery::Element(b"x".to_vec()), + search_options(2, VectorSearchMode::Approximate), + ) + .expect("element search"); + + assert_eq!(hits[0].element, b"x"); + assert_eq!(hits[1].element, b"a"); + }); +} + +#[test] +fn test_vsim_stable_tie_breaks_by_raw_element_bytes() { + with_redis(|redis| { + let x = populate_search_vectors(redis); + let hits = redis + .vsim( + b"search", + VectorQuery::Vector(x), + search_options(4, VectorSearchMode::Approximate), + ) + .expect("search"); + + assert_eq!(hits[1].element, b"a"); + assert_eq!(hits[2].element, b"b"); + assert_eq!(hits[1].score, hits[2].score); + }); +} + +#[test] +fn test_vsim_truth_matches_approximate_in_phase_one() { + with_redis(|redis| { + let x = populate_search_vectors(redis); + let approximate = redis + .vsim( + b"search", + VectorQuery::Vector(x.clone()), + search_options(4, VectorSearchMode::Approximate), + ) + .expect("approximate search"); + let truth = redis + .vsim( + b"search", + VectorQuery::Vector(x), + search_options(4, VectorSearchMode::Truth), + ) + .expect("truth search"); + + assert_eq!(truth, approximate); + }); +} + +#[test] +fn test_vsim_missing_key_is_empty_and_missing_ele_is_error() { + with_redis(|redis| { + let x = CanonicalVector::from_values(&[1.0, 0.0]).expect("x"); + assert!( + redis + .vsim( + b"missing", + VectorQuery::Vector(x), + search_options(3, VectorSearchMode::Approximate), + ) + .expect("missing search") + .is_empty() + ); + + populate_search_vectors(redis); + assert!( + redis + .vsim( + b"search", + VectorQuery::Element(b"missing".to_vec()), + search_options(3, VectorSearchMode::Approximate), + ) + .is_err() + ); + }); +} + +#[test] +fn test_vsim_rejects_query_dimension_mismatch() { + with_redis(|redis| { + let x = populate_search_vectors(redis); + let wrong_dimension = CanonicalVector::from_values(&[1.0, 0.0, 0.0]).expect("3d"); + assert!( + redis + .vsim( + b"search", + VectorQuery::Vector(wrong_dimension), + search_options(3, VectorSearchMode::Approximate), + ) + .is_err() + ); + assert!( + redis + .vsim( + b"search", + VectorQuery::Vector(x), + search_options(0, VectorSearchMode::Approximate), + ) + .is_err() + ); + }); +} From 015252005f1f7eaa2eab11c07a7ad41096db2196 Mon Sep 17 00:00:00 2001 From: happy-v587 Date: Sun, 19 Jul 2026 23:44:09 +0800 Subject: [PATCH 05/32] feat(storage): route and clean up vector sets --- src/storage/src/data_compaction_filter.rs | 32 ++++- src/storage/src/meta_compaction_filter.rs | 36 +++-- src/storage/src/redis.rs | 2 +- src/storage/src/redis_strings.rs | 4 +- src/storage/src/storage_impl.rs | 50 +++++++ src/storage/tests/redis_vector_test.rs | 166 ++++++++++++++++++++++ 6 files changed, 276 insertions(+), 14 deletions(-) diff --git a/src/storage/src/data_compaction_filter.rs b/src/storage/src/data_compaction_filter.rs index 325f6d86..b087c423 100644 --- a/src/storage/src/data_compaction_filter.rs +++ b/src/storage/src/data_compaction_filter.rs @@ -263,9 +263,11 @@ impl DataCompactionFilter { DataType::List => ParsedListsMetaValue::new(raw) .ok() .map(|m| (m.version(), m.etime())), - DataType::Hash | DataType::Set | DataType::ZSet => ParsedBaseMetaValue::new(raw) - .ok() - .map(|m| (m.version(), m.etime())), + DataType::Hash | DataType::Set | DataType::ZSet | DataType::VectorSet => { + ParsedBaseMetaValue::new(raw) + .ok() + .map(|m| (m.version(), m.etime())) + } _ => None, } } @@ -427,7 +429,7 @@ mod tests { meta_value.set_etime(etime); db.put(&meta_key, meta_value.encode()).unwrap(); } - DataType::Hash | DataType::Set | DataType::ZSet => { + DataType::Hash | DataType::Set | DataType::ZSet | DataType::VectorSet => { let mut meta_value = BaseMetaValue::new(bytes::Bytes::copy_from_slice(&1u64.to_le_bytes())); meta_value.inner.data_type = data_type; @@ -552,6 +554,28 @@ mod tests { assert!(matches!(decision, CompactionDecision::Keep)); } + #[test] + fn test_keeps_data_if_meta_is_valid_for_vectorset() { + let path = unique_test_db_path(); + let (db_cell, db) = setup_db_for_filter_test(&path); + + put_meta(&db, b"vector_key", DataType::VectorSet, 1, 0); + + std::thread::sleep(std::time::Duration::from_millis(10)); + + let mut factory = DataCompactionFilterFactory::new(db_cell, DataType::VectorSet); + let context = CompactionFilterContext { + is_full_compaction: false, + is_manual_compaction: false, + }; + let mut filter = factory.create(context); + + let data_key = encode_data_key(b"vector_key", 1); + let decision = filter.filter(0, &data_key, b""); + + assert!(matches!(decision, CompactionDecision::Keep)); + } + #[test] fn test_keeps_data_if_meta_is_valid_for_list() { let path = unique_test_db_path(); diff --git a/src/storage/src/meta_compaction_filter.rs b/src/storage/src/meta_compaction_filter.rs index 62a9efce..8a5a7a12 100644 --- a/src/storage/src/meta_compaction_filter.rs +++ b/src/storage/src/meta_compaction_filter.rs @@ -94,14 +94,16 @@ impl CompactionFilter for MetaCompactionFilter { || (m.count() == 0 && m.version() < cur_time) }, ), - DataType::Hash | DataType::Set | DataType::ZSet => parse_and_check( - value, - |v| ParsedBaseMetaValue::new(v), - |m| { - (m.etime() != 0 && m.etime() < cur_time) - || (m.count() == 0 && m.version() < cur_time) - }, - ), + DataType::Hash | DataType::Set | DataType::ZSet | DataType::VectorSet => { + parse_and_check( + value, + |v| ParsedBaseMetaValue::new(v), + |m| { + (m.etime() != 0 && m.etime() < cur_time) + || (m.count() == 0 && m.version() < cur_time) + }, + ) + } _ => CompactionDecision::Keep, } } @@ -359,6 +361,24 @@ mod tests { assert!(matches!(decision, CompactionDecision::Remove)); } + #[test] + fn test_vectorset_meta_value_expired() { + let mut filter = MetaCompactionFilter; + let key = BaseKey::new(b"test_vector_key"); + let encoded_key = key.encode().unwrap(); + + let mut vector_meta = + BaseMetaValue::new(bytes::Bytes::copy_from_slice(&1u64.to_le_bytes())); + vector_meta.inner.data_type = DataType::VectorSet; + let past_time = Utc::now().timestamp_micros() as u64 - 3_600_000_000; + vector_meta.set_etime(past_time); + vector_meta.set_version(past_time); + let encoded_value = vector_meta.encode(); + + let decision = filter.filter(0, &encoded_key, &encoded_value); + assert!(matches!(decision, CompactionDecision::Remove)); + } + #[test] fn test_invalid_key() { let mut filter = MetaCompactionFilter; diff --git a/src/storage/src/redis.rs b/src/storage/src/redis.rs index 1ab28f7a..d384caed 100644 --- a/src/storage/src/redis.rs +++ b/src/storage/src/redis.rs @@ -915,7 +915,7 @@ impl Redis { } Ok(etime < now) } - DataType::Hash | DataType::Set | DataType::ZSet => { + DataType::Hash | DataType::Set | DataType::ZSet | DataType::VectorSet => { // | type(1B) | count(8B) | version(8B) | reserve(16B) | ctime(8B) | etime(8B) | let count_offset = TYPE_LENGTH; let count_bytes = &val_raw[count_offset..count_offset + 8]; diff --git a/src/storage/src/redis_strings.rs b/src/storage/src/redis_strings.rs index fdfdc1ae..daa49000 100644 --- a/src/storage/src/redis_strings.rs +++ b/src/storage/src/redis_strings.rs @@ -106,7 +106,7 @@ impl Redis { batch.put(ColumnFamilyIndex::MetaCF, &base_meta_key, parsed.value())?; batch.commit()?; } - DataType::Hash | DataType::Set | DataType::ZSet => { + DataType::Hash | DataType::Set | DataType::ZSet | DataType::VectorSet => { let mut parsed = crate::format_base_meta_value::ParsedBaseMetaValue::new(&value[..])?; parsed.set_etime(etime); @@ -2045,6 +2045,7 @@ impl Redis { ColumnFamilyIndex::ListsDataCF, ColumnFamilyIndex::ZsetsDataCF, ColumnFamilyIndex::ZsetsScoreCF, + ColumnFamilyIndex::VectorDataCF, ] { if let Some(cf) = self.get_cf_handle(cf_index) { // Prefix-scan data CF and delete all derived keys @@ -2143,6 +2144,7 @@ impl Redis { ColumnFamilyIndex::ListsDataCF, ColumnFamilyIndex::ZsetsDataCF, ColumnFamilyIndex::ZsetsScoreCF, + ColumnFamilyIndex::VectorDataCF, ]; for cf_index in all_cf_indexes { diff --git a/src/storage/src/storage_impl.rs b/src/storage/src/storage_impl.rs index 45a86f2c..9f2e1d3a 100644 --- a/src/storage/src/storage_impl.rs +++ b/src/storage/src/storage_impl.rs @@ -19,6 +19,7 @@ use crate::error::{Error, InvalidArgumentSnafu, Result}; use crate::format_zset_score_key::ZsetScoreMember; use crate::slot_indexer::key_to_slot_id; use crate::storage::Storage; +use crate::{CanonicalVector, VectorHit, VectorQuery, VectorSearchOptions}; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum BeforeOrAfter { @@ -752,6 +753,55 @@ impl Storage { Ok(None) // No keys found in any instance } + // Vector Set Commands Implementation + + pub fn vadd(&self, key: &[u8], element: &[u8], vector: &CanonicalVector) -> Result { + let slot_id = key_to_slot_id(key); + let instance_id = self.slot_indexer.get_instance_id(slot_id); + self.insts[instance_id].vadd(key, element, vector) + } + + pub fn vsim( + &self, + key: &[u8], + query: VectorQuery, + options: VectorSearchOptions, + ) -> Result> { + let slot_id = key_to_slot_id(key); + let instance_id = self.slot_indexer.get_instance_id(slot_id); + self.insts[instance_id].vsim(key, query, options) + } + + pub fn vrem(&self, key: &[u8], element: &[u8]) -> Result { + let slot_id = key_to_slot_id(key); + let instance_id = self.slot_indexer.get_instance_id(slot_id); + self.insts[instance_id].vrem(key, element) + } + + pub fn vcard(&self, key: &[u8]) -> Result { + let slot_id = key_to_slot_id(key); + let instance_id = self.slot_indexer.get_instance_id(slot_id); + self.insts[instance_id].vcard(key) + } + + pub fn vdim(&self, key: &[u8]) -> Result { + let slot_id = key_to_slot_id(key); + let instance_id = self.slot_indexer.get_instance_id(slot_id); + self.insts[instance_id].vdim(key) + } + + pub fn vemb(&self, key: &[u8], element: &[u8]) -> Result>> { + let slot_id = key_to_slot_id(key); + let instance_id = self.slot_indexer.get_instance_id(slot_id); + self.insts[instance_id].vemb(key, element) + } + + pub fn vismember(&self, key: &[u8], element: &[u8]) -> Result { + let slot_id = key_to_slot_id(key); + let instance_id = self.slot_indexer.get_instance_id(slot_id); + self.insts[instance_id].vismember(key, element) + } + // Sets Commands Implementation // Add the specified members to the set stored at key. Specified members that diff --git a/src/storage/tests/redis_vector_test.rs b/src/storage/tests/redis_vector_test.rs index eb693477..89676164 100644 --- a/src/storage/tests/redis_vector_test.rs +++ b/src/storage/tests/redis_vector_test.rs @@ -25,6 +25,7 @@ use storage::{ BaseMetaKey, BgTaskHandler, CanonicalVector, ColumnFamilyIndex, Redis, StorageOptions, VectorQuery, VectorSearchMode, VectorSearchOptions, safe_cleanup_test_db, unique_test_db_path, }; +use storage::{slot_indexer::key_to_slot_id, storage::Storage}; fn open_redis(path: &PathBuf) -> Redis { let storage_options = Arc::new(StorageOptions::default()); @@ -64,6 +65,14 @@ fn search_options(count: usize, mode: VectorSearchMode) -> VectorSearchOptions { VectorSearchOptions { count, mode } } +fn count_cf_entries(redis: &Redis, cf_index: ColumnFamilyIndex) -> usize { + let db = redis.db.as_ref().expect("db is initialized"); + let cf = redis.get_cf_handle(cf_index).expect("column family exists"); + db.iterator_cf(&cf, IteratorMode::Start) + .map(|entry| entry.expect("read column family entry")) + .count() +} + #[test] fn test_vadd_create_update_and_dimension_guard() { with_redis(|redis| { @@ -365,3 +374,160 @@ fn test_vsim_rejects_query_dimension_mismatch() { ); }); } + +#[tokio::test] +async fn test_storage_routes_all_members_of_one_vectorset_to_one_instance() { + let temp = tempfile::tempdir().expect("temp dir"); + let mut storage = Storage::new(3, 0); + let _receiver = storage + .open(Arc::new(StorageOptions::default()), temp.path()) + .expect("open storage"); + let key = b"routed-vectors"; + let x = CanonicalVector::from_values(&[1.0, 0.0]).expect("x"); + let y = CanonicalVector::from_values(&[0.0, 1.0]).expect("y"); + + assert!(storage.vadd(key, b"x", &x).expect("insert x")); + assert!(storage.vadd(key, b"y", &y).expect("insert y")); + assert_eq!(storage.vcard(key).expect("card"), 2); + assert_eq!(storage.vdim(key).expect("dimension"), 2); + assert!(storage.vismember(key, b"x").expect("membership")); + assert_eq!( + storage.vemb(key, b"x").expect("embedding"), + Some(vec![1.0, 0.0]) + ); + assert_eq!( + storage + .vsim( + key, + VectorQuery::Vector(x), + search_options(1, VectorSearchMode::Approximate), + ) + .expect("search")[0] + .element, + b"x" + ); + assert!(storage.vrem(key, b"y").expect("remove y")); + + let slot_id = key_to_slot_id(key); + let selected = storage.slot_indexer.get_instance_id(slot_id); + let meta_key = BaseMetaKey::new(key).encode().expect("meta key"); + for (instance_id, redis) in storage.insts.iter().enumerate() { + let db = redis.db.as_ref().expect("db is initialized"); + let meta_cf = redis + .get_cf_handle(ColumnFamilyIndex::MetaCF) + .expect("MetaCF exists"); + assert_eq!( + db.get_cf(&meta_cf, &meta_key) + .expect("read routed meta") + .is_some(), + instance_id == selected + ); + assert_eq!( + count_cf_entries(redis, ColumnFamilyIndex::VectorDataCF), + usize::from(instance_id == selected) + ); + } + + storage.shutdown().await; +} + +#[test] +fn test_type_returns_vectorset() { + with_redis(|redis| { + let vector = CanonicalVector::from_values(&[1.0, 0.0]).expect("vector"); + redis.vadd(b"vectors", b"member", &vector).expect("insert"); + assert_eq!( + storage::data_type_to_string(redis.get_key_type(b"vectors").expect("key type")), + "vectorset" + ); + }); +} + +#[tokio::test] +async fn test_expired_vectorset_reads_as_missing() { + let temp = tempfile::tempdir().expect("temp dir"); + let mut storage = Storage::new(1, 0); + let _receiver = storage + .open(Arc::new(StorageOptions::default()), temp.path()) + .expect("open storage"); + let key = b"expiring-vectors"; + let vector = CanonicalVector::from_values(&[1.0, 0.0]).expect("vector"); + storage.vadd(key, b"member", &vector).expect("insert"); + + assert!(storage.expire(key, 60).expect("expire")); + assert!(storage.persist(key).expect("persist")); + assert!(storage.expireat(key, 1).expect("expire in the past")); + assert_eq!(storage.vcard(key).expect("expired card"), 0); + assert!(storage.vdim(key).is_err()); + assert_eq!(storage.vemb(key, b"member").expect("expired emb"), None); + assert!(!storage.vismember(key, b"member").expect("expired member")); + assert!( + storage + .vsim( + key, + VectorQuery::Vector(vector), + search_options(1, VectorSearchMode::Approximate), + ) + .expect("expired search") + .is_empty() + ); + assert_eq!(storage.key_type(key).expect("expired type"), "none"); + + storage.shutdown().await; +} + +#[test] +fn test_del_removes_vector_meta_and_members() { + with_redis(|redis| { + let vector = CanonicalVector::from_values(&[1.0, 0.0]).expect("vector"); + redis.vadd(b"vectors", b"a", &vector).expect("insert a"); + redis.vadd(b"vectors", b"b", &vector).expect("insert b"); + assert_eq!(count_cf_entries(redis, ColumnFamilyIndex::VectorDataCF), 2); + + assert!(redis.del_key(b"vectors").expect("delete vector set")); + assert_eq!(redis.vcard(b"vectors").expect("missing card"), 0); + assert_eq!(count_cf_entries(redis, ColumnFamilyIndex::VectorDataCF), 0); + }); +} + +#[test] +fn test_flushdb_removes_vector_meta_and_members() { + with_redis(|redis| { + let vector = CanonicalVector::from_values(&[1.0, 0.0]).expect("vector"); + redis.vadd(b"vectors", b"a", &vector).expect("insert a"); + redis.vadd(b"vectors", b"b", &vector).expect("insert b"); + + redis.flush_db().expect("flush db"); + assert_eq!(count_cf_entries(redis, ColumnFamilyIndex::MetaCF), 0); + assert_eq!(count_cf_entries(redis, ColumnFamilyIndex::VectorDataCF), 0); + }); +} + +#[tokio::test] +async fn test_vector_storage_rejects_cluster_mode() { + let temp = tempfile::tempdir().expect("temp dir"); + let mut storage = Storage::new(1, 0); + let _receiver = storage + .open(Arc::new(StorageOptions::default()), temp.path()) + .expect("open storage"); + storage.set_append_log_fn(Arc::new(|_| panic!("vector API must not append Raft log"))); + let vector = CanonicalVector::from_values(&[1.0, 0.0]).expect("vector"); + + assert!(storage.vadd(b"vectors", b"member", &vector).is_err()); + assert!( + storage + .vsim( + b"vectors", + VectorQuery::Vector(vector), + search_options(1, VectorSearchMode::Approximate), + ) + .is_err() + ); + assert!(storage.vrem(b"vectors", b"member").is_err()); + assert!(storage.vcard(b"vectors").is_err()); + assert!(storage.vdim(b"vectors").is_err()); + assert!(storage.vemb(b"vectors", b"member").is_err()); + assert!(storage.vismember(b"vectors", b"member").is_err()); + + storage.shutdown().await; +} From f55e74e4d61ef24a95b0c92119dc6cb1a1199d0b Mon Sep 17 00:00:00 2001 From: happy-v587 Date: Mon, 20 Jul 2026 00:24:33 +0800 Subject: [PATCH 06/32] feat(cmd): add redis vector set commands --- src/cmd/src/lib.rs | 1 + src/cmd/src/table.rs | 16 + src/cmd/src/vector.rs | 726 ++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 743 insertions(+) create mode 100644 src/cmd/src/vector.rs diff --git a/src/cmd/src/lib.rs b/src/cmd/src/lib.rs index 8a28f3e0..22ecfb18 100644 --- a/src/cmd/src/lib.rs +++ b/src/cmd/src/lib.rs @@ -90,6 +90,7 @@ pub mod sunionstore; pub mod table; pub mod ttl; pub mod type_cmd; +pub mod vector; pub mod zadd; pub mod zcard; pub mod zcount; diff --git a/src/cmd/src/table.rs b/src/cmd/src/table.rs index 85058d69..69f3f7bd 100644 --- a/src/cmd/src/table.rs +++ b/src/cmd/src/table.rs @@ -165,6 +165,14 @@ pub fn create_command_table(requirepass_provider: RequirepassProvider) -> CmdTab crate::zscan::ZscanCmd, crate::zscore::ZscoreCmd, crate::zunionstore::ZunionstoreCmd, + // Vector Set commands + crate::vector::VAddCmd, + crate::vector::VSimCmd, + crate::vector::VRemCmd, + crate::vector::VCardCmd, + crate::vector::VDimCmd, + crate::vector::VEmbCmd, + crate::vector::VIsMemberCmd, // connection commands crate::ping::PingCmd, ); @@ -258,6 +266,14 @@ mod tests { ); } + #[test] + fn vector_commands_are_registered() { + let table = create_command_table(Arc::new(|| None)); + for name in ["vadd", "vsim", "vrem", "vcard", "vdim", "vemb", "vismember"] { + assert!(table.contains_key(name), "{name} should be registered"); + } + } + #[test] fn hello_bare_with_requirepass_returns_noauth() { let table = create_command_table(Arc::new(|| Some("secret".to_string()))); diff --git a/src/cmd/src/vector.rs b/src/cmd/src/vector.rs new file mode 100644 index 00000000..aa567775 --- /dev/null +++ b/src/cmd/src/vector.rs @@ -0,0 +1,726 @@ +// Copyright (c) 2024-present, arana-db Community. All rights reserved. +// +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use bytes::Bytes; +use client::Client; +use resp::RespData; +use storage::{ + CanonicalVector, VectorQuery, VectorSearchMode, VectorSearchOptions, error::Error, + storage::Storage, +}; + +use crate::{AclCategory, Cmd, CmdFlags, CmdMeta, impl_cmd_clone_box, impl_cmd_meta}; + +const ERR_INVALID_VECTOR: &str = "ERR invalid vector specification"; +const ERR_VECTOR_DIMENSION: &str = "ERR vector dimension mismatch"; +const ERR_DEFAULT_Q8: &str = + "ERR default Q8 quantization is not supported in Phase 1; specify NOQUANT"; +const ERR_VADD_Q8: &str = "ERR VADD option Q8 is not supported yet"; +const ERR_VADD_BIN: &str = "ERR VADD option BIN is not supported yet"; +const ERR_VEMB_RAW: &str = "ERR VEMB option RAW is not supported yet"; +const ERR_ELEMENT_NOT_FOUND: &str = "ERR element not found in set"; +const ERR_KEY_NOT_FOUND: &str = "ERR key does not exist"; +const ERR_STORAGE: &str = "ERR storage error"; + +type ParseResult = std::result::Result; + +#[derive(Debug)] +struct ParsedVAdd { + vector: CanonicalVector, + element: Vec, +} + +#[derive(Debug)] +struct ParsedVSim { + query: VectorQuery, + options: VectorSearchOptions, + with_scores: bool, +} + +fn parse_positive_usize(raw: &[u8]) -> Option { + let value = std::str::from_utf8(raw).ok()?.parse::().ok()?; + (value > 0).then_some(value) +} + +fn parse_vector_values( + argv: &[Vec], + dimension_index: usize, +) -> ParseResult<(CanonicalVector, usize)> { + let dimension = argv + .get(dimension_index) + .and_then(|raw| parse_positive_usize(raw)) + .ok_or(ERR_INVALID_VECTOR)?; + let values_start = dimension_index + 1; + let values_end = values_start + .checked_add(dimension) + .ok_or(ERR_INVALID_VECTOR)?; + let raw_values = argv + .get(values_start..values_end) + .ok_or(ERR_INVALID_VECTOR)?; + let values = raw_values + .iter() + .map(|raw| { + std::str::from_utf8(raw) + .ok() + .and_then(|value| value.parse::().ok()) + .filter(|value| value.is_finite()) + .ok_or(ERR_INVALID_VECTOR) + }) + .collect::>>()?; + let vector = CanonicalVector::from_values(&values).map_err(|_| ERR_INVALID_VECTOR)?; + Ok((vector, values_end)) +} + +fn parse_direct_vector( + argv: &[Vec], + kind_index: usize, +) -> ParseResult<(CanonicalVector, usize)> { + let kind = argv.get(kind_index).ok_or(ERR_INVALID_VECTOR)?; + if kind.eq_ignore_ascii_case(b"FP32") { + let raw = argv.get(kind_index + 1).ok_or(ERR_INVALID_VECTOR)?; + let vector = CanonicalVector::from_fp32_le(raw).map_err(|_| ERR_INVALID_VECTOR)?; + Ok((vector, kind_index + 2)) + } else if kind.eq_ignore_ascii_case(b"VALUES") { + parse_vector_values(argv, kind_index + 1) + } else { + Err(ERR_INVALID_VECTOR) + } +} + +fn parse_vadd(argv: &[Vec]) -> ParseResult { + let (vector, element_index) = parse_direct_vector(argv, 2)?; + let element = argv.get(element_index).cloned().ok_or(ERR_INVALID_VECTOR)?; + let quantization = &argv[element_index + 1..]; + match quantization { + [] => Err(ERR_DEFAULT_Q8), + [option] if option.eq_ignore_ascii_case(b"NOQUANT") => Ok(ParsedVAdd { vector, element }), + [option] if option.eq_ignore_ascii_case(b"Q8") => Err(ERR_VADD_Q8), + [option] if option.eq_ignore_ascii_case(b"BIN") => Err(ERR_VADD_BIN), + _ => Err(ERR_INVALID_VECTOR), + } +} + +fn parse_vsim(argv: &[Vec]) -> ParseResult { + let query_kind = argv.get(2).ok_or(ERR_INVALID_VECTOR)?; + let (query, mut option_index) = if query_kind.eq_ignore_ascii_case(b"ELE") { + let element = argv.get(3).cloned().ok_or(ERR_INVALID_VECTOR)?; + (VectorQuery::Element(element), 4) + } else { + let (vector, next) = parse_direct_vector(argv, 2)?; + (VectorQuery::Vector(vector), next) + }; + + let mut count = 10; + let mut mode = VectorSearchMode::Approximate; + let mut with_scores = false; + let mut count_seen = false; + let mut truth_seen = false; + + while option_index < argv.len() { + let option = &argv[option_index]; + if option.eq_ignore_ascii_case(b"WITHSCORES") { + if with_scores { + return Err(ERR_INVALID_VECTOR); + } + with_scores = true; + option_index += 1; + } else if option.eq_ignore_ascii_case(b"COUNT") { + if count_seen { + return Err(ERR_INVALID_VECTOR); + } + count = argv + .get(option_index + 1) + .and_then(|raw| parse_positive_usize(raw)) + .ok_or(ERR_INVALID_VECTOR)?; + count_seen = true; + option_index += 2; + } else if option.eq_ignore_ascii_case(b"TRUTH") { + if truth_seen { + return Err(ERR_INVALID_VECTOR); + } + mode = VectorSearchMode::Truth; + truth_seen = true; + option_index += 1; + } else { + return Err(ERR_INVALID_VECTOR); + } + } + + Ok(ParsedVSim { + query, + options: VectorSearchOptions { count, mode }, + with_scores, + }) +} + +fn parse_vemb(argv: &[Vec]) -> ParseResult> { + let element = argv.get(2).cloned().ok_or(ERR_INVALID_VECTOR)?; + match &argv[3..] { + [] => Ok(element), + [option] if option.eq_ignore_ascii_case(b"RAW") => Err(ERR_VEMB_RAW), + _ => Err(ERR_INVALID_VECTOR), + } +} + +fn error_reply(message: impl Into) -> RespData { + RespData::Error(message.into().into()) +} + +#[derive(Clone, Copy)] +enum MissingError { + Key, + Element, +} + +fn storage_error_reply(error: Error, missing: MissingError) -> RespData { + match &error { + Error::RedisErr { message, .. } => error_reply(message.clone()), + Error::InvalidArgument { message, .. } if message.contains("dimension mismatch") => { + error_reply(ERR_VECTOR_DIMENSION) + } + Error::KeyNotFound { .. } => match missing { + MissingError::Key => error_reply(ERR_KEY_NOT_FOUND), + MissingError::Element => error_reply(ERR_ELEMENT_NOT_FOUND), + }, + _ => { + log::error!("vector storage command failed: {error}"); + error_reply(ERR_STORAGE) + } + } +} + +fn set_command_key(client: &Client) -> bool { + let argv = client.argv(); + let Some(key) = argv.get(1) else { + client.set_reply(error_reply(ERR_INVALID_VECTOR)); + return false; + }; + client.set_key(key); + true +} + +fn integer_reply(value: u64) -> RespData { + match i64::try_from(value) { + Ok(value) => RespData::Integer(value), + Err(error) => { + log::error!("vector integer reply overflow: {error}"); + error_reply(ERR_STORAGE) + } + } +} + +macro_rules! define_vector_command { + ($type:ident, $name:literal, $arity:expr, $flags:expr, $acl:expr) => { + #[derive(Clone, Default)] + pub struct $type { + meta: CmdMeta, + } + + impl $type { + pub fn new() -> Self { + Self { + meta: CmdMeta { + name: $name.to_string(), + arity: $arity, + flags: $flags, + acl_category: $acl, + ..Default::default() + }, + } + } + } + }; +} + +define_vector_command!( + VAddCmd, + "vadd", + -5, + CmdFlags::WRITE | CmdFlags::FAST | CmdFlags::MODULE_NO_CLUSTER, + AclCategory::KEYSPACE | AclCategory::WRITE +); +define_vector_command!( + VSimCmd, + "vsim", + -4, + CmdFlags::READONLY | CmdFlags::MODULE_NO_CLUSTER, + AclCategory::KEYSPACE | AclCategory::READ | AclCategory::SLOW +); +define_vector_command!( + VRemCmd, + "vrem", + 3, + CmdFlags::WRITE | CmdFlags::FAST | CmdFlags::MODULE_NO_CLUSTER, + AclCategory::KEYSPACE | AclCategory::WRITE +); +define_vector_command!( + VCardCmd, + "vcard", + 2, + CmdFlags::READONLY | CmdFlags::FAST | CmdFlags::MODULE_NO_CLUSTER, + AclCategory::KEYSPACE | AclCategory::READ +); +define_vector_command!( + VDimCmd, + "vdim", + 2, + CmdFlags::READONLY | CmdFlags::FAST | CmdFlags::MODULE_NO_CLUSTER, + AclCategory::KEYSPACE | AclCategory::READ +); +define_vector_command!( + VEmbCmd, + "vemb", + -3, + CmdFlags::READONLY | CmdFlags::FAST | CmdFlags::MODULE_NO_CLUSTER, + AclCategory::KEYSPACE | AclCategory::READ +); +define_vector_command!( + VIsMemberCmd, + "vismember", + 3, + CmdFlags::READONLY | CmdFlags::FAST | CmdFlags::MODULE_NO_CLUSTER, + AclCategory::KEYSPACE | AclCategory::READ +); + +impl Cmd for VAddCmd { + impl_cmd_meta!(); + impl_cmd_clone_box!(); + + fn do_initial(&self, client: &Client) -> bool { + set_command_key(client) + } + + fn do_cmd(&self, client: &Client, storage: Arc) { + let parsed = match parse_vadd(&client.argv()) { + Ok(parsed) => parsed, + Err(message) => { + client.set_reply(error_reply(message)); + return; + } + }; + let reply = match storage.vadd(&client.key(), &parsed.element, &parsed.vector) { + Ok(inserted) => RespData::Integer(i64::from(inserted)), + Err(error) => storage_error_reply(error, MissingError::Key), + }; + client.set_reply(reply); + } +} + +impl Cmd for VSimCmd { + impl_cmd_meta!(); + impl_cmd_clone_box!(); + + fn do_initial(&self, client: &Client) -> bool { + set_command_key(client) + } + + fn do_cmd(&self, client: &Client, storage: Arc) { + let parsed = match parse_vsim(&client.argv()) { + Ok(parsed) => parsed, + Err(message) => { + client.set_reply(error_reply(message)); + return; + } + }; + let reply = match storage.vsim(&client.key(), parsed.query, parsed.options) { + Ok(hits) if parsed.with_scores => RespData::Map( + hits.into_iter() + .map(|hit| { + ( + RespData::BulkString(Some(Bytes::from(hit.element))), + RespData::Double(hit.score), + ) + }) + .collect(), + ), + Ok(hits) => RespData::Array(Some( + hits.into_iter() + .map(|hit| RespData::BulkString(Some(Bytes::from(hit.element)))) + .collect(), + )), + Err(error) => storage_error_reply(error, MissingError::Element), + }; + client.set_reply(reply); + } +} + +impl Cmd for VRemCmd { + impl_cmd_meta!(); + impl_cmd_clone_box!(); + + fn do_initial(&self, client: &Client) -> bool { + set_command_key(client) + } + + fn do_cmd(&self, client: &Client, storage: Arc) { + let argv = client.argv(); + let reply = match storage.vrem(&client.key(), &argv[2]) { + Ok(removed) => RespData::Integer(i64::from(removed)), + Err(error) => storage_error_reply(error, MissingError::Key), + }; + client.set_reply(reply); + } +} + +impl Cmd for VCardCmd { + impl_cmd_meta!(); + impl_cmd_clone_box!(); + + fn do_initial(&self, client: &Client) -> bool { + set_command_key(client) + } + + fn do_cmd(&self, client: &Client, storage: Arc) { + let reply = match storage.vcard(&client.key()) { + Ok(count) => integer_reply(count), + Err(error) => storage_error_reply(error, MissingError::Key), + }; + client.set_reply(reply); + } +} + +impl Cmd for VDimCmd { + impl_cmd_meta!(); + impl_cmd_clone_box!(); + + fn do_initial(&self, client: &Client) -> bool { + set_command_key(client) + } + + fn do_cmd(&self, client: &Client, storage: Arc) { + let reply = match storage.vdim(&client.key()) { + Ok(dimension) => integer_reply(u64::from(dimension)), + Err(error) => storage_error_reply(error, MissingError::Key), + }; + client.set_reply(reply); + } +} + +impl Cmd for VEmbCmd { + impl_cmd_meta!(); + impl_cmd_clone_box!(); + + fn do_initial(&self, client: &Client) -> bool { + set_command_key(client) + } + + fn do_cmd(&self, client: &Client, storage: Arc) { + let element = match parse_vemb(&client.argv()) { + Ok(element) => element, + Err(message) => { + client.set_reply(error_reply(message)); + return; + } + }; + let reply = match storage.vemb(&client.key(), &element) { + Ok(Some(values)) => { + RespData::Array(Some(values.into_iter().map(RespData::Double).collect())) + } + Ok(None) => RespData::BulkString(None), + Err(error) => storage_error_reply(error, MissingError::Key), + }; + client.set_reply(reply); + } +} + +impl Cmd for VIsMemberCmd { + impl_cmd_meta!(); + impl_cmd_clone_box!(); + + fn do_initial(&self, client: &Client) -> bool { + set_command_key(client) + } + + fn do_cmd(&self, client: &Client, storage: Arc) { + let argv = client.argv(); + let reply = match storage.vismember(&client.key(), &argv[2]) { + Ok(is_member) => RespData::Integer(i64::from(is_member)), + Err(error) => storage_error_reply(error, MissingError::Key), + }; + client.set_reply(reply); + } +} + +#[allow(clippy::unwrap_used)] +#[cfg(test)] +mod tests { + use storage::{VectorQuery, VectorSearchMode}; + + use super::*; + + fn fp32(values: &[f32]) -> Vec { + values + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect() + } + + #[test] + fn parses_supported_vadd_shapes() { + let blob = fp32(&[3.0, 4.0]); + let parsed = parse_vadd(&[ + b"vadd".to_vec(), + b"key\0raw".to_vec(), + b"FP32".to_vec(), + blob, + b"\0element".to_vec(), + b"NOQUANT".to_vec(), + ]) + .expect("FP32 VADD"); + assert_eq!(parsed.vector.dimension(), 2); + assert_eq!(parsed.element, b"\0element"); + + let parsed = parse_vadd(&[ + b"vadd".to_vec(), + b"key".to_vec(), + b"VALUES".to_vec(), + b"2".to_vec(), + b"3".to_vec(), + b"4".to_vec(), + b"element".to_vec(), + b"noquant".to_vec(), + ]) + .expect("VALUES VADD"); + assert_eq!(parsed.vector.dimension(), 2); + assert_eq!(parsed.element, b"element"); + } + + #[test] + fn rejects_unsupported_or_invalid_vadd_shapes() { + let malformed_fp32 = vec![ + b"vadd".to_vec(), + b"key".to_vec(), + b"FP32".to_vec(), + vec![1, 2, 3], + b"element".to_vec(), + b"NOQUANT".to_vec(), + ]; + assert_eq!(parse_vadd(&malformed_fp32).unwrap_err(), ERR_INVALID_VECTOR); + + let invalid_values = [ + vec![ + b"vadd".to_vec(), + b"key".to_vec(), + b"VALUES".to_vec(), + b"0".to_vec(), + b"element".to_vec(), + b"NOQUANT".to_vec(), + ], + vec![ + b"vadd".to_vec(), + b"key".to_vec(), + b"VALUES".to_vec(), + b"2".to_vec(), + b"1".to_vec(), + b"element".to_vec(), + b"NOQUANT".to_vec(), + ], + vec![ + b"vadd".to_vec(), + b"key".to_vec(), + b"VALUES".to_vec(), + b"1".to_vec(), + b"not-a-float".to_vec(), + b"element".to_vec(), + b"NOQUANT".to_vec(), + ], + ]; + for argv in invalid_values { + assert_eq!(parse_vadd(&argv).unwrap_err(), ERR_INVALID_VECTOR); + } + + let base = vec![ + b"vadd".to_vec(), + b"key".to_vec(), + b"FP32".to_vec(), + fp32(&[1.0]), + b"element".to_vec(), + ]; + assert_eq!(parse_vadd(&base).unwrap_err(), ERR_DEFAULT_Q8); + + let mut q8 = base.clone(); + q8.push(b"Q8".to_vec()); + assert_eq!(parse_vadd(&q8).unwrap_err(), ERR_VADD_Q8); + + let mut bin = base.clone(); + bin.push(b"BIN".to_vec()); + assert_eq!(parse_vadd(&bin).unwrap_err(), ERR_VADD_BIN); + + let mut trailing = base; + trailing.extend([b"NOQUANT".to_vec(), b"extra".to_vec()]); + assert_eq!(parse_vadd(&trailing).unwrap_err(), ERR_INVALID_VECTOR); + } + + #[test] + fn parses_supported_vsim_shapes_and_options() { + let ele = parse_vsim(&[ + b"vsim".to_vec(), + b"key".to_vec(), + b"ELE".to_vec(), + b"\0element".to_vec(), + b"WITHSCORES".to_vec(), + b"COUNT".to_vec(), + b"3".to_vec(), + b"TRUTH".to_vec(), + ]) + .expect("ELE VSIM"); + assert_eq!(ele.query, VectorQuery::Element(b"\0element".to_vec())); + assert_eq!(ele.options.count, 3); + assert_eq!(ele.options.mode, VectorSearchMode::Truth); + assert!(ele.with_scores); + + let direct = parse_vsim(&[ + b"vsim".to_vec(), + b"key".to_vec(), + b"FP32".to_vec(), + fp32(&[1.0, 0.0]), + ]) + .expect("FP32 VSIM"); + assert!(matches!(direct.query, VectorQuery::Vector(_))); + assert_eq!(direct.options.count, 10); + assert_eq!(direct.options.mode, VectorSearchMode::Approximate); + + let values = parse_vsim(&[ + b"vsim".to_vec(), + b"key".to_vec(), + b"VALUES".to_vec(), + b"2".to_vec(), + b"1".to_vec(), + b"0".to_vec(), + b"COUNT".to_vec(), + b"1".to_vec(), + ]) + .expect("VALUES VSIM"); + assert!(matches!(values.query, VectorQuery::Vector(_))); + assert_eq!(values.options.count, 1); + } + + #[test] + fn rejects_invalid_vsim_options_and_vectors() { + let cases = [ + vec![ + b"vsim".to_vec(), + b"key".to_vec(), + b"FP32".to_vec(), + vec![1, 2, 3], + ], + vec![ + b"vsim".to_vec(), + b"key".to_vec(), + b"VALUES".to_vec(), + b"2".to_vec(), + b"1".to_vec(), + ], + vec![ + b"vsim".to_vec(), + b"key".to_vec(), + b"VALUES".to_vec(), + b"1".to_vec(), + b"NaN".to_vec(), + ], + vec![ + b"vsim".to_vec(), + b"key".to_vec(), + b"ELE".to_vec(), + b"member".to_vec(), + b"COUNT".to_vec(), + b"0".to_vec(), + ], + vec![ + b"vsim".to_vec(), + b"key".to_vec(), + b"ELE".to_vec(), + b"member".to_vec(), + b"COUNT".to_vec(), + b"invalid".to_vec(), + ], + vec![ + b"vsim".to_vec(), + b"key".to_vec(), + b"ELE".to_vec(), + b"member".to_vec(), + b"COUNT".to_vec(), + b"1".to_vec(), + b"COUNT".to_vec(), + b"2".to_vec(), + ], + vec![ + b"vsim".to_vec(), + b"key".to_vec(), + b"ELE".to_vec(), + b"member".to_vec(), + b"WITHSCORES".to_vec(), + b"WITHSCORES".to_vec(), + ], + vec![ + b"vsim".to_vec(), + b"key".to_vec(), + b"ELE".to_vec(), + b"member".to_vec(), + b"TRUTH".to_vec(), + b"TRUTH".to_vec(), + ], + vec![ + b"vsim".to_vec(), + b"key".to_vec(), + b"ELE".to_vec(), + b"member".to_vec(), + b"unknown".to_vec(), + ], + ]; + for argv in cases { + assert_eq!(parse_vsim(&argv).unwrap_err(), ERR_INVALID_VECTOR); + } + } + + #[test] + fn rejects_vemb_raw_and_unknown_trailing_options() { + assert_eq!( + parse_vemb(&[ + b"vemb".to_vec(), + b"key".to_vec(), + b"member".to_vec(), + b"RAW".to_vec(), + ]) + .unwrap_err(), + ERR_VEMB_RAW + ); + assert_eq!( + parse_vemb(&[ + b"vemb".to_vec(), + b"key".to_vec(), + b"member".to_vec(), + b"unknown".to_vec(), + ]) + .unwrap_err(), + ERR_INVALID_VECTOR + ); + } + + #[test] + fn vector_command_metadata_matches_redis_shapes() { + assert_eq!(VAddCmd::new().meta().arity, -5); + assert_eq!(VSimCmd::new().meta().arity, -4); + assert_eq!(VRemCmd::new().meta().arity, 3); + assert_eq!(VCardCmd::new().meta().arity, 2); + assert_eq!(VDimCmd::new().meta().arity, 2); + assert_eq!(VEmbCmd::new().meta().arity, -3); + assert_eq!(VIsMemberCmd::new().meta().arity, 3); + } +} From 53ee9b0d33de1efa89dabe96acfb0a05fd4061bc Mon Sep 17 00:00:00 2001 From: happy-v587 Date: Mon, 20 Jul 2026 00:29:05 +0800 Subject: [PATCH 07/32] fix(resp): downgrade vector scores for resp2 --- src/resp/src/encode.rs | 254 ++++++++++++++++++++---------------- src/resp/src/negotiation.rs | 16 ++- 2 files changed, 154 insertions(+), 116 deletions(-) diff --git a/src/resp/src/encode.rs b/src/resp/src/encode.rs index 5260bc23..17ec1ae6 100644 --- a/src/resp/src/encode.rs +++ b/src/resp/src/encode.rs @@ -23,6 +23,7 @@ use bytes::{Bytes, BytesMut}; use crate::{ CRLF, error::RespError, + negotiation::ProtocolNegotiator, types::{RespData, RespVersion}, }; @@ -198,6 +199,107 @@ impl RespEncoder { let _ = write!(self.buffer, "*{len}"); self.append_crlf() } + + fn encode_resp_data_inner(&mut self, data: &RespData) -> &mut Self { + match data { + RespData::SimpleString(bytes) => { + self.buffer.extend_from_slice(b"+"); + self.buffer.extend_from_slice(bytes); + self.append_crlf() + } + RespData::Error(bytes) => { + self.buffer.extend_from_slice(b"-"); + self.buffer.extend_from_slice(bytes); + self.append_crlf() + } + RespData::Integer(num) => self.append_integer(*num), + RespData::BulkString(Some(bytes)) => self.append_bulk_string(bytes), + RespData::BulkString(None) => self.set_bulk_string_len(-1), + RespData::Array(Some(array)) => { + self.append_array_len(array.len() as i64); + for item in array { + self.encode_resp_data_inner(item); + } + self + } + RespData::Array(None) => self.set_array_len(-1), + RespData::Null => { + self.buffer.extend_from_slice(b"_"); + self.append_crlf() + } + RespData::Boolean(value) => { + self.buffer.extend_from_slice(b"#"); + self.buffer + .extend_from_slice(if *value { b"t" } else { b"f" }); + self.append_crlf() + } + RespData::Double(value) => { + if value.is_nan() { + self.buffer.extend_from_slice(b",nan"); + } else if value.is_infinite() { + self.buffer.extend_from_slice(if value.is_sign_negative() { + b",-inf" + } else { + b",inf" + }); + } else { + let _ = write!(self.buffer, ",{value}"); + } + self.append_crlf() + } + RespData::BigNumber(bytes) => { + self.buffer.extend_from_slice(b"("); + self.buffer.extend_from_slice(bytes); + self.append_crlf() + } + RespData::BulkError(bytes) => { + let _ = write!(self.buffer, "!{}", bytes.len()); + self.append_crlf(); + self.buffer.extend_from_slice(bytes); + self.append_crlf() + } + RespData::VerbatimString { format, data } => { + if format.len() != 3 { + panic!( + "RESP3 VerbatimString format must be exactly 3 bytes, got {}", + format.len() + ); + } + let total_len = format.len() + 1 + data.len(); + let _ = write!(self.buffer, "={total_len}"); + self.append_crlf(); + self.buffer.extend_from_slice(format); + self.buffer.extend_from_slice(b":"); + self.buffer.extend_from_slice(data); + self.append_crlf() + } + RespData::Map(pairs) => { + let _ = write!(self.buffer, "%{}", pairs.len()); + self.append_crlf(); + for (key, value) in pairs { + self.encode_resp_data_inner(key); + self.encode_resp_data_inner(value); + } + self + } + RespData::Set(items) => { + let _ = write!(self.buffer, "~{}", items.len()); + self.append_crlf(); + for item in items { + self.encode_resp_data_inner(item); + } + self + } + RespData::Push(items) => { + let _ = write!(self.buffer, ">{}", items.len()); + self.append_crlf(); + for item in items { + self.encode_resp_data_inner(item); + } + self + } + } + } } impl RespEncode for RespEncoder { @@ -362,103 +464,11 @@ impl RespEncode for RespEncoder { } fn encode_resp_data(&mut self, data: &RespData) -> &mut Self { - match data { - RespData::SimpleString(bytes) => { - self.buffer.extend_from_slice(b"+"); - self.buffer.extend_from_slice(bytes); - self.append_crlf() - } - RespData::Error(bytes) => { - self.buffer.extend_from_slice(b"-"); - self.buffer.extend_from_slice(bytes); - self.append_crlf() - } - RespData::Integer(num) => self.append_integer(*num), - RespData::BulkString(Some(bytes)) => self.append_bulk_string(bytes), - RespData::BulkString(None) | RespData::Array(None) if self.is_resp3() => { - self.append_null() - } - RespData::BulkString(None) => self.set_bulk_string_len(-1), - RespData::Array(Some(array)) => { - self.append_array_len(array.len() as i64); - for item in array { - self.encode_resp_data(item); - } - self - } - RespData::Array(None) => self.set_array_len(-1), - // RESP3 types - RespData::Null => self.append_null(), - RespData::Boolean(b) => { - self.buffer.extend_from_slice(b"#"); - self.buffer.extend_from_slice(if *b { b"t" } else { b"f" }); - self.append_crlf() - } - RespData::Double(d) => { - if d.is_nan() { - self.buffer.extend_from_slice(b",nan"); - } else if d.is_infinite() { - self.buffer.extend_from_slice(if d.is_sign_negative() { - b",-inf" - } else { - b",inf" - }); - } else { - let _ = write!(self.buffer, ",{}", d); - } - self.append_crlf() - } - RespData::BigNumber(bytes) => { - self.buffer.extend_from_slice(b"("); - self.buffer.extend_from_slice(bytes); - self.append_crlf() - } - RespData::BulkError(bytes) => { - let _ = write!(self.buffer, "!{}", bytes.len()); - self.append_crlf(); - self.buffer.extend_from_slice(bytes); - self.append_crlf() - } - RespData::VerbatimString { format, data } => { - if format.len() != 3 { - panic!( - "RESP3 VerbatimString format must be exactly 3 bytes, got {}", - format.len() - ); - } - let total_len = format.len() + 1 + data.len(); // format + ':' + data - let _ = write!(self.buffer, "={total_len}"); - self.append_crlf(); - self.buffer.extend_from_slice(format); - self.buffer.extend_from_slice(b":"); - self.buffer.extend_from_slice(data); - self.append_crlf() - } - RespData::Map(pairs) => { - let _ = write!(self.buffer, "%{}", pairs.len()); - self.append_crlf(); - for (key, value) in pairs { - self.encode_resp_data(key); - self.encode_resp_data(value); - } - self - } - RespData::Set(items) => { - let _ = write!(self.buffer, "~{}", items.len()); - self.append_crlf(); - for item in items { - self.encode_resp_data(item); - } - self - } - RespData::Push(items) => { - let _ = write!(self.buffer, ">{}", items.len()); - self.append_crlf(); - for item in items { - self.encode_resp_data(item); - } - self - } + if self.is_resp3() { + self.encode_resp_data_inner(data) + } else { + let normalized = ProtocolNegotiator::convert_to_resp2(data); + self.encode_resp_data_inner(&normalized) } } @@ -557,6 +567,39 @@ mod tests { use super::*; use crate::types::{RespData, RespVersion}; + fn vsim_withscores_response() -> RespData { + RespData::Map(vec![ + ( + RespData::BulkString(Some(Bytes::from_static(b"a"))), + RespData::Double(1.0), + ), + ( + RespData::BulkString(Some(Bytes::from_static(b"b"))), + RespData::Double(0.5), + ), + ]) + } + + #[test] + fn encode_vsim_withscores_resp3() { + let mut encoder = RespEncoder::new(RespVersion::RESP3); + encoder.encode_resp_data(&vsim_withscores_response()); + assert_eq!( + encoder.get_response(), + Bytes::from_static(b"%2\r\n$1\r\na\r\n,1\r\n$1\r\nb\r\n,0.5\r\n") + ); + } + + #[test] + fn encode_vsim_withscores_resp2() { + let mut encoder = RespEncoder::new(RespVersion::RESP2); + encoder.encode_resp_data(&vsim_withscores_response()); + assert_eq!( + encoder.get_response(), + Bytes::from_static(b"*4\r\n$1\r\na\r\n$1\r\n1\r\n$1\r\nb\r\n$3\r\n0.5\r\n") + ); + } + #[test] fn test_encode_resp3_null() { let mut encoder = RespEncoder::new(RespVersion::RESP3); @@ -744,14 +787,9 @@ mod tests { #[test] fn test_resp3_boolean_with_resp2_encoder() { - // Test current behavior: RESP3 Boolean encodes as RESP3 format even with RESP2 encoder - // This produces invalid RESP2 output (#t\r\n is not valid RESP2) - // TODO: Either fail with error or auto-convert Boolean(true) -> Integer(1) ":1\r\n" let mut encoder = RespEncoder::new(RespVersion::RESP2); encoder.encode_resp_data(&RespData::Boolean(true)); - let result = encoder.get_response(); - // Current implementation produces RESP3 format regardless of encoder version - assert_eq!(result, Bytes::from("#t\r\n")); + assert_eq!(encoder.get_response(), Bytes::from(":1\r\n")); } #[test] @@ -762,18 +800,14 @@ mod tests { } #[test] - fn test_resp3_types_behavior() { - // Document current behavior: RESP3 types encode regardless of version - // Future improvement: Add version-aware encoding or fail early + fn test_resp3_types_downgrade_for_resp2() { let mut encoder = RespEncoder::new(RespVersion::RESP2); - // Double #[allow(clippy::approx_constant)] encoder.clear().encode_resp_data(&RespData::Double(3.14)); - assert_eq!(encoder.get_response(), Bytes::from(",3.14\r\n")); + assert_eq!(encoder.get_response(), Bytes::from("$4\r\n3.14\r\n")); - // Map encoder.clear().encode_resp_data(&RespData::Map(vec![])); - assert_eq!(encoder.get_response(), Bytes::from("%0\r\n")); + assert_eq!(encoder.get_response(), Bytes::from("*0\r\n")); } } diff --git a/src/resp/src/negotiation.rs b/src/resp/src/negotiation.rs index e79f4751..aec17f56 100644 --- a/src/resp/src/negotiation.rs +++ b/src/resp/src/negotiation.rs @@ -415,14 +415,18 @@ mod tests { ), ( RespData::BulkString(Some(Bytes::from("key2"))), - RespData::Integer(42), + RespData::Double(0.5), ), ]); let converted = ProtocolNegotiator::convert_to_resp2(&map_data); - if let RespData::Array(Some(items)) = converted { - assert_eq!(items.len(), 4); // 2 key-value pairs = 4 items - } else { - panic!("Expected array conversion for map"); - } + assert_eq!( + converted, + RespData::Array(Some(vec![ + RespData::BulkString(Some(Bytes::from("key1"))), + RespData::BulkString(Some(Bytes::from("value1"))), + RespData::BulkString(Some(Bytes::from("key2"))), + RespData::BulkString(Some(Bytes::from("0.5"))), + ])) + ); } } From e02d043ce8d070a933e87c10eb69377bad257759 Mon Sep 17 00:00:00 2001 From: happy-v587 Date: Mon, 20 Jul 2026 01:02:25 +0800 Subject: [PATCH 08/32] fix(resp): encode nulls with native protocol type --- src/resp/src/encode.rs | 3 +++ src/resp/tests/integration_tests.rs | 9 +++++---- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/src/resp/src/encode.rs b/src/resp/src/encode.rs index 17ec1ae6..8ab2e5cd 100644 --- a/src/resp/src/encode.rs +++ b/src/resp/src/encode.rs @@ -214,6 +214,9 @@ impl RespEncoder { } RespData::Integer(num) => self.append_integer(*num), RespData::BulkString(Some(bytes)) => self.append_bulk_string(bytes), + RespData::BulkString(None) | RespData::Array(None) if self.is_resp3() => { + self.append_null() + } RespData::BulkString(None) => self.set_bulk_string_len(-1), RespData::Array(Some(array)) => { self.append_array_len(array.len() as i64); diff --git a/src/resp/tests/integration_tests.rs b/src/resp/tests/integration_tests.rs index 1d1ce1a2..3cbbc454 100644 --- a/src/resp/tests/integration_tests.rs +++ b/src/resp/tests/integration_tests.rs @@ -191,6 +191,11 @@ fn test_resp3_backward_compatibility() { ]; for original_data in resp2_data { + let expected_data = match &original_data { + RespData::BulkString(None) => RespData::Null, + _ => original_data.clone(), + }; + // Encode with RESP3 encoder let mut encoder = RespEncoder::new(RespVersion::RESP3); encoder.encode_resp_data(&original_data); @@ -201,10 +206,6 @@ fn test_resp3_backward_compatibility() { let result = parser.parse(encoded); if let RespParseResult::Complete(parsed_data) = result { - let expected_data = match original_data { - RespData::BulkString(None) => RespData::Null, - data => data, - }; assert_eq!(parsed_data, expected_data); } else { panic!( From e0d72be73038fd9aa371d25f7503383b43691a40 Mon Sep 17 00:00:00 2001 From: happy-v587 Date: Mon, 20 Jul 2026 01:21:59 +0800 Subject: [PATCH 09/32] test: cover redis vector set commands --- .../plans/2026-07-19-redis-vector-set.md | 1059 +++++++++++++++++ tests/python/test_vector_set_commands.py | 257 ++++ 2 files changed, 1316 insertions(+) create mode 100644 docs/superpowers/plans/2026-07-19-redis-vector-set.md create mode 100644 tests/python/test_vector_set_commands.py diff --git a/docs/superpowers/plans/2026-07-19-redis-vector-set.md b/docs/superpowers/plans/2026-07-19-redis-vector-set.md new file mode 100644 index 00000000..b6133fc6 --- /dev/null +++ b/docs/superpowers/plans/2026-07-19-redis-vector-set.md @@ -0,0 +1,1059 @@ +# Redis Vector Set Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** 在 Kiwi 中实现 Redis 8 Vector Set 的 standalone Phase 1：持久化 FP32 向量，提供精确 FLAT 相似度查询，并支持 `VADD`、`VSIM`、`VREM`、`VCARD`、`VDIM`、`VEMB`、`VISMEMBER`。 + +**Architecture:** 沿用 Kiwi 现有 Hash/Set/ZSet 的复合类型模型。一个 VectorSet key 只按 user key 路由到一个 RocksDB instance；`MetaCF` 保存类型、数量、生命周期 version 和维度，新增 `VectorDataCF` 保存 `key + version + element` 对应的 canonical vector。写命令通过同一个 RocksDB `WriteBatch` 原子更新 Meta 和成员；`VSIM` 在该 instance 内按 generation prefix 扫描并使用有界 Top-K heap 返回精确结果。 + +**Tech Stack:** Rust 2021、Tokio 双 runtime、RocksDB Column Family / WriteBatch / Snapshot、RESP2/RESP3、Python redis-py 集成测试。 + +## Global Constraints + +- 基线为 `origin/main` 的 `cdada8b`；忽略尚未合入的 error-catalog/error-model 分支，错误处理遵循当前主分支模式。 +- 首版只支持 standalone；所有 Vector Set storage API 在 Raft append hook 已安装时返回 `ERR Vector Set is not supported in cluster mode`。 +- 首版只支持 cosine、canonical FP32 little-endian 和显式 `NOQUANT`；默认 Q8、显式 `Q8`、`BIN`、VEMB `RAW` 均返回明确的 unsupported 错误。 +- 首版只实现 FLAT；`TRUTH` 与普通 `VSIM` 都走同一个精确引擎，但保留不同的搜索模式枚举。 +- 首版不实现 `VINFO`、`INFO VECTOR`、HNSW、FAISS/IVF、`storage_incarnation`、O(1) `DEL`、Raft logical mutation、ReadIndex、snapshot barrier 和滚动升级门禁。 +- 复用 `MemberDataKey` 的 `key + version + element` 布局，不引入第二套 Vector member key codec。 +- Meta 保持现有统一 envelope；reserve 的前 8 字节固定为 `format(1) + encoding(1) + metric(1) + flags(1) + dimension(4 LE)`，其余 8 字节写零。 +- 所有新增 `.rs` 文件复制现有 Apache 2.0 license header；生产代码不得使用 `unwrap()`。 +- 不修改已有未跟踪的 `tests/python/test_vector_basic.py`；该文件验证的是旧 FT/HASH 原型，不属于 Vector Set。 +- 每个任务先写失败测试，再做最小实现；最终执行 `make fmt && make lint && make build && make test`。 + +--- + +## File Map + +| 文件 | 职责 | +|---|---| +| `src/storage/src/format_base_value.rs` | 追加 `DataType::VectorSet = 7`、字符串和 tag 映射 | +| `src/storage/src/redis.rs` | 追加 `VectorDataCF = 6`，创建 CF、配置 compaction filter、加入 batch handles | +| `src/storage/src/vector.rs` | 定义 canonical vector、查询参数、命中结果和数值计算 | +| `src/storage/src/format_vector.rs` | 编解码 VectorSet meta 和 VectorDataCF value | +| `src/storage/src/redis_vectors.rs` | 单个 RocksDB instance 上的 VADD/VREM/point read/FLAT scan | +| `src/storage/src/storage_impl.rs` | 按 user key 路由所有 Vector Set 操作 | +| `src/storage/src/data_compaction_filter.rs` | 让 VectorDataCF 复用现有 version/TTL 垃圾清理 | +| `src/storage/src/meta_compaction_filter.rs` | 让过期或空 VectorSet meta 复用现有清理规则 | +| `src/storage/src/redis_strings.rs` | 将 VectorDataCF 纳入 DEL 与 FLUSHDB 的物理清理 | +| `src/storage/src/batch.rs` | 将 CF index 6 映射到 VectorDataCF | +| `src/storage/src/storage.rs` | Raft apply 的 CF index 6 解码；即使首版拒绝集群也保持 schema 完整 | +| `src/conf/src/raft_type.rs` | 追加跨 crate 使用的 VectorDataCF index | +| `src/storage/src/logindex/types.rs` | 追加 VectorDataCF 的 log-index 元数据 | +| `src/raft/src/lib.rs` | 追加 CF 名称和一致性断言 | +| `src/cmd/src/vector.rs` | 七个命令、参数解析、错误和 RESP reply 构造 | +| `src/cmd/src/lib.rs` | 导出 vector 命令模块 | +| `src/cmd/src/table.rs` | 注册七个命令 | +| `src/resp/src/encode.rs` | RESP2 下递归降级 Map/Double，RESP3 保持原生类型 | +| `src/storage/tests/redis_vector_test.rs` | storage 行为、生命周期、多 instance 和 FLAT 排序测试 | +| `tests/python/test_vector_set_commands.py` | 真实 server 的 RESP2/RESP3 命令兼容测试 | + +--- + +### Task 1: Add the VectorSet data type and VectorDataCF plumbing + +**Files:** + +- Modify: `src/storage/src/format_base_value.rs` +- Modify: `src/storage/src/redis.rs` +- Modify: `src/storage/src/batch.rs` +- Modify: `src/storage/src/storage.rs` +- Modify: `src/conf/src/raft_type.rs` +- Modify: `src/storage/src/logindex/types.rs` +- Modify: `src/raft/src/lib.rs` +- Modify: `src/storage/tests/redis_basic_test.rs` + +**Interfaces:** + +- Produces: `DataType::VectorSet = 7` +- Produces: `ColumnFamilyIndex::VectorDataCF = 6` +- Produces: RocksDB CF name `vector_data_cf` +- Preserves: all existing DataType and CF numeric values + +- [x] **Step 1: Extend the CF expectations in the existing storage test** + +Update `test_open_redis` and `test_column_family_index` in `src/storage/tests/redis_basic_test.rs` to require seven handles and the exact final mapping: + +```rust +assert_eq!(redis.handles.len(), 7); + +let expected_cf_names = [ + "default", + "hash_data_cf", + "set_data_cf", + "list_data_cf", + "zset_data_cf", + "zset_score_cf", + "vector_data_cf", +]; + +assert_eq!(ColumnFamilyIndex::VectorDataCF as usize, 6); +assert_eq!(ColumnFamilyIndex::COUNT, 7); +``` + +- [x] **Step 2: Run the focused test and confirm the expected failure** + +Run: + +```bash +cargo test -p storage --test redis_basic_test test_open_redis -- --exact +``` + +Expected: compile failure because `ColumnFamilyIndex::VectorDataCF` does not exist, or assertion failure because only six CFs are opened. + +- [x] **Step 3: Append DataType without renumbering existing values** + +Apply this exact public shape in `src/storage/src/format_base_value.rs`: + +```rust +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum DataType { + String = 0, + Hash = 1, + Set = 2, + List = 3, + ZSet = 4, + None = 5, + All = 6, + VectorSet = 7, +} + +pub const DATA_TYPE_STRINGS: [&str; 8] = [ + "string", "hash", "set", "list", "zset", "none", "all", "vectorset", +]; +pub const DATA_TYPE_TAG: [char; 8] = ['k', 'h', 's', 'l', 'z', 'n', 'a', 'v']; +``` + +Update both exhaustive matches: + +```rust +DataType::Hash | DataType::Set | DataType::ZSet | DataType::VectorSet => { + Ok(BASE_META_VALUE_LENGTH) +} +``` + +```rust +7 => Ok(DataType::VectorSet), +``` + +- [x] **Step 4: Append VectorDataCF to every schema mapping** + +Use the following invariant in all listed files: + +```rust +VectorDataCF = 6 +ColumnFamilyIndex::COUNT = 7 +VectorDataCF.name() = "vector_data_cf" +VectorDataCF.data_type() = Some(DataType::VectorSet) +``` + +Update these exact mapping sites: + +- `src/storage/src/redis.rs`: enum, `COUNT`, `name`, `data_type`, `CF_CONFIGS`, compaction-filter CF list, and `create_rocks_batch` handle list. +- `src/storage/src/batch.rs`: `cf_index_to_usize(ColumnFamilyIndex::VectorDataCF) => 6`. +- `src/storage/src/storage.rs`: `entry.cf_idx == 6` maps to `ColumnFamilyIndex::VectorDataCF`. +- `src/conf/src/raft_type.rs`: append variant and `from_u32(6)`. +- `src/storage/src/logindex/types.rs`: count becomes 7 and both CF-name arrays append `vector_data_cf`. +- `src/raft/src/lib.rs`: append `vector_data_cf` to `CF_NAMES` and `test_cf_names_match_storage` variants. + +Do not insert the new value in the middle of any enum or array. + +- [x] **Step 5: Run the schema tests** + +Run: + +```bash +cargo test -p storage --test redis_basic_test test_open_redis +cargo test -p storage --test redis_basic_test test_column_family_index +cargo test -p raft test_cf_names_match_storage +``` + +Expected: all tests PASS and a newly opened database reports seven CF handles in index order. + +- [x] **Step 6: Commit the schema slice** + +```bash +git add src/storage/src/format_base_value.rs src/storage/src/redis.rs src/storage/src/batch.rs src/storage/src/storage.rs src/conf/src/raft_type.rs src/storage/src/logindex/types.rs src/raft/src/lib.rs src/storage/tests/redis_basic_test.rs +git commit -m "feat(storage): add vector data column family" +``` + +--- + +### Task 2: Implement canonical vector, meta, and value codecs + +**Files:** + +- Create: `src/storage/src/vector.rs` +- Create: `src/storage/src/format_vector.rs` +- Modify: `src/storage/src/lib.rs` + +**Interfaces:** + +- Produces: `CanonicalVector::from_fp32_le(&[u8]) -> Result` +- Produces: `CanonicalVector::from_values(&[f32]) -> Result` +- Produces: `CanonicalVector::score(&CanonicalVector) -> Result` +- Produces: `CanonicalVector::restore() -> Vec` +- Produces: `VectorMeta::{new, encode, decode, count, set_count, version, dimension, is_stale}` +- Produces: `VectorDataValue::{from_canonical, encode, decode}` +- Produces: `VectorQuery`, `VectorSearchMode`, `VectorSearchOptions`, `VectorHit` + +- [x] **Step 1: Write codec and numerical unit tests before exporting modules** + +Place unit tests in the two new files. Cover these exact cases: + +```rust +#[test] +fn canonical_vector_normalizes_and_restores_values() { + let vector = CanonicalVector::from_values(&[3.0, 4.0]).expect("valid vector"); + assert_eq!(vector.dimension(), 2); + assert!((vector.original_l2() - 5.0).abs() < 1e-6); + let restored = vector.restore(); + assert!((restored[0] - 3.0).abs() < 1e-6); + assert!((restored[1] - 4.0).abs() < 1e-6); +} + +#[test] +fn canonical_vector_rejects_invalid_inputs() { + assert!(CanonicalVector::from_values(&[]).is_err()); + assert!(CanonicalVector::from_values(&[0.0, 0.0]).is_err()); + assert!(CanonicalVector::from_values(&[f32::NAN]).is_err()); + assert!(CanonicalVector::from_fp32_le(&[0, 1, 2]).is_err()); +} + +#[test] +fn cosine_score_maps_to_redis_range() { + let x = CanonicalVector::from_values(&[1.0, 0.0]).expect("valid x"); + let same = CanonicalVector::from_values(&[1.0, 0.0]).expect("valid same"); + let opposite = CanonicalVector::from_values(&[-1.0, 0.0]).expect("valid opposite"); + assert!((x.score(&same).expect("score") - 1.0).abs() < 1e-12); + assert!(x.score(&opposite).expect("score").abs() < 1e-12); +} +``` + +Codec tests must round-trip an empty binary element independently through `MemberDataKey`, a two-dimensional vector value, and a `VectorMeta` with `count=2`, `version=42`, `dimension=2`. + +- [x] **Step 2: Run tests and confirm the modules are missing** + +Run: + +```bash +cargo test -p storage vector::tests +cargo test -p storage format_vector::tests +``` + +Expected: compile failure until the modules and types are implemented and exported. + +- [x] **Step 3: Implement the public vector types** + +Define the following exact public API in `src/storage/src/vector.rs`: + +```rust +#[derive(Debug, Clone, PartialEq)] +pub struct CanonicalVector { + dimension: u32, + original_l2: f32, + normalized: Vec, +} + +#[derive(Debug, Clone, PartialEq)] +pub enum VectorQuery { + Element(Vec), + Vector(CanonicalVector), +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum VectorSearchMode { + Approximate, + Truth, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct VectorSearchOptions { + pub count: usize, + pub mode: VectorSearchMode, +} + +#[derive(Debug, Clone, PartialEq)] +pub struct VectorHit { + pub element: Vec, + pub score: f64, +} +``` + +Numerical rules: + +- FP32 blob length must be non-zero and divisible by four. +- Parse every component with `f32::from_le_bytes`; reject non-finite values. +- Accumulate `norm²` in `f64`; reject zero or non-finite norm. +- Store normalized components as `f32`, store original L2 as `f32`. +- Compute dot product in `f64`, clamp it to `[-1, 1]`, then return `(dot + 1.0) / 2.0` clamped to `[0, 1]`. +- Reject score comparisons when dimensions differ. + +- [x] **Step 4: Implement fixed codecs** + +Use this exact metadata layout in `src/storage/src/format_vector.rs`: + +```text +| type=7 | count | version | format | encoding | metric | flags | dimension | zero reserve | ctime | etime | +| 1B | 8B | 8B | 1B | 1B | 1B | 1B | 4B LE | 8B | 8B | 8B | +``` + +Constants: + +```rust +pub const VECTOR_META_FORMAT: u8 = 1; +pub const VECTOR_ENCODING_FP32_LE: u8 = 1; +pub const VECTOR_METRIC_COSINE: u8 = 1; +pub const VECTOR_VALUE_MAGIC: u8 = 0x56; +pub const VECTOR_VALUE_FORMAT: u8 = 1; +``` + +Use this value layout: + +```text +| magic=0x56 | format=1 | dimension | original_l2 | normalized FP32 payload | +| 1B | 1B | 4B LE | 4B LE | dimension * 4B | +``` + +Decoders must check exact length, magic, format, non-zero dimension, finite positive `original_l2`, finite payload components, and payload dimension. Return `InvalidFormatSnafu` on malformed persisted bytes; never panic. + +- [x] **Step 5: Export the new modules and rerun tests** + +Add to `src/storage/src/lib.rs`: + +```rust +mod format_vector; +pub mod vector; + +pub use vector::{ + CanonicalVector, VectorHit, VectorQuery, VectorSearchMode, VectorSearchOptions, +}; +``` + +Run: + +```bash +cargo test -p storage vector::tests +cargo test -p storage format_vector::tests +``` + +Expected: all canonicalization and codec tests PASS. + +- [x] **Step 6: Commit the codec slice** + +```bash +git add src/storage/src/lib.rs src/storage/src/vector.rs src/storage/src/format_vector.rs +git commit -m "feat(storage): add vector codecs" +``` + +--- + +### Task 3: Implement atomic VectorSet mutations and point reads + +**Files:** + +- Create: `src/storage/src/redis_vectors.rs` +- Modify: `src/storage/src/lib.rs` +- Create: `src/storage/tests/redis_vector_test.rs` + +**Interfaces:** + +- Consumes: `CanonicalVector`, `VectorMeta`, `VectorDataValue`, `MemberDataKey` +- Produces: `Redis::vadd(key, element, vector) -> Result` +- Produces: `Redis::vrem(key, element) -> Result` +- Produces: `Redis::vcard(key) -> Result` +- Produces: `Redis::vdim(key) -> Result` +- Produces: `Redis::vemb(key, element) -> Result>>` +- Produces: `Redis::vismember(key, element) -> Result` +- Produces: `Redis::is_cluster_mode() -> bool` + +- [x] **Step 1: Write storage behavior tests** + +Create the test file using `unique_test_db_path`, `safe_cleanup_test_db`, `StorageOptions`, `BgTaskHandler`, and `LockMgr`, matching existing storage test setup. Add these independently named tests: + +- `test_vadd_create_update_and_dimension_guard` +- `test_vadd_is_binary_safe_and_accepts_empty_element` +- `test_vcard_vdim_vemb_and_vismember_missing_semantics` +- `test_vrem_deletes_last_member_and_meta` +- `test_vector_commands_return_wrongtype_for_string_key` +- `test_vector_meta_and_member_are_committed_together` + +The core assertions are: + +```rust +let a = CanonicalVector::from_values(&[1.0, 0.0]).expect("valid vector"); +let b = CanonicalVector::from_values(&[0.0, 1.0]).expect("valid vector"); + +assert!(redis.vadd(b"vectors", b"a", &a).expect("insert a")); +assert!(redis.vadd(b"vectors", b"b", &b).expect("insert b")); +assert!(!redis.vadd(b"vectors", b"a", &b).expect("update a")); +assert_eq!(redis.vcard(b"vectors").expect("card"), 2); +assert_eq!(redis.vdim(b"vectors").expect("dim"), 2); +assert_eq!(redis.vemb(b"vectors", b"a").expect("emb"), Some(vec![0.0, 1.0])); +assert!(redis.vismember(b"vectors", b"a").expect("member")); +assert!(redis.vadd(b"vectors", b"\x00binary", &a).expect("binary member")); +assert!(redis.vadd(b"empty-element", b"", &a).expect("empty member")); +``` + +For dimension mismatch, assert `is_err()` and confirm that `VCARD` and the old member value remain unchanged. For last-member removal, assert `get_key_type(key)` returns a missing-key error after `VREM`. + +- [x] **Step 2: Run the behavior tests and confirm missing methods** + +Run: + +```bash +cargo test -p storage --test redis_vector_test +``` + +Expected: compile failure because the six `Redis` methods are not implemented. + +- [x] **Step 3: Implement one locked read-modify-write path for VADD** + +Add this private module declaration to `src/storage/src/lib.rs`: + +```rust +mod redis_vectors; +``` + +Implementation sequence inside `Redis::vadd`: + +1. Reject cluster mode by checking `self.append_log_fn.get().is_some()`. +2. Acquire `ScopeRecordLock` using the user key. +3. Read `MetaCF[BaseMetaKey::new(key)]`. +4. Missing or stale meta creates `VectorMeta::new(1, dimension)` with a new timestamp-based version and does not increment it again. +5. Live non-VectorSet meta returns the same WRONGTYPE text used by `check_type_state`. +6. Live VectorSet requires equal dimension. +7. Point-read `VectorDataCF[MemberDataKey(key, version, element)]` to distinguish insert from update. +8. Inserting into an existing live VectorSet increments count; creating a new VectorSet keeps the initial count at one; updating an existing element preserves count. +9. Build one batch containing the member put and meta put; commit once. +10. Return `true` for insert and `false` for update. + +The public signature must be: + +```rust +pub fn vadd( + &self, + key: &[u8], + element: &[u8], + vector: &CanonicalVector, +) -> Result; +``` + +- [x] **Step 4: Implement VREM and point-read commands** + +Use `VectorMeta::decode` and the same `MemberDataKey` generation for every operation. + +VREM rules: + +- Missing/stale key or missing element returns `Ok(false)`. +- Wrong type returns WRONGTYPE. +- If count is greater than one, delete member and write `count - 1` meta in one batch. +- If count is one, delete member and delete MetaCF in one batch. + +Read rules: + +- `vcard`: missing/stale returns `0`. +- `vdim`: missing/stale returns `KeyNotFoundSnafu`. +- `vemb`: missing/stale member or key returns `None`; live value returns `CanonicalVector::restore()`. +- `vismember`: missing/stale member or key returns `false`. +- All live wrong-type keys return WRONGTYPE. + +- [x] **Step 5: Run storage tests** + +Run: + +```bash +RUST_TEST_THREADS=1 cargo test -p storage --test redis_vector_test +cargo test -p storage vector::tests format_vector::tests +``` + +Expected: all tests PASS, including rollback-visible behavior after dimension mismatch. + +- [x] **Step 6: Commit point operations** + +```bash +git add src/storage/src/lib.rs src/storage/src/redis_vectors.rs src/storage/tests/redis_vector_test.rs +git commit -m "feat(storage): persist vector set members" +``` + +--- + +### Task 4: Implement exact FLAT search with stable Top-K ordering + +**Files:** + +- Modify: `src/storage/src/redis_vectors.rs` +- Modify: `src/storage/src/vector.rs` +- Modify: `src/storage/tests/redis_vector_test.rs` + +**Interfaces:** + +- Consumes: `VectorQuery`, `VectorSearchOptions` +- Produces: `Redis::vsim(key, query, options) -> Result>` +- Ordering contract: score descending, then raw element bytes ascending +- Complexity contract: `O(N × DIM + N log K)` time and `O(K + DIM)` extra memory + +- [x] **Step 1: Add failing FLAT search tests** + +Add these tests: + +- `test_vsim_direct_vector_returns_exact_top_k` +- `test_vsim_ele_uses_stored_member_as_query` +- `test_vsim_stable_tie_breaks_by_raw_element_bytes` +- `test_vsim_truth_matches_approximate_in_phase_one` +- `test_vsim_missing_key_is_empty_and_missing_ele_is_error` +- `test_vsim_rejects_query_dimension_mismatch` + +Use this deterministic dataset: + +```rust +let x = CanonicalVector::from_values(&[1.0, 0.0]).expect("x"); +let y = CanonicalVector::from_values(&[0.0, 1.0]).expect("y"); +let neg_x = CanonicalVector::from_values(&[-1.0, 0.0]).expect("negative x"); +redis.vadd(b"search", b"b", &y).expect("insert b"); +redis.vadd(b"search", b"a", &y).expect("insert a"); +redis.vadd(b"search", b"x", &x).expect("insert x"); +redis.vadd(b"search", b"neg", &neg_x).expect("insert neg"); +``` + +For query `x` and `COUNT 3`, require `x` first, then `a`, then `b`; `a` and `b` tie on score and raw bytes decide their order. + +- [x] **Step 2: Run the focused tests and confirm VSim is missing** + +```bash +cargo test -p storage --test redis_vector_test test_vsim +``` + +Expected: compile failure because `Redis::vsim` is not defined. + +- [x] **Step 3: Add a heap item with total ordering** + +Define a private `HeapHit` that compares worse hits as the heap maximum, so a heap capped at K can replace its worst element. Equality and ordering must use both `score.total_cmp()` and raw element bytes; do not use `partial_cmp().unwrap()`. + +Required result ordering after draining the heap: + +```rust +hits.sort_by(|left, right| { + right + .score + .total_cmp(&left.score) + .then_with(|| left.element.cmp(&right.element)) +}); +``` + +- [x] **Step 4: Implement one RocksDB snapshot scan** + +The exact search sequence is: + +1. Return an empty vector when the key is missing or stale. +2. Return WRONGTYPE for a live non-VectorSet key. +3. Decode meta and resolve an `Element` query with a point read from the same RocksDB snapshot. +4. Reject missing query element and dimension mismatch. +5. Build `MemberDataKey::new(key, version, b"").encode_seek_key()`. +6. Iterate `VectorDataCF` forward from the prefix using snapshot read options. +7. Stop at the first key not starting with the prefix. +8. Strip prefix and the 16-byte suffix reserve to recover raw element bytes. +9. Decode each vector value and compute score. +10. Keep at most `options.count` heap entries. +11. Drain and sort using the stable ordering above. + +Reject `options.count == 0` before creating the iterator. + +- [x] **Step 5: Run FLAT and regression tests** + +```bash +RUST_TEST_THREADS=1 cargo test -p storage --test redis_vector_test test_vsim +RUST_TEST_THREADS=1 cargo test -p storage --test redis_set_test +``` + +Expected: Vector search tests PASS and existing Set prefix-iteration tests remain green. + +- [x] **Step 6: Commit the search engine** + +```bash +git add src/storage/src/vector.rs src/storage/src/redis_vectors.rs src/storage/tests/redis_vector_test.rs +git commit -m "feat(storage): add flat vector similarity search" +``` + +--- + +### Task 5: Wire Storage routing, TYPE, TTL cleanup, DEL, and FLUSHDB + +**Files:** + +- Modify: `src/storage/src/storage_impl.rs` +- Modify: `src/storage/src/data_compaction_filter.rs` +- Modify: `src/storage/src/meta_compaction_filter.rs` +- Modify: `src/storage/src/redis.rs` +- Modify: `src/storage/src/redis_strings.rs` +- Modify: `src/storage/tests/redis_vector_test.rs` + +**Interfaces:** + +- Produces: the same seven operations on `Storage`, each routed exactly once by user key +- Preserves: `TYPE key -> vectorset` +- Preserves: generic EXPIRE/PERSIST behavior through the common Meta envelope +- Preserves: generic DEL and FLUSHDB physically delete VectorDataCF entries + +- [x] **Step 1: Add facade and lifecycle tests** + +Add tests using `Storage::new(3, 0)` for: + +- `test_storage_routes_all_members_of_one_vectorset_to_one_instance` +- `test_type_returns_vectorset` +- `test_expired_vectorset_reads_as_missing` +- `test_del_removes_vector_meta_and_members` +- `test_flushdb_removes_vector_meta_and_members` +- `test_vector_storage_rejects_cluster_mode` + +The multi-instance test must compute `key_to_slot_id(key)` and confirm only the selected instance contains the MetaCF key; it must not scan or fan out for VSIM. + +- [x] **Step 2: Run tests and confirm missing Storage methods** + +```bash +cargo test -p storage --test redis_vector_test test_storage_routes +cargo test -p storage --test redis_vector_test test_del_removes_vector +``` + +Expected: compile failure because the Vector methods exist only on `Redis`. + +- [x] **Step 3: Add exact one-key routing methods** + +Add these signatures to `src/storage/src/storage_impl.rs`: + +```rust +pub fn vadd(&self, key: &[u8], element: &[u8], vector: &CanonicalVector) -> Result; +pub fn vsim( + &self, + key: &[u8], + query: VectorQuery, + options: VectorSearchOptions, +) -> Result>; +pub fn vrem(&self, key: &[u8], element: &[u8]) -> Result; +pub fn vcard(&self, key: &[u8]) -> Result; +pub fn vdim(&self, key: &[u8]) -> Result; +pub fn vemb(&self, key: &[u8], element: &[u8]) -> Result>>; +pub fn vismember(&self, key: &[u8], element: &[u8]) -> Result; +``` + +Every method must perform the same two routing statements and then call its matching Redis method. For example, `Storage::vadd` is: + +```rust +let slot_id = key_to_slot_id(key); +let instance_id = self.slot_indexer.get_instance_id(slot_id); +self.insts[instance_id].vadd(key, element, vector) +``` + +Apply the identical prefix to `vsim`, `vrem`, `vcard`, `vdim`, `vemb`, and `vismember`, forwarding their declared arguments unchanged. There is no element routing and no instance fan-out. + +- [x] **Step 4: Extend common lifecycle handling** + +Apply these exact additions: + +- `src/storage/src/redis.rs::is_stale_static`: include `DataType::VectorSet` in the common count/version meta branch. +- `src/storage/src/redis_strings.rs::set_key_etime`: parse `DataType::VectorSet` with `ParsedBaseMetaValue`, identical to Hash/Set/ZSet. +- `src/storage/src/meta_compaction_filter.rs`: parse `DataType::VectorSet` with `ParsedBaseMetaValue`. +- `src/storage/src/data_compaction_filter.rs::parse_meta_value` and its test-only meta parser: parse VectorSet with `ParsedBaseMetaValue`. +- `Redis::del_key`: add `ColumnFamilyIndex::VectorDataCF` to the composite data CF scan list. +- `Redis::flush_db`: add `ColumnFamilyIndex::VectorDataCF` to `all_cf_indexes`. + +Do not introduce a wildcard in the schema matches above; each VectorSet/VectorDataCF branch must be explicit. + +- [x] **Step 5: Run lifecycle and existing TTL tests** + +```bash +RUST_TEST_THREADS=1 cargo test -p storage --test redis_vector_test +RUST_TEST_THREADS=1 cargo test -p storage --test ttl_test +RUST_TEST_THREADS=1 cargo test -p storage --test redis_basic_test +``` + +Expected: all tests PASS; expired vector data becomes invisible immediately and eligible for compaction, while DEL/FLUSHDB remove physical VectorDataCF entries. + +- [x] **Step 6: Commit routing and lifecycle support** + +```bash +git add src/storage/src/storage_impl.rs src/storage/src/data_compaction_filter.rs src/storage/src/meta_compaction_filter.rs src/storage/src/redis.rs src/storage/src/redis_strings.rs src/storage/tests/redis_vector_test.rs +git commit -m "feat(storage): route and clean up vector sets" +``` + +--- + +### Task 6: Add Vector Set command parsing and command-table registration + +**Files:** + +- Create: `src/cmd/src/vector.rs` +- Modify: `src/cmd/src/lib.rs` +- Modify: `src/cmd/src/table.rs` + +**Interfaces:** + +- Produces: `VAddCmd`, `VSimCmd`, `VRemCmd`, `VCardCmd`, `VDimCmd`, `VEmbCmd`, `VIsMemberCmd` +- Consumes: all seven `Storage` APIs from Task 5 +- Produces: binary-safe argv parsing and `RespData` replies + +- [x] **Step 1: Write parser and metadata unit tests in the new module** + +Cover all supported command shapes: + +```text +VADD key FP32 blob element NOQUANT +VADD key VALUES num value [value ...] element NOQUANT +VSIM key ELE element [WITHSCORES] [COUNT num] [TRUTH] +VSIM key FP32 blob [WITHSCORES] [COUNT num] [TRUTH] +VSIM key VALUES num value [value ...] [WITHSCORES] [COUNT num] [TRUTH] +VREM key element +VCARD key +VDIM key +VEMB key element +VISMEMBER key element +``` + +Add explicit failure tests for malformed FP32 length, invalid VALUES dimension/count, invalid float, missing `NOQUANT`, `Q8`, `BIN`, duplicate VSIM options, zero/invalid COUNT, `VEMB RAW`, and trailing unknown options. + +Metadata assertions: + +```rust +assert_eq!(VAddCmd::new().meta().arity, -5); +assert_eq!(VSimCmd::new().meta().arity, -4); +assert_eq!(VRemCmd::new().meta().arity, 3); +assert_eq!(VCardCmd::new().meta().arity, 2); +assert_eq!(VDimCmd::new().meta().arity, 2); +assert_eq!(VEmbCmd::new().meta().arity, -3); +assert_eq!(VIsMemberCmd::new().meta().arity, 3); +``` + +- [x] **Step 2: Run command tests and confirm the module is not registered** + +```bash +cargo test -p cmd vector::tests +``` + +Expected: compile failure until the new module, command structs, and parsers exist. + +- [x] **Step 3: Implement shared parsers without UTF-8 assumptions for keys/elements/blobs** + +Only option keywords and numeric VALUES tokens are decoded as ASCII/UTF-8. Keep key, element, and FP32 blob as raw bytes. + +Define private parsed forms: + +```rust +struct ParsedVAdd { + vector: CanonicalVector, + element: Vec, +} + +struct ParsedVSim { + query: VectorQuery, + options: VectorSearchOptions, + with_scores: bool, +} +``` + +Error strings must be exact and already include the Redis error class: + +```text +ERR invalid vector specification +ERR vector dimension mismatch +ERR default Q8 quantization is not supported in Phase 1; specify NOQUANT +ERR VADD option Q8 is not supported yet +ERR VADD option BIN is not supported yet +ERR VEMB option RAW is not supported yet +ERR element not found in set +ERR key does not exist +WRONGTYPE Operation against a key holding the wrong kind of value +``` + +Add one local reply helper that avoids double-prefixing current storage error strings: + +```rust +fn error_reply(message: impl Into) -> RespData { + RespData::Error(message.into().into()) +} +``` + +Map internal non-Redis storage failures to `ERR storage error`; log the detailed error with `log::error!`. + +- [x] **Step 4: Implement command replies** + +Reply contracts: + +- VADD/VREM/VISMEMBER: integer `1` or `0`. +- VCARD/VDIM: integer after checked `u64/u32 -> i64` conversion. +- VEMB: null bulk for missing key/member; otherwise array of `RespData::Double`. +- VSIM without scores: array of bulk-string elements. +- VSIM with scores: `RespData::Map` of bulk-string element to `RespData::Double`; Task 7 performs RESP2 downgrade at the encoder boundary. + +Flags and ACL categories: + +- VADD/VREM: `WRITE | FAST | MODULE_NO_CLUSTER`, `KEYSPACE | WRITE`. +- VSIM: `READONLY | MODULE_NO_CLUSTER`, `KEYSPACE | READ | SLOW`. +- VCARD/VDIM/VEMB/VISMEMBER: `READONLY | FAST | MODULE_NO_CLUSTER`, `KEYSPACE | READ`. + +- [x] **Step 5: Export and register all seven commands** + +Add `pub mod vector;` to `src/cmd/src/lib.rs`. Append these types to `register_cmd!` in `src/cmd/src/table.rs`: + +```rust +crate::vector::VAddCmd, +crate::vector::VSimCmd, +crate::vector::VRemCmd, +crate::vector::VCardCmd, +crate::vector::VDimCmd, +crate::vector::VEmbCmd, +crate::vector::VIsMemberCmd, +``` + +Add a table test asserting all lowercase names resolve. + +- [x] **Step 6: Run command and runtime dispatch tests** + +```bash +cargo test -p cmd vector::tests +cargo test -p cmd table +cargo test -p runtime handle_execute_command +``` + +Expected: all tests PASS and the generic `StorageCommand::Execute` path finds every Vector command without a new runtime message variant. + +- [x] **Step 7: Commit the command layer** + +```bash +git add src/cmd/src/vector.rs src/cmd/src/lib.rs src/cmd/src/table.rs +git commit -m "feat(cmd): add redis vector set commands" +``` + +--- + +### Task 7: Make VSIM WITHSCORES correct in RESP2 and RESP3 + +**Files:** + +- Modify: `src/resp/src/encode.rs` +- Modify: `src/resp/src/negotiation.rs` + +**Interfaces:** + +- Consumes: command-layer `RespData::Map` and `RespData::Double` +- Produces: RESP3 `%` map with `,` doubles +- Produces: RESP2 flat array with bulk-string scores + +- [x] **Step 1: Add encoder regression tests** + +Add tests that encode this response under both protocol versions: + +```rust +let response = RespData::Map(vec![ + ( + RespData::BulkString(Some(Bytes::from_static(b"a"))), + RespData::Double(1.0), + ), + ( + RespData::BulkString(Some(Bytes::from_static(b"b"))), + RespData::Double(0.5), + ), +]); +``` + +Expected RESP3 bytes: + +```text +%2\r\n$1\r\na\r\n,1\r\n$1\r\nb\r\n,0.5\r\n +``` + +Expected RESP2 bytes: + +```text +*4\r\n$1\r\na\r\n$1\r\n1\r\n$1\r\nb\r\n$3\r\n0.5\r\n +``` + +- [x] **Step 2: Run the test and confirm RESP2 currently emits RESP3 types** + +```bash +cargo test -p resp encode_vsim_withscores +``` + +Expected: RESP2 assertion FAIL because `RespEncoder` currently stores the version but does not use it when encoding Map/Double. + +- [x] **Step 3: Normalize once at the encoder boundary** + +Refactor `RespEncoder` so the public trait method converts the complete response once, then calls a private recursive encoder: + +```rust +fn encode_resp_data(&mut self, data: &RespData) -> &mut Self { + let normalized = if self.is_resp3() { + data.clone() + } else { + ProtocolNegotiator::convert_to_resp2(data) + }; + self.encode_resp_data_inner(&normalized) +} +``` + +Move the existing recursive match into `encode_resp_data_inner`, and make all nested array/map/set/push recursion call `encode_resp_data_inner` rather than invoking the version-normalizing method again. + +Keep `ProtocolNegotiator::convert_to_resp2` recursive for Array, Map, Set, Push, Null, Boolean, Double, BigNumber, BulkError, and VerbatimString. + +- [x] **Step 4: Run all RESP tests** + +```bash +cargo test -p resp +``` + +Expected: all RESP tests PASS; existing RESP3 HELLO behavior remains unchanged and RESP2 never emits RESP3-only prefixes. + +- [x] **Step 5: Commit protocol-aware encoding** + +```bash +git add src/resp/src/encode.rs src/resp/src/negotiation.rs +git commit -m "fix(resp): downgrade vector scores for resp2" +``` + +--- + +### Task 8: Add end-to-end compatibility tests and run the Kiwi verification pipeline + +**Files:** + +- Create: `tests/python/test_vector_set_commands.py` +- Modify: `docs/superpowers/plans/2026-07-19-redis-vector-set.md` only to check completed boxes during execution + +**Interfaces:** + +- Verifies: command parsing through TCP, dual runtime dispatch, storage persistence, RESP2/RESP3 encoding, TYPE/DEL/EXPIRE interaction +- Does not consume: `tests/python/test_vector_basic.py`, FT.CREATE, FT.SEARCH, HSET vector fields + +- [x] **Step 1: Write binary-safe Python integration tests** + +Create a local fixture in the new test file using: + +```python +import os +import struct + +import pytest +import redis + + +@pytest.fixture(params=[2, 3]) +def vector_client(request): + client = redis.Redis( + host=os.getenv("KIWI_HOST", "127.0.0.1"), + port=int(os.getenv("KIWI_PORT", "7379")), + decode_responses=False, + protocol=request.param, + ) + client.ping() + yield client + for key in client.scan_iter(match=b"test_vset:*"): + client.delete(key) + client.close() +``` + +Add tests for: + +- VALUES create/update return values and dimension mismatch. +- FP32 little-endian blob via `struct.pack("<2f", 1.0, 0.0)`. +- Binary key and element including `b"\x00"`; empty element. +- VCARD, VDIM, VEMB, VISMEMBER, VREM and last-member key deletion. +- VSIM direct vector, ELE query, COUNT, WITHSCORES and TRUTH. +- Stable tie order. +- TYPE returns `b"vectorset"`. +- WRONGTYPE against an existing string key. +- Missing-key semantics for every read command. +- EXPIRE makes the key immediately behave as missing. +- DEL removes members and allows same-name recreation with a new dimension. +- Missing NOQUANT, Q8, BIN, RAW and malformed input errors. + +For WITHSCORES, assert redis-py returns a list under RESP2 and a dict under RESP3, with the same element order and numerical scores. + +- [x] **Step 2: Start standalone Kiwi and run only Vector integration tests** + +Terminal 1: + +```bash +make standalone +``` + +Terminal 2: + +```bash +make -C tests install-deps +KIWI_PORT=7379 pytest -q tests/python/test_vector_set_commands.py +``` + +Expected: all RESP2 and RESP3 parameterized cases PASS. + +- [x] **Step 3: Run focused Rust regression suites** + +```bash +RUST_TEST_THREADS=1 cargo test -p storage --test redis_vector_test +RUST_TEST_THREADS=1 cargo test -p storage --test redis_hash_test +RUST_TEST_THREADS=1 cargo test -p storage --test redis_set_test +RUST_TEST_THREADS=1 cargo test -p storage --test ttl_test +cargo test -p cmd vector::tests +cargo test -p resp +``` + +Expected: all commands exit 0. + +- [x] **Step 4: Run repository-standard verification** + +Run exactly in this order: + +```bash +make fmt +make lint +make build +make test +``` + +Expected: all commands exit 0. The first RocksDB build may be long; retain sccache and do not disable it repository-wide. + +- [x] **Step 5: Inspect the final diff for scope and generated files** + +```bash +git status --short +git diff --stat origin/main...HEAD +git diff --check origin/main...HEAD +``` + +Expected: + +- No changes to `tests/python/test_vector_basic.py`. +- No FT.CREATE/FT.SEARCH implementation. +- No HNSW, Raft mutation, snapshot, or O(1) deletion code. +- No `__pycache__`, RocksDB data directory, build artifacts, or scratch notes staged. +- `git diff --check` produces no output. + +- [x] **Step 6: Commit end-to-end tests** + +```bash +git add tests/python/test_vector_set_commands.py docs/superpowers/plans/2026-07-19-redis-vector-set.md +git commit -m "test: cover redis vector set commands" +``` + +--- + +## Completion Criteria + +The Phase 1 implementation is complete only when all statements below are true: + +- A VectorSet key and all its elements live in one RocksDB instance selected only from the user key. +- Meta and member mutations commit in one batch; failed validation leaves both unchanged. +- `TYPE` reports `vectorset`; generic TTL, DEL and FLUSHDB semantics work. +- VADD accepts FP32 and VALUES only with explicit NOQUANT and is binary-safe for elements. +- VSIM returns exact, deterministic Top-K results without collecting and sorting every hit. +- RESP2 and RESP3 WITHSCORES replies use their native compatible shapes. +- Cluster mode rejects all seven Vector Set operations before creating a batch or snapshot. +- Focused Rust/Python tests and the complete Kiwi verification pipeline pass. +- The final diff contains no FT prototype extraction, HNSW, Raft/snapshot design implementation, or unrelated cleanup. + +## Deferred Follow-up Plans + +Create separate reviewed designs and plans before implementing any of these independent projects: + +1. Redis-compatible VINFO and INFO VECTOR observability. +2. Raft logical VectorSet mutations, leader-linearizable reads and snapshots. +3. O(1) generation deletion with background compaction cleanup. +4. HNSW derived index, rebuild/recovery and rolling-upgrade capability gates. diff --git a/tests/python/test_vector_set_commands.py b/tests/python/test_vector_set_commands.py new file mode 100644 index 00000000..2cbb3ae4 --- /dev/null +++ b/tests/python/test_vector_set_commands.py @@ -0,0 +1,257 @@ +# Copyright (c) 2024-present, arana-db Community. All rights reserved. +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import re +import struct +import time + +import pytest +import redis + + +@pytest.fixture(params=[2, 3]) +def vector_client(request): + protocol = request.param + prefix = f"test_vset:p{protocol}:".encode() + keys = [ + prefix + b"values", + prefix + b"binary:\x00key", + prefix + b"search", + prefix + b"missing", + prefix + b"expired", + prefix + b"recreated", + prefix + b"errors", + prefix + b"string", + prefix + b"malformed", + ] + client = redis.Redis( + host=os.getenv("KIWI_HOST", "127.0.0.1"), + port=int(os.getenv("KIWI_PORT", "7379")), + decode_responses=False, + protocol=protocol, + ) + client.ping() + client.delete(*keys) + yield client, protocol, prefix + client.delete(*keys) + client.close() + + +def vadd_values(client, key, values, element, quantization=b"NOQUANT"): + return client.execute_command( + b"VADD", + key, + b"VALUES", + len(values), + *values, + element, + quantization, + ) + + +def assert_response_error(client, expected, *command): + with pytest.raises(redis.ResponseError, match=re.escape(expected)): + client.execute_command(*command) + + +def assert_vector(values, expected): + assert [float(value) for value in values] == pytest.approx(expected) + + +def test_values_create_update_and_point_commands(vector_client): + client, _protocol, prefix = vector_client + key = prefix + b"values" + + assert vadd_values(client, key, [1, 0], b"member") == 1 + assert vadd_values(client, key, [0.5, 0.5], b"member") == 0 + assert client.execute_command(b"VCARD", key) == 1 + assert client.execute_command(b"VDIM", key) == 2 + assert client.execute_command(b"VISMEMBER", key, b"member") == 1 + assert_vector(client.execute_command(b"VEMB", key, b"member"), [0.5, 0.5]) + assert client.type(key) == b"vectorset" + + assert_response_error( + client, + "vector dimension mismatch", + b"VADD", + key, + b"VALUES", + 3, + 1, + 0, + 0, + b"other", + b"NOQUANT", + ) + + +def test_fp32_binary_members_and_last_member_removal(vector_client): + client, _protocol, prefix = vector_client + key = prefix + b"binary:\x00key" + blob = struct.pack("<2f", 1.0, 0.0) + + assert client.execute_command(b"VADD", key, b"FP32", blob, b"", b"NOQUANT") == 1 + assert vadd_values(client, key, [0, 1], b"\x00member") == 1 + assert client.execute_command(b"VCARD", key) == 2 + assert_vector(client.execute_command(b"VEMB", key, b""), [1.0, 0.0]) + assert client.execute_command(b"VISMEMBER", key, b"\x00member") == 1 + + assert client.execute_command(b"VREM", key, b"") == 1 + assert client.execute_command(b"VREM", key, b"\x00member") == 1 + assert client.execute_command(b"VREM", key, b"\x00member") == 0 + assert client.type(key) == b"none" + + +def test_vsim_queries_scores_truth_and_stable_ties(vector_client): + client, protocol, prefix = vector_client + key = prefix + b"search" + vadd_values(client, key, [1, 0], b"a") + vadd_values(client, key, [1, 0], b"b") + vadd_values(client, key, [0, 1], b"c") + + direct = client.execute_command( + b"VSIM", key, b"VALUES", 2, 1, 0, b"COUNT", 2, b"TRUTH" + ) + assert direct == [b"a", b"b"] + + by_element = client.execute_command( + b"VSIM", key, b"ELE", b"a", b"COUNT", 3, b"TRUTH" + ) + assert by_element == [b"a", b"b", b"c"] + + scores = client.execute_command( + b"VSIM", + key, + b"FP32", + struct.pack("<2f", 1.0, 0.0), + b"WITHSCORES", + b"COUNT", + 2, + b"TRUTH", + ) + if protocol == 2: + assert scores[::2] == [b"a", b"b"] + assert [float(value) for value in scores[1::2]] == pytest.approx([1.0, 1.0]) + else: + assert list(scores) == [b"a", b"b"] + assert list(scores.values()) == pytest.approx([1.0, 1.0]) + + +def test_missing_key_and_missing_element_semantics(vector_client): + client, _protocol, prefix = vector_client + key = prefix + b"missing" + + assert client.execute_command(b"VCARD", key) == 0 + assert client.execute_command(b"VEMB", key, b"member") is None + assert client.execute_command(b"VISMEMBER", key, b"member") == 0 + assert client.execute_command(b"VSIM", key, b"VALUES", 2, 1, 0) == [] + assert_response_error(client, "key does not exist", b"VDIM", key) + + vadd_values(client, key, [1, 0], b"present") + assert_response_error( + client, "element not found in set", b"VSIM", key, b"ELE", b"absent" + ) + + +def test_expire_del_and_recreate_with_new_dimension(vector_client): + client, _protocol, prefix = vector_client + expired = prefix + b"expired" + recreated = prefix + b"recreated" + + vadd_values(client, expired, [1, 0], b"member") + assert client.expire(expired, 1) + time.sleep(1.1) + assert client.type(expired) == b"none" + assert client.execute_command(b"VCARD", expired) == 0 + assert client.execute_command(b"VEMB", expired, b"member") is None + + vadd_values(client, recreated, [1, 0], b"member") + assert client.delete(recreated) == 1 + assert client.execute_command(b"VISMEMBER", recreated, b"member") == 0 + assert vadd_values(client, recreated, [1, 0, 0], b"new-member") == 1 + assert client.execute_command(b"VDIM", recreated) == 3 + + +def test_wrongtype_and_unsupported_options(vector_client): + client, _protocol, prefix = vector_client + key = prefix + b"errors" + string_key = prefix + b"string" + client.set(string_key, b"value") + + assert_response_error( + client, + "WRONGTYPE Operation against a key holding the wrong kind of value", + b"VADD", + string_key, + b"VALUES", + 2, + 1, + 0, + b"member", + b"NOQUANT", + ) + + base = (b"VADD", key, b"VALUES", 2, 1, 0, b"member") + assert_response_error(client, "default Q8 quantization is not supported", *base) + assert_response_error(client, "VADD option Q8 is not supported yet", *base, b"Q8") + assert_response_error(client, "VADD option BIN is not supported yet", *base, b"BIN") + + vadd_values(client, key, [1, 0], b"member") + assert_response_error( + client, "VEMB option RAW is not supported yet", b"VEMB", key, b"member", b"RAW" + ) + + +def test_malformed_vectors_and_options(vector_client): + client, _protocol, prefix = vector_client + key = prefix + b"malformed" + + assert_response_error( + client, + "invalid vector specification", + b"VADD", + key, + b"FP32", + b"abc", + b"member", + b"NOQUANT", + ) + assert_response_error( + client, + "invalid vector specification", + b"VADD", + key, + b"VALUES", + 2, + 1, + b"not-a-float", + b"member", + b"NOQUANT", + ) + assert_response_error( + client, + "invalid vector specification", + b"VSIM", + key, + b"VALUES", + 2, + 1, + 0, + b"COUNT", + 0, + ) From dfd841fd4d49d110b2ed79553daf45761276ec82 Mon Sep 17 00:00:00 2001 From: happy-v587 Date: Mon, 20 Jul 2026 11:37:37 +0800 Subject: [PATCH 10/32] test: skip vector set commands when server is unavailable --- tests/python/test_vector_set_commands.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/python/test_vector_set_commands.py b/tests/python/test_vector_set_commands.py index 2cbb3ae4..c03b59fc 100644 --- a/tests/python/test_vector_set_commands.py +++ b/tests/python/test_vector_set_commands.py @@ -45,7 +45,11 @@ def vector_client(request): decode_responses=False, protocol=protocol, ) - client.ping() + try: + client.ping() + except redis.ConnectionError: + client.close() + pytest.skip("Kiwi server is not running") client.delete(*keys) yield client, protocol, prefix client.delete(*keys) From ede10cd81f950c0a49212eb20d1edfaf41c57e2e Mon Sep 17 00:00:00 2001 From: happy-v587 Date: Thu, 23 Jul 2026 21:21:03 +0800 Subject: [PATCH 11/32] test: use unique test db path helpers in vector storage tests --- src/storage/tests/redis_vector_test.rs | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/src/storage/tests/redis_vector_test.rs b/src/storage/tests/redis_vector_test.rs index 89676164..6d2cacb9 100644 --- a/src/storage/tests/redis_vector_test.rs +++ b/src/storage/tests/redis_vector_test.rs @@ -377,10 +377,11 @@ fn test_vsim_rejects_query_dimension_mismatch() { #[tokio::test] async fn test_storage_routes_all_members_of_one_vectorset_to_one_instance() { - let temp = tempfile::tempdir().expect("temp dir"); + let test_db_path = unique_test_db_path(); + safe_cleanup_test_db(&test_db_path); let mut storage = Storage::new(3, 0); let _receiver = storage - .open(Arc::new(StorageOptions::default()), temp.path()) + .open(Arc::new(StorageOptions::default()), &test_db_path) .expect("open storage"); let key = b"routed-vectors"; let x = CanonicalVector::from_values(&[1.0, 0.0]).expect("x"); @@ -429,6 +430,7 @@ async fn test_storage_routes_all_members_of_one_vectorset_to_one_instance() { } storage.shutdown().await; + safe_cleanup_test_db(&test_db_path); } #[test] @@ -445,10 +447,11 @@ fn test_type_returns_vectorset() { #[tokio::test] async fn test_expired_vectorset_reads_as_missing() { - let temp = tempfile::tempdir().expect("temp dir"); + let test_db_path = unique_test_db_path(); + safe_cleanup_test_db(&test_db_path); let mut storage = Storage::new(1, 0); let _receiver = storage - .open(Arc::new(StorageOptions::default()), temp.path()) + .open(Arc::new(StorageOptions::default()), &test_db_path) .expect("open storage"); let key = b"expiring-vectors"; let vector = CanonicalVector::from_values(&[1.0, 0.0]).expect("vector"); @@ -474,6 +477,7 @@ async fn test_expired_vectorset_reads_as_missing() { assert_eq!(storage.key_type(key).expect("expired type"), "none"); storage.shutdown().await; + safe_cleanup_test_db(&test_db_path); } #[test] @@ -505,10 +509,11 @@ fn test_flushdb_removes_vector_meta_and_members() { #[tokio::test] async fn test_vector_storage_rejects_cluster_mode() { - let temp = tempfile::tempdir().expect("temp dir"); + let test_db_path = unique_test_db_path(); + safe_cleanup_test_db(&test_db_path); let mut storage = Storage::new(1, 0); let _receiver = storage - .open(Arc::new(StorageOptions::default()), temp.path()) + .open(Arc::new(StorageOptions::default()), &test_db_path) .expect("open storage"); storage.set_append_log_fn(Arc::new(|_| panic!("vector API must not append Raft log"))); let vector = CanonicalVector::from_values(&[1.0, 0.0]).expect("vector"); @@ -530,4 +535,5 @@ async fn test_vector_storage_rejects_cluster_mode() { assert!(storage.vismember(b"vectors", b"member").is_err()); storage.shutdown().await; + safe_cleanup_test_db(&test_db_path); } From 049f9c0c67dd4f3cfc75e50ca72e9b9d7465a2c1 Mon Sep 17 00:00:00 2001 From: happy-v587 Date: Sat, 25 Jul 2026 20:59:27 +0800 Subject: [PATCH 12/32] refactor(storage): adapt vector set to engine removal --- src/storage/src/redis_strings.rs | 2 +- src/storage/tests/redis_vector_test.rs | 8 +++----- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/src/storage/src/redis_strings.rs b/src/storage/src/redis_strings.rs index daa49000..a86b7bd1 100644 --- a/src/storage/src/redis_strings.rs +++ b/src/storage/src/redis_strings.rs @@ -2105,7 +2105,7 @@ impl Redis { let meta = ParsedListsMetaValue::new(&value_bytes[..])?; !meta.is_stale() && meta.count() > 0 } - DataType::Hash | DataType::Set | DataType::ZSet => { + DataType::Hash | DataType::Set | DataType::ZSet | DataType::VectorSet => { let meta = ParsedBaseMetaValue::new(&value_bytes[..])?; !meta.is_stale() && meta.count() > 0 } diff --git a/src/storage/tests/redis_vector_test.rs b/src/storage/tests/redis_vector_test.rs index 6d2cacb9..2a498938 100644 --- a/src/storage/tests/redis_vector_test.rs +++ b/src/storage/tests/redis_vector_test.rs @@ -45,7 +45,6 @@ fn with_redis(test: impl FnOnce(&Redis)) { test(&redis); - redis.set_need_close(true); drop(redis); safe_cleanup_test_db(&path); } @@ -66,7 +65,7 @@ fn search_options(count: usize, mode: VectorSearchMode) -> VectorSearchOptions { } fn count_cf_entries(redis: &Redis, cf_index: ColumnFamilyIndex) -> usize { - let db = redis.db.as_ref().expect("db is initialized"); + let db = redis.db().expect("db is initialized"); let cf = redis.get_cf_handle(cf_index).expect("column family exists"); db.iterator_cf(&cf, IteratorMode::Start) .map(|entry| entry.expect("read column family entry")) @@ -205,7 +204,7 @@ fn test_vector_meta_and_member_are_committed_together() { assert!(redis.vadd(b"vectors", b"member", &vector).expect("insert")); { - let db = redis.db.as_ref().expect("db is initialized"); + let db = redis.db().expect("db is initialized"); let meta_cf = redis .get_cf_handle(ColumnFamilyIndex::MetaCF) .expect("MetaCF exists"); @@ -234,7 +233,6 @@ fn test_vector_meta_and_member_are_committed_together() { assert_eq!(redis.vcard(b"vectors").expect("card"), 1); assert!(redis.vismember(b"vectors", b"member").expect("membership")); - redis.set_need_close(true); drop(redis); safe_cleanup_test_db(&path); } @@ -413,7 +411,7 @@ async fn test_storage_routes_all_members_of_one_vectorset_to_one_instance() { let selected = storage.slot_indexer.get_instance_id(slot_id); let meta_key = BaseMetaKey::new(key).encode().expect("meta key"); for (instance_id, redis) in storage.insts.iter().enumerate() { - let db = redis.db.as_ref().expect("db is initialized"); + let db = redis.db().expect("db is initialized"); let meta_cf = redis .get_cf_handle(ColumnFamilyIndex::MetaCF) .expect("MetaCF exists"); From 05148c07bd6952600b5c6fc6a201f499d752b4c8 Mon Sep 17 00:00:00 2001 From: happy-v587 Date: Sat, 25 Jul 2026 23:49:08 +0800 Subject: [PATCH 13/32] docs: split cargo test filters and specify monotonic vector generation --- docs/superpowers/plans/2026-07-19-redis-vector-set.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/docs/superpowers/plans/2026-07-19-redis-vector-set.md b/docs/superpowers/plans/2026-07-19-redis-vector-set.md index b6133fc6..e3e9e1c7 100644 --- a/docs/superpowers/plans/2026-07-19-redis-vector-set.md +++ b/docs/superpowers/plans/2026-07-19-redis-vector-set.md @@ -420,7 +420,7 @@ Implementation sequence inside `Redis::vadd`: 1. Reject cluster mode by checking `self.append_log_fn.get().is_some()`. 2. Acquire `ScopeRecordLock` using the user key. 3. Read `MetaCF[BaseMetaKey::new(key)]`. -4. Missing or stale meta creates `VectorMeta::new(1, dimension)` with a new timestamp-based version and does not increment it again. +4. Missing or stale meta creates `VectorMeta::new_after(1, dimension, previous_generation)` whose version is the current timestamp clamped above the previous generation (monotonic, so a recreated VectorSet can never address stale `VectorDataCF` rows), and does not increment it again. 5. Live non-VectorSet meta returns the same WRONGTYPE text used by `check_type_state`. 6. Live VectorSet requires equal dimension. 7. Point-read `VectorDataCF[MemberDataKey(key, version, element)]` to distinguish insert from update. @@ -464,7 +464,8 @@ Run: ```bash RUST_TEST_THREADS=1 cargo test -p storage --test redis_vector_test -cargo test -p storage vector::tests format_vector::tests +cargo test -p storage vector::tests +cargo test -p storage format_vector::tests ``` Expected: all tests PASS, including rollback-visible behavior after dimension mismatch. From 20a5b8b54ab2c0a26c7b6f91f15f51dd5d4b05c4 Mon Sep 17 00:00:00 2001 From: happy-v587 Date: Sun, 26 Jul 2026 00:07:08 +0800 Subject: [PATCH 14/32] fix(vector): prevent stale members and protocol regressions --- src/net/src/executor_ext.rs | 11 +- src/net/tests/storage_command_e2e_tests.rs | 45 +++++++- src/resp/src/encode.rs | 123 +++++++++++++-------- src/resp/tests/resp2_encoding.rs | 93 ++++++++++++++++ src/storage/src/format_vector.rs | 12 +- src/storage/src/redis_vectors.rs | 8 +- src/storage/tests/redis_vector_test.rs | 43 +++++++ 7 files changed, 278 insertions(+), 57 deletions(-) create mode 100644 src/resp/tests/resp2_encoding.rs diff --git a/src/net/src/executor_ext.rs b/src/net/src/executor_ext.rs index 37b03989..c4e62abf 100644 --- a/src/net/src/executor_ext.rs +++ b/src/net/src/executor_ext.rs @@ -60,9 +60,16 @@ impl CmdExecutorNetworkExt for CmdExecutor { return Ok(()); } - // Cluster-mode leader gate: reject writes on non-leaders before any - // command-specific setup runs. + // Cluster-mode gates run before any command-specific setup: reject + // unsupported module commands on every node, then writes on followers. if let Some(gate) = exec.leader_gate.as_ref() { + if exec.cmd.has_flag(CmdFlags::MODULE_NO_CLUSTER) { + exec.client.set_reply(RespData::Error( + "ERR Vector Set is not supported in cluster mode".into(), + )); + return Ok(()); + } + if exec.cmd.has_flag(CmdFlags::WRITE) && !gate.is_leader() { // Simplified redirect: Kiwi returns "MOVED " (no hash slot, // unlike Redis Cluster's "MOVED "). Clients are diff --git a/src/net/tests/storage_command_e2e_tests.rs b/src/net/tests/storage_command_e2e_tests.rs index 164dfbdb..c9a81c36 100644 --- a/src/net/tests/storage_command_e2e_tests.rs +++ b/src/net/tests/storage_command_e2e_tests.rs @@ -33,6 +33,7 @@ use std::time::Duration; use bytes::Bytes; use net::{ServerTrait, network_server::NetworkServer, storage_client::StorageClient}; +use raft::leader_gate::LeaderGate; use resp::{ Parse, RespData, RespEncode, RespParse, RespParseResult, RespVersion, encode::RespEncoder, }; @@ -51,10 +52,29 @@ struct TestServer { storage_client: Arc, } +struct FollowerGate; + +impl LeaderGate for FollowerGate { + fn is_leader(&self) -> bool { + false + } + + fn leader_resp_addr(&self) -> Option { + Some("127.0.0.1:7380".to_string()) + } +} + impl TestServer { /// Start the network and storage runtimes, open a real storage DB, and bind /// a `NetworkServer` to an ephemeral port. async fn start(requirepass: Option) -> Self { + Self::start_with_leader_gate(requirepass, None).await + } + + async fn start_with_leader_gate( + requirepass: Option, + leader_gate: Option>, + ) -> Self { let db_path = unique_test_db_path(); safe_cleanup_test_db(&db_path); @@ -104,7 +124,7 @@ impl TestServer { cmd_table, executor, requirepass, - None, + leader_gate, ) .expect("network server"), ); @@ -403,6 +423,29 @@ async fn storage_command_e2e_generic_storage_commands_use_storage_path() { server.shutdown().await; } +#[tokio::test] +async fn storage_command_e2e_no_cluster_commands_precede_follower_redirect() { + let server = TestServer::start_with_leader_gate(None, Some(Arc::new(FollowerGate))).await; + let mut stream = tokio::net::TcpStream::connect(server.addr) + .await + .expect("connect to server"); + let expected = RespData::Error(Bytes::from_static( + b"ERR Vector Set is not supported in cluster mode", + )); + + let reply = send_command( + &mut stream, + &["VADD", "vectors", "VALUES", "2", "1", "member"], + ) + .await; + assert_eq!(reply, expected); + + let reply = send_command(&mut stream, &["VREM", "vectors", "member"]).await; + assert_eq!(reply, expected); + + server.shutdown().await; +} + // Regression for issue #349: GET/MGET must return stored bytes unchanged over RESP. #[tokio::test] async fn storage_command_e2e_get_and_mget_preserve_binary_values() { diff --git a/src/resp/src/encode.rs b/src/resp/src/encode.rs index 8ab2e5cd..9e6277cd 100644 --- a/src/resp/src/encode.rs +++ b/src/resp/src/encode.rs @@ -23,7 +23,6 @@ use bytes::{Bytes, BytesMut}; use crate::{ CRLF, error::RespError, - negotiation::ProtocolNegotiator, types::{RespData, RespVersion}, }; @@ -226,59 +225,82 @@ impl RespEncoder { self } RespData::Array(None) => self.set_array_len(-1), - RespData::Null => { - self.buffer.extend_from_slice(b"_"); - self.append_crlf() - } + RespData::Null => self.append_null(), RespData::Boolean(value) => { - self.buffer.extend_from_slice(b"#"); - self.buffer - .extend_from_slice(if *value { b"t" } else { b"f" }); - self.append_crlf() + if self.is_resp3() { + self.buffer.extend_from_slice(b"#"); + self.buffer + .extend_from_slice(if *value { b"t" } else { b"f" }); + self.append_crlf() + } else { + self.append_integer(if *value { 1 } else { 0 }) + } } RespData::Double(value) => { - if value.is_nan() { - self.buffer.extend_from_slice(b",nan"); - } else if value.is_infinite() { - self.buffer.extend_from_slice(if value.is_sign_negative() { - b",-inf" + if self.is_resp3() { + if value.is_nan() { + self.buffer.extend_from_slice(b",nan"); + } else if value.is_infinite() { + self.buffer.extend_from_slice(if value.is_sign_negative() { + b",-inf" + } else { + b",inf" + }); } else { - b",inf" - }); + let _ = write!(self.buffer, ",{value}"); + } + self.append_crlf() } else { - let _ = write!(self.buffer, ",{value}"); + self.append_bulk_string(value.to_string().as_bytes()) } - self.append_crlf() } RespData::BigNumber(bytes) => { - self.buffer.extend_from_slice(b"("); - self.buffer.extend_from_slice(bytes); - self.append_crlf() + if self.is_resp3() { + self.buffer.extend_from_slice(b"("); + self.buffer.extend_from_slice(bytes); + self.append_crlf() + } else { + self.append_bulk_string(bytes) + } } RespData::BulkError(bytes) => { - let _ = write!(self.buffer, "!{}", bytes.len()); - self.append_crlf(); - self.buffer.extend_from_slice(bytes); - self.append_crlf() + if self.is_resp3() { + let _ = write!(self.buffer, "!{}", bytes.len()); + self.append_crlf(); + self.buffer.extend_from_slice(bytes); + self.append_crlf() + } else { + self.buffer.extend_from_slice(b"-"); + self.buffer.extend_from_slice(bytes); + self.append_crlf() + } } RespData::VerbatimString { format, data } => { - if format.len() != 3 { - panic!( - "RESP3 VerbatimString format must be exactly 3 bytes, got {}", - format.len() - ); + if self.is_resp3() { + if format.len() != 3 { + panic!( + "RESP3 VerbatimString format must be exactly 3 bytes, got {}", + format.len() + ); + } + let total_len = format.len() + 1 + data.len(); + let _ = write!(self.buffer, "={total_len}"); + self.append_crlf(); + self.buffer.extend_from_slice(format); + self.buffer.extend_from_slice(b":"); + self.buffer.extend_from_slice(data); + self.append_crlf() + } else { + self.append_bulk_string(data) } - let total_len = format.len() + 1 + data.len(); - let _ = write!(self.buffer, "={total_len}"); - self.append_crlf(); - self.buffer.extend_from_slice(format); - self.buffer.extend_from_slice(b":"); - self.buffer.extend_from_slice(data); - self.append_crlf() } RespData::Map(pairs) => { - let _ = write!(self.buffer, "%{}", pairs.len()); - self.append_crlf(); + if self.is_resp3() { + let _ = write!(self.buffer, "%{}", pairs.len()); + self.append_crlf(); + } else { + self.append_array_len((pairs.len() * 2) as i64); + } for (key, value) in pairs { self.encode_resp_data_inner(key); self.encode_resp_data_inner(value); @@ -286,16 +308,24 @@ impl RespEncoder { self } RespData::Set(items) => { - let _ = write!(self.buffer, "~{}", items.len()); - self.append_crlf(); + if self.is_resp3() { + let _ = write!(self.buffer, "~{}", items.len()); + self.append_crlf(); + } else { + self.append_array_len(items.len() as i64); + } for item in items { self.encode_resp_data_inner(item); } self } RespData::Push(items) => { - let _ = write!(self.buffer, ">{}", items.len()); - self.append_crlf(); + if self.is_resp3() { + let _ = write!(self.buffer, ">{}", items.len()); + self.append_crlf(); + } else { + self.append_array_len(items.len() as i64); + } for item in items { self.encode_resp_data_inner(item); } @@ -467,12 +497,7 @@ impl RespEncode for RespEncoder { } fn encode_resp_data(&mut self, data: &RespData) -> &mut Self { - if self.is_resp3() { - self.encode_resp_data_inner(data) - } else { - let normalized = ProtocolNegotiator::convert_to_resp2(data); - self.encode_resp_data_inner(&normalized) - } + self.encode_resp_data_inner(data) } fn append_null(&mut self) -> &mut Self { diff --git a/src/resp/tests/resp2_encoding.rs b/src/resp/tests/resp2_encoding.rs new file mode 100644 index 00000000..0055e8b6 --- /dev/null +++ b/src/resp/tests/resp2_encoding.rs @@ -0,0 +1,93 @@ +// Copyright (c) 2024-present, arana-db Community. All rights reserved. +// +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::{ + alloc::{GlobalAlloc, Layout, System}, + sync::atomic::{AtomicBool, AtomicUsize, Ordering}, +}; + +use bytes::Bytes; +use resp::{ + RespData, RespVersion, + encode::{RespEncode, RespEncoder}, +}; + +struct CountingAllocator; + +static COUNT_ALLOCATIONS: AtomicBool = AtomicBool::new(false); +static ALLOCATION_COUNT: AtomicUsize = AtomicUsize::new(0); + +unsafe impl GlobalAlloc for CountingAllocator { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + if COUNT_ALLOCATIONS.load(Ordering::Relaxed) { + ALLOCATION_COUNT.fetch_add(1, Ordering::Relaxed); + } + unsafe { System.alloc(layout) } + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + unsafe { System.dealloc(ptr, layout) } + } + + unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + if COUNT_ALLOCATIONS.load(Ordering::Relaxed) { + ALLOCATION_COUNT.fetch_add(1, Ordering::Relaxed); + } + unsafe { System.realloc(ptr, layout, new_size) } + } +} + +#[global_allocator] +static ALLOCATOR: CountingAllocator = CountingAllocator; + +fn count_allocations(f: impl FnOnce()) -> usize { + ALLOCATION_COUNT.store(0, Ordering::Relaxed); + COUNT_ALLOCATIONS.store(true, Ordering::Relaxed); + f(); + COUNT_ALLOCATIONS.store(false, Ordering::Relaxed); + ALLOCATION_COUNT.load(Ordering::Relaxed) +} + +#[test] +fn resp2_arrays_encode_without_normalization_and_downgrade_nested_resp3_values() { + let data = RespData::Array(Some(vec![ + RespData::Integer(1), + RespData::BulkString(Some(Bytes::from_static(b"value"))), + ])); + let mut encoder = RespEncoder::new(RespVersion::RESP2); + + // Warm the output buffer so the measured encode does not need to grow it. + encoder.encode_resp_data(&data).clear(); + + let allocations = count_allocations(|| { + encoder.encode_resp_data(&data); + }); + + assert_eq!( + allocations, 0, + "RESP2-compatible aggregates should encode without a normalized copy" + ); + assert_eq!( + encoder.get_response(), + Bytes::from_static(b"*2\r\n:1\r\n$5\r\nvalue\r\n") + ); + + let data = RespData::Array(Some(vec![RespData::Boolean(true)])); + encoder.clear().encode_resp_data(&data); + + assert_eq!(encoder.get_response(), Bytes::from_static(b"*1\r\n:1\r\n")); +} diff --git a/src/storage/src/format_vector.rs b/src/storage/src/format_vector.rs index d59e3e89..20b23bb5 100644 --- a/src/storage/src/format_vector.rs +++ b/src/storage/src/format_vector.rs @@ -45,11 +45,15 @@ pub(crate) struct VectorMeta { } impl VectorMeta { - pub(crate) fn new(count: u64, dimension: u32) -> Self { + pub(crate) fn new_after(count: u64, dimension: u32, previous_version: u64) -> Self { let now = Utc::now().timestamp_micros() as u64; + let version = match previous_version >= now { + true => previous_version + 1, + false => now, + }; Self { count, - version: now, + version, dimension, ctime: now, etime: 0, @@ -306,7 +310,7 @@ mod tests { #[test] fn vector_meta_round_trips() { - let mut meta = VectorMeta::new(2, 2); + let mut meta = VectorMeta::new_after(2, 2, 0); meta.version = 42; let encoded = meta.encode(); let decoded = VectorMeta::decode(&encoded).expect("decode vector meta"); @@ -334,7 +338,7 @@ mod tests { non_finite_payload[10..14].copy_from_slice(&f32::NAN.to_le_bytes()); assert!(VectorDataValue::decode(&non_finite_payload).is_err()); - let encoded_meta = VectorMeta::new(2, 2).encode(); + let encoded_meta = VectorMeta::new_after(2, 2, 0).encode(); assert!(VectorMeta::decode(&encoded_meta[..encoded_meta.len() - 1]).is_err()); let mut bad_meta_format = encoded_meta; diff --git a/src/storage/src/redis_vectors.rs b/src/storage/src/redis_vectors.rs index 67441a65..83db5853 100644 --- a/src/storage/src/redis_vectors.rs +++ b/src/storage/src/redis_vectors.rs @@ -119,6 +119,12 @@ impl Redis { .map(|value| self.parse_vector_meta(value)) .transpose()? .flatten(); + let previous_generation = stored_meta + .as_deref() + .filter(|value| value.first().copied() == Some(DataType::VectorSet as u8)) + .map(VectorMeta::decode) + .transpose()? + .map_or(0, |meta| meta.version()); let is_new_set = live_meta.is_none(); let mut meta = match live_meta { @@ -135,7 +141,7 @@ impl Redis { } meta } - None => VectorMeta::new(1, vector.dimension()), + None => VectorMeta::new_after(1, vector.dimension(), previous_generation), }; let member_key = MemberDataKey::new(key, meta.version(), element).encode()?; diff --git a/src/storage/tests/redis_vector_test.rs b/src/storage/tests/redis_vector_test.rs index 2a498938..be7fbdf8 100644 --- a/src/storage/tests/redis_vector_test.rs +++ b/src/storage/tests/redis_vector_test.rs @@ -443,6 +443,49 @@ fn test_type_returns_vectorset() { }); } +#[test] +fn test_vadd_rebuilds_expired_vectorset_with_newer_generation() { + with_redis(|redis| { + let key = b"expiring-vectors"; + let vector = CanonicalVector::from_values(&[1.0, 0.0]).expect("vector"); + redis.vadd(key, b"old", &vector).expect("insert old member"); + + let db = redis.db().expect("db is initialized"); + let meta_cf = redis + .get_cf_handle(ColumnFamilyIndex::MetaCF) + .expect("MetaCF exists"); + let meta_key = BaseMetaKey::new(key).encode().expect("meta key"); + let mut meta = db + .get_cf(&meta_cf, &meta_key) + .expect("read vector meta") + .expect("vector meta exists"); + let previous_generation = u64::MAX - 1; + meta[9..17].copy_from_slice(&previous_generation.to_le_bytes()); + let etime_offset = meta.len() - size_of::(); + meta[etime_offset..].copy_from_slice(&1_u64.to_le_bytes()); + db.put_cf(&meta_cf, &meta_key, &meta) + .expect("store expired vector meta"); + + assert_eq!(redis.vcard(key).expect("expired card"), 0); + assert_eq!(count_cf_entries(redis, ColumnFamilyIndex::VectorDataCF), 1); + + redis + .vadd(key, b"new", &vector) + .expect("rebuild vector set"); + + assert_eq!(redis.vcard(key).expect("rebuilt card"), 1); + assert!(!redis.vismember(key, b"old").expect("old membership")); + assert!(redis.vismember(key, b"new").expect("new membership")); + let rebuilt_meta = db + .get_cf(&meta_cf, &meta_key) + .expect("read rebuilt vector meta") + .expect("rebuilt vector meta exists"); + let rebuilt_generation = + u64::from_le_bytes(rebuilt_meta[9..17].try_into().expect("generation bytes")); + assert!(rebuilt_generation > previous_generation); + }); +} + #[tokio::test] async fn test_expired_vectorset_reads_as_missing() { let test_db_path = unique_test_db_path(); From 5e0e972ff0f64abd85c94f88dfc5bae9b143f70d Mon Sep 17 00:00:00 2001 From: happy-v587 Date: Sun, 26 Jul 2026 14:46:53 +0800 Subject: [PATCH 15/32] feat(vector): enable raft cluster mode and add metric/search abstractions - Remove MODULE_NO_CLUSTER guard so vector commands work in Raft Group clusters (binlog replication); keep follower redirect for writes. - Split src/cmd/src/vector.rs into one module per command. - Introduce SimilarityMetric (currently Cosine) persisted in VectorMeta. - Introduce VectorSearchEngine (currently Flat) for future HNSW. - Add RocksDB encoding layout comments for vector meta/data values. - Update the Redis Vector Set superpower plan to reflect cluster support. Co-Authored-By: Claude --- .../plans/2026-07-19-redis-vector-set.md | 31 ++- src/cmd/src/lib.rs | 1 - src/cmd/src/{vector.rs => vector/mod.rs} | 245 ++---------------- src/cmd/src/vector/vadd.rs | 58 +++++ src/cmd/src/vector/vcard.rs | 50 ++++ src/cmd/src/vector/vdim.rs | 50 ++++ src/cmd/src/vector/vemb.rs | 61 +++++ src/cmd/src/vector/vismember.rs | 52 ++++ src/cmd/src/vector/vrem.rs | 52 ++++ src/cmd/src/vector/vsim.rs | 73 ++++++ src/net/src/executor_ext.rs | 10 +- src/net/tests/storage_command_e2e_tests.rs | 6 +- src/storage/src/format_member_data_key.rs | 3 + src/storage/src/format_vector.rs | 65 ++++- src/storage/src/lib.rs | 5 +- src/storage/src/redis_vectors.rs | 136 +++------- src/storage/src/vector.rs | 92 +++++++ src/storage/tests/redis_vector_test.rs | 64 +++-- 18 files changed, 677 insertions(+), 377 deletions(-) rename src/cmd/src/{vector.rs => vector/mod.rs} (70%) create mode 100644 src/cmd/src/vector/vadd.rs create mode 100644 src/cmd/src/vector/vcard.rs create mode 100644 src/cmd/src/vector/vdim.rs create mode 100644 src/cmd/src/vector/vemb.rs create mode 100644 src/cmd/src/vector/vismember.rs create mode 100644 src/cmd/src/vector/vrem.rs create mode 100644 src/cmd/src/vector/vsim.rs diff --git a/docs/superpowers/plans/2026-07-19-redis-vector-set.md b/docs/superpowers/plans/2026-07-19-redis-vector-set.md index e3e9e1c7..e9cd4077 100644 --- a/docs/superpowers/plans/2026-07-19-redis-vector-set.md +++ b/docs/superpowers/plans/2026-07-19-redis-vector-set.md @@ -11,7 +11,7 @@ ## Global Constraints - 基线为 `origin/main` 的 `cdada8b`；忽略尚未合入的 error-catalog/error-model 分支，错误处理遵循当前主分支模式。 -- 首版只支持 standalone；所有 Vector Set storage API 在 Raft append hook 已安装时返回 `ERR Vector Set is not supported in cluster mode`。 +- 首版支持 standalone 与 Raft Group cluster 模式；暂不支持 Redis Cluster 的槽位路由。 - 首版只支持 cosine、canonical FP32 little-endian 和显式 `NOQUANT`；默认 Q8、显式 `Q8`、`BIN`、VEMB `RAW` 均返回明确的 unsupported 错误。 - 首版只实现 FLAT；`TRUTH` 与普通 `VSIM` 都走同一个精确引擎，但保留不同的搜索模式枚举。 - 首版不实现 `VINFO`、`INFO VECTOR`、HNSW、FAISS/IVF、`storage_incarnation`、O(1) `DEL`、Raft logical mutation、ReadIndex、snapshot barrier 和滚动升级门禁。 @@ -41,7 +41,14 @@ | `src/conf/src/raft_type.rs` | 追加跨 crate 使用的 VectorDataCF index | | `src/storage/src/logindex/types.rs` | 追加 VectorDataCF 的 log-index 元数据 | | `src/raft/src/lib.rs` | 追加 CF 名称和一致性断言 | -| `src/cmd/src/vector.rs` | 七个命令、参数解析、错误和 RESP reply 构造 | +| `src/cmd/src/vector/mod.rs` | 共享解析器、错误/回复辅助函数与命令注册测试 | +| `src/cmd/src/vector/vadd.rs` | `VAddCmd` 实现 | +| `src/cmd/src/vector/vsim.rs` | `VSimCmd` 实现 | +| `src/cmd/src/vector/vrem.rs` | `VRemCmd` 实现 | +| `src/cmd/src/vector/vcard.rs` | `VCardCmd` 实现 | +| `src/cmd/src/vector/vdim.rs` | `VDimCmd` 实现 | +| `src/cmd/src/vector/vemb.rs` | `VEmbCmd` 实现 | +| `src/cmd/src/vector/vismember.rs` | `VIsMemberCmd` 实现 | | `src/cmd/src/lib.rs` | 导出 vector 命令模块 | | `src/cmd/src/table.rs` | 注册七个命令 | | `src/resp/src/encode.rs` | RESP2 下递归降级 Map/Double，RESP3 保持原生类型 | @@ -417,8 +424,7 @@ mod redis_vectors; Implementation sequence inside `Redis::vadd`: -1. Reject cluster mode by checking `self.append_log_fn.get().is_some()`. -2. Acquire `ScopeRecordLock` using the user key. +1. Acquire `ScopeRecordLock` using the user key. 3. Read `MetaCF[BaseMetaKey::new(key)]`. 4. Missing or stale meta creates `VectorMeta::new_after(1, dimension, previous_generation)` whose version is the current timestamp clamped above the previous generation (monotonic, so a recreated VectorSet can never address stale `VectorDataCF` rows), and does not increment it again. 5. Live non-VectorSet meta returns the same WRONGTYPE text used by `check_type_state`. @@ -787,9 +793,9 @@ Reply contracts: Flags and ACL categories: -- VADD/VREM: `WRITE | FAST | MODULE_NO_CLUSTER`, `KEYSPACE | WRITE`. -- VSIM: `READONLY | MODULE_NO_CLUSTER`, `KEYSPACE | READ | SLOW`. -- VCARD/VDIM/VEMB/VISMEMBER: `READONLY | FAST | MODULE_NO_CLUSTER`, `KEYSPACE | READ`. +- VADD/VREM: `WRITE | FAST`, `KEYSPACE | WRITE`. +- VSIM: `READONLY`, `KEYSPACE | READ | SLOW`. +- VCARD/VDIM/VEMB/VISMEMBER: `READONLY | FAST`, `KEYSPACE | READ`. - [x] **Step 5: Export and register all seven commands** @@ -1046,7 +1052,7 @@ The Phase 1 implementation is complete only when all statements below are true: - VADD accepts FP32 and VALUES only with explicit NOQUANT and is binary-safe for elements. - VSIM returns exact, deterministic Top-K results without collecting and sorting every hit. - RESP2 and RESP3 WITHSCORES replies use their native compatible shapes. -- Cluster mode rejects all seven Vector Set operations before creating a batch or snapshot. +- Raft Group cluster 模式支持所有七个命令的写入和读取（通过 binlog 复制到 follower）；Redis Cluster 槽位路由不在本版范围内。 - Focused Rust/Python tests and the complete Kiwi verification pipeline pass. - The final diff contains no FT prototype extraction, HNSW, Raft/snapshot design implementation, or unrelated cleanup. @@ -1054,7 +1060,8 @@ The Phase 1 implementation is complete only when all statements below are true: Create separate reviewed designs and plans before implementing any of these independent projects: -1. Redis-compatible VINFO and INFO VECTOR observability. -2. Raft logical VectorSet mutations, leader-linearizable reads and snapshots. -3. O(1) generation deletion with background compaction cleanup. -4. HNSW derived index, rebuild/recovery and rolling-upgrade capability gates. +1. Redis Cluster 槽位路由与跨 slot 查询。 +2. Redis-compatible VINFO and INFO VECTOR observability. +3. Raft logical VectorSet mutations, leader-linearizable reads and snapshots. +4. O(1) generation deletion with background compaction cleanup. +5. HNSW derived index, rebuild/recovery and rolling-upgrade capability gates. diff --git a/src/cmd/src/lib.rs b/src/cmd/src/lib.rs index 22ecfb18..3b8cc7f1 100644 --- a/src/cmd/src/lib.rs +++ b/src/cmd/src/lib.rs @@ -139,7 +139,6 @@ bitflags! { const NO_AUTH = 1 << 10; // Skip ACL checks const MAY_REPLICATE = 1 << 11; // May replicate even if writes are disabled const PROTECTED = 1 << 12; // Don't accept in scripts - const MODULE_NO_CLUSTER = 1 << 13; // No cluster mode support const NO_MULTI = 1 << 14; // Cannot be pipelined const EXCLUSIVE = 1 << 15; // May change Storage pointer const RAFT = 1 << 16; diff --git a/src/cmd/src/vector.rs b/src/cmd/src/vector/mod.rs similarity index 70% rename from src/cmd/src/vector.rs rename to src/cmd/src/vector/mod.rs index aa567775..a5d5dd34 100644 --- a/src/cmd/src/vector.rs +++ b/src/cmd/src/vector/mod.rs @@ -15,17 +15,9 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::sync::Arc; - -use bytes::Bytes; use client::Client; use resp::RespData; -use storage::{ - CanonicalVector, VectorQuery, VectorSearchMode, VectorSearchOptions, error::Error, - storage::Storage, -}; - -use crate::{AclCategory, Cmd, CmdFlags, CmdMeta, impl_cmd_clone_box, impl_cmd_meta}; +use storage::{CanonicalVector, VectorQuery, VectorSearchMode, VectorSearchOptions, error::Error}; const ERR_INVALID_VECTOR: &str = "ERR invalid vector specification"; const ERR_VECTOR_DIMENSION: &str = "ERR vector dimension mismatch"; @@ -178,17 +170,17 @@ fn parse_vemb(argv: &[Vec]) -> ParseResult> { } } -fn error_reply(message: impl Into) -> RespData { +pub(crate) fn error_reply(message: impl Into) -> RespData { RespData::Error(message.into().into()) } #[derive(Clone, Copy)] -enum MissingError { +pub(crate) enum MissingError { Key, Element, } -fn storage_error_reply(error: Error, missing: MissingError) -> RespData { +pub(crate) fn storage_error_reply(error: Error, missing: MissingError) -> RespData { match &error { Error::RedisErr { message, .. } => error_reply(message.clone()), Error::InvalidArgument { message, .. } if message.contains("dimension mismatch") => { @@ -205,7 +197,7 @@ fn storage_error_reply(error: Error, missing: MissingError) -> RespData { } } -fn set_command_key(client: &Client) -> bool { +pub(crate) fn set_command_key(client: &Client) -> bool { let argv = client.argv(); let Some(key) = argv.get(1) else { client.set_reply(error_reply(ERR_INVALID_VECTOR)); @@ -215,7 +207,7 @@ fn set_command_key(client: &Client) -> bool { true } -fn integer_reply(value: u64) -> RespData { +pub(crate) fn integer_reply(value: u64) -> RespData { match i64::try_from(value) { Ok(value) => RespData::Integer(value), Err(error) => { @@ -225,6 +217,7 @@ fn integer_reply(value: u64) -> RespData { } } +#[macro_export] macro_rules! define_vector_command { ($type:ident, $name:literal, $arity:expr, $flags:expr, $acl:expr) => { #[derive(Clone, Default)] @@ -248,214 +241,21 @@ macro_rules! define_vector_command { }; } -define_vector_command!( - VAddCmd, - "vadd", - -5, - CmdFlags::WRITE | CmdFlags::FAST | CmdFlags::MODULE_NO_CLUSTER, - AclCategory::KEYSPACE | AclCategory::WRITE -); -define_vector_command!( - VSimCmd, - "vsim", - -4, - CmdFlags::READONLY | CmdFlags::MODULE_NO_CLUSTER, - AclCategory::KEYSPACE | AclCategory::READ | AclCategory::SLOW -); -define_vector_command!( - VRemCmd, - "vrem", - 3, - CmdFlags::WRITE | CmdFlags::FAST | CmdFlags::MODULE_NO_CLUSTER, - AclCategory::KEYSPACE | AclCategory::WRITE -); -define_vector_command!( - VCardCmd, - "vcard", - 2, - CmdFlags::READONLY | CmdFlags::FAST | CmdFlags::MODULE_NO_CLUSTER, - AclCategory::KEYSPACE | AclCategory::READ -); -define_vector_command!( - VDimCmd, - "vdim", - 2, - CmdFlags::READONLY | CmdFlags::FAST | CmdFlags::MODULE_NO_CLUSTER, - AclCategory::KEYSPACE | AclCategory::READ -); -define_vector_command!( - VEmbCmd, - "vemb", - -3, - CmdFlags::READONLY | CmdFlags::FAST | CmdFlags::MODULE_NO_CLUSTER, - AclCategory::KEYSPACE | AclCategory::READ -); -define_vector_command!( - VIsMemberCmd, - "vismember", - 3, - CmdFlags::READONLY | CmdFlags::FAST | CmdFlags::MODULE_NO_CLUSTER, - AclCategory::KEYSPACE | AclCategory::READ -); - -impl Cmd for VAddCmd { - impl_cmd_meta!(); - impl_cmd_clone_box!(); - - fn do_initial(&self, client: &Client) -> bool { - set_command_key(client) - } - - fn do_cmd(&self, client: &Client, storage: Arc) { - let parsed = match parse_vadd(&client.argv()) { - Ok(parsed) => parsed, - Err(message) => { - client.set_reply(error_reply(message)); - return; - } - }; - let reply = match storage.vadd(&client.key(), &parsed.element, &parsed.vector) { - Ok(inserted) => RespData::Integer(i64::from(inserted)), - Err(error) => storage_error_reply(error, MissingError::Key), - }; - client.set_reply(reply); - } -} - -impl Cmd for VSimCmd { - impl_cmd_meta!(); - impl_cmd_clone_box!(); - - fn do_initial(&self, client: &Client) -> bool { - set_command_key(client) - } - - fn do_cmd(&self, client: &Client, storage: Arc) { - let parsed = match parse_vsim(&client.argv()) { - Ok(parsed) => parsed, - Err(message) => { - client.set_reply(error_reply(message)); - return; - } - }; - let reply = match storage.vsim(&client.key(), parsed.query, parsed.options) { - Ok(hits) if parsed.with_scores => RespData::Map( - hits.into_iter() - .map(|hit| { - ( - RespData::BulkString(Some(Bytes::from(hit.element))), - RespData::Double(hit.score), - ) - }) - .collect(), - ), - Ok(hits) => RespData::Array(Some( - hits.into_iter() - .map(|hit| RespData::BulkString(Some(Bytes::from(hit.element)))) - .collect(), - )), - Err(error) => storage_error_reply(error, MissingError::Element), - }; - client.set_reply(reply); - } -} - -impl Cmd for VRemCmd { - impl_cmd_meta!(); - impl_cmd_clone_box!(); - - fn do_initial(&self, client: &Client) -> bool { - set_command_key(client) - } - - fn do_cmd(&self, client: &Client, storage: Arc) { - let argv = client.argv(); - let reply = match storage.vrem(&client.key(), &argv[2]) { - Ok(removed) => RespData::Integer(i64::from(removed)), - Err(error) => storage_error_reply(error, MissingError::Key), - }; - client.set_reply(reply); - } -} - -impl Cmd for VCardCmd { - impl_cmd_meta!(); - impl_cmd_clone_box!(); - - fn do_initial(&self, client: &Client) -> bool { - set_command_key(client) - } - - fn do_cmd(&self, client: &Client, storage: Arc) { - let reply = match storage.vcard(&client.key()) { - Ok(count) => integer_reply(count), - Err(error) => storage_error_reply(error, MissingError::Key), - }; - client.set_reply(reply); - } -} - -impl Cmd for VDimCmd { - impl_cmd_meta!(); - impl_cmd_clone_box!(); - - fn do_initial(&self, client: &Client) -> bool { - set_command_key(client) - } - - fn do_cmd(&self, client: &Client, storage: Arc) { - let reply = match storage.vdim(&client.key()) { - Ok(dimension) => integer_reply(u64::from(dimension)), - Err(error) => storage_error_reply(error, MissingError::Key), - }; - client.set_reply(reply); - } -} - -impl Cmd for VEmbCmd { - impl_cmd_meta!(); - impl_cmd_clone_box!(); - - fn do_initial(&self, client: &Client) -> bool { - set_command_key(client) - } - - fn do_cmd(&self, client: &Client, storage: Arc) { - let element = match parse_vemb(&client.argv()) { - Ok(element) => element, - Err(message) => { - client.set_reply(error_reply(message)); - return; - } - }; - let reply = match storage.vemb(&client.key(), &element) { - Ok(Some(values)) => { - RespData::Array(Some(values.into_iter().map(RespData::Double).collect())) - } - Ok(None) => RespData::BulkString(None), - Err(error) => storage_error_reply(error, MissingError::Key), - }; - client.set_reply(reply); - } -} - -impl Cmd for VIsMemberCmd { - impl_cmd_meta!(); - impl_cmd_clone_box!(); - - fn do_initial(&self, client: &Client) -> bool { - set_command_key(client) - } - - fn do_cmd(&self, client: &Client, storage: Arc) { - let argv = client.argv(); - let reply = match storage.vismember(&client.key(), &argv[2]) { - Ok(is_member) => RespData::Integer(i64::from(is_member)), - Err(error) => storage_error_reply(error, MissingError::Key), - }; - client.set_reply(reply); - } -} +pub mod vadd; +pub mod vcard; +pub mod vdim; +pub mod vemb; +pub mod vismember; +pub mod vrem; +pub mod vsim; + +pub use vadd::VAddCmd; +pub use vcard::VCardCmd; +pub use vdim::VDimCmd; +pub use vemb::VEmbCmd; +pub use vismember::VIsMemberCmd; +pub use vrem::VRemCmd; +pub use vsim::VSimCmd; #[allow(clippy::unwrap_used)] #[cfg(test)] @@ -463,6 +263,7 @@ mod tests { use storage::{VectorQuery, VectorSearchMode}; use super::*; + use crate::Cmd; fn fp32(values: &[f32]) -> Vec { values diff --git a/src/cmd/src/vector/vadd.rs b/src/cmd/src/vector/vadd.rs new file mode 100644 index 00000000..e4c69ac5 --- /dev/null +++ b/src/cmd/src/vector/vadd.rs @@ -0,0 +1,58 @@ +// Copyright (c) 2024-present, arana-db Community. All rights reserved. +// +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use client::Client; +use resp::RespData; +use storage::storage::Storage; + +use crate::{AclCategory, Cmd, CmdFlags, CmdMeta, impl_cmd_clone_box, impl_cmd_meta}; + +use super::{MissingError, error_reply, parse_vadd, storage_error_reply}; + +crate::define_vector_command!( + VAddCmd, + "vadd", + -5, + CmdFlags::WRITE | CmdFlags::FAST, + AclCategory::KEYSPACE | AclCategory::WRITE +); + +impl Cmd for VAddCmd { + impl_cmd_meta!(); + impl_cmd_clone_box!(); + + fn do_initial(&self, client: &Client) -> bool { + super::set_command_key(client) + } + + fn do_cmd(&self, client: &Client, storage: Arc) { + let parsed = match parse_vadd(&client.argv()) { + Ok(parsed) => parsed, + Err(message) => { + client.set_reply(error_reply(message)); + return; + } + }; + let reply = match storage.vadd(&client.key(), &parsed.element, &parsed.vector) { + Ok(inserted) => RespData::Integer(i64::from(inserted)), + Err(error) => storage_error_reply(error, MissingError::Key), + }; + client.set_reply(reply); + } +} diff --git a/src/cmd/src/vector/vcard.rs b/src/cmd/src/vector/vcard.rs new file mode 100644 index 00000000..ae11347f --- /dev/null +++ b/src/cmd/src/vector/vcard.rs @@ -0,0 +1,50 @@ +// Copyright (c) 2024-present, arana-db Community. All rights reserved. +// +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use client::Client; +use storage::storage::Storage; + +use crate::{AclCategory, Cmd, CmdFlags, CmdMeta, impl_cmd_clone_box, impl_cmd_meta}; + +use super::{MissingError, integer_reply, storage_error_reply}; + +crate::define_vector_command!( + VCardCmd, + "vcard", + 2, + CmdFlags::READONLY | CmdFlags::FAST, + AclCategory::KEYSPACE | AclCategory::READ +); + +impl Cmd for VCardCmd { + impl_cmd_meta!(); + impl_cmd_clone_box!(); + + fn do_initial(&self, client: &Client) -> bool { + super::set_command_key(client) + } + + fn do_cmd(&self, client: &Client, storage: Arc) { + let reply = match storage.vcard(&client.key()) { + Ok(count) => integer_reply(count), + Err(error) => storage_error_reply(error, MissingError::Key), + }; + client.set_reply(reply); + } +} diff --git a/src/cmd/src/vector/vdim.rs b/src/cmd/src/vector/vdim.rs new file mode 100644 index 00000000..fb1db93b --- /dev/null +++ b/src/cmd/src/vector/vdim.rs @@ -0,0 +1,50 @@ +// Copyright (c) 2024-present, arana-db Community. All rights reserved. +// +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use client::Client; +use storage::storage::Storage; + +use crate::{AclCategory, Cmd, CmdFlags, CmdMeta, impl_cmd_clone_box, impl_cmd_meta}; + +use super::{MissingError, integer_reply, storage_error_reply}; + +crate::define_vector_command!( + VDimCmd, + "vdim", + 2, + CmdFlags::READONLY | CmdFlags::FAST, + AclCategory::KEYSPACE | AclCategory::READ +); + +impl Cmd for VDimCmd { + impl_cmd_meta!(); + impl_cmd_clone_box!(); + + fn do_initial(&self, client: &Client) -> bool { + super::set_command_key(client) + } + + fn do_cmd(&self, client: &Client, storage: Arc) { + let reply = match storage.vdim(&client.key()) { + Ok(dimension) => integer_reply(u64::from(dimension)), + Err(error) => storage_error_reply(error, MissingError::Key), + }; + client.set_reply(reply); + } +} diff --git a/src/cmd/src/vector/vemb.rs b/src/cmd/src/vector/vemb.rs new file mode 100644 index 00000000..b8715cb4 --- /dev/null +++ b/src/cmd/src/vector/vemb.rs @@ -0,0 +1,61 @@ +// Copyright (c) 2024-present, arana-db Community. All rights reserved. +// +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use client::Client; +use resp::RespData; +use storage::storage::Storage; + +use crate::{AclCategory, Cmd, CmdFlags, CmdMeta, impl_cmd_clone_box, impl_cmd_meta}; + +use super::{MissingError, error_reply, parse_vemb, storage_error_reply}; + +crate::define_vector_command!( + VEmbCmd, + "vemb", + -3, + CmdFlags::READONLY | CmdFlags::FAST, + AclCategory::KEYSPACE | AclCategory::READ +); + +impl Cmd for VEmbCmd { + impl_cmd_meta!(); + impl_cmd_clone_box!(); + + fn do_initial(&self, client: &Client) -> bool { + super::set_command_key(client) + } + + fn do_cmd(&self, client: &Client, storage: Arc) { + let element = match parse_vemb(&client.argv()) { + Ok(element) => element, + Err(message) => { + client.set_reply(error_reply(message)); + return; + } + }; + let reply = match storage.vemb(&client.key(), &element) { + Ok(Some(values)) => { + RespData::Array(Some(values.into_iter().map(RespData::Double).collect())) + } + Ok(None) => RespData::BulkString(None), + Err(error) => storage_error_reply(error, MissingError::Key), + }; + client.set_reply(reply); + } +} diff --git a/src/cmd/src/vector/vismember.rs b/src/cmd/src/vector/vismember.rs new file mode 100644 index 00000000..0457d0ae --- /dev/null +++ b/src/cmd/src/vector/vismember.rs @@ -0,0 +1,52 @@ +// Copyright (c) 2024-present, arana-db Community. All rights reserved. +// +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use client::Client; +use resp::RespData; +use storage::storage::Storage; + +use crate::{AclCategory, Cmd, CmdFlags, CmdMeta, impl_cmd_clone_box, impl_cmd_meta}; + +use super::{MissingError, storage_error_reply}; + +crate::define_vector_command!( + VIsMemberCmd, + "vismember", + 3, + CmdFlags::READONLY | CmdFlags::FAST, + AclCategory::KEYSPACE | AclCategory::READ +); + +impl Cmd for VIsMemberCmd { + impl_cmd_meta!(); + impl_cmd_clone_box!(); + + fn do_initial(&self, client: &Client) -> bool { + super::set_command_key(client) + } + + fn do_cmd(&self, client: &Client, storage: Arc) { + let argv = client.argv(); + let reply = match storage.vismember(&client.key(), &argv[2]) { + Ok(is_member) => RespData::Integer(i64::from(is_member)), + Err(error) => storage_error_reply(error, MissingError::Key), + }; + client.set_reply(reply); + } +} diff --git a/src/cmd/src/vector/vrem.rs b/src/cmd/src/vector/vrem.rs new file mode 100644 index 00000000..1752738f --- /dev/null +++ b/src/cmd/src/vector/vrem.rs @@ -0,0 +1,52 @@ +// Copyright (c) 2024-present, arana-db Community. All rights reserved. +// +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use client::Client; +use resp::RespData; +use storage::storage::Storage; + +use crate::{AclCategory, Cmd, CmdFlags, CmdMeta, impl_cmd_clone_box, impl_cmd_meta}; + +use super::{MissingError, storage_error_reply}; + +crate::define_vector_command!( + VRemCmd, + "vrem", + 3, + CmdFlags::WRITE | CmdFlags::FAST, + AclCategory::KEYSPACE | AclCategory::WRITE +); + +impl Cmd for VRemCmd { + impl_cmd_meta!(); + impl_cmd_clone_box!(); + + fn do_initial(&self, client: &Client) -> bool { + super::set_command_key(client) + } + + fn do_cmd(&self, client: &Client, storage: Arc) { + let argv = client.argv(); + let reply = match storage.vrem(&client.key(), &argv[2]) { + Ok(removed) => RespData::Integer(i64::from(removed)), + Err(error) => storage_error_reply(error, MissingError::Key), + }; + client.set_reply(reply); + } +} diff --git a/src/cmd/src/vector/vsim.rs b/src/cmd/src/vector/vsim.rs new file mode 100644 index 00000000..a0c39ddb --- /dev/null +++ b/src/cmd/src/vector/vsim.rs @@ -0,0 +1,73 @@ +// Copyright (c) 2024-present, arana-db Community. All rights reserved. +// +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use bytes::Bytes; +use client::Client; +use resp::RespData; +use storage::storage::Storage; + +use crate::{AclCategory, Cmd, CmdFlags, CmdMeta, impl_cmd_clone_box, impl_cmd_meta}; + +use super::{MissingError, error_reply, parse_vsim, storage_error_reply}; + +crate::define_vector_command!( + VSimCmd, + "vsim", + -4, + CmdFlags::READONLY, + AclCategory::KEYSPACE | AclCategory::READ | AclCategory::SLOW +); + +impl Cmd for VSimCmd { + impl_cmd_meta!(); + impl_cmd_clone_box!(); + + fn do_initial(&self, client: &Client) -> bool { + super::set_command_key(client) + } + + fn do_cmd(&self, client: &Client, storage: Arc) { + let parsed = match parse_vsim(&client.argv()) { + Ok(parsed) => parsed, + Err(message) => { + client.set_reply(error_reply(message)); + return; + } + }; + let reply = match storage.vsim(&client.key(), parsed.query, parsed.options) { + Ok(hits) if parsed.with_scores => RespData::Map( + hits.into_iter() + .map(|hit| { + ( + RespData::BulkString(Some(Bytes::from(hit.element))), + RespData::Double(hit.score), + ) + }) + .collect(), + ), + Ok(hits) => RespData::Array(Some( + hits.into_iter() + .map(|hit| RespData::BulkString(Some(Bytes::from(hit.element)))) + .collect(), + )), + Err(error) => storage_error_reply(error, MissingError::Element), + }; + client.set_reply(reply); + } +} diff --git a/src/net/src/executor_ext.rs b/src/net/src/executor_ext.rs index c4e62abf..ea6bbbd6 100644 --- a/src/net/src/executor_ext.rs +++ b/src/net/src/executor_ext.rs @@ -60,16 +60,8 @@ impl CmdExecutorNetworkExt for CmdExecutor { return Ok(()); } - // Cluster-mode gates run before any command-specific setup: reject - // unsupported module commands on every node, then writes on followers. + // Cluster-mode gate: redirect write commands on followers to the leader. if let Some(gate) = exec.leader_gate.as_ref() { - if exec.cmd.has_flag(CmdFlags::MODULE_NO_CLUSTER) { - exec.client.set_reply(RespData::Error( - "ERR Vector Set is not supported in cluster mode".into(), - )); - return Ok(()); - } - if exec.cmd.has_flag(CmdFlags::WRITE) && !gate.is_leader() { // Simplified redirect: Kiwi returns "MOVED " (no hash slot, // unlike Redis Cluster's "MOVED "). Clients are diff --git a/src/net/tests/storage_command_e2e_tests.rs b/src/net/tests/storage_command_e2e_tests.rs index c9a81c36..e8bd88ee 100644 --- a/src/net/tests/storage_command_e2e_tests.rs +++ b/src/net/tests/storage_command_e2e_tests.rs @@ -424,14 +424,12 @@ async fn storage_command_e2e_generic_storage_commands_use_storage_path() { } #[tokio::test] -async fn storage_command_e2e_no_cluster_commands_precede_follower_redirect() { +async fn storage_command_e2e_vector_write_commands_redirect_on_follower() { let server = TestServer::start_with_leader_gate(None, Some(Arc::new(FollowerGate))).await; let mut stream = tokio::net::TcpStream::connect(server.addr) .await .expect("connect to server"); - let expected = RespData::Error(Bytes::from_static( - b"ERR Vector Set is not supported in cluster mode", - )); + let expected = RespData::Error(Bytes::from_static(b"MOVED 127.0.0.1:7380")); let reply = send_command( &mut stream, diff --git a/src/storage/src/format_member_data_key.rs b/src/storage/src/format_member_data_key.rs index 457b2f38..add8768e 100644 --- a/src/storage/src/format_member_data_key.rs +++ b/src/storage/src/format_member_data_key.rs @@ -28,6 +28,9 @@ use crate::{ // used for Hash/Set/Zset's member data key. format: // | reserve1 | key | version | data | reserve2 | // | 8B | | 8B | | 16B | +// +// For vector sets, `data` is the element name and the value stored at this key +// is a VectorDataValue (see format_vector.rs). #[derive(Debug, Clone)] pub struct MemberDataKey { diff --git a/src/storage/src/format_vector.rs b/src/storage/src/format_vector.rs index 20b23bb5..909d9e51 100644 --- a/src/storage/src/format_vector.rs +++ b/src/storage/src/format_vector.rs @@ -32,14 +32,64 @@ pub const VECTOR_METRIC_COSINE: u8 = 1; pub const VECTOR_VALUE_MAGIC: u8 = 0x56; pub const VECTOR_VALUE_FORMAT: u8 = 1; +// Vector set meta value layout stored in MetaCF: +// +// | data_type | count | version | format | encoding | metric | flags | dimension | zero_reserve | ctime | etime | +// | 1B | 8B | 8B | 1B | 1B | 1B | 1B | 4B | 8B | 8B | 8B | +// +// `data_type` is DataType::VectorSet, `encoding` is VECTOR_ENCODING_FP32_LE, +// and `metric` is the similarity metric used for VSIM (e.g. cosine). + +// Vector member data value layout stored in VectorDataCF: +// +// | magic | format | dimension | original_l2 | normalized_components ... | +// | 1B | 1B | 4B | 4B | 4B * dimension | +// +// `magic` is VECTOR_VALUE_MAGIC and `original_l2` preserves the pre-normalization +// L2 norm so VEMB can reconstruct the original FP32 vector. + const VECTOR_META_ZERO_RESERVE_LENGTH: usize = 8; const VECTOR_VALUE_HEADER_LENGTH: usize = 10; +/// Similarity metric used to compare vectors in a vector set. +/// +/// The metric is persisted in `VectorMeta` so all future VSIM queries against +/// the set use the same formula that was established at creation time. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SimilarityMetric { + Cosine, +} + +impl SimilarityMetric { + pub fn from_u8(value: u8) -> Result { + match value { + VECTOR_METRIC_COSINE => Ok(Self::Cosine), + _ => InvalidFormatSnafu { + message: format!("unsupported vector metric: {value}"), + } + .fail(), + } + } + + pub const fn to_u8(self) -> u8 { + match self { + Self::Cosine => VECTOR_METRIC_COSINE, + } + } + + pub fn score(&self, left: &CanonicalVector, right: &CanonicalVector) -> Result { + match self { + Self::Cosine => left.cosine_score(right), + } + } +} + #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) struct VectorMeta { count: u64, pub(crate) version: u64, dimension: u32, + metric: SimilarityMetric, ctime: u64, etime: u64, } @@ -55,11 +105,16 @@ impl VectorMeta { count, version, dimension, + metric: SimilarityMetric::Cosine, ctime: now, etime: 0, } } + pub(crate) fn metric(&self) -> SimilarityMetric { + self.metric + } + pub(crate) fn encode(&self) -> BytesMut { let mut output = BytesMut::with_capacity(BASE_META_VALUE_LENGTH); output.put_u8(DataType::VectorSet as u8); @@ -67,7 +122,7 @@ impl VectorMeta { output.put_u64_le(self.version); output.put_u8(VECTOR_META_FORMAT); output.put_u8(VECTOR_ENCODING_FP32_LE); - output.put_u8(VECTOR_METRIC_COSINE); + output.put_u8(self.metric.to_u8()); output.put_u8(0); output.put_u32_le(self.dimension); output.put_bytes(0, VECTOR_META_ZERO_RESERVE_LENGTH); @@ -121,12 +176,7 @@ impl VectorMeta { message: format!("unsupported vector encoding: {encoding}") } ); - ensure!( - metric == VECTOR_METRIC_COSINE, - InvalidFormatSnafu { - message: format!("unsupported vector metric: {metric}") - } - ); + let metric = SimilarityMetric::from_u8(metric)?; ensure!( flags == 0 && zero_reserve.iter().all(|byte| *byte == 0), InvalidFormatSnafu { @@ -144,6 +194,7 @@ impl VectorMeta { count, version, dimension, + metric, ctime, etime, }) diff --git a/src/storage/src/lib.rs b/src/storage/src/lib.rs index 4ca63142..ab6d3c3f 100644 --- a/src/storage/src/lib.rs +++ b/src/storage/src/lib.rs @@ -83,4 +83,7 @@ pub use statistics::KeyStatistics; pub use storage::{BgTask, BgTaskHandler}; pub use storage_impl::BeforeOrAfter; pub use util::{safe_cleanup_test_db, unique_test_db_path}; -pub use vector::{CanonicalVector, VectorHit, VectorQuery, VectorSearchMode, VectorSearchOptions}; +pub use vector::{ + CanonicalVector, VectorHit, VectorQuery, VectorSearchEngine, VectorSearchMode, + VectorSearchOptions, +}; diff --git a/src/storage/src/redis_vectors.rs b/src/storage/src/redis_vectors.rs index 83db5853..6ea7fd89 100644 --- a/src/storage/src/redis_vectors.rs +++ b/src/storage/src/redis_vectors.rs @@ -15,18 +15,15 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::{cmp::Ordering, collections::BinaryHeap}; - use kstd::lock_mgr::ScopeRecordLock; use rocksdb::{Direction, IteratorMode, ReadOptions}; use snafu::{OptionExt, ResultExt}; use crate::{ CanonicalVector, ColumnFamilyIndex, DataType, Redis, Result, TypeCheckState, VectorHit, - VectorQuery, VectorSearchOptions, + VectorQuery, VectorSearchEngine, VectorSearchOptions, error::{ - InvalidArgumentSnafu, InvalidFormatSnafu, KeyNotFoundSnafu, OptionNoneSnafu, RedisErrSnafu, - RocksSnafu, + InvalidArgumentSnafu, InvalidFormatSnafu, KeyNotFoundSnafu, OptionNoneSnafu, RocksSnafu, }, format_base_key::BaseMetaKey, format_member_data_key::MemberDataKey, @@ -34,50 +31,7 @@ use crate::{ storage_define::SUFFIX_RESERVE_LENGTH, }; -#[derive(Debug)] -struct HeapHit { - element: Vec, - score: f64, -} - -impl PartialEq for HeapHit { - fn eq(&self, other: &Self) -> bool { - self.score.total_cmp(&other.score) == Ordering::Equal && self.element == other.element - } -} - -impl Eq for HeapHit {} - -impl PartialOrd for HeapHit { - fn partial_cmp(&self, other: &Self) -> Option { - Some(self.cmp(other)) - } -} - -impl Ord for HeapHit { - fn cmp(&self, other: &Self) -> Ordering { - other - .score - .total_cmp(&self.score) - .then_with(|| self.element.cmp(&other.element)) - } -} - impl Redis { - pub fn is_cluster_mode(&self) -> bool { - self.append_log_fn.get().is_some() - } - - fn ensure_vector_standalone(&self) -> Result<()> { - if self.is_cluster_mode() { - return RedisErrSnafu { - message: "ERR Vector Set is not supported in cluster mode".to_string(), - } - .fail(); - } - Ok(()) - } - fn parse_vector_meta(&self, value: &[u8]) -> Result> { if value.is_empty() { return Ok(None); @@ -95,7 +49,6 @@ impl Redis { } pub fn vadd(&self, key: &[u8], element: &[u8], vector: &CanonicalVector) -> Result { - self.ensure_vector_standalone()?; let db = self.db.as_ref().context(OptionNoneSnafu { message: "db is not initialized".to_string(), })?; @@ -173,7 +126,6 @@ impl Redis { } pub fn vrem(&self, key: &[u8], element: &[u8]) -> Result { - self.ensure_vector_standalone()?; let db = self.db.as_ref().context(OptionNoneSnafu { message: "db is not initialized".to_string(), })?; @@ -221,12 +173,10 @@ impl Redis { } pub fn vcard(&self, key: &[u8]) -> Result { - self.ensure_vector_standalone()?; Ok(self.read_vector_meta(key)?.map_or(0, |meta| meta.count())) } pub fn vdim(&self, key: &[u8]) -> Result { - self.ensure_vector_standalone()?; match self.read_vector_meta(key)? { Some(meta) => Ok(meta.dimension()), None => KeyNotFoundSnafu { @@ -237,7 +187,6 @@ impl Redis { } pub fn vemb(&self, key: &[u8], element: &[u8]) -> Result>> { - self.ensure_vector_standalone()?; let db = self.db.as_ref().context(OptionNoneSnafu { message: "db is not initialized".to_string(), })?; @@ -287,7 +236,6 @@ impl Redis { } pub fn vismember(&self, key: &[u8], element: &[u8]) -> Result { - self.ensure_vector_standalone()?; let db = self.db.as_ref().context(OptionNoneSnafu { message: "db is not initialized".to_string(), })?; @@ -328,7 +276,6 @@ impl Redis { query: VectorQuery, options: VectorSearchOptions, ) -> Result> { - self.ensure_vector_standalone()?; if options.count == 0 { return InvalidArgumentSnafu { message: "vector search count must be greater than zero".to_string(), @@ -399,60 +346,39 @@ impl Redis { scan_options, IteratorMode::From(&prefix, Direction::Forward), ); - let mut heap = BinaryHeap::new(); - - for entry in iterator { - let (encoded_key, encoded_value) = entry.context(RocksSnafu)?; - if !encoded_key.starts_with(&prefix) { - break; - } - if encoded_key.len() < prefix.len() + SUFFIX_RESERVE_LENGTH { - return InvalidFormatSnafu { - message: "vector member key is shorter than its generation prefix".to_string(), + let engine = VectorSearchEngine::Flat; + let candidates = iterator + .take_while(|result| match result { + Ok((encoded_key, _)) => encoded_key.starts_with(&prefix), + Err(_) => true, + }) + .map(|entry| { + let (encoded_key, encoded_value) = entry.context(RocksSnafu)?; + if encoded_key.len() < prefix.len() + SUFFIX_RESERVE_LENGTH { + return InvalidFormatSnafu { + message: "vector member key is shorter than its generation prefix" + .to_string(), + } + .fail(); } - .fail(); - } - let element_end = encoded_key.len() - SUFFIX_RESERVE_LENGTH; - let element = encoded_key[prefix.len()..element_end].to_vec(); - let value = VectorDataValue::decode(&encoded_value)?; - if value.dimension() != meta.dimension() { - return InvalidFormatSnafu { - message: format!( - "vector member dimension {} does not match meta dimension {}", - value.dimension(), - meta.dimension() - ), + let element_end = encoded_key.len() - SUFFIX_RESERVE_LENGTH; + let element = encoded_key[prefix.len()..element_end].to_vec(); + let value = VectorDataValue::decode(&encoded_value)?; + if value.dimension() != meta.dimension() { + return InvalidFormatSnafu { + message: format!( + "vector member dimension {} does not match meta dimension {}", + value.dimension(), + meta.dimension() + ), + } + .fail(); } - .fail(); - } - let hit = HeapHit { - element, - score: query_vector.score(value.canonical())?, - }; + Ok((element, value.canonical().clone())) + }); - if heap.len() < options.count { - heap.push(hit); - } else if heap.peek().is_some_and(|worst| hit < *worst) { - heap.pop(); - heap.push(hit); - } - } - - let mut hits = heap - .into_iter() - .map(|hit| VectorHit { - element: hit.element, - score: hit.score, - }) - .collect::>(); - hits.sort_by(|left, right| { - right - .score - .total_cmp(&left.score) - .then_with(|| left.element.cmp(&right.element)) - }); - Ok(hits) + engine.search(&query_vector, &meta.metric(), options.count, candidates) } fn read_vector_meta(&self, key: &[u8]) -> Result> { diff --git a/src/storage/src/vector.rs b/src/storage/src/vector.rs index 11dfb5a0..6df41973 100644 --- a/src/storage/src/vector.rs +++ b/src/storage/src/vector.rs @@ -15,9 +15,12 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::{cmp::Ordering, collections::BinaryHeap}; + use snafu::ensure; use crate::error::{InvalidArgumentSnafu, Result}; +use crate::format_vector::SimilarityMetric; #[derive(Debug, Clone, PartialEq)] pub struct CanonicalVector { @@ -120,6 +123,10 @@ impl CanonicalVector { } pub fn score(&self, other: &Self) -> Result { + self.cosine_score(other) + } + + pub fn cosine_score(&self, other: &Self) -> Result { ensure!( self.dimension == other.dimension, InvalidArgumentSnafu { @@ -172,6 +179,91 @@ pub struct VectorHit { pub score: f64, } +#[derive(Debug, Clone, PartialEq)] +struct ScoredCandidate { + element: Vec, + score: f64, +} + +impl PartialOrd for ScoredCandidate { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for ScoredCandidate { + fn cmp(&self, other: &Self) -> Ordering { + other + .score + .total_cmp(&self.score) + .then_with(|| self.element.cmp(&other.element)) + } +} + +impl Eq for ScoredCandidate {} + +/// Search strategy for vector-similarity queries. +/// +/// Currently only exhaustive flat search is implemented. Future variants can +/// hold pre-built approximate indices such as HNSW. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum VectorSearchEngine { + Flat, +} + +impl VectorSearchEngine { + /// Search `candidates` and return up to `count` best hits according to + /// `metric`. + pub fn search( + &self, + query: &CanonicalVector, + metric: &SimilarityMetric, + count: usize, + candidates: impl Iterator, CanonicalVector)>>, + ) -> Result> { + match self { + Self::Flat => Self::flat_search(query, metric, count, candidates), + } + } + + fn flat_search( + query: &CanonicalVector, + metric: &SimilarityMetric, + count: usize, + candidates: impl Iterator, CanonicalVector)>>, + ) -> Result> { + let mut heap = BinaryHeap::new(); + + for candidate in candidates { + let (element, vector) = candidate?; + let score = metric.score(query, &vector)?; + let item = ScoredCandidate { element, score }; + + if heap.len() < count { + heap.push(item); + } else if heap.peek().is_some_and(|worst| item < *worst) { + heap.pop(); + heap.push(item); + } + } + + let mut hits = heap + .into_iter() + .map(|item| VectorHit { + element: item.element, + score: item.score, + }) + .collect::>(); + hits.sort_by(|left, right| { + right + .score + .total_cmp(&left.score) + .then_with(|| left.element.cmp(&right.element)) + }); + Ok(hits) + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/src/storage/tests/redis_vector_test.rs b/src/storage/tests/redis_vector_test.rs index be7fbdf8..9aea4894 100644 --- a/src/storage/tests/redis_vector_test.rs +++ b/src/storage/tests/redis_vector_test.rs @@ -17,7 +17,7 @@ #![allow(clippy::unwrap_used)] -use std::{path::PathBuf, sync::Arc}; +use std::{path::PathBuf, sync::Arc, sync::Mutex}; use kstd::lock_mgr::LockMgr; use rocksdb::{IteratorMode, ReadOptions}; @@ -549,32 +549,64 @@ fn test_flushdb_removes_vector_meta_and_members() { } #[tokio::test] -async fn test_vector_storage_rejects_cluster_mode() { +async fn test_vector_storage_proposes_binlog_in_cluster_mode() { let test_db_path = unique_test_db_path(); safe_cleanup_test_db(&test_db_path); let mut storage = Storage::new(1, 0); let _receiver = storage .open(Arc::new(StorageOptions::default()), &test_db_path) .expect("open storage"); - storage.set_append_log_fn(Arc::new(|_| panic!("vector API must not append Raft log"))); + + // Wrap Storage in an Arc so the append_log_fn callback can apply binlogs + // back to the local instance. A Weak reference is used to avoid creating a + // strong-reference cycle between the callback and Storage. + let storage_arc = Arc::new(storage); + let storage_weak = Arc::downgrade(&storage_arc); + let captured = Arc::new(Mutex::new(None)); + let captured_clone = captured.clone(); + storage_arc.set_append_log_fn(Arc::new(move |binlog| { + *captured_clone.lock().expect("lock captured binlog") = Some(binlog.clone()); + let storage = storage_weak + .upgrade() + .ok_or("storage dropped before binlog apply")?; + storage + .on_binlog_write(&binlog, 1) + .map_err(|error| error.to_string())?; + Ok(conf::raft_type::BinlogResponse::ok()) + })); + let vector = CanonicalVector::from_values(&[1.0, 0.0]).expect("vector"); + assert!( + storage_arc + .vadd(b"vectors", b"member", &vector) + .expect("vadd in cluster mode"), + "vadd should insert a new member" + ); + { + let binlog = captured.lock().expect("lock captured binlog").take(); + let binlog = binlog.expect("vadd should propose a binlog"); + assert!( + !binlog.entries.is_empty(), + "vadd binlog should contain vector writes" + ); + } - assert!(storage.vadd(b"vectors", b"member", &vector).is_err()); assert!( - storage - .vsim( - b"vectors", - VectorQuery::Vector(vector), - search_options(1, VectorSearchMode::Approximate), - ) - .is_err() + storage_arc + .vrem(b"vectors", b"member") + .expect("vrem in cluster mode"), + "vrem should remove the member" + ); + let binlog = captured.lock().expect("lock captured binlog").take(); + let binlog = binlog.expect("vrem should propose a binlog"); + assert!( + !binlog.entries.is_empty(), + "vrem binlog should contain vector deletes" ); - assert!(storage.vrem(b"vectors", b"member").is_err()); - assert!(storage.vcard(b"vectors").is_err()); - assert!(storage.vdim(b"vectors").is_err()); - assert!(storage.vemb(b"vectors", b"member").is_err()); - assert!(storage.vismember(b"vectors", b"member").is_err()); + // Weak reference, so the strong count is 1. + let mut storage = Arc::try_unwrap(storage_arc) + .unwrap_or_else(|_| panic!("storage should not be shared after test operations")); storage.shutdown().await; safe_cleanup_test_db(&test_db_path); } From 1ffceef241234705434333bf0bbf137ce85b29be Mon Sep 17 00:00:00 2001 From: happy-v587 Date: Sun, 26 Jul 2026 17:03:54 +0800 Subject: [PATCH 16/32] fix(cmd): keep MODULE_NO_CLUSTER flag definition Restore the flag constant so the bit remains reserved, even though vector commands no longer use it. Co-Authored-By: Claude --- src/cmd/src/lib.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/src/cmd/src/lib.rs b/src/cmd/src/lib.rs index 3b8cc7f1..22ecfb18 100644 --- a/src/cmd/src/lib.rs +++ b/src/cmd/src/lib.rs @@ -139,6 +139,7 @@ bitflags! { const NO_AUTH = 1 << 10; // Skip ACL checks const MAY_REPLICATE = 1 << 11; // May replicate even if writes are disabled const PROTECTED = 1 << 12; // Don't accept in scripts + const MODULE_NO_CLUSTER = 1 << 13; // No cluster mode support const NO_MULTI = 1 << 14; // Cannot be pipelined const EXCLUSIVE = 1 << 15; // May change Storage pointer const RAFT = 1 << 16; From 4ce0aed7047813507d247aa92396597d99447445 Mon Sep 17 00:00:00 2001 From: happy-v587 Date: Sun, 26 Jul 2026 21:27:19 +0800 Subject: [PATCH 17/32] refactor(vector,resp): address CodeRabbit review comments on PR #356 - Document reversed score ordering in HeapHit::cmp. - Add regression test for unsupported vector metric byte. - Expose VectorMeta accessors and use decode/encode in tests. - Delegate RESP encoding to version-aware append_* helpers. Co-Authored-By: Claude --- src/resp/src/encode.rs | 243 +++++++++---------------- src/storage/src/format_vector.rs | 21 ++- src/storage/src/lib.rs | 2 +- src/storage/src/redis_vectors.rs | 4 + src/storage/tests/redis_vector_test.rs | 34 ++-- 5 files changed, 130 insertions(+), 174 deletions(-) diff --git a/src/resp/src/encode.rs b/src/resp/src/encode.rs index 9e6277cd..994a8da6 100644 --- a/src/resp/src/encode.rs +++ b/src/resp/src/encode.rs @@ -135,11 +135,11 @@ pub trait RespEncode { fn append_double(&mut self, value: f64) -> &mut Self; - fn append_big_number(&mut self, value: &str) -> &mut Self; + fn append_big_number(&mut self, value: &[u8]) -> &mut Self; fn append_bulk_error(&mut self, value: &[u8]) -> &mut Self; - fn append_verbatim_string(&mut self, format: &str, data: &[u8]) -> &mut Self; + fn append_verbatim_string(&mut self, format: &[u8], data: &[u8]) -> &mut Self; fn append_map(&mut self, pairs: &[(RespData, RespData)]) -> &mut Self; @@ -226,111 +226,14 @@ impl RespEncoder { } RespData::Array(None) => self.set_array_len(-1), RespData::Null => self.append_null(), - RespData::Boolean(value) => { - if self.is_resp3() { - self.buffer.extend_from_slice(b"#"); - self.buffer - .extend_from_slice(if *value { b"t" } else { b"f" }); - self.append_crlf() - } else { - self.append_integer(if *value { 1 } else { 0 }) - } - } - RespData::Double(value) => { - if self.is_resp3() { - if value.is_nan() { - self.buffer.extend_from_slice(b",nan"); - } else if value.is_infinite() { - self.buffer.extend_from_slice(if value.is_sign_negative() { - b",-inf" - } else { - b",inf" - }); - } else { - let _ = write!(self.buffer, ",{value}"); - } - self.append_crlf() - } else { - self.append_bulk_string(value.to_string().as_bytes()) - } - } - RespData::BigNumber(bytes) => { - if self.is_resp3() { - self.buffer.extend_from_slice(b"("); - self.buffer.extend_from_slice(bytes); - self.append_crlf() - } else { - self.append_bulk_string(bytes) - } - } - RespData::BulkError(bytes) => { - if self.is_resp3() { - let _ = write!(self.buffer, "!{}", bytes.len()); - self.append_crlf(); - self.buffer.extend_from_slice(bytes); - self.append_crlf() - } else { - self.buffer.extend_from_slice(b"-"); - self.buffer.extend_from_slice(bytes); - self.append_crlf() - } - } - RespData::VerbatimString { format, data } => { - if self.is_resp3() { - if format.len() != 3 { - panic!( - "RESP3 VerbatimString format must be exactly 3 bytes, got {}", - format.len() - ); - } - let total_len = format.len() + 1 + data.len(); - let _ = write!(self.buffer, "={total_len}"); - self.append_crlf(); - self.buffer.extend_from_slice(format); - self.buffer.extend_from_slice(b":"); - self.buffer.extend_from_slice(data); - self.append_crlf() - } else { - self.append_bulk_string(data) - } - } - RespData::Map(pairs) => { - if self.is_resp3() { - let _ = write!(self.buffer, "%{}", pairs.len()); - self.append_crlf(); - } else { - self.append_array_len((pairs.len() * 2) as i64); - } - for (key, value) in pairs { - self.encode_resp_data_inner(key); - self.encode_resp_data_inner(value); - } - self - } - RespData::Set(items) => { - if self.is_resp3() { - let _ = write!(self.buffer, "~{}", items.len()); - self.append_crlf(); - } else { - self.append_array_len(items.len() as i64); - } - for item in items { - self.encode_resp_data_inner(item); - } - self - } - RespData::Push(items) => { - if self.is_resp3() { - let _ = write!(self.buffer, ">{}", items.len()); - self.append_crlf(); - } else { - self.append_array_len(items.len() as i64); - } - for item in items { - self.encode_resp_data_inner(item); - } - self - } + RespData::Boolean(value) => self.append_boolean(*value), + RespData::Double(value) => self.append_double(*value), + RespData::BigNumber(bytes) => self.append_big_number(bytes), + RespData::BulkError(bytes) => self.append_bulk_error(bytes), + RespData::VerbatimString { format, data } => self.append_verbatim_string(format, data), + RespData::Map(pairs) => self.append_map(pairs), + RespData::Set(items) => self.append_set(items), + RespData::Push(items) => self.append_push(items), } } } @@ -510,80 +413,114 @@ impl RespEncode for RespEncoder { } fn append_boolean(&mut self, value: bool) -> &mut Self { - self.buffer.extend_from_slice(b"#"); - self.buffer - .extend_from_slice(if value { b"t" } else { b"f" }); - self.append_crlf() + if self.is_resp3() { + self.buffer.extend_from_slice(b"#"); + self.buffer + .extend_from_slice(if value { b"t" } else { b"f" }); + self.append_crlf() + } else { + self.append_integer(if value { 1 } else { 0 }) + } } fn append_double(&mut self, value: f64) -> &mut Self { - if value.is_nan() { - self.buffer.extend_from_slice(b",nan"); - } else if value.is_infinite() { - self.buffer.extend_from_slice(if value.is_sign_negative() { - b",-inf" + if self.is_resp3() { + if value.is_nan() { + self.buffer.extend_from_slice(b",nan"); + } else if value.is_infinite() { + self.buffer.extend_from_slice(if value.is_sign_negative() { + b",-inf" + } else { + b",inf" + }); } else { - b",inf" - }); + let _ = write!(self.buffer, ",{value}"); + } + self.append_crlf() } else { - let _ = write!(self.buffer, ",{}", value); + self.append_bulk_string(value.to_string().as_bytes()) } - self.append_crlf() } - fn append_big_number(&mut self, value: &str) -> &mut Self { - self.buffer.extend_from_slice(b"("); - self.buffer.extend_from_slice(value.as_bytes()); - self.append_crlf() + fn append_big_number(&mut self, value: &[u8]) -> &mut Self { + if self.is_resp3() { + self.buffer.extend_from_slice(b"("); + self.buffer.extend_from_slice(value); + self.append_crlf() + } else { + self.append_bulk_string(value) + } } fn append_bulk_error(&mut self, value: &[u8]) -> &mut Self { - let _ = write!(self.buffer, "!{}", value.len()); - self.append_crlf(); - self.buffer.extend_from_slice(value); - self.append_crlf() + if self.is_resp3() { + let _ = write!(self.buffer, "!{}", value.len()); + self.append_crlf(); + self.buffer.extend_from_slice(value); + self.append_crlf() + } else { + self.buffer.extend_from_slice(b"-"); + self.buffer.extend_from_slice(value); + self.append_crlf() + } } - fn append_verbatim_string(&mut self, format: &str, data: &[u8]) -> &mut Self { - if format.len() != 3 { - panic!( - "RESP3 VerbatimString format must be exactly 3 bytes, got {}", - format.len() - ); + fn append_verbatim_string(&mut self, format: &[u8], data: &[u8]) -> &mut Self { + if self.is_resp3() { + if format.len() != 3 { + panic!( + "RESP3 VerbatimString format must be exactly 3 bytes, got {}", + format.len() + ); + } + let total_len = format.len() + 1 + data.len(); // format + ':' + data + let _ = write!(self.buffer, "={total_len}"); + self.append_crlf(); + self.buffer.extend_from_slice(format); + self.buffer.extend_from_slice(b":"); + self.buffer.extend_from_slice(data); + self.append_crlf() + } else { + self.append_bulk_string(data) } - let total_len = format.len() + 1 + data.len(); // format + ':' + data - let _ = write!(self.buffer, "={total_len}"); - self.append_crlf(); - self.buffer.extend_from_slice(format.as_bytes()); - self.buffer.extend_from_slice(b":"); - self.buffer.extend_from_slice(data); - self.append_crlf() } fn append_map(&mut self, pairs: &[(RespData, RespData)]) -> &mut Self { - let _ = write!(self.buffer, "%{}", pairs.len()); - self.append_crlf(); + if self.is_resp3() { + let _ = write!(self.buffer, "%{}", pairs.len()); + self.append_crlf(); + } else { + self.append_array_len((pairs.len() * 2) as i64); + } for (key, value) in pairs { - self.encode_resp_data(key); - self.encode_resp_data(value); + self.encode_resp_data_inner(key); + self.encode_resp_data_inner(value); } self } fn append_set(&mut self, items: &[RespData]) -> &mut Self { - let _ = write!(self.buffer, "~{}", items.len()); - self.append_crlf(); + if self.is_resp3() { + let _ = write!(self.buffer, "~{}", items.len()); + self.append_crlf(); + } else { + self.append_array_len(items.len() as i64); + } for item in items { - self.encode_resp_data(item); + self.encode_resp_data_inner(item); } self } fn append_push(&mut self, items: &[RespData]) -> &mut Self { - let _ = write!(self.buffer, ">{}", items.len()); - self.append_crlf(); + if self.is_resp3() { + let _ = write!(self.buffer, ">{}", items.len()); + self.append_crlf(); + } else { + self.append_array_len(items.len() as i64); + } for item in items { - self.encode_resp_data(item); + self.encode_resp_data_inner(item); } self } @@ -773,7 +710,7 @@ mod tests { encoder.clear().append_double(2.718); assert_eq!(encoder.get_response(), Bytes::from(",2.718\r\n")); - encoder.clear().append_big_number("999999999999999999999"); + encoder.clear().append_big_number(b"999999999999999999999"); assert_eq!( encoder.get_response(), Bytes::from("(999999999999999999999\r\n") @@ -789,7 +726,7 @@ mod tests { encoder .clear() - .append_verbatim_string("txt", b"Hello World"); + .append_verbatim_string(b"txt", b"Hello World"); assert_eq!( encoder.get_response(), Bytes::from("=15\r\ntxt:Hello World\r\n") diff --git a/src/storage/src/format_vector.rs b/src/storage/src/format_vector.rs index 20b23bb5..754fb3f0 100644 --- a/src/storage/src/format_vector.rs +++ b/src/storage/src/format_vector.rs @@ -36,7 +36,7 @@ const VECTOR_META_ZERO_RESERVE_LENGTH: usize = 8; const VECTOR_VALUE_HEADER_LENGTH: usize = 10; #[derive(Debug, Clone, PartialEq, Eq)] -pub(crate) struct VectorMeta { +pub struct VectorMeta { count: u64, pub(crate) version: u64, dimension: u32, @@ -60,7 +60,7 @@ impl VectorMeta { } } - pub(crate) fn encode(&self) -> BytesMut { + pub fn encode(&self) -> BytesMut { let mut output = BytesMut::with_capacity(BASE_META_VALUE_LENGTH); output.put_u8(DataType::VectorSet as u8); output.put_u64_le(self.count); @@ -76,7 +76,7 @@ impl VectorMeta { output } - pub(crate) fn decode(value: &[u8]) -> Result { + pub fn decode(value: &[u8]) -> Result { ensure!( value.len() == BASE_META_VALUE_LENGTH, InvalidFormatSnafu { @@ -157,10 +157,14 @@ impl VectorMeta { self.count = count; } - pub(crate) fn version(&self) -> u64 { + pub fn version(&self) -> u64 { self.version } + pub fn set_version(&mut self, version: u64) { + self.version = version; + } + pub(crate) fn dimension(&self) -> u32 { self.dimension } @@ -168,6 +172,10 @@ impl VectorMeta { pub(crate) fn is_stale(&self) -> bool { self.etime != 0 && self.etime < Utc::now().timestamp_micros() as u64 } + + pub fn set_etime(&mut self, etime: u64) { + self.etime = etime; + } } #[derive(Debug, Clone, PartialEq)] @@ -344,5 +352,10 @@ mod tests { let mut bad_meta_format = encoded_meta; bad_meta_format[17] = 0; assert!(VectorMeta::decode(&bad_meta_format).is_err()); + + let mut bad_metric = bad_meta_format; + bad_metric[17] = VECTOR_META_FORMAT; + bad_metric[19] = 0xFF; + assert!(VectorMeta::decode(&bad_metric).is_err()); } } diff --git a/src/storage/src/lib.rs b/src/storage/src/lib.rs index 4ca63142..a5adf2f3 100644 --- a/src/storage/src/lib.rs +++ b/src/storage/src/lib.rs @@ -19,7 +19,7 @@ extern crate core; mod format_base_data_value; mod format_member_data_key; -mod format_vector; +pub mod format_vector; pub mod vector; mod data_compaction_filter; diff --git a/src/storage/src/redis_vectors.rs b/src/storage/src/redis_vectors.rs index 83db5853..6104604c 100644 --- a/src/storage/src/redis_vectors.rs +++ b/src/storage/src/redis_vectors.rs @@ -55,6 +55,10 @@ impl PartialOrd for HeapHit { } impl Ord for HeapHit { + // Score is reversed so `BinaryHeap` (a max-heap) surfaces the *worst* + // candidate via `peek()`, enabling top-`count` retention by evicting it. + // The element tie-break is intentionally NOT reversed; it must match the + // final `sort_by` ordering used when draining the heap. fn cmp(&self, other: &Self) -> Ordering { other .score diff --git a/src/storage/tests/redis_vector_test.rs b/src/storage/tests/redis_vector_test.rs index be7fbdf8..403438df 100644 --- a/src/storage/tests/redis_vector_test.rs +++ b/src/storage/tests/redis_vector_test.rs @@ -23,7 +23,8 @@ use kstd::lock_mgr::LockMgr; use rocksdb::{IteratorMode, ReadOptions}; use storage::{ BaseMetaKey, BgTaskHandler, CanonicalVector, ColumnFamilyIndex, Redis, StorageOptions, - VectorQuery, VectorSearchMode, VectorSearchOptions, safe_cleanup_test_db, unique_test_db_path, + VectorQuery, VectorSearchMode, VectorSearchOptions, format_vector::VectorMeta, + safe_cleanup_test_db, unique_test_db_path, }; use storage::{slot_indexer::key_to_slot_id, storage::Storage}; @@ -455,15 +456,16 @@ fn test_vadd_rebuilds_expired_vectorset_with_newer_generation() { .get_cf_handle(ColumnFamilyIndex::MetaCF) .expect("MetaCF exists"); let meta_key = BaseMetaKey::new(key).encode().expect("meta key"); - let mut meta = db - .get_cf(&meta_cf, &meta_key) - .expect("read vector meta") - .expect("vector meta exists"); + let mut meta = VectorMeta::decode( + &db.get_cf(&meta_cf, &meta_key) + .expect("read vector meta") + .expect("vector meta exists"), + ) + .expect("decode vector meta"); let previous_generation = u64::MAX - 1; - meta[9..17].copy_from_slice(&previous_generation.to_le_bytes()); - let etime_offset = meta.len() - size_of::(); - meta[etime_offset..].copy_from_slice(&1_u64.to_le_bytes()); - db.put_cf(&meta_cf, &meta_key, &meta) + meta.set_version(previous_generation); + meta.set_etime(1); + db.put_cf(&meta_cf, &meta_key, &meta.encode()) .expect("store expired vector meta"); assert_eq!(redis.vcard(key).expect("expired card"), 0); @@ -476,13 +478,13 @@ fn test_vadd_rebuilds_expired_vectorset_with_newer_generation() { assert_eq!(redis.vcard(key).expect("rebuilt card"), 1); assert!(!redis.vismember(key, b"old").expect("old membership")); assert!(redis.vismember(key, b"new").expect("new membership")); - let rebuilt_meta = db - .get_cf(&meta_cf, &meta_key) - .expect("read rebuilt vector meta") - .expect("rebuilt vector meta exists"); - let rebuilt_generation = - u64::from_le_bytes(rebuilt_meta[9..17].try_into().expect("generation bytes")); - assert!(rebuilt_generation > previous_generation); + let rebuilt_meta = VectorMeta::decode( + &db.get_cf(&meta_cf, &meta_key) + .expect("read rebuilt vector meta") + .expect("rebuilt vector meta exists"), + ) + .expect("decode rebuilt vector meta"); + assert!(rebuilt_meta.version() > previous_generation); }); } From ea4eb2aaf7c4fd2190310534bc9898dd86b6d09b Mon Sep 17 00:00:00 2001 From: happy-v587 Date: Mon, 27 Jul 2026 11:15:49 +0800 Subject: [PATCH 18/32] refactor(resp): inline encode_resp_data_inner into encode_resp_data - merge the private encode_resp_data_inner into the RespEncode trait method; version handling stays inline at the encoding points - translate the vector-set plan doc to Chinese and update Step 3 to match the merged encoder design - move vector value layout comments next to their struct definitions --- .../plans/2026-07-19-redis-vector-set.md | 683 +++++++++--------- src/resp/src/encode.rs | 82 +-- src/storage/src/format_vector.rs | 30 +- 3 files changed, 399 insertions(+), 396 deletions(-) diff --git a/docs/superpowers/plans/2026-07-19-redis-vector-set.md b/docs/superpowers/plans/2026-07-19-redis-vector-set.md index e9cd4077..f99cd141 100644 --- a/docs/superpowers/plans/2026-07-19-redis-vector-set.md +++ b/docs/superpowers/plans/2026-07-19-redis-vector-set.md @@ -1,34 +1,34 @@ -# Redis Vector Set Implementation Plan +# Redis Vector Set 实施计划 -> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. +> **致智能体执行者:** 必备子技能:使用 superpowers:subagent-driven-development(推荐)或 superpowers:executing-plans 按任务逐步实现本计划。各步骤使用复选框(`- [ ]`)语法进行跟踪。 -**Goal:** 在 Kiwi 中实现 Redis 8 Vector Set 的 standalone Phase 1：持久化 FP32 向量，提供精确 FLAT 相似度查询，并支持 `VADD`、`VSIM`、`VREM`、`VCARD`、`VDIM`、`VEMB`、`VISMEMBER`。 +**目标:** 在 Kiwi 中实现 Redis 8 Vector Set 的 standalone Phase 1:持久化 FP32 向量,提供精确 FLAT 相似度查询,并支持 `VADD`、`VSIM`、`VREM`、`VCARD`、`VDIM`、`VEMB`、`VISMEMBER`。 -**Architecture:** 沿用 Kiwi 现有 Hash/Set/ZSet 的复合类型模型。一个 VectorSet key 只按 user key 路由到一个 RocksDB instance；`MetaCF` 保存类型、数量、生命周期 version 和维度，新增 `VectorDataCF` 保存 `key + version + element` 对应的 canonical vector。写命令通过同一个 RocksDB `WriteBatch` 原子更新 Meta 和成员；`VSIM` 在该 instance 内按 generation prefix 扫描并使用有界 Top-K heap 返回精确结果。 +**架构:** 沿用 Kiwi 现有 Hash/Set/ZSet 的复合类型模型。一个 VectorSet key 只按 user key 路由到一个 RocksDB instance;`MetaCF` 保存类型、数量、生命周期 version 和维度,新增 `VectorDataCF` 保存 `key + version + element` 对应的 canonical vector。写命令通过同一个 RocksDB `WriteBatch` 原子更新 Meta 和成员;`VSIM` 在该 instance 内按 generation prefix 扫描并使用有界 Top-K heap 返回精确结果。 -**Tech Stack:** Rust 2021、Tokio 双 runtime、RocksDB Column Family / WriteBatch / Snapshot、RESP2/RESP3、Python redis-py 集成测试。 +**技术栈:** Rust 2021、Tokio 双 runtime、RocksDB Column Family / WriteBatch / Snapshot、RESP2/RESP3、Python redis-py 集成测试。 -## Global Constraints +## 全局约束 -- 基线为 `origin/main` 的 `cdada8b`；忽略尚未合入的 error-catalog/error-model 分支，错误处理遵循当前主分支模式。 -- 首版支持 standalone 与 Raft Group cluster 模式；暂不支持 Redis Cluster 的槽位路由。 -- 首版只支持 cosine、canonical FP32 little-endian 和显式 `NOQUANT`；默认 Q8、显式 `Q8`、`BIN`、VEMB `RAW` 均返回明确的 unsupported 错误。 -- 首版只实现 FLAT；`TRUTH` 与普通 `VSIM` 都走同一个精确引擎，但保留不同的搜索模式枚举。 +- 基线为 `origin/main` 的 `cdada8b`;忽略尚未合入的 error-catalog/error-model 分支,错误处理遵循当前主分支模式。 +- 首版支持 standalone 与 Raft Group cluster 模式;暂不支持 Redis Cluster 的槽位路由。 +- 首版只支持 cosine、canonical FP32 little-endian 和显式 `NOQUANT`;默认 Q8、显式 `Q8`、`BIN`、VEMB `RAW` 均返回明确的 unsupported 错误。 +- 首版只实现 FLAT;`TRUTH` 与普通 `VSIM` 都走同一个精确引擎,但保留不同的搜索模式枚举。 - 首版不实现 `VINFO`、`INFO VECTOR`、HNSW、FAISS/IVF、`storage_incarnation`、O(1) `DEL`、Raft logical mutation、ReadIndex、snapshot barrier 和滚动升级门禁。 -- 复用 `MemberDataKey` 的 `key + version + element` 布局，不引入第二套 Vector member key codec。 -- Meta 保持现有统一 envelope；reserve 的前 8 字节固定为 `format(1) + encoding(1) + metric(1) + flags(1) + dimension(4 LE)`，其余 8 字节写零。 -- 所有新增 `.rs` 文件复制现有 Apache 2.0 license header；生产代码不得使用 `unwrap()`。 -- 不修改已有未跟踪的 `tests/python/test_vector_basic.py`；该文件验证的是旧 FT/HASH 原型，不属于 Vector Set。 -- 每个任务先写失败测试，再做最小实现；最终执行 `make fmt && make lint && make build && make test`。 +- 复用 `MemberDataKey` 的 `key + version + element` 布局,不引入第二套 Vector member key codec。 +- Meta 保持现有统一 envelope;reserve 的前 8 字节固定为 `format(1) + encoding(1) + metric(1) + flags(1) + dimension(4 LE)`,其余 8 字节写零。 +- 所有新增 `.rs` 文件复制现有 Apache 2.0 license header;生产代码不得使用 `unwrap()`。 +- 不修改已有未跟踪的 `tests/python/test_vector_basic.py`;该文件验证的是旧 FT/HASH 原型,不属于 Vector Set。 +- 每个任务先写失败测试,再做最小实现;最终执行 `make fmt && make lint && make build && make test`。 --- -## File Map +## 文件地图 | 文件 | 职责 | |---|---| | `src/storage/src/format_base_value.rs` | 追加 `DataType::VectorSet = 7`、字符串和 tag 映射 | -| `src/storage/src/redis.rs` | 追加 `VectorDataCF = 6`，创建 CF、配置 compaction filter、加入 batch handles | +| `src/storage/src/redis.rs` | 追加 `VectorDataCF = 6`,创建 CF、配置 compaction filter、加入 batch handles | | `src/storage/src/vector.rs` | 定义 canonical vector、查询参数、命中结果和数值计算 | | `src/storage/src/format_vector.rs` | 编解码 VectorSet meta 和 VectorDataCF value | | `src/storage/src/redis_vectors.rs` | 单个 RocksDB instance 上的 VADD/VREM/point read/FLAT scan | @@ -37,7 +37,7 @@ | `src/storage/src/meta_compaction_filter.rs` | 让过期或空 VectorSet meta 复用现有清理规则 | | `src/storage/src/redis_strings.rs` | 将 VectorDataCF 纳入 DEL 与 FLUSHDB 的物理清理 | | `src/storage/src/batch.rs` | 将 CF index 6 映射到 VectorDataCF | -| `src/storage/src/storage.rs` | Raft apply 的 CF index 6 解码；即使首版拒绝集群也保持 schema 完整 | +| `src/storage/src/storage.rs` | Raft apply 的 CF index 6 解码;即使首版拒绝集群也保持 schema 完整 | | `src/conf/src/raft_type.rs` | 追加跨 crate 使用的 VectorDataCF index | | `src/storage/src/logindex/types.rs` | 追加 VectorDataCF 的 log-index 元数据 | | `src/raft/src/lib.rs` | 追加 CF 名称和一致性断言 | @@ -51,35 +51,35 @@ | `src/cmd/src/vector/vismember.rs` | `VIsMemberCmd` 实现 | | `src/cmd/src/lib.rs` | 导出 vector 命令模块 | | `src/cmd/src/table.rs` | 注册七个命令 | -| `src/resp/src/encode.rs` | RESP2 下递归降级 Map/Double，RESP3 保持原生类型 | +| `src/resp/src/encode.rs` | RESP2 下递归降级 Map/Double,RESP3 保持原生类型 | | `src/storage/tests/redis_vector_test.rs` | storage 行为、生命周期、多 instance 和 FLAT 排序测试 | | `tests/python/test_vector_set_commands.py` | 真实 server 的 RESP2/RESP3 命令兼容测试 | --- -### Task 1: Add the VectorSet data type and VectorDataCF plumbing +### 任务 1:添加 VectorSet 数据类型和 VectorDataCF 管线 -**Files:** +**文件:** -- Modify: `src/storage/src/format_base_value.rs` -- Modify: `src/storage/src/redis.rs` -- Modify: `src/storage/src/batch.rs` -- Modify: `src/storage/src/storage.rs` -- Modify: `src/conf/src/raft_type.rs` -- Modify: `src/storage/src/logindex/types.rs` -- Modify: `src/raft/src/lib.rs` -- Modify: `src/storage/tests/redis_basic_test.rs` +- 修改:`src/storage/src/format_base_value.rs` +- 修改:`src/storage/src/redis.rs` +- 修改:`src/storage/src/batch.rs` +- 修改:`src/storage/src/storage.rs` +- 修改:`src/conf/src/raft_type.rs` +- 修改:`src/storage/src/logindex/types.rs` +- 修改:`src/raft/src/lib.rs` +- 修改:`src/storage/tests/redis_basic_test.rs` -**Interfaces:** +**接口:** -- Produces: `DataType::VectorSet = 7` -- Produces: `ColumnFamilyIndex::VectorDataCF = 6` -- Produces: RocksDB CF name `vector_data_cf` -- Preserves: all existing DataType and CF numeric values +- 产出:`DataType::VectorSet = 7` +- 产出:`ColumnFamilyIndex::VectorDataCF = 6` +- 产出:RocksDB CF 名称 `vector_data_cf` +- 保持:所有既有 DataType 和 CF 数值不变 -- [x] **Step 1: Extend the CF expectations in the existing storage test** +- [x] **步骤 1:扩展现有 storage 测试中的 CF 预期** -Update `test_open_redis` and `test_column_family_index` in `src/storage/tests/redis_basic_test.rs` to require seven handles and the exact final mapping: +更新 `src/storage/tests/redis_basic_test.rs` 中的 `test_open_redis` 和 `test_column_family_index`,要求七个 handle 以及如下的精确最终映射: ```rust assert_eq!(redis.handles.len(), 7); @@ -98,19 +98,19 @@ assert_eq!(ColumnFamilyIndex::VectorDataCF as usize, 6); assert_eq!(ColumnFamilyIndex::COUNT, 7); ``` -- [x] **Step 2: Run the focused test and confirm the expected failure** +- [x] **步骤 2:运行聚焦测试并确认预期失败** -Run: +运行: ```bash cargo test -p storage --test redis_basic_test test_open_redis -- --exact ``` -Expected: compile failure because `ColumnFamilyIndex::VectorDataCF` does not exist, or assertion failure because only six CFs are opened. +预期:编译失败,因为 `ColumnFamilyIndex::VectorDataCF` 尚不存在;或者断言失败,因为只打开了六个 CF。 -- [x] **Step 3: Append DataType without renumbering existing values** +- [x] **步骤 3:追加 DataType,不重排既有数值** -Apply this exact public shape in `src/storage/src/format_base_value.rs`: +在 `src/storage/src/format_base_value.rs` 中应用如下精确的公开形态: ```rust #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -131,7 +131,7 @@ pub const DATA_TYPE_STRINGS: [&str; 8] = [ pub const DATA_TYPE_TAG: [char; 8] = ['k', 'h', 's', 'l', 'z', 'n', 'a', 'v']; ``` -Update both exhaustive matches: +更新两处穷尽式 match: ```rust DataType::Hash | DataType::Set | DataType::ZSet | DataType::VectorSet => { @@ -143,9 +143,9 @@ DataType::Hash | DataType::Set | DataType::ZSet | DataType::VectorSet => { 7 => Ok(DataType::VectorSet), ``` -- [x] **Step 4: Append VectorDataCF to every schema mapping** +- [x] **步骤 4:将 VectorDataCF 追加到每一处 schema 映射** -Use the following invariant in all listed files: +在所有列出的文件中使用以下不变量: ```rust VectorDataCF = 6 @@ -154,20 +154,20 @@ VectorDataCF.name() = "vector_data_cf" VectorDataCF.data_type() = Some(DataType::VectorSet) ``` -Update these exact mapping sites: +更新以下精确的映射位置: -- `src/storage/src/redis.rs`: enum, `COUNT`, `name`, `data_type`, `CF_CONFIGS`, compaction-filter CF list, and `create_rocks_batch` handle list. -- `src/storage/src/batch.rs`: `cf_index_to_usize(ColumnFamilyIndex::VectorDataCF) => 6`. -- `src/storage/src/storage.rs`: `entry.cf_idx == 6` maps to `ColumnFamilyIndex::VectorDataCF`. -- `src/conf/src/raft_type.rs`: append variant and `from_u32(6)`. -- `src/storage/src/logindex/types.rs`: count becomes 7 and both CF-name arrays append `vector_data_cf`. -- `src/raft/src/lib.rs`: append `vector_data_cf` to `CF_NAMES` and `test_cf_names_match_storage` variants. +- `src/storage/src/redis.rs`:enum、`COUNT`、`name`、`data_type`、`CF_CONFIGS`、compaction-filter CF 列表,以及 `create_rocks_batch` 的 handle 列表。 +- `src/storage/src/batch.rs`:`cf_index_to_usize(ColumnFamilyIndex::VectorDataCF) => 6`。 +- `src/storage/src/storage.rs`:`entry.cf_idx == 6` 映射到 `ColumnFamilyIndex::VectorDataCF`。 +- `src/conf/src/raft_type.rs`:追加变体和 `from_u32(6)`。 +- `src/storage/src/logindex/types.rs`:count 变为 7,两个 CF 名称数组都追加 `vector_data_cf`。 +- `src/raft/src/lib.rs`:在 `CF_NAMES` 和 `test_cf_names_match_storage` 的变体中追加 `vector_data_cf`。 -Do not insert the new value in the middle of any enum or array. +不要将新值插入任何 enum 或数组的中间位置。 -- [x] **Step 5: Run the schema tests** +- [x] **步骤 5:运行 schema 测试** -Run: +运行: ```bash cargo test -p storage --test redis_basic_test test_open_redis @@ -175,9 +175,9 @@ cargo test -p storage --test redis_basic_test test_column_family_index cargo test -p raft test_cf_names_match_storage ``` -Expected: all tests PASS and a newly opened database reports seven CF handles in index order. +预期:所有测试通过,新打开的数据库按索引顺序报告七个 CF handle。 -- [x] **Step 6: Commit the schema slice** +- [x] **步骤 6:提交 schema 切片** ```bash git add src/storage/src/format_base_value.rs src/storage/src/redis.rs src/storage/src/batch.rs src/storage/src/storage.rs src/conf/src/raft_type.rs src/storage/src/logindex/types.rs src/raft/src/lib.rs src/storage/tests/redis_basic_test.rs @@ -186,27 +186,27 @@ git commit -m "feat(storage): add vector data column family" --- -### Task 2: Implement canonical vector, meta, and value codecs +### 任务 2:实现 canonical vector、meta 和 value 编解码器 -**Files:** +**文件:** -- Create: `src/storage/src/vector.rs` -- Create: `src/storage/src/format_vector.rs` -- Modify: `src/storage/src/lib.rs` +- 新建:`src/storage/src/vector.rs` +- 新建:`src/storage/src/format_vector.rs` +- 修改:`src/storage/src/lib.rs` -**Interfaces:** +**接口:** -- Produces: `CanonicalVector::from_fp32_le(&[u8]) -> Result` -- Produces: `CanonicalVector::from_values(&[f32]) -> Result` -- Produces: `CanonicalVector::score(&CanonicalVector) -> Result` -- Produces: `CanonicalVector::restore() -> Vec` -- Produces: `VectorMeta::{new, encode, decode, count, set_count, version, dimension, is_stale}` -- Produces: `VectorDataValue::{from_canonical, encode, decode}` -- Produces: `VectorQuery`, `VectorSearchMode`, `VectorSearchOptions`, `VectorHit` +- 产出:`CanonicalVector::from_fp32_le(&[u8]) -> Result` +- 产出:`CanonicalVector::from_values(&[f32]) -> Result` +- 产出:`CanonicalVector::score(&CanonicalVector) -> Result` +- 产出:`CanonicalVector::restore() -> Vec` +- 产出:`VectorMeta::{new, encode, decode, count, set_count, version, dimension, is_stale}` +- 产出:`VectorDataValue::{from_canonical, encode, decode}` +- 产出:`VectorQuery`、`VectorSearchMode`、`VectorSearchOptions`、`VectorHit` -- [x] **Step 1: Write codec and numerical unit tests before exporting modules** +- [x] **步骤 1:在导出模块之前先写编解码和数值单元测试** -Place unit tests in the two new files. Cover these exact cases: +将单元测试放在两个新文件中。覆盖以下精确的用例: ```rust #[test] @@ -237,22 +237,22 @@ fn cosine_score_maps_to_redis_range() { } ``` -Codec tests must round-trip an empty binary element independently through `MemberDataKey`, a two-dimensional vector value, and a `VectorMeta` with `count=2`, `version=42`, `dimension=2`. +编解码测试必须独立地对空二进制 element 通过 `MemberDataKey` 往返、对二维 vector value 往返,以及对 `count=2`、`version=42`、`dimension=2` 的 `VectorMeta` 往返。 -- [x] **Step 2: Run tests and confirm the modules are missing** +- [x] **步骤 2:运行测试并确认模块尚不存在** -Run: +运行: ```bash cargo test -p storage vector::tests cargo test -p storage format_vector::tests ``` -Expected: compile failure until the modules and types are implemented and exported. +预期:在模块和类型实现并导出之前编译失败。 -- [x] **Step 3: Implement the public vector types** +- [x] **步骤 3:实现公开的 vector 类型** -Define the following exact public API in `src/storage/src/vector.rs`: +在 `src/storage/src/vector.rs` 中定义如下精确的公开 API: ```rust #[derive(Debug, Clone, PartialEq)] @@ -287,25 +287,25 @@ pub struct VectorHit { } ``` -Numerical rules: +数值规则: -- FP32 blob length must be non-zero and divisible by four. -- Parse every component with `f32::from_le_bytes`; reject non-finite values. -- Accumulate `norm²` in `f64`; reject zero or non-finite norm. -- Store normalized components as `f32`, store original L2 as `f32`. -- Compute dot product in `f64`, clamp it to `[-1, 1]`, then return `(dot + 1.0) / 2.0` clamped to `[0, 1]`. -- Reject score comparisons when dimensions differ. +- FP32 blob 长度必须非零且能被四整除。 +- 每个分量用 `f32::from_le_bytes` 解析;拒绝非有限值。 +- 用 `f64` 累加 `norm²`;拒绝零或非有限的 norm。 +- 归一化分量按 `f32` 存储,原始 L2 按 `f32` 存储。 +- 用 `f64` 计算点积,将其钳制到 `[-1, 1]`,然后返回 `(dot + 1.0) / 2.0` 并钳制到 `[0, 1]`。 +- 维度不同时拒绝分数比较。 -- [x] **Step 4: Implement fixed codecs** +- [x] **步骤 4:实现固定编解码器** -Use this exact metadata layout in `src/storage/src/format_vector.rs`: +在 `src/storage/src/format_vector.rs` 中使用如下精确的 metadata 布局: ```text | type=7 | count | version | format | encoding | metric | flags | dimension | zero reserve | ctime | etime | | 1B | 8B | 8B | 1B | 1B | 1B | 1B | 4B LE | 8B | 8B | 8B | ``` -Constants: +常量: ```rust pub const VECTOR_META_FORMAT: u8 = 1; @@ -315,18 +315,18 @@ pub const VECTOR_VALUE_MAGIC: u8 = 0x56; pub const VECTOR_VALUE_FORMAT: u8 = 1; ``` -Use this value layout: +使用如下 value 布局: ```text | magic=0x56 | format=1 | dimension | original_l2 | normalized FP32 payload | | 1B | 1B | 4B LE | 4B LE | dimension * 4B | ``` -Decoders must check exact length, magic, format, non-zero dimension, finite positive `original_l2`, finite payload components, and payload dimension. Return `InvalidFormatSnafu` on malformed persisted bytes; never panic. +解码器必须检查精确长度、magic、format、非零 dimension、有限且为正的 `original_l2`、有限的 payload 分量,以及 payload 维度。对格式错误的持久化字节返回 `InvalidFormatSnafu`;永不 panic。 -- [x] **Step 5: Export the new modules and rerun tests** +- [x] **步骤 5:导出新模块并重新运行测试** -Add to `src/storage/src/lib.rs`: +在 `src/storage/src/lib.rs` 中添加: ```rust mod format_vector; @@ -337,16 +337,16 @@ pub use vector::{ }; ``` -Run: +运行: ```bash cargo test -p storage vector::tests cargo test -p storage format_vector::tests ``` -Expected: all canonicalization and codec tests PASS. +预期:所有 canonicalization 和编解码测试通过。 -- [x] **Step 6: Commit the codec slice** +- [x] **步骤 6:提交编解码器切片** ```bash git add src/storage/src/lib.rs src/storage/src/vector.rs src/storage/src/format_vector.rs @@ -355,28 +355,28 @@ git commit -m "feat(storage): add vector codecs" --- -### Task 3: Implement atomic VectorSet mutations and point reads +### 任务 3:实现原子的 VectorSet 变更和点读 -**Files:** +**文件:** -- Create: `src/storage/src/redis_vectors.rs` -- Modify: `src/storage/src/lib.rs` -- Create: `src/storage/tests/redis_vector_test.rs` +- 新建:`src/storage/src/redis_vectors.rs` +- 修改:`src/storage/src/lib.rs` +- 新建:`src/storage/tests/redis_vector_test.rs` -**Interfaces:** +**接口:** -- Consumes: `CanonicalVector`, `VectorMeta`, `VectorDataValue`, `MemberDataKey` -- Produces: `Redis::vadd(key, element, vector) -> Result` -- Produces: `Redis::vrem(key, element) -> Result` -- Produces: `Redis::vcard(key) -> Result` -- Produces: `Redis::vdim(key) -> Result` -- Produces: `Redis::vemb(key, element) -> Result>>` -- Produces: `Redis::vismember(key, element) -> Result` -- Produces: `Redis::is_cluster_mode() -> bool` +- 消费:`CanonicalVector`、`VectorMeta`、`VectorDataValue`、`MemberDataKey` +- 产出:`Redis::vadd(key, element, vector) -> Result` +- 产出:`Redis::vrem(key, element) -> Result` +- 产出:`Redis::vcard(key) -> Result` +- 产出:`Redis::vdim(key) -> Result` +- 产出:`Redis::vemb(key, element) -> Result>>` +- 产出:`Redis::vismember(key, element) -> Result` +- 产出:`Redis::is_cluster_mode() -> bool` -- [x] **Step 1: Write storage behavior tests** +- [x] **步骤 1:编写 storage 行为测试** -Create the test file using `unique_test_db_path`, `safe_cleanup_test_db`, `StorageOptions`, `BgTaskHandler`, and `LockMgr`, matching existing storage test setup. Add these independently named tests: +使用 `unique_test_db_path`、`safe_cleanup_test_db`、`StorageOptions`、`BgTaskHandler` 和 `LockMgr` 创建测试文件,与现有 storage 测试的搭建方式保持一致。添加以下独立命名的测试: - `test_vadd_create_update_and_dimension_guard` - `test_vadd_is_binary_safe_and_accepts_empty_element` @@ -385,7 +385,7 @@ Create the test file using `unique_test_db_path`, `safe_cleanup_test_db`, `Stora - `test_vector_commands_return_wrongtype_for_string_key` - `test_vector_meta_and_member_are_committed_together` -The core assertions are: +核心断言如下: ```rust let a = CanonicalVector::from_values(&[1.0, 0.0]).expect("valid vector"); @@ -402,39 +402,39 @@ assert!(redis.vadd(b"vectors", b"\x00binary", &a).expect("binary member")); assert!(redis.vadd(b"empty-element", b"", &a).expect("empty member")); ``` -For dimension mismatch, assert `is_err()` and confirm that `VCARD` and the old member value remain unchanged. For last-member removal, assert `get_key_type(key)` returns a missing-key error after `VREM`. +对于维度不匹配,断言 `is_err()` 并确认 `VCARD` 和旧成员值保持不变。对于最后一个成员的删除,断言 `VREM` 之后 `get_key_type(key)` 返回 key 不存在的错误。 -- [x] **Step 2: Run the behavior tests and confirm missing methods** +- [x] **步骤 2:运行行为测试并确认方法缺失** -Run: +运行: ```bash cargo test -p storage --test redis_vector_test ``` -Expected: compile failure because the six `Redis` methods are not implemented. +预期:编译失败,因为六个 `Redis` 方法尚未实现。 -- [x] **Step 3: Implement one locked read-modify-write path for VADD** +- [x] **步骤 3:为 VADD 实现一条加锁的读-改-写路径** -Add this private module declaration to `src/storage/src/lib.rs`: +在 `src/storage/src/lib.rs` 中添加这个私有模块声明: ```rust mod redis_vectors; ``` -Implementation sequence inside `Redis::vadd`: +`Redis::vadd` 内部的实现顺序: -1. Acquire `ScopeRecordLock` using the user key. -3. Read `MetaCF[BaseMetaKey::new(key)]`. -4. Missing or stale meta creates `VectorMeta::new_after(1, dimension, previous_generation)` whose version is the current timestamp clamped above the previous generation (monotonic, so a recreated VectorSet can never address stale `VectorDataCF` rows), and does not increment it again. -5. Live non-VectorSet meta returns the same WRONGTYPE text used by `check_type_state`. -6. Live VectorSet requires equal dimension. -7. Point-read `VectorDataCF[MemberDataKey(key, version, element)]` to distinguish insert from update. -8. Inserting into an existing live VectorSet increments count; creating a new VectorSet keeps the initial count at one; updating an existing element preserves count. -9. Build one batch containing the member put and meta put; commit once. -10. Return `true` for insert and `false` for update. +1. 使用 user key 获取 `ScopeRecordLock`。 +3. 读取 `MetaCF[BaseMetaKey::new(key)]`。 +4. meta 缺失或过期时创建 `VectorMeta::new_after(1, dimension, previous_generation)`,其 version 为钳制在前一代 generation 之上的当前时间戳(单调递增,因此重建的 VectorSet 永远不可能寻址到过期的 `VectorDataCF` 行),且不再二次递增。 +5. 存活的非 VectorSet meta 返回与 `check_type_state` 相同的 WRONGTYPE 文本。 +6. 存活的 VectorSet 要求维度相等。 +7. 点读 `VectorDataCF[MemberDataKey(key, version, element)]` 以区分插入与更新。 +8. 向存活的既有 VectorSet 插入时递增 count;创建新 VectorSet 时保持初始 count 为一;更新既有 element 时保持 count 不变。 +9. 构建一个同时包含 member put 和 meta put 的 batch;提交一次。 +10. 插入返回 `true`,更新返回 `false`。 -The public signature must be: +公开签名必须是: ```rust pub fn vadd( @@ -445,28 +445,28 @@ pub fn vadd( ) -> Result; ``` -- [x] **Step 4: Implement VREM and point-read commands** +- [x] **步骤 4:实现 VREM 和点读命令** -Use `VectorMeta::decode` and the same `MemberDataKey` generation for every operation. +每个操作都使用 `VectorMeta::decode` 和相同的 `MemberDataKey` generation。 -VREM rules: +VREM 规则: -- Missing/stale key or missing element returns `Ok(false)`. -- Wrong type returns WRONGTYPE. -- If count is greater than one, delete member and write `count - 1` meta in one batch. -- If count is one, delete member and delete MetaCF in one batch. +- key 缺失/过期或 element 缺失时返回 `Ok(false)`。 +- 类型错误时返回 WRONGTYPE。 +- count 大于一时,在一个 batch 中删除 member 并写入 `count - 1` 的 meta。 +- count 等于一时,在一个 batch 中删除 member 并删除 MetaCF。 -Read rules: +读取规则: -- `vcard`: missing/stale returns `0`. -- `vdim`: missing/stale returns `KeyNotFoundSnafu`. -- `vemb`: missing/stale member or key returns `None`; live value returns `CanonicalVector::restore()`. -- `vismember`: missing/stale member or key returns `false`. -- All live wrong-type keys return WRONGTYPE. +- `vcard`:缺失/过期返回 `0`。 +- `vdim`:缺失/过期返回 `KeyNotFoundSnafu`。 +- `vemb`:member 或 key 缺失/过期返回 `None`;存活值返回 `CanonicalVector::restore()`。 +- `vismember`:member 或 key 缺失/过期返回 `false`。 +- 所有存活的类型错误 key 都返回 WRONGTYPE。 -- [x] **Step 5: Run storage tests** +- [x] **步骤 5:运行 storage 测试** -Run: +运行: ```bash RUST_TEST_THREADS=1 cargo test -p storage --test redis_vector_test @@ -474,9 +474,9 @@ cargo test -p storage vector::tests cargo test -p storage format_vector::tests ``` -Expected: all tests PASS, including rollback-visible behavior after dimension mismatch. +预期:所有测试通过,包括维度不匹配后可见的回滚行为。 -- [x] **Step 6: Commit point operations** +- [x] **步骤 6:提交点操作** ```bash git add src/storage/src/lib.rs src/storage/src/redis_vectors.rs src/storage/tests/redis_vector_test.rs @@ -485,24 +485,24 @@ git commit -m "feat(storage): persist vector set members" --- -### Task 4: Implement exact FLAT search with stable Top-K ordering +### 任务 4:实现带稳定 Top-K 排序的精确 FLAT 搜索 -**Files:** +**文件:** -- Modify: `src/storage/src/redis_vectors.rs` -- Modify: `src/storage/src/vector.rs` -- Modify: `src/storage/tests/redis_vector_test.rs` +- 修改:`src/storage/src/redis_vectors.rs` +- 修改:`src/storage/src/vector.rs` +- 修改:`src/storage/tests/redis_vector_test.rs` -**Interfaces:** +**接口:** -- Consumes: `VectorQuery`, `VectorSearchOptions` -- Produces: `Redis::vsim(key, query, options) -> Result>` -- Ordering contract: score descending, then raw element bytes ascending -- Complexity contract: `O(N × DIM + N log K)` time and `O(K + DIM)` extra memory +- 消费:`VectorQuery`、`VectorSearchOptions` +- 产出:`Redis::vsim(key, query, options) -> Result>` +- 排序约定:分数降序,然后按原始 element 字节升序 +- 复杂度约定:`O(N × DIM + N log K)` 时间,`O(K + DIM)` 额外内存 -- [x] **Step 1: Add failing FLAT search tests** +- [x] **步骤 1:添加失败的 FLAT 搜索测试** -Add these tests: +添加以下测试: - `test_vsim_direct_vector_returns_exact_top_k` - `test_vsim_ele_uses_stored_member_as_query` @@ -511,7 +511,7 @@ Add these tests: - `test_vsim_missing_key_is_empty_and_missing_ele_is_error` - `test_vsim_rejects_query_dimension_mismatch` -Use this deterministic dataset: +使用这个确定性数据集: ```rust let x = CanonicalVector::from_values(&[1.0, 0.0]).expect("x"); @@ -523,21 +523,21 @@ redis.vadd(b"search", b"x", &x).expect("insert x"); redis.vadd(b"search", b"neg", &neg_x).expect("insert neg"); ``` -For query `x` and `COUNT 3`, require `x` first, then `a`, then `b`; `a` and `b` tie on score and raw bytes decide their order. +对于查询 `x` 和 `COUNT 3`,要求 `x` 在前,然后是 `a`,然后是 `b`;`a` 和 `b` 分数相同,由原始字节决定其顺序。 -- [x] **Step 2: Run the focused tests and confirm VSim is missing** +- [x] **步骤 2:运行聚焦测试并确认 VSim 缺失** ```bash cargo test -p storage --test redis_vector_test test_vsim ``` -Expected: compile failure because `Redis::vsim` is not defined. +预期:编译失败,因为 `Redis::vsim` 尚未定义。 -- [x] **Step 3: Add a heap item with total ordering** +- [x] **步骤 3:添加具有全序关系的堆元素** -Define a private `HeapHit` that compares worse hits as the heap maximum, so a heap capped at K can replace its worst element. Equality and ordering must use both `score.total_cmp()` and raw element bytes; do not use `partial_cmp().unwrap()`. +定义一个私有的 `HeapHit`,将更差的命中比较为堆顶最大值,这样容量为 K 的堆可以替换其最差元素。相等性和排序必须同时使用 `score.total_cmp()` 和原始 element 字节;不得使用 `partial_cmp().unwrap()`。 -Required result ordering after draining the heap: +排空堆之后要求的结果排序: ```rust hits.sort_by(|left, right| { @@ -548,34 +548,34 @@ hits.sort_by(|left, right| { }); ``` -- [x] **Step 4: Implement one RocksDB snapshot scan** +- [x] **步骤 4:实现一次 RocksDB 快照扫描** -The exact search sequence is: +精确的搜索顺序是: -1. Return an empty vector when the key is missing or stale. -2. Return WRONGTYPE for a live non-VectorSet key. -3. Decode meta and resolve an `Element` query with a point read from the same RocksDB snapshot. -4. Reject missing query element and dimension mismatch. -5. Build `MemberDataKey::new(key, version, b"").encode_seek_key()`. -6. Iterate `VectorDataCF` forward from the prefix using snapshot read options. -7. Stop at the first key not starting with the prefix. -8. Strip prefix and the 16-byte suffix reserve to recover raw element bytes. -9. Decode each vector value and compute score. -10. Keep at most `options.count` heap entries. -11. Drain and sort using the stable ordering above. +1. key 缺失或过期时返回空 vector。 +2. 存活的非 VectorSet key 返回 WRONGTYPE。 +3. 解码 meta,并从同一个 RocksDB 快照中通过点读解析 `Element` 查询。 +4. 拒绝缺失的查询 element 和维度不匹配。 +5. 构建 `MemberDataKey::new(key, version, b"").encode_seek_key()`。 +6. 使用快照读选项从该前缀开始向前迭代 `VectorDataCF`。 +7. 在第一个不以该前缀开头的 key 处停止。 +8. 剥掉前缀和 16 字节的后缀 reserve,恢复原始 element 字节。 +9. 解码每个 vector value 并计算分数。 +10. 堆中最多保留 `options.count` 个条目。 +11. 排空并使用上述稳定排序。 -Reject `options.count == 0` before creating the iterator. +在创建迭代器之前拒绝 `options.count == 0`。 -- [x] **Step 5: Run FLAT and regression tests** +- [x] **步骤 5:运行 FLAT 和回归测试** ```bash RUST_TEST_THREADS=1 cargo test -p storage --test redis_vector_test test_vsim RUST_TEST_THREADS=1 cargo test -p storage --test redis_set_test ``` -Expected: Vector search tests PASS and existing Set prefix-iteration tests remain green. +预期:Vector 搜索测试通过,且现有 Set 前缀迭代测试保持绿色。 -- [x] **Step 6: Commit the search engine** +- [x] **步骤 6:提交搜索引擎** ```bash git add src/storage/src/vector.rs src/storage/src/redis_vectors.rs src/storage/tests/redis_vector_test.rs @@ -584,27 +584,27 @@ git commit -m "feat(storage): add flat vector similarity search" --- -### Task 5: Wire Storage routing, TYPE, TTL cleanup, DEL, and FLUSHDB +### 任务 5:接通 Storage 路由、TYPE、TTL 清理、DEL 和 FLUSHDB -**Files:** +**文件:** -- Modify: `src/storage/src/storage_impl.rs` -- Modify: `src/storage/src/data_compaction_filter.rs` -- Modify: `src/storage/src/meta_compaction_filter.rs` -- Modify: `src/storage/src/redis.rs` -- Modify: `src/storage/src/redis_strings.rs` -- Modify: `src/storage/tests/redis_vector_test.rs` +- 修改:`src/storage/src/storage_impl.rs` +- 修改:`src/storage/src/data_compaction_filter.rs` +- 修改:`src/storage/src/meta_compaction_filter.rs` +- 修改:`src/storage/src/redis.rs` +- 修改:`src/storage/src/redis_strings.rs` +- 修改:`src/storage/tests/redis_vector_test.rs` -**Interfaces:** +**接口:** -- Produces: the same seven operations on `Storage`, each routed exactly once by user key -- Preserves: `TYPE key -> vectorset` -- Preserves: generic EXPIRE/PERSIST behavior through the common Meta envelope -- Preserves: generic DEL and FLUSHDB physically delete VectorDataCF entries +- 产出:`Storage` 上相同的七个操作,每个都按 user key 恰好路由一次 +- 保持:`TYPE key -> vectorset` +- 保持:通过通用 Meta envelope 的通用 EXPIRE/PERSIST 行为 +- 保持:通用 DEL 和 FLUSHDB 物理删除 VectorDataCF 条目 -- [x] **Step 1: Add facade and lifecycle tests** +- [x] **步骤 1:添加门面和生命周期测试** -Add tests using `Storage::new(3, 0)` for: +使用 `Storage::new(3, 0)` 添加以下测试: - `test_storage_routes_all_members_of_one_vectorset_to_one_instance` - `test_type_returns_vectorset` @@ -613,20 +613,20 @@ Add tests using `Storage::new(3, 0)` for: - `test_flushdb_removes_vector_meta_and_members` - `test_vector_storage_rejects_cluster_mode` -The multi-instance test must compute `key_to_slot_id(key)` and confirm only the selected instance contains the MetaCF key; it must not scan or fan out for VSIM. +多 instance 测试必须计算 `key_to_slot_id(key)`,并确认只有被选中的 instance 包含该 MetaCF key;不得为 VSIM 做扫描或扇出。 -- [x] **Step 2: Run tests and confirm missing Storage methods** +- [x] **步骤 2:运行测试并确认 Storage 方法缺失** ```bash cargo test -p storage --test redis_vector_test test_storage_routes cargo test -p storage --test redis_vector_test test_del_removes_vector ``` -Expected: compile failure because the Vector methods exist only on `Redis`. +预期:编译失败,因为 Vector 方法只存在于 `Redis` 上。 -- [x] **Step 3: Add exact one-key routing methods** +- [x] **步骤 3:添加精确的单 key 路由方法** -Add these signatures to `src/storage/src/storage_impl.rs`: +在 `src/storage/src/storage_impl.rs` 中添加以下签名: ```rust pub fn vadd(&self, key: &[u8], element: &[u8], vector: &CanonicalVector) -> Result; @@ -643,7 +643,7 @@ pub fn vemb(&self, key: &[u8], element: &[u8]) -> Result>>; pub fn vismember(&self, key: &[u8], element: &[u8]) -> Result; ``` -Every method must perform the same two routing statements and then call its matching Redis method. For example, `Storage::vadd` is: +每个方法都必须执行相同的两条路由语句,然后调用对应的 Redis 方法。例如,`Storage::vadd` 是: ```rust let slot_id = key_to_slot_id(key); @@ -651,22 +651,22 @@ let instance_id = self.slot_indexer.get_instance_id(slot_id); self.insts[instance_id].vadd(key, element, vector) ``` -Apply the identical prefix to `vsim`, `vrem`, `vcard`, `vdim`, `vemb`, and `vismember`, forwarding their declared arguments unchanged. There is no element routing and no instance fan-out. +对 `vsim`、`vrem`、`vcard`、`vdim`、`vemb` 和 `vismember` 应用相同的前缀,按声明的原样转发参数。没有 element 路由,也没有 instance 扇出。 -- [x] **Step 4: Extend common lifecycle handling** +- [x] **步骤 4:扩展通用生命周期处理** -Apply these exact additions: +应用以下精确的增补: -- `src/storage/src/redis.rs::is_stale_static`: include `DataType::VectorSet` in the common count/version meta branch. -- `src/storage/src/redis_strings.rs::set_key_etime`: parse `DataType::VectorSet` with `ParsedBaseMetaValue`, identical to Hash/Set/ZSet. -- `src/storage/src/meta_compaction_filter.rs`: parse `DataType::VectorSet` with `ParsedBaseMetaValue`. -- `src/storage/src/data_compaction_filter.rs::parse_meta_value` and its test-only meta parser: parse VectorSet with `ParsedBaseMetaValue`. -- `Redis::del_key`: add `ColumnFamilyIndex::VectorDataCF` to the composite data CF scan list. -- `Redis::flush_db`: add `ColumnFamilyIndex::VectorDataCF` to `all_cf_indexes`. +- `src/storage/src/redis.rs::is_stale_static`:在通用的 count/version meta 分支中包含 `DataType::VectorSet`。 +- `src/storage/src/redis_strings.rs::set_key_etime`:用 `ParsedBaseMetaValue` 解析 `DataType::VectorSet`,与 Hash/Set/ZSet 完全一致。 +- `src/storage/src/meta_compaction_filter.rs`:用 `ParsedBaseMetaValue` 解析 `DataType::VectorSet`。 +- `src/storage/src/data_compaction_filter.rs::parse_meta_value` 及其仅用于测试的 meta 解析器:用 `ParsedBaseMetaValue` 解析 VectorSet。 +- `Redis::del_key`:将 `ColumnFamilyIndex::VectorDataCF` 加入复合数据 CF 扫描列表。 +- `Redis::flush_db`:将 `ColumnFamilyIndex::VectorDataCF` 加入 `all_cf_indexes`。 -Do not introduce a wildcard in the schema matches above; each VectorSet/VectorDataCF branch must be explicit. +不要在上述 schema match 中引入通配符;每个 VectorSet/VectorDataCF 分支都必须显式写出。 -- [x] **Step 5: Run lifecycle and existing TTL tests** +- [x] **步骤 5:运行生命周期和现有 TTL 测试** ```bash RUST_TEST_THREADS=1 cargo test -p storage --test redis_vector_test @@ -674,9 +674,9 @@ RUST_TEST_THREADS=1 cargo test -p storage --test ttl_test RUST_TEST_THREADS=1 cargo test -p storage --test redis_basic_test ``` -Expected: all tests PASS; expired vector data becomes invisible immediately and eligible for compaction, while DEL/FLUSHDB remove physical VectorDataCF entries. +预期:所有测试通过;过期的 vector 数据立即不可见并具备被 compaction 清理的资格,而 DEL/FLUSHDB 会移除物理的 VectorDataCF 条目。 -- [x] **Step 6: Commit routing and lifecycle support** +- [x] **步骤 6:提交路由和生命周期支持** ```bash git add src/storage/src/storage_impl.rs src/storage/src/data_compaction_filter.rs src/storage/src/meta_compaction_filter.rs src/storage/src/redis.rs src/storage/src/redis_strings.rs src/storage/tests/redis_vector_test.rs @@ -685,23 +685,23 @@ git commit -m "feat(storage): route and clean up vector sets" --- -### Task 6: Add Vector Set command parsing and command-table registration +### 任务 6:添加 Vector Set 命令解析和命令表注册 -**Files:** +**文件:** -- Create: `src/cmd/src/vector.rs` -- Modify: `src/cmd/src/lib.rs` -- Modify: `src/cmd/src/table.rs` +- 新建:`src/cmd/src/vector.rs` +- 修改:`src/cmd/src/lib.rs` +- 修改:`src/cmd/src/table.rs` -**Interfaces:** +**接口:** -- Produces: `VAddCmd`, `VSimCmd`, `VRemCmd`, `VCardCmd`, `VDimCmd`, `VEmbCmd`, `VIsMemberCmd` -- Consumes: all seven `Storage` APIs from Task 5 -- Produces: binary-safe argv parsing and `RespData` replies +- 产出:`VAddCmd`、`VSimCmd`、`VRemCmd`、`VCardCmd`、`VDimCmd`、`VEmbCmd`、`VIsMemberCmd` +- 消费:任务 5 中的全部七个 `Storage` API +- 产出:二进制安全的 argv 解析和 `RespData` 回复 -- [x] **Step 1: Write parser and metadata unit tests in the new module** +- [x] **步骤 1:在新模块中编写解析器和元数据单元测试** -Cover all supported command shapes: +覆盖所有受支持的命令形态: ```text VADD key FP32 blob element NOQUANT @@ -716,9 +716,9 @@ VEMB key element VISMEMBER key element ``` -Add explicit failure tests for malformed FP32 length, invalid VALUES dimension/count, invalid float, missing `NOQUANT`, `Q8`, `BIN`, duplicate VSIM options, zero/invalid COUNT, `VEMB RAW`, and trailing unknown options. +为以下情况添加显式的失败测试:FP32 长度错误、VALUES 维度/数量非法、浮点数非法、缺少 `NOQUANT`、`Q8`、`BIN`、重复的 VSIM 选项、COUNT 为零/非法、`VEMB RAW`,以及结尾出现未知选项。 -Metadata assertions: +元数据断言: ```rust assert_eq!(VAddCmd::new().meta().arity, -5); @@ -730,19 +730,19 @@ assert_eq!(VEmbCmd::new().meta().arity, -3); assert_eq!(VIsMemberCmd::new().meta().arity, 3); ``` -- [x] **Step 2: Run command tests and confirm the module is not registered** +- [x] **步骤 2:运行命令测试并确认模块尚未注册** ```bash cargo test -p cmd vector::tests ``` -Expected: compile failure until the new module, command structs, and parsers exist. +预期:在新模块、命令结构体和解析器存在之前编译失败。 -- [x] **Step 3: Implement shared parsers without UTF-8 assumptions for keys/elements/blobs** +- [x] **步骤 3:实现共享解析器,对 key/element/blob 不做 UTF-8 假设** -Only option keywords and numeric VALUES tokens are decoded as ASCII/UTF-8. Keep key, element, and FP32 blob as raw bytes. +只有选项关键字和 VALUES 的数字 token 按 ASCII/UTF-8 解码。key、element 和 FP32 blob 保持为原始字节。 -Define private parsed forms: +定义私有的解析结果形式: ```rust struct ParsedVAdd { @@ -757,7 +757,7 @@ struct ParsedVSim { } ``` -Error strings must be exact and already include the Redis error class: +错误字符串必须精确匹配,且已包含 Redis 错误类别前缀: ```text ERR invalid vector specification @@ -771,7 +771,7 @@ ERR key does not exist WRONGTYPE Operation against a key holding the wrong kind of value ``` -Add one local reply helper that avoids double-prefixing current storage error strings: +添加一个本地回复辅助函数,避免对现有 storage 错误字符串重复加前缀: ```rust fn error_reply(message: impl Into) -> RespData { @@ -779,27 +779,27 @@ fn error_reply(message: impl Into) -> RespData { } ``` -Map internal non-Redis storage failures to `ERR storage error`; log the detailed error with `log::error!`. +将内部的非 Redis storage 失败映射为 `ERR storage error`;用 `log::error!` 记录详细错误。 -- [x] **Step 4: Implement command replies** +- [x] **步骤 4:实现命令回复** -Reply contracts: +回复约定: -- VADD/VREM/VISMEMBER: integer `1` or `0`. -- VCARD/VDIM: integer after checked `u64/u32 -> i64` conversion. -- VEMB: null bulk for missing key/member; otherwise array of `RespData::Double`. -- VSIM without scores: array of bulk-string elements. -- VSIM with scores: `RespData::Map` of bulk-string element to `RespData::Double`; Task 7 performs RESP2 downgrade at the encoder boundary. +- VADD/VREM/VISMEMBER:整数 `1` 或 `0`。 +- VCARD/VDIM:经检查的 `u64/u32 -> i64` 转换后的整数。 +- VEMB:key/member 缺失时返回 null bulk;否则返回 `RespData::Double` 数组。 +- 不带分数的 VSIM:bulk-string element 数组。 +- 带分数的 VSIM:bulk-string element 到 `RespData::Double` 的 `RespData::Map`;RESP2 降级由任务 7 在编码器边界执行。 -Flags and ACL categories: +标志和 ACL 类别: -- VADD/VREM: `WRITE | FAST`, `KEYSPACE | WRITE`. -- VSIM: `READONLY`, `KEYSPACE | READ | SLOW`. -- VCARD/VDIM/VEMB/VISMEMBER: `READONLY | FAST`, `KEYSPACE | READ`. +- VADD/VREM:`WRITE | FAST`,`KEYSPACE | WRITE`。 +- VSIM:`READONLY`,`KEYSPACE | READ | SLOW`。 +- VCARD/VDIM/VEMB/VISMEMBER:`READONLY | FAST`,`KEYSPACE | READ`。 -- [x] **Step 5: Export and register all seven commands** +- [x] **步骤 5:导出并注册全部七个命令** -Add `pub mod vector;` to `src/cmd/src/lib.rs`. Append these types to `register_cmd!` in `src/cmd/src/table.rs`: +在 `src/cmd/src/lib.rs` 中添加 `pub mod vector;`。在 `src/cmd/src/table.rs` 的 `register_cmd!` 中追加以下类型: ```rust crate::vector::VAddCmd, @@ -811,9 +811,9 @@ crate::vector::VEmbCmd, crate::vector::VIsMemberCmd, ``` -Add a table test asserting all lowercase names resolve. +添加一个表测试,断言所有小写名称都能解析。 -- [x] **Step 6: Run command and runtime dispatch tests** +- [x] **步骤 6:运行命令和 runtime 分发测试** ```bash cargo test -p cmd vector::tests @@ -821,9 +821,9 @@ cargo test -p cmd table cargo test -p runtime handle_execute_command ``` -Expected: all tests PASS and the generic `StorageCommand::Execute` path finds every Vector command without a new runtime message variant. +预期:所有测试通过,且通用的 `StorageCommand::Execute` 路径无需新增 runtime 消息变体即可找到每个 Vector 命令。 -- [x] **Step 7: Commit the command layer** +- [x] **步骤 7:提交命令层** ```bash git add src/cmd/src/vector.rs src/cmd/src/lib.rs src/cmd/src/table.rs @@ -832,22 +832,22 @@ git commit -m "feat(cmd): add redis vector set commands" --- -### Task 7: Make VSIM WITHSCORES correct in RESP2 and RESP3 +### 任务 7:让 VSIM WITHSCORES 在 RESP2 和 RESP3 下都正确 -**Files:** +**文件:** -- Modify: `src/resp/src/encode.rs` -- Modify: `src/resp/src/negotiation.rs` +- 修改:`src/resp/src/encode.rs` +- 修改:`src/resp/src/negotiation.rs` -**Interfaces:** +**接口:** -- Consumes: command-layer `RespData::Map` and `RespData::Double` -- Produces: RESP3 `%` map with `,` doubles -- Produces: RESP2 flat array with bulk-string scores +- 消费:命令层的 `RespData::Map` 和 `RespData::Double` +- 产出:RESP3 的 `%` map 与 `,` double +- 产出:RESP2 的扁平数组,分数为 bulk-string -- [x] **Step 1: Add encoder regression tests** +- [x] **步骤 1:添加编码器回归测试** -Add tests that encode this response under both protocol versions: +添加在两种协议版本下编码以下响应的测试: ```rust let response = RespData::Map(vec![ @@ -862,54 +862,63 @@ let response = RespData::Map(vec![ ]); ``` -Expected RESP3 bytes: +预期的 RESP3 字节: ```text %2\r\n$1\r\na\r\n,1\r\n$1\r\nb\r\n,0.5\r\n ``` -Expected RESP2 bytes: +预期的 RESP2 字节: ```text *4\r\n$1\r\na\r\n$1\r\n1\r\n$1\r\nb\r\n$3\r\n0.5\r\n ``` -- [x] **Step 2: Run the test and confirm RESP2 currently emits RESP3 types** +- [x] **步骤 2:运行测试并确认 RESP2 当前会输出 RESP3 类型** ```bash cargo test -p resp encode_vsim_withscores ``` -Expected: RESP2 assertion FAIL because `RespEncoder` currently stores the version but does not use it when encoding Map/Double. +预期:RESP2 断言失败,因为 `RespEncoder` 目前只是存储了版本号,但在编码 Map/Double 时并未使用它。 -- [x] **Step 3: Normalize once at the encoder boundary** +- [x] **步骤 3:用版本感知分支直接编码** -Refactor `RespEncoder` so the public trait method converts the complete response once, then calls a private recursive encoder: +递归 match 直接位于 trait 方法 `encode_resp_data` 中;版本差异在编码点处内联处理(null,以及经由 `append_map`/`append_set`/`append_push` 的 map/set/push 前缀): ```rust fn encode_resp_data(&mut self, data: &RespData) -> &mut Self { - let normalized = if self.is_resp3() { - data.clone() - } else { - ProtocolNegotiator::convert_to_resp2(data) - }; - self.encode_resp_data_inner(&normalized) + match data { + // ... + RespData::BulkString(None) | RespData::Array(None) if self.is_resp3() => { + self.append_null() + } + RespData::BulkString(None) => self.set_bulk_string_len(-1), + RespData::Array(Some(array)) => { + self.append_array_len(array.len() as i64); + for item in array { + self.encode_resp_data(item); + } + self + } + // ... + } } ``` -Move the existing recursive match into `encode_resp_data_inner`, and make all nested array/map/set/push recursion call `encode_resp_data_inner` rather than invoking the version-normalizing method again. +所有嵌套的 array/map/set/push 递归都调用 `encode_resp_data` 自身;不存在单独的内部/归一化方法。 -Keep `ProtocolNegotiator::convert_to_resp2` recursive for Array, Map, Set, Push, Null, Boolean, Double, BigNumber, BulkError, and VerbatimString. +`ProtocolNegotiator::convert_to_resp2` 保持可用,并对 Array、Map、Set、Push、Null、Boolean、Double、BigNumber、BulkError 和 VerbatimString 递归。 -- [x] **Step 4: Run all RESP tests** +- [x] **步骤 4:运行所有 RESP 测试** ```bash cargo test -p resp ``` -Expected: all RESP tests PASS; existing RESP3 HELLO behavior remains unchanged and RESP2 never emits RESP3-only prefixes. +预期:所有 RESP 测试通过;现有 RESP3 HELLO 行为保持不变,且 RESP2 永不输出仅属于 RESP3 的前缀。 -- [x] **Step 5: Commit protocol-aware encoding** +- [x] **步骤 5:提交协议感知编码** ```bash git add src/resp/src/encode.rs src/resp/src/negotiation.rs @@ -918,21 +927,21 @@ git commit -m "fix(resp): downgrade vector scores for resp2" --- -### Task 8: Add end-to-end compatibility tests and run the Kiwi verification pipeline +### 任务 8:添加端到端兼容性测试并运行 Kiwi 验证流水线 -**Files:** +**文件:** -- Create: `tests/python/test_vector_set_commands.py` -- Modify: `docs/superpowers/plans/2026-07-19-redis-vector-set.md` only to check completed boxes during execution +- 新建:`tests/python/test_vector_set_commands.py` +- 修改:`docs/superpowers/plans/2026-07-19-redis-vector-set.md`,仅在执行过程中勾选已完成的复选框 -**Interfaces:** +**接口:** -- Verifies: command parsing through TCP, dual runtime dispatch, storage persistence, RESP2/RESP3 encoding, TYPE/DEL/EXPIRE interaction -- Does not consume: `tests/python/test_vector_basic.py`, FT.CREATE, FT.SEARCH, HSET vector fields +- 验证:经由 TCP 的命令解析、双 runtime 分发、storage 持久化、RESP2/RESP3 编码、TYPE/DEL/EXPIRE 交互 +- 不消费:`tests/python/test_vector_basic.py`、FT.CREATE、FT.SEARCH、HSET vector 字段 -- [x] **Step 1: Write binary-safe Python integration tests** +- [x] **步骤 1:编写二进制安全的 Python 集成测试** -Create a local fixture in the new test file using: +在新测试文件中使用以下代码创建本地 fixture: ```python import os @@ -957,41 +966,41 @@ def vector_client(request): client.close() ``` -Add tests for: +添加以下测试: -- VALUES create/update return values and dimension mismatch. -- FP32 little-endian blob via `struct.pack("<2f", 1.0, 0.0)`. -- Binary key and element including `b"\x00"`; empty element. -- VCARD, VDIM, VEMB, VISMEMBER, VREM and last-member key deletion. -- VSIM direct vector, ELE query, COUNT, WITHSCORES and TRUTH. -- Stable tie order. -- TYPE returns `b"vectorset"`. -- WRONGTYPE against an existing string key. -- Missing-key semantics for every read command. -- EXPIRE makes the key immediately behave as missing. -- DEL removes members and allows same-name recreation with a new dimension. -- Missing NOQUANT, Q8, BIN, RAW and malformed input errors. +- VALUES 创建/更新的返回值,以及维度不匹配。 +- 通过 `struct.pack("<2f", 1.0, 0.0)` 构造的 FP32 little-endian blob。 +- 包含 `b"\x00"` 的二进制 key 和 element;空 element。 +- VCARD、VDIM、VEMB、VISMEMBER、VREM,以及最后一个成员删除后 key 被删除。 +- VSIM 直接向量查询、ELE 查询、COUNT、WITHSCORES 和 TRUTH。 +- 稳定的同分顺序。 +- TYPE 返回 `b"vectorset"`。 +- 对既有 string key 的 WRONGTYPE。 +- 每个读命令的缺失 key 语义。 +- EXPIRE 使 key 立即表现为缺失。 +- DEL 移除成员,并允许以新维度重建同名 key。 +- 缺少 NOQUANT、Q8、BIN、RAW 以及格式错误输入的报错。 -For WITHSCORES, assert redis-py returns a list under RESP2 and a dict under RESP3, with the same element order and numerical scores. +对于 WITHSCORES,断言 redis-py 在 RESP2 下返回 list、在 RESP3 下返回 dict,且 element 顺序和数值分数相同。 -- [x] **Step 2: Start standalone Kiwi and run only Vector integration tests** +- [x] **步骤 2:启动 standalone Kiwi 并只运行 Vector 集成测试** -Terminal 1: +终端 1: ```bash make standalone ``` -Terminal 2: +终端 2: ```bash make -C tests install-deps KIWI_PORT=7379 pytest -q tests/python/test_vector_set_commands.py ``` -Expected: all RESP2 and RESP3 parameterized cases PASS. +预期:所有 RESP2 和 RESP3 参数化用例通过。 -- [x] **Step 3: Run focused Rust regression suites** +- [x] **步骤 3:运行聚焦的 Rust 回归套件** ```bash RUST_TEST_THREADS=1 cargo test -p storage --test redis_vector_test @@ -1002,11 +1011,11 @@ cargo test -p cmd vector::tests cargo test -p resp ``` -Expected: all commands exit 0. +预期:所有命令退出码为 0。 -- [x] **Step 4: Run repository-standard verification** +- [x] **步骤 4:运行仓库标准验证** -Run exactly in this order: +严格按以下顺序运行: ```bash make fmt @@ -1015,9 +1024,9 @@ make build make test ``` -Expected: all commands exit 0. The first RocksDB build may be long; retain sccache and do not disable it repository-wide. +预期:所有命令退出码为 0。首次 RocksDB 构建可能耗时较长;保留 sccache,不要在仓库范围内禁用它。 -- [x] **Step 5: Inspect the final diff for scope and generated files** +- [x] **步骤 5:检查最终 diff 的范围和生成物** ```bash git status --short @@ -1025,15 +1034,15 @@ git diff --stat origin/main...HEAD git diff --check origin/main...HEAD ``` -Expected: +预期: -- No changes to `tests/python/test_vector_basic.py`. -- No FT.CREATE/FT.SEARCH implementation. -- No HNSW, Raft mutation, snapshot, or O(1) deletion code. -- No `__pycache__`, RocksDB data directory, build artifacts, or scratch notes staged. -- `git diff --check` produces no output. +- 没有改动 `tests/python/test_vector_basic.py`。 +- 没有 FT.CREATE/FT.SEARCH 实现。 +- 没有 HNSW、Raft mutation、snapshot 或 O(1) 删除代码。 +- 没有 `__pycache__`、RocksDB 数据目录、构建产物或临时笔记被加入暂存区。 +- `git diff --check` 无输出。 -- [x] **Step 6: Commit end-to-end tests** +- [x] **步骤 6:提交端到端测试** ```bash git add tests/python/test_vector_set_commands.py docs/superpowers/plans/2026-07-19-redis-vector-set.md @@ -1042,26 +1051,26 @@ git commit -m "test: cover redis vector set commands" --- -## Completion Criteria +## 完成标准 -The Phase 1 implementation is complete only when all statements below are true: +只有当以下所有陈述都为真时,Phase 1 实现才算完成: -- A VectorSet key and all its elements live in one RocksDB instance selected only from the user key. -- Meta and member mutations commit in one batch; failed validation leaves both unchanged. -- `TYPE` reports `vectorset`; generic TTL, DEL and FLUSHDB semantics work. -- VADD accepts FP32 and VALUES only with explicit NOQUANT and is binary-safe for elements. -- VSIM returns exact, deterministic Top-K results without collecting and sorting every hit. -- RESP2 and RESP3 WITHSCORES replies use their native compatible shapes. -- Raft Group cluster 模式支持所有七个命令的写入和读取（通过 binlog 复制到 follower）；Redis Cluster 槽位路由不在本版范围内。 -- Focused Rust/Python tests and the complete Kiwi verification pipeline pass. -- The final diff contains no FT prototype extraction, HNSW, Raft/snapshot design implementation, or unrelated cleanup. +- 一个 VectorSet key 及其全部 element 只存在于由 user key 唯一选定的一个 RocksDB instance 中。 +- Meta 和 member 的变更在一个 batch 中提交;校验失败时两者都保持不变。 +- `TYPE` 报告 `vectorset`;通用的 TTL、DEL 和 FLUSHDB 语义正常工作。 +- VADD 只在显式指定 NOQUANT 时接受 FP32 和 VALUES,且对 element 二进制安全。 +- VSIM 返回精确、确定性的 Top-K 结果,且不收集并排序所有命中。 +- RESP2 和 RESP3 的 WITHSCORES 回复使用各自原生兼容的形态。 +- Raft Group cluster 模式支持所有七个命令的写入和读取(通过 binlog 复制到 follower);Redis Cluster 槽位路由不在本版范围内。 +- 聚焦的 Rust/Python 测试和完整的 Kiwi 验证流水线全部通过。 +- 最终 diff 不包含 FT 原型抽取、HNSW、Raft/snapshot 设计实现或无关清理。 -## Deferred Follow-up Plans +## 推迟的后续计划 -Create separate reviewed designs and plans before implementing any of these independent projects: +在实现以下任何独立项目之前,先创建单独的、经过评审的设计和计划: 1. Redis Cluster 槽位路由与跨 slot 查询。 -2. Redis-compatible VINFO and INFO VECTOR observability. -3. Raft logical VectorSet mutations, leader-linearizable reads and snapshots. -4. O(1) generation deletion with background compaction cleanup. -5. HNSW derived index, rebuild/recovery and rolling-upgrade capability gates. +2. Redis 兼容的 VINFO 和 INFO VECTOR 可观测性。 +3. Raft 逻辑的 VectorSet 变更、领导者线性化读取和快照。 +4. O(1) generation 删除与后台 compaction 清理。 +5. HNSW 派生索引、重建/恢复以及滚动升级能力门禁。 diff --git a/src/resp/src/encode.rs b/src/resp/src/encode.rs index 994a8da6..fcd6bdc2 100644 --- a/src/resp/src/encode.rs +++ b/src/resp/src/encode.rs @@ -198,44 +198,6 @@ impl RespEncoder { let _ = write!(self.buffer, "*{len}"); self.append_crlf() } - - fn encode_resp_data_inner(&mut self, data: &RespData) -> &mut Self { - match data { - RespData::SimpleString(bytes) => { - self.buffer.extend_from_slice(b"+"); - self.buffer.extend_from_slice(bytes); - self.append_crlf() - } - RespData::Error(bytes) => { - self.buffer.extend_from_slice(b"-"); - self.buffer.extend_from_slice(bytes); - self.append_crlf() - } - RespData::Integer(num) => self.append_integer(*num), - RespData::BulkString(Some(bytes)) => self.append_bulk_string(bytes), - RespData::BulkString(None) | RespData::Array(None) if self.is_resp3() => { - self.append_null() - } - RespData::BulkString(None) => self.set_bulk_string_len(-1), - RespData::Array(Some(array)) => { - self.append_array_len(array.len() as i64); - for item in array { - self.encode_resp_data_inner(item); - } - self - } - RespData::Array(None) => self.set_array_len(-1), - RespData::Null => self.append_null(), - RespData::Boolean(value) => self.append_boolean(*value), - RespData::Double(value) => self.append_double(*value), - RespData::BigNumber(bytes) => self.append_big_number(bytes), - RespData::BulkError(bytes) => self.append_bulk_error(bytes), - RespData::VerbatimString { format, data } => self.append_verbatim_string(format, data), - RespData::Map(pairs) => self.append_map(pairs), - RespData::Set(items) => self.append_set(items), - RespData::Push(items) => self.append_push(items), - } - } } impl RespEncode for RespEncoder { @@ -400,7 +362,41 @@ impl RespEncode for RespEncoder { } fn encode_resp_data(&mut self, data: &RespData) -> &mut Self { - self.encode_resp_data_inner(data) + match data { + RespData::SimpleString(bytes) => { + self.buffer.extend_from_slice(b"+"); + self.buffer.extend_from_slice(bytes); + self.append_crlf() + } + RespData::Error(bytes) => { + self.buffer.extend_from_slice(b"-"); + self.buffer.extend_from_slice(bytes); + self.append_crlf() + } + RespData::Integer(num) => self.append_integer(*num), + RespData::BulkString(Some(bytes)) => self.append_bulk_string(bytes), + RespData::BulkString(None) | RespData::Array(None) if self.is_resp3() => { + self.append_null() + } + RespData::BulkString(None) => self.set_bulk_string_len(-1), + RespData::Array(Some(array)) => { + self.append_array_len(array.len() as i64); + for item in array { + self.encode_resp_data(item); + } + self + } + RespData::Array(None) => self.set_array_len(-1), + RespData::Null => self.append_null(), + RespData::Boolean(value) => self.append_boolean(*value), + RespData::Double(value) => self.append_double(*value), + RespData::BigNumber(bytes) => self.append_big_number(bytes), + RespData::BulkError(bytes) => self.append_bulk_error(bytes), + RespData::VerbatimString { format, data } => self.append_verbatim_string(format, data), + RespData::Map(pairs) => self.append_map(pairs), + RespData::Set(items) => self.append_set(items), + RespData::Push(items) => self.append_push(items), + } } fn append_null(&mut self) -> &mut Self { @@ -493,8 +489,8 @@ impl RespEncode for RespEncoder { self.append_array_len((pairs.len() * 2) as i64); } for (key, value) in pairs { - self.encode_resp_data_inner(key); - self.encode_resp_data_inner(value); + self.encode_resp_data(key); + self.encode_resp_data(value); } self } @@ -507,7 +503,7 @@ impl RespEncode for RespEncoder { self.append_array_len(items.len() as i64); } for item in items { - self.encode_resp_data_inner(item); + self.encode_resp_data(item); } self } @@ -520,7 +516,7 @@ impl RespEncode for RespEncoder { self.append_array_len(items.len() as i64); } for item in items { - self.encode_resp_data_inner(item); + self.encode_resp_data(item); } self } diff --git a/src/storage/src/format_vector.rs b/src/storage/src/format_vector.rs index 584c08b4..0832576c 100644 --- a/src/storage/src/format_vector.rs +++ b/src/storage/src/format_vector.rs @@ -32,22 +32,6 @@ pub const VECTOR_METRIC_COSINE: u8 = 1; pub const VECTOR_VALUE_MAGIC: u8 = 0x56; pub const VECTOR_VALUE_FORMAT: u8 = 1; -// Vector set meta value layout stored in MetaCF: -// -// | data_type | count | version | format | encoding | metric | flags | dimension | zero_reserve | ctime | etime | -// | 1B | 8B | 8B | 1B | 1B | 1B | 1B | 4B | 8B | 8B | 8B | -// -// `data_type` is DataType::VectorSet, `encoding` is VECTOR_ENCODING_FP32_LE, -// and `metric` is the similarity metric used for VSIM (e.g. cosine). - -// Vector member data value layout stored in VectorDataCF: -// -// | magic | format | dimension | original_l2 | normalized_components ... | -// | 1B | 1B | 4B | 4B | 4B * dimension | -// -// `magic` is VECTOR_VALUE_MAGIC and `original_l2` preserves the pre-normalization -// L2 norm so VEMB can reconstruct the original FP32 vector. - const VECTOR_META_ZERO_RESERVE_LENGTH: usize = 8; const VECTOR_VALUE_HEADER_LENGTH: usize = 10; @@ -84,6 +68,13 @@ impl SimilarityMetric { } } +// Vector set meta value layout stored in MetaCF: +// +// | data_type | count | version | format | encoding | metric | flags | dimension | zero_reserve | ctime | etime | +// | 1B | 8B | 8B | 1B | 1B | 1B | 1B | 4B | 8B | 8B | 8B | +// +// `data_type` is DataType::VectorSet, `encoding` is VECTOR_ENCODING_FP32_LE, +// and `metric` is the similarity metric used for VSIM (e.g. cosine). #[derive(Debug, Clone, PartialEq, Eq)] pub struct VectorMeta { count: u64, @@ -229,6 +220,13 @@ impl VectorMeta { } } +// Vector member data value layout stored in VectorDataCF: +// +// | magic | format | dimension | original_l2 | normalized_components ... | +// | 1B | 1B | 4B | 4B | 4B * dimension | +// +// `magic` is VECTOR_VALUE_MAGIC and `original_l2` preserves the pre-normalization +// L2 norm so VEMB can reconstruct the original FP32 vector. #[derive(Debug, Clone, PartialEq)] pub(crate) struct VectorDataValue { canonical: CanonicalVector, From 5a3535648d1d073cb61b0d7e01858523ce871635 Mon Sep 17 00:00:00 2001 From: happy-v587 Date: Mon, 27 Jul 2026 21:32:34 +0800 Subject: [PATCH 19/32] feat(vector): add quantization framework (NOQUANT/BIN/Q8) - CanonicalVector now holds quantized data (VectorData: Fp32/Binary/Int8) with to_quantized() as the single conversion entry point - score() dispatches per quantization: hamming similarity for BIN, cosine on the dequantized FP32 form otherwise - VectorDataValue format v3: quant byte, reserved flags byte (for future SETATTR attributes), per-quantization payload layouts - VectorMeta carries the set-level quantization; vadd/vsim convert members and queries to the set's quantization - unify vector meta reads via decode_vector_meta/read_vector_meta_opt --- src/storage/src/format_vector.rs | 234 +++++++++++++++++++++++------- src/storage/src/lib.rs | 4 +- src/storage/src/redis_vectors.rs | 123 +++++++--------- src/storage/src/vector.rs | 239 +++++++++++++++++++++++++++++-- 4 files changed, 466 insertions(+), 134 deletions(-) diff --git a/src/storage/src/format_vector.rs b/src/storage/src/format_vector.rs index 0832576c..c72fd0c2 100644 --- a/src/storage/src/format_vector.rs +++ b/src/storage/src/format_vector.rs @@ -23,17 +23,23 @@ use crate::{ DataType, error::{InvalidFormatSnafu, Result}, storage_define::BASE_META_VALUE_LENGTH, - vector::CanonicalVector, + vector::{CanonicalVector, QuantizationType, VectorData}, }; pub const VECTOR_META_FORMAT: u8 = 1; -pub const VECTOR_ENCODING_FP32_LE: u8 = 1; + +pub const VECTOR_ENCODING_NOQUANT: u8 = 0; +pub const VECTOR_ENCODING_BIN: u8 = 1; +pub const VECTOR_ENCODING_Q8: u8 = 2; + pub const VECTOR_METRIC_COSINE: u8 = 1; + pub const VECTOR_VALUE_MAGIC: u8 = 0x56; pub const VECTOR_VALUE_FORMAT: u8 = 1; const VECTOR_META_ZERO_RESERVE_LENGTH: usize = 8; -const VECTOR_VALUE_HEADER_LENGTH: usize = 10; +const VECTOR_VALUE_HEADER_LENGTH: usize = 12; +const VECTOR_Q8_PARAMS_LENGTH: usize = 8; /// Similarity metric used to compare vectors in a vector set. /// @@ -44,6 +50,28 @@ pub enum SimilarityMetric { Cosine, } +impl QuantizationType { + pub fn from_u8(value: u8) -> Result { + match value { + VECTOR_ENCODING_NOQUANT => Ok(Self::None), + VECTOR_ENCODING_BIN => Ok(Self::Binary), + VECTOR_ENCODING_Q8 => Ok(Self::Int8), + _ => InvalidFormatSnafu { + message: format!("unsupported quantization type: {value}"), + } + .fail(), + } + } + + pub const fn to_u8(self) -> u8 { + match self { + Self::None => VECTOR_ENCODING_NOQUANT, + Self::Binary => VECTOR_ENCODING_BIN, + Self::Int8 => VECTOR_ENCODING_Q8, + } + } +} + impl SimilarityMetric { pub fn from_u8(value: u8) -> Result { match value { @@ -63,30 +91,36 @@ impl SimilarityMetric { pub fn score(&self, left: &CanonicalVector, right: &CanonicalVector) -> Result { match self { - Self::Cosine => left.cosine_score(right), + Self::Cosine => left.score(right), } } } // Vector set meta value layout stored in MetaCF: // -// | data_type | count | version | format | encoding | metric | flags | dimension | zero_reserve | ctime | etime | -// | 1B | 8B | 8B | 1B | 1B | 1B | 1B | 4B | 8B | 8B | 8B | +// | data_type | count | version | format | quant | metric | flags | dimension | zero_reserve | ctime | etime | +// | 1B | 8B | 8B | 1B | 1B | 1B | 1B | 4B | 8B | 8B | 8B | // -// `data_type` is DataType::VectorSet, `encoding` is VECTOR_ENCODING_FP32_LE, +// `data_type` is DataType::VectorSet, `quant` is the quantization type (NOQUANT/BIN/Q8), // and `metric` is the similarity metric used for VSIM (e.g. cosine). #[derive(Debug, Clone, PartialEq, Eq)] pub struct VectorMeta { count: u64, pub(crate) version: u64, dimension: u32, + quantization: QuantizationType, metric: SimilarityMetric, ctime: u64, etime: u64, } impl VectorMeta { - pub(crate) fn new_after(count: u64, dimension: u32, previous_version: u64) -> Self { + pub(crate) fn new_after( + count: u64, + dimension: u32, + quantization: QuantizationType, + previous_version: u64, + ) -> Self { let now = Utc::now().timestamp_micros() as u64; let version = match previous_version >= now { true => previous_version + 1, @@ -96,12 +130,17 @@ impl VectorMeta { count, version, dimension, + quantization, metric: SimilarityMetric::Cosine, ctime: now, etime: 0, } } + pub(crate) fn quantization(&self) -> QuantizationType { + self.quantization + } + pub(crate) fn metric(&self) -> SimilarityMetric { self.metric } @@ -112,7 +151,7 @@ impl VectorMeta { output.put_u64_le(self.count); output.put_u64_le(self.version); output.put_u8(VECTOR_META_FORMAT); - output.put_u8(VECTOR_ENCODING_FP32_LE); + output.put_u8(self.quantization.to_u8()); output.put_u8(self.metric.to_u8()); output.put_u8(0); output.put_u32_le(self.dimension); @@ -146,7 +185,7 @@ impl VectorMeta { let count = reader.get_u64_le(); let version = reader.get_u64_le(); let format = reader.get_u8(); - let encoding = reader.get_u8(); + let quant = reader.get_u8(); let metric = reader.get_u8(); let flags = reader.get_u8(); let dimension = reader.get_u32_le(); @@ -161,12 +200,7 @@ impl VectorMeta { message: format!("unsupported vector meta format: {format}") } ); - ensure!( - encoding == VECTOR_ENCODING_FP32_LE, - InvalidFormatSnafu { - message: format!("unsupported vector encoding: {encoding}") - } - ); + let quantization = QuantizationType::from_u8(quant)?; let metric = SimilarityMetric::from_u8(metric)?; ensure!( flags == 0 && zero_reserve.iter().all(|byte| *byte == 0), @@ -185,6 +219,7 @@ impl VectorMeta { count, version, dimension, + quantization, metric, ctime, etime, @@ -222,11 +257,17 @@ impl VectorMeta { // Vector member data value layout stored in VectorDataCF: // -// | magic | format | dimension | original_l2 | normalized_components ... | -// | 1B | 1B | 4B | 4B | 4B * dimension | +// | magic | format | quant | flags | dimension | original_l2 | [quant_params] | payload | +// | 1B | 1B | 1B | 1B | 4B | 4B | 0B or 8B | varies | // // `magic` is VECTOR_VALUE_MAGIC and `original_l2` preserves the pre-normalization -// L2 norm so VEMB can reconstruct the original FP32 vector. +// L2 norm so VEMB can reconstruct the original FP32 vector. `flags` is reserved +// for optional sections (e.g. bit 0 = trailing SETATTR attributes JSON) and +// must be zero until such a section is implemented. The payload layout +// depends on `quant`: +// NOQUANT: 4B * dimension FP32 components, no quant_params. +// BIN: ceil(dimension / 8) bitmap bytes, no quant_params. +// Q8: 1B * dimension INT8 codes, quant_params = min FP32 + max FP32. #[derive(Debug, Clone, PartialEq)] pub(crate) struct VectorDataValue { canonical: CanonicalVector, @@ -240,15 +281,35 @@ impl VectorDataValue { } pub(crate) fn encode(&self) -> BytesMut { - let mut output = BytesMut::with_capacity( - VECTOR_VALUE_HEADER_LENGTH + size_of_val(self.canonical.normalized()), - ); + let canonical = &self.canonical; + let payload_length = match canonical.data() { + VectorData::Fp32(values) => values.len() * size_of::(), + VectorData::Binary(bits) => bits.len(), + VectorData::Int8 { values, .. } => VECTOR_Q8_PARAMS_LENGTH + values.len(), + }; + let mut output = BytesMut::with_capacity(VECTOR_VALUE_HEADER_LENGTH + payload_length); output.put_u8(VECTOR_VALUE_MAGIC); output.put_u8(VECTOR_VALUE_FORMAT); - output.put_u32_le(self.canonical.dimension()); - output.put_f32_le(self.canonical.original_l2()); - for component in self.canonical.normalized() { - output.put_f32_le(*component); + output.put_u8(canonical.quantization().to_u8()); + output.put_u8(0); // flags: reserved, no optional sections yet + output.put_u32_le(canonical.dimension()); + output.put_f32_le(canonical.original_l2()); + match canonical.data() { + VectorData::Fp32(values) => { + for component in values { + output.put_f32_le(*component); + } + } + VectorData::Binary(bits) => { + output.put_slice(bits); + } + VectorData::Int8 { values, min, max } => { + output.put_f32_le(*min); + output.put_f32_le(*max); + for value in values { + output.put_i8(*value); + } + } } output } @@ -268,6 +329,8 @@ impl VectorDataValue { let mut reader = value; let magic = reader.get_u8(); let format = reader.get_u8(); + let quantization = QuantizationType::from_u8(reader.get_u8())?; + let flags = reader.get_u8(); let dimension = reader.get_u32_le(); let original_l2 = reader.get_f32_le(); @@ -283,6 +346,12 @@ impl VectorDataValue { message: format!("unsupported vector value format: {format}") } ); + ensure!( + flags == 0, + InvalidFormatSnafu { + message: format!("unsupported vector value flags: {flags:#04x}") + } + ); ensure!( dimension != 0, InvalidFormatSnafu { @@ -296,31 +365,75 @@ impl VectorDataValue { } ); - let payload_length = (dimension as usize).checked_mul(size_of::()); - ensure!( - payload_length.is_some_and(|length| reader.len() == length), - InvalidFormatSnafu { - message: format!( - "invalid vector payload length: {} for dimension {}", - reader.len(), - dimension - ) + let dimension = dimension as usize; + let data = match quantization { + QuantizationType::None => { + ensure!( + reader.len() == dimension * size_of::(), + InvalidFormatSnafu { + message: format!( + "invalid vector payload length: {} for dimension {}", + reader.len(), + dimension + ) + } + ); + let normalized = reader + .chunks_exact(size_of::()) + .map(|chunk| f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]])) + .collect::>(); + ensure!( + normalized.iter().all(|component| component.is_finite()), + InvalidFormatSnafu { + message: "vector payload components must be finite".to_string() + } + ); + VectorData::Fp32(normalized) } - ); - - let normalized = reader - .chunks_exact(size_of::()) - .map(|chunk| f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]])) - .collect::>(); - ensure!( - normalized.iter().all(|component| component.is_finite()), - InvalidFormatSnafu { - message: "vector payload components must be finite".to_string() + QuantizationType::Binary => { + ensure!( + reader.len() == dimension.div_ceil(8), + InvalidFormatSnafu { + message: format!( + "invalid binary vector payload length: {} for dimension {}", + reader.len(), + dimension + ) + } + ); + VectorData::Binary(reader.to_vec()) } - ); + QuantizationType::Int8 => { + ensure!( + reader.len() == VECTOR_Q8_PARAMS_LENGTH + dimension, + InvalidFormatSnafu { + message: format!( + "invalid q8 vector payload length: {} for dimension {}", + reader.len(), + dimension + ) + } + ); + let min = reader.get_f32_le(); + let max = reader.get_f32_le(); + ensure!( + min.is_finite() && max.is_finite(), + InvalidFormatSnafu { + message: "q8 quantization params must be finite".to_string() + } + ); + let values = reader.iter().map(|byte| *byte as i8).collect(); + VectorData::Int8 { values, min, max } + } + }; Ok(Self { - canonical: CanonicalVector::from_normalized_parts(dimension, original_l2, normalized), + canonical: CanonicalVector::from_parts( + dimension as u32, + original_l2, + quantization, + data, + ), }) } @@ -365,9 +478,20 @@ mod tests { assert_eq!(decoded.canonical(), &canonical); } + #[test] + fn vector_data_value_round_trips_quantized() { + let canonical = CanonicalVector::from_values(&[3.0, 4.0, -1.0, 2.0]).expect("valid vector"); + for quantization in [QuantizationType::Binary, QuantizationType::Int8] { + let quantized = canonical.to_quantized(quantization).expect("quantize"); + let encoded = VectorDataValue::from_canonical(&quantized).encode(); + let decoded = VectorDataValue::decode(&encoded).expect("decode vector value"); + assert_eq!(decoded.canonical(), &quantized); + } + } + #[test] fn vector_meta_round_trips() { - let mut meta = VectorMeta::new_after(2, 2, 0); + let mut meta = VectorMeta::new_after(2, 2, QuantizationType::None, 0); meta.version = 42; let encoded = meta.encode(); let decoded = VectorMeta::decode(&encoded).expect("decode vector meta"); @@ -387,15 +511,23 @@ mod tests { bad_magic[0] = 0; assert!(VectorDataValue::decode(&bad_magic).is_err()); + let mut bad_quant = encoded_value.clone(); + bad_quant[2] = 0xFF; + assert!(VectorDataValue::decode(&bad_quant).is_err()); + + let mut bad_flags = encoded_value.clone(); + bad_flags[3] = 0xFF; + assert!(VectorDataValue::decode(&bad_flags).is_err()); + let mut zero_dimension = encoded_value.clone(); - zero_dimension[2..6].copy_from_slice(&0_u32.to_le_bytes()); + zero_dimension[4..8].copy_from_slice(&0_u32.to_le_bytes()); assert!(VectorDataValue::decode(&zero_dimension).is_err()); let mut non_finite_payload = encoded_value; - non_finite_payload[10..14].copy_from_slice(&f32::NAN.to_le_bytes()); + non_finite_payload[12..16].copy_from_slice(&f32::NAN.to_le_bytes()); assert!(VectorDataValue::decode(&non_finite_payload).is_err()); - let encoded_meta = VectorMeta::new_after(2, 2, 0).encode(); + let encoded_meta = VectorMeta::new_after(2, 2, QuantizationType::None, 0).encode(); assert!(VectorMeta::decode(&encoded_meta[..encoded_meta.len() - 1]).is_err()); let mut bad_meta_format = encoded_meta; diff --git a/src/storage/src/lib.rs b/src/storage/src/lib.rs index 984c27c8..b146db29 100644 --- a/src/storage/src/lib.rs +++ b/src/storage/src/lib.rs @@ -84,6 +84,6 @@ pub use storage::{BgTask, BgTaskHandler}; pub use storage_impl::BeforeOrAfter; pub use util::{safe_cleanup_test_db, unique_test_db_path}; pub use vector::{ - CanonicalVector, VectorHit, VectorQuery, VectorSearchEngine, VectorSearchMode, - VectorSearchOptions, + CanonicalVector, QuantizationType, VectorHit, VectorQuery, VectorSearchEngine, + VectorSearchMode, VectorSearchOptions, }; diff --git a/src/storage/src/redis_vectors.rs b/src/storage/src/redis_vectors.rs index 6ea7fd89..1f407b98 100644 --- a/src/storage/src/redis_vectors.rs +++ b/src/storage/src/redis_vectors.rs @@ -32,14 +32,16 @@ use crate::{ }; impl Redis { - fn parse_vector_meta(&self, value: &[u8]) -> Result> { + /// Decode raw meta bytes into a `VectorMeta` without liveness filtering: + /// stale or emptied sets are still returned so callers can inspect their + /// version. Returns an error when the key holds another live data type. + fn decode_vector_meta(&self, value: &[u8]) -> Result> { if value.is_empty() { return Ok(None); } if value[0] == DataType::VectorSet as u8 { - let meta = VectorMeta::decode(value)?; - return Ok((!meta.is_stale() && meta.count() != 0).then_some(meta)); + return VectorMeta::decode(value).map(Some); } match self.check_type_state(value, DataType::VectorSet)? { @@ -48,6 +50,14 @@ impl Redis { } } + /// Decode raw meta bytes into a live `VectorMeta`, treating stale or + /// emptied sets as absent. + fn parse_vector_meta(&self, value: &[u8]) -> Result> { + Ok(self + .decode_vector_meta(value)? + .filter(|meta| !meta.is_stale() && meta.count() != 0)) + } + pub fn vadd(&self, key: &[u8], element: &[u8], vector: &CanonicalVector) -> Result { let db = self.db.as_ref().context(OptionNoneSnafu { message: "db is not initialized".to_string(), @@ -66,18 +76,15 @@ impl Redis { let lock_key = String::from_utf8_lossy(key); let _lock = ScopeRecordLock::new(self.lock_mgr.as_ref(), &lock_key); let meta_key = BaseMetaKey::new(key).encode()?; - let stored_meta = db.get_cf(&meta_cf, &meta_key).context(RocksSnafu)?; - let live_meta = stored_meta - .as_deref() - .map(|value| self.parse_vector_meta(value)) - .transpose()? - .flatten(); - let previous_generation = stored_meta - .as_deref() - .filter(|value| value.first().copied() == Some(DataType::VectorSet as u8)) - .map(VectorMeta::decode) - .transpose()? - .map_or(0, |meta| meta.version()); + let stored_raw = db.get_cf(&meta_cf, &meta_key).context(RocksSnafu)?; + let stored_meta = match stored_raw.as_deref() { + Some(value) => self.decode_vector_meta(value)?, + None => None, + }; + // Keep the previous version even for stale or emptied sets so a + // recreated set always gets a fresh, monotonically increasing generation. + let previous_generation = stored_meta.as_ref().map_or(0, VectorMeta::version); + let live_meta = stored_meta.filter(|meta| !meta.is_stale() && meta.count() != 0); let is_new_set = live_meta.is_none(); let mut meta = match live_meta { @@ -94,7 +101,12 @@ impl Redis { } meta } - None => VectorMeta::new_after(1, vector.dimension(), previous_generation), + None => VectorMeta::new_after( + 1, + vector.dimension(), + vector.quantization(), + previous_generation, + ), }; let member_key = MemberDataKey::new(key, meta.version(), element).encode()?; @@ -116,7 +128,10 @@ impl Redis { meta.set_count(count); } - let member_value = VectorDataValue::from_canonical(vector).encode(); + // Quantization is a per-set property: store the member in the set's + // quantization regardless of the form the client supplied. + let vector = vector.to_quantized(meta.quantization())?; + let member_value = VectorDataValue::from_canonical(&vector).encode(); let meta_value = meta.encode(); let mut batch = self.create_batch()?; batch.put(ColumnFamilyIndex::VectorDataCF, &member_key, &member_value)?; @@ -129,11 +144,6 @@ impl Redis { let db = self.db.as_ref().context(OptionNoneSnafu { message: "db is not initialized".to_string(), })?; - let meta_cf = self - .get_cf_handle(ColumnFamilyIndex::MetaCF) - .context(OptionNoneSnafu { - message: "MetaCF is not initialized".to_string(), - })?; let vector_cf = self .get_cf_handle(ColumnFamilyIndex::VectorDataCF) .context(OptionNoneSnafu { @@ -142,11 +152,7 @@ impl Redis { let lock_key = String::from_utf8_lossy(key); let _lock = ScopeRecordLock::new(self.lock_mgr.as_ref(), &lock_key); - let meta_key = BaseMetaKey::new(key).encode()?; - let Some(meta_raw) = db.get_cf(&meta_cf, &meta_key).context(RocksSnafu)? else { - return Ok(false); - }; - let Some(mut meta) = self.parse_vector_meta(&meta_raw)? else { + let Some(mut meta) = self.read_vector_meta(key)? else { return Ok(false); }; @@ -159,6 +165,7 @@ impl Redis { return Ok(false); } + let meta_key = BaseMetaKey::new(key).encode()?; let mut batch = self.create_batch()?; batch.delete(ColumnFamilyIndex::VectorDataCF, &member_key)?; if meta.count() > 1 { @@ -190,11 +197,6 @@ impl Redis { let db = self.db.as_ref().context(OptionNoneSnafu { message: "db is not initialized".to_string(), })?; - let meta_cf = self - .get_cf_handle(ColumnFamilyIndex::MetaCF) - .context(OptionNoneSnafu { - message: "MetaCF is not initialized".to_string(), - })?; let vector_cf = self .get_cf_handle(ColumnFamilyIndex::VectorDataCF) .context(OptionNoneSnafu { @@ -204,14 +206,7 @@ impl Redis { let mut read_options = ReadOptions::default(); read_options.set_snapshot(&snapshot); - let meta_key = BaseMetaKey::new(key).encode()?; - let Some(meta_raw) = db - .get_cf_opt(&meta_cf, &meta_key, &read_options) - .context(RocksSnafu)? - else { - return Ok(None); - }; - let Some(meta) = self.parse_vector_meta(&meta_raw)? else { + let Some(meta) = self.read_vector_meta_opt(key, Some(&read_options))? else { return Ok(None); }; let member_key = MemberDataKey::new(key, meta.version(), element).encode()?; @@ -239,11 +234,6 @@ impl Redis { let db = self.db.as_ref().context(OptionNoneSnafu { message: "db is not initialized".to_string(), })?; - let meta_cf = self - .get_cf_handle(ColumnFamilyIndex::MetaCF) - .context(OptionNoneSnafu { - message: "MetaCF is not initialized".to_string(), - })?; let vector_cf = self .get_cf_handle(ColumnFamilyIndex::VectorDataCF) .context(OptionNoneSnafu { @@ -253,14 +243,7 @@ impl Redis { let mut read_options = ReadOptions::default(); read_options.set_snapshot(&snapshot); - let meta_key = BaseMetaKey::new(key).encode()?; - let Some(meta_raw) = db - .get_cf_opt(&meta_cf, &meta_key, &read_options) - .context(RocksSnafu)? - else { - return Ok(false); - }; - let Some(meta) = self.parse_vector_meta(&meta_raw)? else { + let Some(meta) = self.read_vector_meta_opt(key, Some(&read_options))? else { return Ok(false); }; let member_key = MemberDataKey::new(key, meta.version(), element).encode()?; @@ -286,11 +269,6 @@ impl Redis { let db = self.db.as_ref().context(OptionNoneSnafu { message: "db is not initialized".to_string(), })?; - let meta_cf = self - .get_cf_handle(ColumnFamilyIndex::MetaCF) - .context(OptionNoneSnafu { - message: "MetaCF is not initialized".to_string(), - })?; let vector_cf = self .get_cf_handle(ColumnFamilyIndex::VectorDataCF) .context(OptionNoneSnafu { @@ -300,14 +278,7 @@ impl Redis { let mut point_read_options = ReadOptions::default(); point_read_options.set_snapshot(&snapshot); - let meta_key = BaseMetaKey::new(key).encode()?; - let Some(meta_raw) = db - .get_cf_opt(&meta_cf, &meta_key, &point_read_options) - .context(RocksSnafu)? - else { - return Ok(Vec::new()); - }; - let Some(meta) = self.parse_vector_meta(&meta_raw)? else { + let Some(meta) = self.read_vector_meta_opt(key, Some(&point_read_options))? else { return Ok(Vec::new()); }; @@ -337,6 +308,9 @@ impl Redis { } .fail(); } + // Score the query in the set's quantization so it is comparable to + // the stored members. + let query_vector = query_vector.to_quantized(meta.quantization())?; let prefix = MemberDataKey::new(key, meta.version(), b"").encode_seek_key()?; let mut scan_options = ReadOptions::default(); @@ -381,7 +355,11 @@ impl Redis { engine.search(&query_vector, &meta.metric(), options.count, candidates) } - fn read_vector_meta(&self, key: &[u8]) -> Result> { + fn read_vector_meta_opt( + &self, + key: &[u8], + read_options: Option<&ReadOptions>, + ) -> Result> { let db = self.db.as_ref().context(OptionNoneSnafu { message: "db is not initialized".to_string(), })?; @@ -391,9 +369,18 @@ impl Redis { message: "MetaCF is not initialized".to_string(), })?; let meta_key = BaseMetaKey::new(key).encode()?; - match db.get_cf(&meta_cf, &meta_key).context(RocksSnafu)? { + let value = match read_options { + Some(opts) => db.get_cf_opt(&meta_cf, &meta_key, opts), + None => db.get_cf(&meta_cf, &meta_key), + } + .context(RocksSnafu)?; + match value { Some(value) => self.parse_vector_meta(&value), None => Ok(None), } } + + fn read_vector_meta(&self, key: &[u8]) -> Result> { + self.read_vector_meta_opt(key, None) + } } diff --git a/src/storage/src/vector.rs b/src/storage/src/vector.rs index e32aae5f..3b0ced99 100644 --- a/src/storage/src/vector.rs +++ b/src/storage/src/vector.rs @@ -22,14 +22,41 @@ use snafu::ensure; use crate::error::{InvalidArgumentSnafu, Result}; use crate::format_vector::SimilarityMetric; +/// Quantization type for vector storage. +/// Matches Redis vector set quantization options. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[repr(u8)] +pub enum QuantizationType { + /// No quantization, store as FP32 (NOQUANT) + None = 0, + /// Binary quantization, 1 bit per component (BIN) + Binary = 1, + /// 8-bit integer quantization (Q8) + Int8 = 2, +} + #[derive(Debug, Clone, PartialEq)] pub struct CanonicalVector { dimension: u32, original_l2: f32, - normalized: Vec, + quantization: QuantizationType, + data: VectorData, +} + +/// Internal representation of quantized vector data. +#[derive(Debug, Clone, PartialEq)] +pub enum VectorData { + /// Unquantized FP32 values (NOQUANT) + Fp32(Vec), + /// Binary packed bits, dimension/8 bytes rounded up (BIN) + Binary(Vec), + /// 8-bit quantized values with min/max range (Q8) + Int8 { values: Vec, min: f32, max: f32 }, } impl CanonicalVector { + /// Create a CanonicalVector from raw little-endian FP32 bytes. + /// The vector is stored as-is (NOQUANT quantization). pub fn from_fp32_le(raw: &[u8]) -> Result { ensure!( !raw.is_empty() && raw.len().is_multiple_of(size_of::()), @@ -46,6 +73,8 @@ impl CanonicalVector { Self::from_values(&values) } + /// Create a CanonicalVector from FP32 values. + /// Validates that all values are finite and non-zero norm. pub fn from_values(values: &[f32]) -> Result { ensure!( !values.is_empty(), @@ -86,7 +115,7 @@ impl CanonicalVector { } ); - let normalized = values + let normalized: Vec = values .iter() .map(|value| (f64::from(*value) / norm) as f32) .collect(); @@ -94,10 +123,82 @@ impl CanonicalVector { Ok(Self { dimension: values.len() as u32, original_l2, - normalized, + quantization: QuantizationType::None, + data: VectorData::Fp32(normalized), }) } + /// Quantize to binary (1 bit per component). + /// Positive values become 1, negative/zero become 0. + pub fn to_binary(&self) -> Result { + let fp32 = self.as_fp32()?; + let byte_count = (self.dimension as usize).div_ceil(8); + let mut bits = vec![0u8; byte_count]; + + for (i, &value) in fp32.iter().enumerate() { + if value > 0.0 { + bits[i / 8] |= 1 << (i % 8); + } + } + + Ok(Self { + dimension: self.dimension, + original_l2: self.original_l2, + quantization: QuantizationType::Binary, + data: VectorData::Binary(bits), + }) + } + + /// Quantize to 8-bit integers using scalar quantization. + /// Maps [min, max] to [-128, 127]. + pub fn to_int8(&self) -> Result { + let fp32 = self.as_fp32()?; + + let min = fp32.iter().cloned().fold(f32::INFINITY, f32::min); + let max = fp32.iter().cloned().fold(f32::NEG_INFINITY, f32::max); + let range = max - min; + + // A normalized vector can still have all components equal (e.g. + // [1/sqrt(n), ...]); every code then dequantizes back to `min`. + let values: Vec = if range > 0.0 { + fp32.iter() + .map(|&v| (((v - min) / range) * 255.0 - 128.0).round() as i8) + .collect() + } else { + vec![-128; fp32.len()] + }; + + Ok(Self { + dimension: self.dimension, + original_l2: self.original_l2, + quantization: QuantizationType::Int8, + data: VectorData::Int8 { values, min, max }, + }) + } + + /// Convert this vector to `target` quantization, going through the FP32 + /// representation when the quantization differs. + pub fn to_quantized(&self, target: QuantizationType) -> Result { + if self.quantization == target { + return Ok(self.clone()); + } + match target { + QuantizationType::None => Ok(Self::from_parts( + self.dimension, + self.original_l2, + QuantizationType::None, + VectorData::Fp32(self.as_fp32()?), + )), + QuantizationType::Binary => self.to_binary(), + QuantizationType::Int8 => self.to_int8(), + } + } + + /// Get the quantization type of this vector. + pub fn quantization(&self) -> QuantizationType { + self.quantization + } + pub fn dimension(&self) -> u32 { self.dimension } @@ -106,24 +207,98 @@ impl CanonicalVector { self.original_l2 } - pub(crate) fn normalized(&self) -> &[f32] { - &self.normalized + /// Get FP32 representation, converting from quantized format if needed. + pub fn as_fp32(&self) -> Result> { + match &self.data { + VectorData::Fp32(values) => Ok(values.clone()), + VectorData::Binary(bits) => { + let mut values = Vec::with_capacity(self.dimension as usize); + for i in 0..self.dimension as usize { + let byte_idx = i / 8; + let bit_idx = i % 8; + let is_set = (bits[byte_idx] >> bit_idx) & 1; + values.push(if is_set == 1 { 1.0 } else { -1.0 }); + } + Ok(values) + } + VectorData::Int8 { values, min, max } => { + let range = max - min; + let fp32: Vec = values + .iter() + .map(|&v| { + let normalized = (f32::from(v) + 128.0) / 255.0; // [0, 1] + min + normalized * range + }) + .collect(); + Ok(fp32) + } + } + } + + /// Get the raw quantized data (for storage encoding). + pub(crate) fn data(&self) -> &VectorData { + &self.data } - pub(crate) fn from_normalized_parts( + /// Create from pre-quantized data (for loading from storage). + pub(crate) fn from_parts( dimension: u32, original_l2: f32, - normalized: Vec, + quantization: QuantizationType, + data: VectorData, ) -> Self { Self { dimension, original_l2, - normalized, + quantization, + data, } } + /// Score against another vector. Both vectors must share the same + /// quantization: callers are expected to convert the query with + /// `to_quantized` first (quantization is a per-set property). pub fn score(&self, other: &Self) -> Result { - self.cosine_score(other) + ensure!( + self.quantization == other.quantization, + InvalidArgumentSnafu { + message: format!( + "vector quantization mismatch: {:?} vs {:?}", + self.quantization, other.quantization + ) + } + ); + match (&self.data, &other.data) { + (VectorData::Binary(left), VectorData::Binary(right)) => { + self.hamming_score(left, right) + } + // NOQUANT and Q8 both score on the (dequantized) FP32 form. + // A native INT8 dot product can replace the Q8 path later. + _ => self.cosine_score(other), + } + } + + /// Hamming similarity for binary vectors, mapped to [0, 1]. + /// Padding bits beyond `dimension` are zero on both sides and therefore + /// never count as mismatches. + fn hamming_score(&self, left: &[u8], right: &[u8]) -> Result { + ensure!( + left.len() == right.len(), + InvalidArgumentSnafu { + message: format!( + "binary vector length mismatch: {} vs {}", + left.len(), + right.len() + ) + } + ); + let mismatches = left + .iter() + .zip(right) + .map(|(a, b)| u64::from((a ^ b).count_ones())) + .sum::(); + let dimension = u64::from(self.dimension); + Ok((dimension.saturating_sub(mismatches)) as f64 / dimension as f64) } pub fn cosine_score(&self, other: &Self) -> Result { @@ -137,10 +312,11 @@ impl CanonicalVector { } ); - let dot = self - .normalized + let self_fp32 = self.as_fp32()?; + let other_fp32 = other.as_fp32()?; + let dot = self_fp32 .iter() - .zip(&other.normalized) + .zip(&other_fp32) .map(|(left, right)| f64::from(*left) * f64::from(*right)) .sum::() .clamp(-1.0, 1.0); @@ -148,7 +324,8 @@ impl CanonicalVector { } pub fn restore(&self) -> Vec { - self.normalized + self.as_fp32() + .expect("failed to restore vector") .iter() .map(|value| f64::from(*value) * f64::from(self.original_l2)) .collect() @@ -298,4 +475,40 @@ mod tests { assert!((x.score(&same).expect("score") - 1.0).abs() < 1e-12); assert!(x.score(&opposite).expect("score").abs() < 1e-12); } + + #[test] + fn binary_score_uses_hamming_similarity() { + let x = CanonicalVector::from_values(&[1.0, -1.0, 1.0, -1.0]).expect("valid x"); + let y = CanonicalVector::from_values(&[1.0, 1.0, 1.0, -1.0]).expect("valid y"); + let x_bin = x + .to_quantized(QuantizationType::Binary) + .expect("quantize x"); + let y_bin = y + .to_quantized(QuantizationType::Binary) + .expect("quantize y"); + // Signs differ in one of four components. + assert!((x_bin.score(&y_bin).expect("score") - 0.75).abs() < 1e-12); + // Mixed quantization is rejected instead of panicking. + assert!(x.score(&y_bin).is_err()); + } + + #[test] + fn int8_round_trip_approximately_restores_values() { + let x = CanonicalVector::from_values(&[0.5, -0.25, 0.1, 0.9]).expect("valid x"); + let quantized = x.to_quantized(QuantizationType::Int8).expect("quantize"); + let restored = quantized.as_fp32().expect("dequantize"); + for (restored, original) in restored.iter().zip(x.as_fp32().expect("fp32")) { + assert!((restored - original).abs() < 0.01); + } + } + + #[test] + fn int8_handles_constant_vector() { + let component = 1.0 / 4.0_f32.sqrt(); + let x = CanonicalVector::from_values(&[component; 4]).expect("valid x"); + let quantized = x.to_quantized(QuantizationType::Int8).expect("quantize"); + for value in quantized.as_fp32().expect("dequantize") { + assert!((value - component).abs() < 1e-6); + } + } } From d65d7fa0f18f29513a5a6f4d43f11abaef2389b6 Mon Sep 17 00:00:00 2001 From: happy-v587 Date: Mon, 27 Jul 2026 21:32:46 +0800 Subject: [PATCH 20/32] feat(vector): parse VADD options and split per-command parsing - move parse_vadd/parse_vsim/parse_vemb and their tests from vector/mod.rs into the respective command modules; mod.rs keeps only the shared helpers - parse_vadd becomes a keyword option loop matching the Redis syntax: NOQUANT/Q8/BIN are wired to quantization (default NOQUANT), while CAS/EF/SETATTR/M/REDUCE are rejected with dedicated errors - add Redis syntax comments to vector command arity definitions --- src/cmd/src/vector/mod.rs | 346 +---------------------- src/cmd/src/vector/vadd.rs | 221 ++++++++++++++- src/cmd/src/vector/vcard.rs | 2 +- src/cmd/src/vector/vdim.rs | 2 +- src/cmd/src/vector/vemb.rs | 45 ++- src/cmd/src/vector/vismember.rs | 2 +- src/cmd/src/vector/vrem.rs | 2 +- src/cmd/src/vector/vsim.rs | 204 ++++++++++++- tests/python/test_vector_set_commands.py | 31 +- 9 files changed, 495 insertions(+), 360 deletions(-) diff --git a/src/cmd/src/vector/mod.rs b/src/cmd/src/vector/mod.rs index a5d5dd34..00c82a0c 100644 --- a/src/cmd/src/vector/mod.rs +++ b/src/cmd/src/vector/mod.rs @@ -17,34 +17,16 @@ use client::Client; use resp::RespData; -use storage::{CanonicalVector, VectorQuery, VectorSearchMode, VectorSearchOptions, error::Error}; +use storage::{CanonicalVector, error::Error}; const ERR_INVALID_VECTOR: &str = "ERR invalid vector specification"; const ERR_VECTOR_DIMENSION: &str = "ERR vector dimension mismatch"; -const ERR_DEFAULT_Q8: &str = - "ERR default Q8 quantization is not supported in Phase 1; specify NOQUANT"; -const ERR_VADD_Q8: &str = "ERR VADD option Q8 is not supported yet"; -const ERR_VADD_BIN: &str = "ERR VADD option BIN is not supported yet"; -const ERR_VEMB_RAW: &str = "ERR VEMB option RAW is not supported yet"; const ERR_ELEMENT_NOT_FOUND: &str = "ERR element not found in set"; const ERR_KEY_NOT_FOUND: &str = "ERR key does not exist"; const ERR_STORAGE: &str = "ERR storage error"; type ParseResult = std::result::Result; -#[derive(Debug)] -struct ParsedVAdd { - vector: CanonicalVector, - element: Vec, -} - -#[derive(Debug)] -struct ParsedVSim { - query: VectorQuery, - options: VectorSearchOptions, - with_scores: bool, -} - fn parse_positive_usize(raw: &[u8]) -> Option { let value = std::str::from_utf8(raw).ok()?.parse::().ok()?; (value > 0).then_some(value) @@ -95,81 +77,6 @@ fn parse_direct_vector( } } -fn parse_vadd(argv: &[Vec]) -> ParseResult { - let (vector, element_index) = parse_direct_vector(argv, 2)?; - let element = argv.get(element_index).cloned().ok_or(ERR_INVALID_VECTOR)?; - let quantization = &argv[element_index + 1..]; - match quantization { - [] => Err(ERR_DEFAULT_Q8), - [option] if option.eq_ignore_ascii_case(b"NOQUANT") => Ok(ParsedVAdd { vector, element }), - [option] if option.eq_ignore_ascii_case(b"Q8") => Err(ERR_VADD_Q8), - [option] if option.eq_ignore_ascii_case(b"BIN") => Err(ERR_VADD_BIN), - _ => Err(ERR_INVALID_VECTOR), - } -} - -fn parse_vsim(argv: &[Vec]) -> ParseResult { - let query_kind = argv.get(2).ok_or(ERR_INVALID_VECTOR)?; - let (query, mut option_index) = if query_kind.eq_ignore_ascii_case(b"ELE") { - let element = argv.get(3).cloned().ok_or(ERR_INVALID_VECTOR)?; - (VectorQuery::Element(element), 4) - } else { - let (vector, next) = parse_direct_vector(argv, 2)?; - (VectorQuery::Vector(vector), next) - }; - - let mut count = 10; - let mut mode = VectorSearchMode::Approximate; - let mut with_scores = false; - let mut count_seen = false; - let mut truth_seen = false; - - while option_index < argv.len() { - let option = &argv[option_index]; - if option.eq_ignore_ascii_case(b"WITHSCORES") { - if with_scores { - return Err(ERR_INVALID_VECTOR); - } - with_scores = true; - option_index += 1; - } else if option.eq_ignore_ascii_case(b"COUNT") { - if count_seen { - return Err(ERR_INVALID_VECTOR); - } - count = argv - .get(option_index + 1) - .and_then(|raw| parse_positive_usize(raw)) - .ok_or(ERR_INVALID_VECTOR)?; - count_seen = true; - option_index += 2; - } else if option.eq_ignore_ascii_case(b"TRUTH") { - if truth_seen { - return Err(ERR_INVALID_VECTOR); - } - mode = VectorSearchMode::Truth; - truth_seen = true; - option_index += 1; - } else { - return Err(ERR_INVALID_VECTOR); - } - } - - Ok(ParsedVSim { - query, - options: VectorSearchOptions { count, mode }, - with_scores, - }) -} - -fn parse_vemb(argv: &[Vec]) -> ParseResult> { - let element = argv.get(2).cloned().ok_or(ERR_INVALID_VECTOR)?; - match &argv[3..] { - [] => Ok(element), - [option] if option.eq_ignore_ascii_case(b"RAW") => Err(ERR_VEMB_RAW), - _ => Err(ERR_INVALID_VECTOR), - } -} - pub(crate) fn error_reply(message: impl Into) -> RespData { RespData::Error(message.into().into()) } @@ -260,260 +167,9 @@ pub use vsim::VSimCmd; #[allow(clippy::unwrap_used)] #[cfg(test)] mod tests { - use storage::{VectorQuery, VectorSearchMode}; - use super::*; use crate::Cmd; - fn fp32(values: &[f32]) -> Vec { - values - .iter() - .flat_map(|value| value.to_le_bytes()) - .collect() - } - - #[test] - fn parses_supported_vadd_shapes() { - let blob = fp32(&[3.0, 4.0]); - let parsed = parse_vadd(&[ - b"vadd".to_vec(), - b"key\0raw".to_vec(), - b"FP32".to_vec(), - blob, - b"\0element".to_vec(), - b"NOQUANT".to_vec(), - ]) - .expect("FP32 VADD"); - assert_eq!(parsed.vector.dimension(), 2); - assert_eq!(parsed.element, b"\0element"); - - let parsed = parse_vadd(&[ - b"vadd".to_vec(), - b"key".to_vec(), - b"VALUES".to_vec(), - b"2".to_vec(), - b"3".to_vec(), - b"4".to_vec(), - b"element".to_vec(), - b"noquant".to_vec(), - ]) - .expect("VALUES VADD"); - assert_eq!(parsed.vector.dimension(), 2); - assert_eq!(parsed.element, b"element"); - } - - #[test] - fn rejects_unsupported_or_invalid_vadd_shapes() { - let malformed_fp32 = vec![ - b"vadd".to_vec(), - b"key".to_vec(), - b"FP32".to_vec(), - vec![1, 2, 3], - b"element".to_vec(), - b"NOQUANT".to_vec(), - ]; - assert_eq!(parse_vadd(&malformed_fp32).unwrap_err(), ERR_INVALID_VECTOR); - - let invalid_values = [ - vec![ - b"vadd".to_vec(), - b"key".to_vec(), - b"VALUES".to_vec(), - b"0".to_vec(), - b"element".to_vec(), - b"NOQUANT".to_vec(), - ], - vec![ - b"vadd".to_vec(), - b"key".to_vec(), - b"VALUES".to_vec(), - b"2".to_vec(), - b"1".to_vec(), - b"element".to_vec(), - b"NOQUANT".to_vec(), - ], - vec![ - b"vadd".to_vec(), - b"key".to_vec(), - b"VALUES".to_vec(), - b"1".to_vec(), - b"not-a-float".to_vec(), - b"element".to_vec(), - b"NOQUANT".to_vec(), - ], - ]; - for argv in invalid_values { - assert_eq!(parse_vadd(&argv).unwrap_err(), ERR_INVALID_VECTOR); - } - - let base = vec![ - b"vadd".to_vec(), - b"key".to_vec(), - b"FP32".to_vec(), - fp32(&[1.0]), - b"element".to_vec(), - ]; - assert_eq!(parse_vadd(&base).unwrap_err(), ERR_DEFAULT_Q8); - - let mut q8 = base.clone(); - q8.push(b"Q8".to_vec()); - assert_eq!(parse_vadd(&q8).unwrap_err(), ERR_VADD_Q8); - - let mut bin = base.clone(); - bin.push(b"BIN".to_vec()); - assert_eq!(parse_vadd(&bin).unwrap_err(), ERR_VADD_BIN); - - let mut trailing = base; - trailing.extend([b"NOQUANT".to_vec(), b"extra".to_vec()]); - assert_eq!(parse_vadd(&trailing).unwrap_err(), ERR_INVALID_VECTOR); - } - - #[test] - fn parses_supported_vsim_shapes_and_options() { - let ele = parse_vsim(&[ - b"vsim".to_vec(), - b"key".to_vec(), - b"ELE".to_vec(), - b"\0element".to_vec(), - b"WITHSCORES".to_vec(), - b"COUNT".to_vec(), - b"3".to_vec(), - b"TRUTH".to_vec(), - ]) - .expect("ELE VSIM"); - assert_eq!(ele.query, VectorQuery::Element(b"\0element".to_vec())); - assert_eq!(ele.options.count, 3); - assert_eq!(ele.options.mode, VectorSearchMode::Truth); - assert!(ele.with_scores); - - let direct = parse_vsim(&[ - b"vsim".to_vec(), - b"key".to_vec(), - b"FP32".to_vec(), - fp32(&[1.0, 0.0]), - ]) - .expect("FP32 VSIM"); - assert!(matches!(direct.query, VectorQuery::Vector(_))); - assert_eq!(direct.options.count, 10); - assert_eq!(direct.options.mode, VectorSearchMode::Approximate); - - let values = parse_vsim(&[ - b"vsim".to_vec(), - b"key".to_vec(), - b"VALUES".to_vec(), - b"2".to_vec(), - b"1".to_vec(), - b"0".to_vec(), - b"COUNT".to_vec(), - b"1".to_vec(), - ]) - .expect("VALUES VSIM"); - assert!(matches!(values.query, VectorQuery::Vector(_))); - assert_eq!(values.options.count, 1); - } - - #[test] - fn rejects_invalid_vsim_options_and_vectors() { - let cases = [ - vec![ - b"vsim".to_vec(), - b"key".to_vec(), - b"FP32".to_vec(), - vec![1, 2, 3], - ], - vec![ - b"vsim".to_vec(), - b"key".to_vec(), - b"VALUES".to_vec(), - b"2".to_vec(), - b"1".to_vec(), - ], - vec![ - b"vsim".to_vec(), - b"key".to_vec(), - b"VALUES".to_vec(), - b"1".to_vec(), - b"NaN".to_vec(), - ], - vec![ - b"vsim".to_vec(), - b"key".to_vec(), - b"ELE".to_vec(), - b"member".to_vec(), - b"COUNT".to_vec(), - b"0".to_vec(), - ], - vec![ - b"vsim".to_vec(), - b"key".to_vec(), - b"ELE".to_vec(), - b"member".to_vec(), - b"COUNT".to_vec(), - b"invalid".to_vec(), - ], - vec![ - b"vsim".to_vec(), - b"key".to_vec(), - b"ELE".to_vec(), - b"member".to_vec(), - b"COUNT".to_vec(), - b"1".to_vec(), - b"COUNT".to_vec(), - b"2".to_vec(), - ], - vec![ - b"vsim".to_vec(), - b"key".to_vec(), - b"ELE".to_vec(), - b"member".to_vec(), - b"WITHSCORES".to_vec(), - b"WITHSCORES".to_vec(), - ], - vec![ - b"vsim".to_vec(), - b"key".to_vec(), - b"ELE".to_vec(), - b"member".to_vec(), - b"TRUTH".to_vec(), - b"TRUTH".to_vec(), - ], - vec![ - b"vsim".to_vec(), - b"key".to_vec(), - b"ELE".to_vec(), - b"member".to_vec(), - b"unknown".to_vec(), - ], - ]; - for argv in cases { - assert_eq!(parse_vsim(&argv).unwrap_err(), ERR_INVALID_VECTOR); - } - } - - #[test] - fn rejects_vemb_raw_and_unknown_trailing_options() { - assert_eq!( - parse_vemb(&[ - b"vemb".to_vec(), - b"key".to_vec(), - b"member".to_vec(), - b"RAW".to_vec(), - ]) - .unwrap_err(), - ERR_VEMB_RAW - ); - assert_eq!( - parse_vemb(&[ - b"vemb".to_vec(), - b"key".to_vec(), - b"member".to_vec(), - b"unknown".to_vec(), - ]) - .unwrap_err(), - ERR_INVALID_VECTOR - ); - } - #[test] fn vector_command_metadata_matches_redis_shapes() { assert_eq!(VAddCmd::new().meta().arity, -5); diff --git a/src/cmd/src/vector/vadd.rs b/src/cmd/src/vector/vadd.rs index e4c69ac5..54aac824 100644 --- a/src/cmd/src/vector/vadd.rs +++ b/src/cmd/src/vector/vadd.rs @@ -19,20 +19,83 @@ use std::sync::Arc; use client::Client; use resp::RespData; -use storage::storage::Storage; +use storage::{CanonicalVector, QuantizationType, storage::Storage}; use crate::{AclCategory, Cmd, CmdFlags, CmdMeta, impl_cmd_clone_box, impl_cmd_meta}; -use super::{MissingError, error_reply, parse_vadd, storage_error_reply}; +use super::{ + ERR_INVALID_VECTOR, MissingError, ParseResult, error_reply, parse_direct_vector, + storage_error_reply, +}; + +const ERR_VADD_REDUCE: &str = "ERR VADD option REDUCE is not supported yet"; +const ERR_VADD_CAS: &str = "ERR VADD option CAS is not supported yet"; +const ERR_VADD_EF: &str = "ERR VADD option EF is not supported yet"; +const ERR_VADD_SETATTR: &str = "ERR VADD option SETATTR is not supported yet"; +const ERR_VADD_M: &str = "ERR VADD option M is not supported yet"; crate::define_vector_command!( VAddCmd, "vadd", - -5, + -5, // VADD key (FP32 | VALUES num) vector element [NOQUANT | Q8 | BIN] CmdFlags::WRITE | CmdFlags::FAST, AclCategory::KEYSPACE | AclCategory::WRITE ); +#[derive(Debug)] +struct ParsedVAdd { + vector: CanonicalVector, + element: Vec, +} + +fn parse_vadd(argv: &[Vec]) -> ParseResult { + // REDUCE must precede the vector spec; reject it explicitly until + // random-projection support lands. + if argv + .get(2) + .is_some_and(|arg| arg.eq_ignore_ascii_case(b"REDUCE")) + { + return Err(ERR_VADD_REDUCE); + } + let (vector, element_index) = parse_direct_vector(argv, 2)?; + let element = argv.get(element_index).cloned().ok_or(ERR_INVALID_VECTOR)?; + + let mut quantization = None; + for option in &argv[element_index + 1..] { + if option.eq_ignore_ascii_case(b"NOQUANT") + || option.eq_ignore_ascii_case(b"Q8") + || option.eq_ignore_ascii_case(b"BIN") + { + if quantization.is_some() { + return Err(ERR_INVALID_VECTOR); + } else if option.eq_ignore_ascii_case(b"NOQUANT") { + quantization = Some(QuantizationType::None); + } else if option.eq_ignore_ascii_case(b"Q8") { + quantization = Some(QuantizationType::Int8); + } else { + quantization = Some(QuantizationType::Binary); + } + } else if option.eq_ignore_ascii_case(b"CAS") { + return Err(ERR_VADD_CAS); + } else if option.eq_ignore_ascii_case(b"EF") { + return Err(ERR_VADD_EF); + } else if option.eq_ignore_ascii_case(b"SETATTR") { + return Err(ERR_VADD_SETATTR); + } else if option.eq_ignore_ascii_case(b"M") { + return Err(ERR_VADD_M); + } else { + return Err(ERR_INVALID_VECTOR); + } + } + + // Fold the requested quantization into the vector here; from this point + // on the vector's own quantization is the single source of truth. + let vector = vector + .to_quantized(quantization.unwrap_or(QuantizationType::None)) + .map_err(|_| ERR_INVALID_VECTOR)?; + Ok(ParsedVAdd { vector, element }) +} + impl Cmd for VAddCmd { impl_cmd_meta!(); impl_cmd_clone_box!(); @@ -56,3 +119,155 @@ impl Cmd for VAddCmd { client.set_reply(reply); } } + +#[allow(clippy::unwrap_used)] +#[cfg(test)] +mod tests { + use super::*; + + fn fp32(values: &[f32]) -> Vec { + values + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect() + } + + #[test] + fn parses_supported_vadd_shapes() { + let blob = fp32(&[3.0, 4.0]); + let parsed = parse_vadd(&[ + b"vadd".to_vec(), + b"key\0raw".to_vec(), + b"FP32".to_vec(), + blob, + b"\0element".to_vec(), + b"NOQUANT".to_vec(), + ]) + .expect("FP32 VADD"); + assert_eq!(parsed.vector.dimension(), 2); + assert_eq!(parsed.element, b"\0element"); + assert_eq!(parsed.vector.quantization(), QuantizationType::None); + + let parsed = parse_vadd(&[ + b"vadd".to_vec(), + b"key".to_vec(), + b"VALUES".to_vec(), + b"2".to_vec(), + b"3".to_vec(), + b"4".to_vec(), + b"element".to_vec(), + b"noquant".to_vec(), + ]) + .expect("VALUES VADD"); + assert_eq!(parsed.vector.dimension(), 2); + assert_eq!(parsed.element, b"element"); + + // The quantization option defaults to NOQUANT and accepts Q8/BIN. + let base = vec![ + b"vadd".to_vec(), + b"key".to_vec(), + b"FP32".to_vec(), + fp32(&[1.0]), + b"element".to_vec(), + ]; + assert_eq!( + parse_vadd(&base) + .expect("default VADD") + .vector + .quantization(), + QuantizationType::None + ); + for (option, expected) in [ + (b"Q8".to_vec(), QuantizationType::Int8), + (b"BIN".to_vec(), QuantizationType::Binary), + ] { + let mut argv = base.clone(); + argv.push(option); + assert_eq!( + parse_vadd(&argv) + .expect("quantized VADD") + .vector + .quantization(), + expected + ); + } + } + + #[test] + fn rejects_unsupported_or_invalid_vadd_shapes() { + let malformed_fp32 = vec![ + b"vadd".to_vec(), + b"key".to_vec(), + b"FP32".to_vec(), + vec![1, 2, 3], + b"element".to_vec(), + b"NOQUANT".to_vec(), + ]; + assert_eq!(parse_vadd(&malformed_fp32).unwrap_err(), ERR_INVALID_VECTOR); + + let invalid_values = [ + vec![ + b"vadd".to_vec(), + b"key".to_vec(), + b"VALUES".to_vec(), + b"0".to_vec(), + b"element".to_vec(), + b"NOQUANT".to_vec(), + ], + vec![ + b"vadd".to_vec(), + b"key".to_vec(), + b"VALUES".to_vec(), + b"2".to_vec(), + b"1".to_vec(), + b"element".to_vec(), + b"NOQUANT".to_vec(), + ], + vec![ + b"vadd".to_vec(), + b"key".to_vec(), + b"VALUES".to_vec(), + b"1".to_vec(), + b"not-a-float".to_vec(), + b"element".to_vec(), + b"NOQUANT".to_vec(), + ], + ]; + for argv in invalid_values { + assert_eq!(parse_vadd(&argv).unwrap_err(), ERR_INVALID_VECTOR); + } + + let base = vec![ + b"vadd".to_vec(), + b"key".to_vec(), + b"FP32".to_vec(), + fp32(&[1.0]), + b"element".to_vec(), + ]; + + // Quantization options are mutually exclusive. + let mut duplicated = base.clone(); + duplicated.extend([b"Q8".to_vec(), b"BIN".to_vec()]); + assert_eq!(parse_vadd(&duplicated).unwrap_err(), ERR_INVALID_VECTOR); + + // Recognized but not yet supported options get dedicated errors. + for (option, expected) in [ + (b"CAS".to_vec(), ERR_VADD_CAS), + (b"EF".to_vec(), ERR_VADD_EF), + (b"SETATTR".to_vec(), ERR_VADD_SETATTR), + (b"M".to_vec(), ERR_VADD_M), + ] { + let mut argv = base.clone(); + argv.push(option); + assert_eq!(parse_vadd(&argv).unwrap_err(), expected); + } + + let mut reduce = vec![b"vadd".to_vec(), b"key".to_vec(), b"REDUCE".to_vec()]; + reduce.extend(base.iter().skip(2).cloned()); + assert_eq!(parse_vadd(&reduce).unwrap_err(), ERR_VADD_REDUCE); + + let mut trailing = base; + trailing.extend([b"NOQUANT".to_vec(), b"extra".to_vec()]); + assert_eq!(parse_vadd(&trailing).unwrap_err(), ERR_INVALID_VECTOR); + } +} diff --git a/src/cmd/src/vector/vcard.rs b/src/cmd/src/vector/vcard.rs index ae11347f..3dda37eb 100644 --- a/src/cmd/src/vector/vcard.rs +++ b/src/cmd/src/vector/vcard.rs @@ -27,7 +27,7 @@ use super::{MissingError, integer_reply, storage_error_reply}; crate::define_vector_command!( VCardCmd, "vcard", - 2, + 2, // VCARD key CmdFlags::READONLY | CmdFlags::FAST, AclCategory::KEYSPACE | AclCategory::READ ); diff --git a/src/cmd/src/vector/vdim.rs b/src/cmd/src/vector/vdim.rs index fb1db93b..b67ea4b1 100644 --- a/src/cmd/src/vector/vdim.rs +++ b/src/cmd/src/vector/vdim.rs @@ -27,7 +27,7 @@ use super::{MissingError, integer_reply, storage_error_reply}; crate::define_vector_command!( VDimCmd, "vdim", - 2, + 2, // VDIM key CmdFlags::READONLY | CmdFlags::FAST, AclCategory::KEYSPACE | AclCategory::READ ); diff --git a/src/cmd/src/vector/vemb.rs b/src/cmd/src/vector/vemb.rs index b8715cb4..f001fb4f 100644 --- a/src/cmd/src/vector/vemb.rs +++ b/src/cmd/src/vector/vemb.rs @@ -23,16 +23,27 @@ use storage::storage::Storage; use crate::{AclCategory, Cmd, CmdFlags, CmdMeta, impl_cmd_clone_box, impl_cmd_meta}; -use super::{MissingError, error_reply, parse_vemb, storage_error_reply}; +use super::{ERR_INVALID_VECTOR, MissingError, ParseResult, error_reply, storage_error_reply}; + +const ERR_VEMB_RAW: &str = "ERR VEMB option RAW is not supported yet"; crate::define_vector_command!( VEmbCmd, "vemb", - -3, + -3, // VEMB key element CmdFlags::READONLY | CmdFlags::FAST, AclCategory::KEYSPACE | AclCategory::READ ); +fn parse_vemb(argv: &[Vec]) -> ParseResult> { + let element = argv.get(2).cloned().ok_or(ERR_INVALID_VECTOR)?; + match &argv[3..] { + [] => Ok(element), + [option] if option.eq_ignore_ascii_case(b"RAW") => Err(ERR_VEMB_RAW), + _ => Err(ERR_INVALID_VECTOR), + } +} + impl Cmd for VEmbCmd { impl_cmd_meta!(); impl_cmd_clone_box!(); @@ -59,3 +70,33 @@ impl Cmd for VEmbCmd { client.set_reply(reply); } } + +#[allow(clippy::unwrap_used)] +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn rejects_vemb_raw_and_unknown_trailing_options() { + assert_eq!( + parse_vemb(&[ + b"vemb".to_vec(), + b"key".to_vec(), + b"member".to_vec(), + b"RAW".to_vec(), + ]) + .unwrap_err(), + ERR_VEMB_RAW + ); + assert_eq!( + parse_vemb(&[ + b"vemb".to_vec(), + b"key".to_vec(), + b"member".to_vec(), + b"unknown".to_vec(), + ]) + .unwrap_err(), + ERR_INVALID_VECTOR + ); + } +} diff --git a/src/cmd/src/vector/vismember.rs b/src/cmd/src/vector/vismember.rs index 0457d0ae..e6a609c3 100644 --- a/src/cmd/src/vector/vismember.rs +++ b/src/cmd/src/vector/vismember.rs @@ -28,7 +28,7 @@ use super::{MissingError, storage_error_reply}; crate::define_vector_command!( VIsMemberCmd, "vismember", - 3, + 3, // VISMEMBER key element CmdFlags::READONLY | CmdFlags::FAST, AclCategory::KEYSPACE | AclCategory::READ ); diff --git a/src/cmd/src/vector/vrem.rs b/src/cmd/src/vector/vrem.rs index 1752738f..0d2db6a6 100644 --- a/src/cmd/src/vector/vrem.rs +++ b/src/cmd/src/vector/vrem.rs @@ -28,7 +28,7 @@ use super::{MissingError, storage_error_reply}; crate::define_vector_command!( VRemCmd, "vrem", - 3, + 3, // VREM key element CmdFlags::WRITE | CmdFlags::FAST, AclCategory::KEYSPACE | AclCategory::WRITE ); diff --git a/src/cmd/src/vector/vsim.rs b/src/cmd/src/vector/vsim.rs index a0c39ddb..e58d696e 100644 --- a/src/cmd/src/vector/vsim.rs +++ b/src/cmd/src/vector/vsim.rs @@ -20,20 +20,83 @@ use std::sync::Arc; use bytes::Bytes; use client::Client; use resp::RespData; -use storage::storage::Storage; +use storage::{VectorQuery, VectorSearchMode, VectorSearchOptions, storage::Storage}; use crate::{AclCategory, Cmd, CmdFlags, CmdMeta, impl_cmd_clone_box, impl_cmd_meta}; -use super::{MissingError, error_reply, parse_vsim, storage_error_reply}; +use super::{ + ERR_INVALID_VECTOR, MissingError, ParseResult, error_reply, parse_direct_vector, + parse_positive_usize, storage_error_reply, +}; crate::define_vector_command!( VSimCmd, "vsim", - -4, + -4, // VSIM key (ELE element | FP32 vector | VALUES num vector) [WITHSCORES] [COUNT count] [TRUTH] CmdFlags::READONLY, AclCategory::KEYSPACE | AclCategory::READ | AclCategory::SLOW ); +#[derive(Debug)] +struct ParsedVSim { + query: VectorQuery, + options: VectorSearchOptions, + with_scores: bool, +} + +fn parse_vsim(argv: &[Vec]) -> ParseResult { + let query_kind = argv.get(2).ok_or(ERR_INVALID_VECTOR)?; + let (query, mut option_index) = if query_kind.eq_ignore_ascii_case(b"ELE") { + let element = argv.get(3).cloned().ok_or(ERR_INVALID_VECTOR)?; + (VectorQuery::Element(element), 4) + } else { + let (vector, next) = parse_direct_vector(argv, 2)?; + (VectorQuery::Vector(vector), next) + }; + + let mut count = 10; + let mut mode = VectorSearchMode::Approximate; + let mut with_scores = false; + let mut count_seen = false; + let mut truth_seen = false; + + while option_index < argv.len() { + let option = &argv[option_index]; + if option.eq_ignore_ascii_case(b"WITHSCORES") { + if with_scores { + return Err(ERR_INVALID_VECTOR); + } + with_scores = true; + option_index += 1; + } else if option.eq_ignore_ascii_case(b"COUNT") { + if count_seen { + return Err(ERR_INVALID_VECTOR); + } + count = argv + .get(option_index + 1) + .and_then(|raw| parse_positive_usize(raw)) + .ok_or(ERR_INVALID_VECTOR)?; + count_seen = true; + option_index += 2; + } else if option.eq_ignore_ascii_case(b"TRUTH") { + if truth_seen { + return Err(ERR_INVALID_VECTOR); + } + mode = VectorSearchMode::Truth; + truth_seen = true; + option_index += 1; + } else { + return Err(ERR_INVALID_VECTOR); + } + } + + Ok(ParsedVSim { + query, + options: VectorSearchOptions { count, mode }, + with_scores, + }) +} + impl Cmd for VSimCmd { impl_cmd_meta!(); impl_cmd_clone_box!(); @@ -71,3 +134,138 @@ impl Cmd for VSimCmd { client.set_reply(reply); } } + +#[allow(clippy::unwrap_used)] +#[cfg(test)] +mod tests { + use super::*; + + fn fp32(values: &[f32]) -> Vec { + values + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect() + } + + #[test] + fn parses_supported_vsim_shapes_and_options() { + let ele = parse_vsim(&[ + b"vsim".to_vec(), + b"key".to_vec(), + b"ELE".to_vec(), + b"\0element".to_vec(), + b"WITHSCORES".to_vec(), + b"COUNT".to_vec(), + b"3".to_vec(), + b"TRUTH".to_vec(), + ]) + .expect("ELE VSIM"); + assert_eq!(ele.query, VectorQuery::Element(b"\0element".to_vec())); + assert_eq!(ele.options.count, 3); + assert_eq!(ele.options.mode, VectorSearchMode::Truth); + assert!(ele.with_scores); + + let direct = parse_vsim(&[ + b"vsim".to_vec(), + b"key".to_vec(), + b"FP32".to_vec(), + fp32(&[1.0, 0.0]), + ]) + .expect("FP32 VSIM"); + assert!(matches!(direct.query, VectorQuery::Vector(_))); + assert_eq!(direct.options.count, 10); + assert_eq!(direct.options.mode, VectorSearchMode::Approximate); + + let values = parse_vsim(&[ + b"vsim".to_vec(), + b"key".to_vec(), + b"VALUES".to_vec(), + b"2".to_vec(), + b"1".to_vec(), + b"0".to_vec(), + b"COUNT".to_vec(), + b"1".to_vec(), + ]) + .expect("VALUES VSIM"); + assert!(matches!(values.query, VectorQuery::Vector(_))); + assert_eq!(values.options.count, 1); + } + + #[test] + fn rejects_invalid_vsim_options_and_vectors() { + let cases = [ + vec![ + b"vsim".to_vec(), + b"key".to_vec(), + b"FP32".to_vec(), + vec![1, 2, 3], + ], + vec![ + b"vsim".to_vec(), + b"key".to_vec(), + b"VALUES".to_vec(), + b"2".to_vec(), + b"1".to_vec(), + ], + vec![ + b"vsim".to_vec(), + b"key".to_vec(), + b"VALUES".to_vec(), + b"1".to_vec(), + b"NaN".to_vec(), + ], + vec![ + b"vsim".to_vec(), + b"key".to_vec(), + b"ELE".to_vec(), + b"member".to_vec(), + b"COUNT".to_vec(), + b"0".to_vec(), + ], + vec![ + b"vsim".to_vec(), + b"key".to_vec(), + b"ELE".to_vec(), + b"member".to_vec(), + b"COUNT".to_vec(), + b"invalid".to_vec(), + ], + vec![ + b"vsim".to_vec(), + b"key".to_vec(), + b"ELE".to_vec(), + b"member".to_vec(), + b"COUNT".to_vec(), + b"1".to_vec(), + b"COUNT".to_vec(), + b"2".to_vec(), + ], + vec![ + b"vsim".to_vec(), + b"key".to_vec(), + b"ELE".to_vec(), + b"member".to_vec(), + b"WITHSCORES".to_vec(), + b"WITHSCORES".to_vec(), + ], + vec![ + b"vsim".to_vec(), + b"key".to_vec(), + b"ELE".to_vec(), + b"member".to_vec(), + b"TRUTH".to_vec(), + b"TRUTH".to_vec(), + ], + vec![ + b"vsim".to_vec(), + b"key".to_vec(), + b"ELE".to_vec(), + b"member".to_vec(), + b"unknown".to_vec(), + ], + ]; + for argv in cases { + assert_eq!(parse_vsim(&argv).unwrap_err(), ERR_INVALID_VECTOR); + } + } +} diff --git a/tests/python/test_vector_set_commands.py b/tests/python/test_vector_set_commands.py index c03b59fc..49580280 100644 --- a/tests/python/test_vector_set_commands.py +++ b/tests/python/test_vector_set_commands.py @@ -36,6 +36,8 @@ def vector_client(request): prefix + b"expired", prefix + b"recreated", prefix + b"errors", + prefix + b"errors:q8", + prefix + b"errors:bin", prefix + b"string", prefix + b"malformed", ] @@ -211,9 +213,32 @@ def test_wrongtype_and_unsupported_options(vector_client): ) base = (b"VADD", key, b"VALUES", 2, 1, 0, b"member") - assert_response_error(client, "default Q8 quantization is not supported", *base) - assert_response_error(client, "VADD option Q8 is not supported yet", *base, b"Q8") - assert_response_error(client, "VADD option BIN is not supported yet", *base, b"BIN") + # The quantization option defaults to NOQUANT and Q8/BIN are accepted. + assert client.execute_command(*base) == 1 + assert vadd_values(client, prefix + b"errors:q8", [1, 0], b"member", b"Q8") == 1 + assert vadd_values(client, prefix + b"errors:bin", [1, 0], b"member", b"BIN") == 1 + + # Recognized but not yet supported options get dedicated errors. + for option in (b"CAS", b"EF", b"SETATTR", b"M"): + assert_response_error( + client, + f"VADD option {option.decode()} is not supported yet", + *base, + option, + ) + assert_response_error( + client, + "VADD option REDUCE is not supported yet", + b"VADD", + key, + b"REDUCE", + 2, + b"VALUES", + 2, + 1, + 0, + b"member", + ) vadd_values(client, key, [1, 0], b"member") assert_response_error( From b9edac1eeb0fa14c57f477d6fb9524709ab7a15a Mon Sep 17 00:00:00 2001 From: happy-v587 Date: Tue, 28 Jul 2026 15:48:17 +0800 Subject: [PATCH 21/32] fix(vector): enforce phase one protocol semantics --- src/cmd/src/vector/vadd.rs | 48 +++---- src/cmd/src/vector/vismember.rs | 2 +- src/cmd/src/vector/vrem.rs | 2 +- src/net/tests/storage_command_e2e_tests.rs | 143 +++++++++++++++++++++ tests/python/test_vector_set_commands.py | 68 +++++++--- 5 files changed, 218 insertions(+), 45 deletions(-) diff --git a/src/cmd/src/vector/vadd.rs b/src/cmd/src/vector/vadd.rs index 54aac824..5e01b435 100644 --- a/src/cmd/src/vector/vadd.rs +++ b/src/cmd/src/vector/vadd.rs @@ -29,6 +29,10 @@ use super::{ }; const ERR_VADD_REDUCE: &str = "ERR VADD option REDUCE is not supported yet"; +const ERR_VADD_DEFAULT_QUANTIZATION: &str = + "ERR default Q8 quantization is not supported in Phase 1; specify NOQUANT"; +const ERR_VADD_Q8: &str = "ERR VADD option Q8 is not supported yet"; +const ERR_VADD_BIN: &str = "ERR VADD option BIN is not supported yet"; const ERR_VADD_CAS: &str = "ERR VADD option CAS is not supported yet"; const ERR_VADD_EF: &str = "ERR VADD option EF is not supported yet"; const ERR_VADD_SETATTR: &str = "ERR VADD option SETATTR is not supported yet"; @@ -62,19 +66,16 @@ fn parse_vadd(argv: &[Vec]) -> ParseResult { let mut quantization = None; for option in &argv[element_index + 1..] { - if option.eq_ignore_ascii_case(b"NOQUANT") - || option.eq_ignore_ascii_case(b"Q8") - || option.eq_ignore_ascii_case(b"BIN") - { + if option.eq_ignore_ascii_case(b"NOQUANT") { if quantization.is_some() { return Err(ERR_INVALID_VECTOR); - } else if option.eq_ignore_ascii_case(b"NOQUANT") { - quantization = Some(QuantizationType::None); - } else if option.eq_ignore_ascii_case(b"Q8") { - quantization = Some(QuantizationType::Int8); } else { - quantization = Some(QuantizationType::Binary); + quantization = Some(QuantizationType::None); } + } else if option.eq_ignore_ascii_case(b"Q8") { + return Err(ERR_VADD_Q8); + } else if option.eq_ignore_ascii_case(b"BIN") { + return Err(ERR_VADD_BIN); } else if option.eq_ignore_ascii_case(b"CAS") { return Err(ERR_VADD_CAS); } else if option.eq_ignore_ascii_case(b"EF") { @@ -91,7 +92,7 @@ fn parse_vadd(argv: &[Vec]) -> ParseResult { // Fold the requested quantization into the vector here; from this point // on the vector's own quantization is the single source of truth. let vector = vector - .to_quantized(quantization.unwrap_or(QuantizationType::None)) + .to_quantized(quantization.ok_or(ERR_VADD_DEFAULT_QUANTIZATION)?) .map_err(|_| ERR_INVALID_VECTOR)?; Ok(ParsedVAdd { vector, element }) } @@ -113,7 +114,7 @@ impl Cmd for VAddCmd { } }; let reply = match storage.vadd(&client.key(), &parsed.element, &parsed.vector) { - Ok(inserted) => RespData::Integer(i64::from(inserted)), + Ok(inserted) => RespData::Boolean(inserted), Err(error) => storage_error_reply(error, MissingError::Key), }; client.set_reply(reply); @@ -162,7 +163,7 @@ mod tests { assert_eq!(parsed.vector.dimension(), 2); assert_eq!(parsed.element, b"element"); - // The quantization option defaults to NOQUANT and accepts Q8/BIN. + // Phase 1 requires explicit NOQUANT and does not support Q8/BIN yet. let base = vec![ b"vadd".to_vec(), b"key".to_vec(), @@ -171,25 +172,16 @@ mod tests { b"element".to_vec(), ]; assert_eq!( - parse_vadd(&base) - .expect("default VADD") - .vector - .quantization(), - QuantizationType::None + parse_vadd(&base).unwrap_err(), + "ERR default Q8 quantization is not supported in Phase 1; specify NOQUANT" ); for (option, expected) in [ - (b"Q8".to_vec(), QuantizationType::Int8), - (b"BIN".to_vec(), QuantizationType::Binary), + (b"Q8".to_vec(), "ERR VADD option Q8 is not supported yet"), + (b"BIN".to_vec(), "ERR VADD option BIN is not supported yet"), ] { let mut argv = base.clone(); argv.push(option); - assert_eq!( - parse_vadd(&argv) - .expect("quantized VADD") - .vector - .quantization(), - expected - ); + assert_eq!(parse_vadd(&argv).unwrap_err(), expected); } } @@ -245,10 +237,10 @@ mod tests { b"element".to_vec(), ]; - // Quantization options are mutually exclusive. + // Q8 is rejected before a following quantization option is considered. let mut duplicated = base.clone(); duplicated.extend([b"Q8".to_vec(), b"BIN".to_vec()]); - assert_eq!(parse_vadd(&duplicated).unwrap_err(), ERR_INVALID_VECTOR); + assert_eq!(parse_vadd(&duplicated).unwrap_err(), ERR_VADD_Q8); // Recognized but not yet supported options get dedicated errors. for (option, expected) in [ diff --git a/src/cmd/src/vector/vismember.rs b/src/cmd/src/vector/vismember.rs index e6a609c3..368241a4 100644 --- a/src/cmd/src/vector/vismember.rs +++ b/src/cmd/src/vector/vismember.rs @@ -44,7 +44,7 @@ impl Cmd for VIsMemberCmd { fn do_cmd(&self, client: &Client, storage: Arc) { let argv = client.argv(); let reply = match storage.vismember(&client.key(), &argv[2]) { - Ok(is_member) => RespData::Integer(i64::from(is_member)), + Ok(is_member) => RespData::Boolean(is_member), Err(error) => storage_error_reply(error, MissingError::Key), }; client.set_reply(reply); diff --git a/src/cmd/src/vector/vrem.rs b/src/cmd/src/vector/vrem.rs index 0d2db6a6..91401876 100644 --- a/src/cmd/src/vector/vrem.rs +++ b/src/cmd/src/vector/vrem.rs @@ -44,7 +44,7 @@ impl Cmd for VRemCmd { fn do_cmd(&self, client: &Client, storage: Arc) { let argv = client.argv(); let reply = match storage.vrem(&client.key(), &argv[2]) { - Ok(removed) => RespData::Integer(i64::from(removed)), + Ok(removed) => RespData::Boolean(removed), Err(error) => storage_error_reply(error, MissingError::Key), }; client.set_reply(reply); diff --git a/src/net/tests/storage_command_e2e_tests.rs b/src/net/tests/storage_command_e2e_tests.rs index e8bd88ee..371a3456 100644 --- a/src/net/tests/storage_command_e2e_tests.rs +++ b/src/net/tests/storage_command_e2e_tests.rs @@ -444,6 +444,149 @@ async fn storage_command_e2e_vector_write_commands_redirect_on_follower() { server.shutdown().await; } +#[tokio::test] +async fn storage_command_e2e_vector_membership_replies_follow_negotiated_protocol() { + let server = TestServer::start(None).await; + + let mut resp2_stream = tokio::net::TcpStream::connect(server.addr) + .await + .expect("connect RESP2 client"); + assert_eq!( + send_command( + &mut resp2_stream, + &[ + "VADD", + "resp2-vectors", + "VALUES", + "2", + "1", + "0", + "member", + "NOQUANT" + ], + ) + .await, + RespData::Integer(1) + ); + assert_eq!( + send_command( + &mut resp2_stream, + &[ + "VADD", + "resp2-vectors", + "VALUES", + "2", + "1", + "0", + "member", + "NOQUANT" + ] + ) + .await, + RespData::Integer(0) + ); + assert_eq!( + send_command(&mut resp2_stream, &["VISMEMBER", "resp2-vectors", "member"]).await, + RespData::Integer(1) + ); + assert_eq!( + send_command(&mut resp2_stream, &["VREM", "resp2-vectors", "member"]).await, + RespData::Integer(1) + ); + assert_eq!( + send_command(&mut resp2_stream, &["VREM", "resp2-vectors", "member"]).await, + RespData::Integer(0) + ); + assert_eq!( + send_command(&mut resp2_stream, &["VISMEMBER", "resp2-vectors", "member"]).await, + RespData::Integer(0) + ); + + let mut resp3_stream = tokio::net::TcpStream::connect(server.addr) + .await + .expect("connect RESP3 client"); + let reply = + send_command_with_version(&mut resp3_stream, &["HELLO", "3"], RespVersion::RESP3).await; + assert!( + matches!(reply, RespData::Map(_)), + "expected RESP3 HELLO map, got {reply:?}" + ); + assert_eq!( + send_command_with_version( + &mut resp3_stream, + &[ + "VADD", + "resp3-vectors", + "VALUES", + "2", + "1", + "0", + "member", + "NOQUANT" + ], + RespVersion::RESP3 + ) + .await, + RespData::Boolean(true) + ); + assert_eq!( + send_command_with_version( + &mut resp3_stream, + &[ + "VADD", + "resp3-vectors", + "VALUES", + "2", + "1", + "0", + "member", + "NOQUANT" + ], + RespVersion::RESP3 + ) + .await, + RespData::Boolean(false) + ); + assert_eq!( + send_command_with_version( + &mut resp3_stream, + &["VISMEMBER", "resp3-vectors", "member"], + RespVersion::RESP3 + ) + .await, + RespData::Boolean(true) + ); + assert_eq!( + send_command_with_version( + &mut resp3_stream, + &["VREM", "resp3-vectors", "member"], + RespVersion::RESP3 + ) + .await, + RespData::Boolean(true) + ); + assert_eq!( + send_command_with_version( + &mut resp3_stream, + &["VREM", "resp3-vectors", "member"], + RespVersion::RESP3 + ) + .await, + RespData::Boolean(false) + ); + assert_eq!( + send_command_with_version( + &mut resp3_stream, + &["VISMEMBER", "resp3-vectors", "member"], + RespVersion::RESP3 + ) + .await, + RespData::Boolean(false) + ); + + server.shutdown().await; +} + // Regression for issue #349: GET/MGET must return stored bytes unchanged over RESP. #[tokio::test] async fn storage_command_e2e_get_and_mget_preserve_binary_values() { diff --git a/tests/python/test_vector_set_commands.py b/tests/python/test_vector_set_commands.py index 49580280..711c0668 100644 --- a/tests/python/test_vector_set_commands.py +++ b/tests/python/test_vector_set_commands.py @@ -79,15 +79,28 @@ def assert_vector(values, expected): assert [float(value) for value in values] == pytest.approx(expected) +def assert_membership_reply(reply, protocol, expected): + if protocol == 2: + assert type(reply) is int + assert reply == int(expected) + else: + assert type(reply) is bool + assert reply is expected + + def test_values_create_update_and_point_commands(vector_client): - client, _protocol, prefix = vector_client + client, protocol, prefix = vector_client key = prefix + b"values" - assert vadd_values(client, key, [1, 0], b"member") == 1 - assert vadd_values(client, key, [0.5, 0.5], b"member") == 0 + assert_membership_reply(vadd_values(client, key, [1, 0], b"member"), protocol, True) + assert_membership_reply( + vadd_values(client, key, [0.5, 0.5], b"member"), protocol, False + ) assert client.execute_command(b"VCARD", key) == 1 assert client.execute_command(b"VDIM", key) == 2 - assert client.execute_command(b"VISMEMBER", key, b"member") == 1 + assert_membership_reply( + client.execute_command(b"VISMEMBER", key, b"member"), protocol, True + ) assert_vector(client.execute_command(b"VEMB", key, b"member"), [0.5, 0.5]) assert client.type(key) == b"vectorset" @@ -107,19 +120,31 @@ def test_values_create_update_and_point_commands(vector_client): def test_fp32_binary_members_and_last_member_removal(vector_client): - client, _protocol, prefix = vector_client + client, protocol, prefix = vector_client key = prefix + b"binary:\x00key" blob = struct.pack("<2f", 1.0, 0.0) - assert client.execute_command(b"VADD", key, b"FP32", blob, b"", b"NOQUANT") == 1 - assert vadd_values(client, key, [0, 1], b"\x00member") == 1 + assert_membership_reply( + client.execute_command(b"VADD", key, b"FP32", blob, b"", b"NOQUANT"), + protocol, + True, + ) + assert_membership_reply( + vadd_values(client, key, [0, 1], b"\x00member"), protocol, True + ) assert client.execute_command(b"VCARD", key) == 2 assert_vector(client.execute_command(b"VEMB", key, b""), [1.0, 0.0]) - assert client.execute_command(b"VISMEMBER", key, b"\x00member") == 1 + assert_membership_reply( + client.execute_command(b"VISMEMBER", key, b"\x00member"), protocol, True + ) - assert client.execute_command(b"VREM", key, b"") == 1 - assert client.execute_command(b"VREM", key, b"\x00member") == 1 - assert client.execute_command(b"VREM", key, b"\x00member") == 0 + assert_membership_reply(client.execute_command(b"VREM", key, b""), protocol, True) + assert_membership_reply( + client.execute_command(b"VREM", key, b"\x00member"), protocol, True + ) + assert_membership_reply( + client.execute_command(b"VREM", key, b"\x00member"), protocol, False + ) assert client.type(key) == b"none" @@ -213,10 +238,23 @@ def test_wrongtype_and_unsupported_options(vector_client): ) base = (b"VADD", key, b"VALUES", 2, 1, 0, b"member") - # The quantization option defaults to NOQUANT and Q8/BIN are accepted. - assert client.execute_command(*base) == 1 - assert vadd_values(client, prefix + b"errors:q8", [1, 0], b"member", b"Q8") == 1 - assert vadd_values(client, prefix + b"errors:bin", [1, 0], b"member", b"BIN") == 1 + assert_response_error( + client, + "default Q8 quantization is not supported in Phase 1; specify NOQUANT", + *base, + ) + assert_response_error( + client, + "VADD option Q8 is not supported yet", + *base, + b"Q8", + ) + assert_response_error( + client, + "VADD option BIN is not supported yet", + *base, + b"BIN", + ) # Recognized but not yet supported options get dedicated errors. for option in (b"CAS", b"EF", b"SETATTR", b"M"): From 9f1e37e3cedd1afbd94dac1c662f82dec0bb47ca Mon Sep 17 00:00:00 2001 From: happy-v587 Date: Fri, 31 Jul 2026 09:47:40 +0800 Subject: [PATCH 22/32] feat: extend vector lifecycle and admin capabilities --- Cargo.lock | 2 + src/cmd/src/admin.rs | 30 +- src/cmd/src/table.rs | 379 ++++++- src/cmd/src/vector/mod.rs | 7 + src/cmd/src/vector/vinfo.rs | 125 +++ src/common/runtime/error.rs | 5 + src/common/runtime/lib.rs | 1 + src/common/runtime/storage_server.rs | 15 +- src/conf/src/config.rs | 59 + src/conf/src/lib.rs | 78 ++ src/conf/src/vector_config.rs | 159 +++ src/net/src/executor_ext.rs | 46 +- src/net/src/lib.rs | 19 +- src/net/tests/storage_command_e2e_tests.rs | 116 ++ src/raft/proto/admin.proto | 13 + src/raft/src/capabilities.rs | 226 ++++ src/raft/src/grpc/admin.rs | 16 +- src/raft/src/leader_gate.rs | 15 + src/raft/src/lib.rs | 1 + src/raft/src/node.rs | 18 + src/raft/src/snapshot_archive.rs | 2 +- src/raft/src/state_machine.rs | 32 +- src/raft/tests/snapshot_roundtrip_test.rs | 8 +- src/server/Cargo.toml | 1 + src/server/src/main.rs | 24 +- src/storage/Cargo.toml | 3 + src/storage/src/batch.rs | 8 + src/storage/src/checkpoint.rs | 126 ++- src/storage/src/data_compaction_filter.rs | 292 ++++- src/storage/src/error.rs | 20 + src/storage/src/format_member_data_key.rs | 51 +- src/storage/src/format_vector.rs | 197 +++- src/storage/src/format_vector_member_key.rs | 508 +++++++++ src/storage/src/lib.rs | 22 +- src/storage/src/options.rs | 4 + src/storage/src/redis.rs | 118 +- src/storage/src/redis_strings.rs | 15 +- src/storage/src/redis_vectors.rs | 434 ++++++-- src/storage/src/storage.rs | 16 + src/storage/src/storage_impl.rs | 15 + src/storage/src/storage_manifest.rs | 243 +++++ src/storage/src/vector.rs | 9 + src/storage/src/vector_fault.rs | 81 ++ src/storage/src/vector_flat.rs | 190 ++++ src/storage/src/vector_metrics.rs | 104 ++ src/storage/src/vector_mutation.rs | 525 +++++++++ src/storage/tests/checkpoint_test.rs | 156 ++- src/storage/tests/redis_vector_test.rs | 1029 +++++++++++++++++- tests/python/conftest.py | 18 +- tests/python/test_mset.py | 45 +- tests/python/test_vector_basic.py | 128 +++ tests/python/test_vector_cluster.py | 706 ++++++++++++ tests/python/test_vector_set_commands.py | 73 ++ tests/python/test_vector_set_differential.py | 341 ++++++ tests/run_python_integration.sh | 12 +- 55 files changed, 6635 insertions(+), 251 deletions(-) create mode 100644 src/cmd/src/vector/vinfo.rs create mode 100644 src/conf/src/vector_config.rs create mode 100644 src/raft/src/capabilities.rs create mode 100644 src/storage/src/format_vector_member_key.rs create mode 100644 src/storage/src/storage_manifest.rs create mode 100644 src/storage/src/vector_fault.rs create mode 100644 src/storage/src/vector_flat.rs create mode 100644 src/storage/src/vector_metrics.rs create mode 100644 src/storage/src/vector_mutation.rs create mode 100644 tests/python/test_vector_basic.py create mode 100644 tests/python/test_vector_cluster.py create mode 100644 tests/python/test_vector_set_differential.py diff --git a/Cargo.lock b/Cargo.lock index 867f7bd5..797bc62e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2512,6 +2512,7 @@ name = "server" version = "0.1.0" dependencies = [ "clap", + "cmd", "conf", "env_logger", "log", @@ -2621,6 +2622,7 @@ dependencies = [ "num_cpus", "once_cell", "parking_lot", + "proptest", "rand 0.8.6", "rust-rocksdb", "serde", diff --git a/src/cmd/src/admin.rs b/src/cmd/src/admin.rs index 356c0214..978e7e2d 100644 --- a/src/cmd/src/admin.rs +++ b/src/cmd/src/admin.rs @@ -26,6 +26,29 @@ use storage::storage::Storage; use crate::{AclCategory, Cmd, CmdFlags, CmdMeta, impl_cmd_clone_box, impl_cmd_meta}; +/// INFO VECTOR section: the Phase 1 index kind plus the FLAT query counters. +/// vector_sets / vector_elements are omitted on purpose: counting them needs +/// a full keyspace scan and INFO must stay O(1). +fn vector_section(storage: &Storage) -> String { + let metrics = storage.vector_metrics(); + format!( + "# Vector\r\n\ + index-kind:flat\r\n\ + vector_flat_queries_total:{}\r\n\ + vector_flat_query_timeouts_total:{}\r\n\ + vector_flat_query_errors_total:{}\r\n\ + vector_search_capacity_rejected_total:{}\r\n\ + vector_flat_query_duration_micros_total:{}\r\n\ + vector_flat_query_duration_count:{}\r\n", + metrics.flat_queries_total, + metrics.flat_query_timeouts_total, + metrics.flat_query_errors_total, + metrics.capacity_rejected_total, + metrics.flat_query_duration_micros_total, + metrics.flat_query_duration_count, + ) +} + /// INFO command - Show server information including cluster status #[derive(Clone, Default)] pub struct InfoCmd { @@ -54,7 +77,7 @@ impl Cmd for InfoCmd { true } - fn do_cmd(&self, client: &Client, _storage: Arc) { + fn do_cmd(&self, client: &Client, storage: Arc) { let section = if client.argv().len() > 1 { String::from_utf8_lossy(&client.argv()[1]).to_lowercase() } else { @@ -69,6 +92,9 @@ impl Cmd for InfoCmd { info.push_str("cluster_enabled:0\r\n"); info.push_str("cluster_state:disabled\r\n"); } + "vector" => { + info.push_str(&vector_section(&storage)); + } "server" | "default" => { info.push_str("# Server\r\n"); info.push_str("redis_version:7.0.0\r\n"); @@ -95,6 +121,8 @@ impl Cmd for InfoCmd { info.push_str("\r\n# Cluster\r\n"); info.push_str("cluster_enabled:0\r\n"); info.push_str("cluster_state:disabled\r\n"); + info.push_str("\r\n"); + info.push_str(&vector_section(&storage)); } } _ => { diff --git a/src/cmd/src/table.rs b/src/cmd/src/table.rs index ecb59f6f..599bba0a 100644 --- a/src/cmd/src/table.rs +++ b/src/cmd/src/table.rs @@ -18,11 +18,128 @@ use std::collections::HashMap; use std::sync::Arc; -use crate::Cmd; +use client::Client; +use resp::RespData; +use storage::storage::Storage; + use crate::auth::RequirepassProvider; +use crate::{Cmd, CmdMeta}; pub type CmdTable = HashMap>; +/// Provider returning whether a gated command family is currently allowed. +pub type GateFlagProvider = Arc bool + Send + Sync>; + +/// Feature gates consulted when a command table is built. Gates are evaluated +/// on every command execution, so a table always reflects the injected flags. +#[derive(Clone)] +pub struct CommandTableGates { + /// Whether the Vector Set commands (VADD/VSIM/...) are enabled. + pub vector_enabled: GateFlagProvider, + /// Whether Vector Set commands are allowed given the cluster state: + /// false in cluster mode until the Raft apply-correctness contract (PR0) + /// lands, unless `vector-cluster-enabled` is set. + pub vector_cluster_allowed: GateFlagProvider, + /// Whether FLUSHDB/FLUSHALL are allowed. Disabled in cluster mode unless + /// `cluster-flush-enabled` is set. + pub cluster_flush_allowed: GateFlagProvider, +} + +impl Default for CommandTableGates { + fn default() -> Self { + Self { + vector_enabled: Arc::new(|| true), + vector_cluster_allowed: Arc::new(|| true), + cluster_flush_allowed: Arc::new(|| true), + } + } +} + +impl CommandTableGates { + /// Build gates from static flags (the common case: values come from config). + pub fn from_flags( + vector_enabled: bool, + vector_cluster_allowed: bool, + cluster_flush_allowed: bool, + ) -> Self { + Self { + vector_enabled: Arc::new(move || vector_enabled), + vector_cluster_allowed: Arc::new(move || vector_cluster_allowed), + cluster_flush_allowed: Arc::new(move || cluster_flush_allowed), + } + } +} + +/// Wraps a command with a deterministic pre-execution gate: when the flag +/// provider returns false the command replies with `disabled_error` and the +/// inner command never runs. +#[derive(Clone)] +struct GatedCmd { + inner: Arc, + allowed: GateFlagProvider, + disabled_error: String, +} + +impl Cmd for GatedCmd { + fn meta(&self) -> &CmdMeta { + self.inner.meta() + } + + fn do_initial(&self, client: &Client) -> bool { + if !(self.allowed)() { + client.set_reply(RespData::Error(self.disabled_error.clone().into())); + return false; + } + self.inner.do_initial(client) + } + + fn do_cmd(&self, client: &Client, storage: Arc) { + self.inner.do_cmd(client, storage); + } + + fn clone_box(&self) -> Box { + Box::new(self.clone()) + } +} + +/// Register `cmds` wrapped in a gate that replies `disabled_error` when +/// `allowed` evaluates to false. +fn register_gated_cmds( + cmd_table: &mut CmdTable, + cmds: Vec>, + allowed: &GateFlagProvider, + disabled_error: impl Fn(&CmdMeta) -> String, +) { + for cmd in cmds { + let meta = cmd.meta().clone(); + let gated = GatedCmd { + inner: cmd, + allowed: Arc::clone(allowed), + disabled_error: disabled_error(&meta), + }; + cmd_table.insert(meta.name, Arc::new(gated)); + } +} + +/// Wrap each command in a gate without registering it, so another gate can be +/// layered on top before insertion into the table. +fn wrap_gated_cmds( + cmds: Vec>, + allowed: &GateFlagProvider, + disabled_error: impl Fn(&CmdMeta) -> String, +) -> Vec> { + cmds.into_iter() + .map(|cmd| { + let meta = cmd.meta().clone(); + Arc::new(GatedCmd { + inner: cmd, + allowed: Arc::clone(allowed), + disabled_error: disabled_error(&meta), + }) as Arc + }) + .collect() +} + #[macro_export] macro_rules! register_cmd { ($cmd_table:expr_2021, $($cmd_struct:ty),+ $(,)?) => { @@ -51,6 +168,13 @@ macro_rules! register_group_cmd { } pub fn create_command_table(requirepass_provider: RequirepassProvider) -> CmdTable { + create_command_table_with_gates(requirepass_provider, CommandTableGates::default()) +} + +pub fn create_command_table_with_gates( + requirepass_provider: RequirepassProvider, + gates: CommandTableGates, +) -> CmdTable { let mut cmd_table: CmdTable = HashMap::new(); register_cmd!( @@ -92,8 +216,6 @@ pub fn create_command_table(requirepass_provider: RequirepassProvider) -> CmdTab crate::type_cmd::TypeCmd, crate::keys::KeysCmd, crate::randomkey::RandomkeyCmd, - crate::flushdb::FlushdbCmd, - crate::flushall::FlushallCmd, // Hash commands crate::hset::HSetCmd, crate::hget::HGetCmd, @@ -165,18 +287,51 @@ pub fn create_command_table(requirepass_provider: RequirepassProvider) -> CmdTab crate::zscan::ZscanCmd, crate::zscore::ZscoreCmd, crate::zunionstore::ZunionstoreCmd, - // Vector Set commands - crate::vector::VAddCmd, - crate::vector::VSimCmd, - crate::vector::VRemCmd, - crate::vector::VCardCmd, - crate::vector::VDimCmd, - crate::vector::VEmbCmd, - crate::vector::VIsMemberCmd, // connection commands crate::ping::PingCmd, ); + // FLUSHDB/FLUSHALL are gated: rejected in cluster mode unless + // `cluster-flush-enabled` restores the legacy behavior. + let flush_cmds: Vec> = vec![ + Arc::new(crate::flushdb::FlushdbCmd::new()), + Arc::new(crate::flushall::FlushallCmd::new()), + ]; + register_gated_cmds( + &mut cmd_table, + flush_cmds, + &gates.cluster_flush_allowed, + |meta| { + format!( + "ERR {} is not supported in cluster mode yet", + meta.name.to_uppercase() + ) + }, + ); + + // Vector Set commands are gated behind `vector-enabled`, and additionally + // rejected in cluster mode until the Raft apply-correctness contract + // (PR0) lands: physical binlog replay cannot re-encode member keys with + // the local storage incarnation, so cluster vector writes would be + // unreadable after a leader failover. `vector-cluster-enabled` restores + // the pre-gate behavior for development. + let vector_cmds: Vec> = vec![ + Arc::new(crate::vector::VAddCmd::new()), + Arc::new(crate::vector::VSimCmd::new()), + Arc::new(crate::vector::VRemCmd::new()), + Arc::new(crate::vector::VCardCmd::new()), + Arc::new(crate::vector::VDimCmd::new()), + Arc::new(crate::vector::VEmbCmd::new()), + Arc::new(crate::vector::VInfoCmd::new()), + Arc::new(crate::vector::VIsMemberCmd::new()), + ]; + let vector_cmds = wrap_gated_cmds(vector_cmds, &gates.vector_cluster_allowed, |_| { + "ERR vector commands are not supported in cluster mode yet".to_string() + }); + register_gated_cmds(&mut cmd_table, vector_cmds, &gates.vector_enabled, |_| { + "ERR vector support is disabled (vector-enabled=false)".to_string() + }); + // AuthCmd and HelloCmd require the requirepass provider for authentication. { let auth_cmd = crate::auth::AuthCmd::new(Arc::clone(&requirepass_provider)); @@ -206,7 +361,9 @@ mod tests { use resp::RespData; use storage::storage::Storage; - use super::create_command_table; + use super::{ + CmdTable, CommandTableGates, create_command_table, create_command_table_with_gates, + }; struct TestStream; @@ -221,6 +378,22 @@ mod tests { } } + fn run_command(table: &CmdTable, name: &str, argv: &[Vec]) -> RespData { + let command = table.get(name).expect("command should be registered"); + let client = Client::new(Box::new(TestStream)); + client.set_cmd_name(name.as_bytes()); + client.set_argv(argv); + command.execute(&client, Arc::new(Storage::new(1, 0))); + client.take_reply() + } + + fn error_text(reply: &RespData) -> String { + match reply { + RespData::Error(e) => String::from_utf8_lossy(e).into_owned(), + other => panic!("expected error reply, got {other:?}"), + } + } + #[test] fn hello_command_returns_resp3_handshake() { let table = create_command_table(Arc::new(|| None)); @@ -269,11 +442,191 @@ mod tests { #[test] fn vector_commands_are_registered() { let table = create_command_table(Arc::new(|| None)); - for name in ["vadd", "vsim", "vrem", "vcard", "vdim", "vemb", "vismember"] { + for name in [ + "vadd", + "vsim", + "vrem", + "vcard", + "vdim", + "vemb", + "vinfo", + "vismember", + ] { assert!(table.contains_key(name), "{name} should be registered"); } } + #[test] + fn vector_commands_are_rejected_when_disabled() { + let table = create_command_table_with_gates( + Arc::new(|| None), + CommandTableGates::from_flags(false, true, true), + ); + let argvs: [(&str, Vec>); 8] = [ + ( + "vadd", + vec![ + b"vadd".to_vec(), + b"k".to_vec(), + b"FP32".to_vec(), + vec![0; 4], + b"e".to_vec(), + ], + ), + ( + "vsim", + vec![ + b"vsim".to_vec(), + b"k".to_vec(), + b"FP32".to_vec(), + vec![0; 4], + ], + ), + ("vrem", vec![b"vrem".to_vec(), b"k".to_vec(), b"e".to_vec()]), + ("vcard", vec![b"vcard".to_vec(), b"k".to_vec()]), + ("vdim", vec![b"vdim".to_vec(), b"k".to_vec()]), + ("vemb", vec![b"vemb".to_vec(), b"k".to_vec(), b"e".to_vec()]), + ("vinfo", vec![b"vinfo".to_vec(), b"k".to_vec()]), + ( + "vismember", + vec![b"vismember".to_vec(), b"k".to_vec(), b"e".to_vec()], + ), + ]; + for (name, argv) in argvs { + let reply = run_command(&table, name, &argv); + assert_eq!( + error_text(&reply), + "ERR vector support is disabled (vector-enabled=false)", + "{name} should be rejected when vector-enabled=false" + ); + } + } + + #[test] + fn vector_commands_pass_gate_when_enabled() { + let table = create_command_table_with_gates( + Arc::new(|| None), + CommandTableGates::from_flags(true, true, true), + ); + // Malformed vector spec: parsing fails before storage is touched, so + // reaching this error proves the command passed the gate. + let reply = run_command( + &table, + "vadd", + &[ + b"vadd".to_vec(), + b"k".to_vec(), + b"VALUES".to_vec(), + b"2".to_vec(), + b"1.0".to_vec(), + b"e".to_vec(), + b"NOQUANT".to_vec(), + ], + ); + assert_eq!(error_text(&reply), "ERR invalid vector specification"); + } + + #[test] + fn vector_commands_are_rejected_when_cluster_gate_disallows() { + let table = create_command_table_with_gates( + Arc::new(|| None), + CommandTableGates::from_flags(true, false, true), + ); + let reply = run_command(&table, "vcard", &[b"vcard".to_vec(), b"k".to_vec()]); + assert_eq!( + error_text(&reply), + "ERR vector commands are not supported in cluster mode yet" + ); + let reply = run_command( + &table, + "vadd", + &[ + b"vadd".to_vec(), + b"k".to_vec(), + b"VALUES".to_vec(), + b"1".to_vec(), + b"1.0".to_vec(), + b"e".to_vec(), + b"NOQUANT".to_vec(), + ], + ); + assert_eq!( + error_text(&reply), + "ERR vector commands are not supported in cluster mode yet" + ); + } + + #[test] + fn info_vector_section_reports_flat_index_and_metrics() { + let table = create_command_table(Arc::new(|| None)); + + let reply = run_command(&table, "info", &[b"info".to_vec(), b"vector".to_vec()]); + let RespData::BulkString(Some(body)) = reply else { + panic!("INFO VECTOR must return a bulk string"); + }; + let body = String::from_utf8(body.to_vec()).expect("utf8 info"); + assert!(body.starts_with("# Vector\r\n")); + assert!(body.contains("index-kind:flat\r\n")); + assert!(body.contains("vector_flat_queries_total:0\r\n")); + assert!(body.contains("vector_flat_query_timeouts_total:0\r\n")); + assert!(body.contains("vector_flat_query_errors_total:0\r\n")); + assert!(body.contains("vector_search_capacity_rejected_total:0\r\n")); + assert!(body.contains("vector_flat_query_duration_micros_total:0\r\n")); + assert!(body.contains("vector_flat_query_duration_count:0\r\n")); + + let reply = run_command(&table, "info", &[b"info".to_vec()]); + let RespData::BulkString(Some(body)) = reply else { + panic!("INFO must return a bulk string"); + }; + let body = String::from_utf8(body.to_vec()).expect("utf8 info"); + assert!( + body.contains("# Vector\r\n"), + "full INFO must include the Vector section" + ); + } + + #[test] + fn flush_commands_are_rejected_when_cluster_gate_disallows() { + let table = create_command_table_with_gates( + Arc::new(|| None), + CommandTableGates::from_flags(true, true, false), + ); + let reply = run_command(&table, "flushdb", &[b"flushdb".to_vec()]); + assert_eq!( + error_text(&reply), + "ERR FLUSHDB is not supported in cluster mode yet" + ); + let reply = run_command(&table, "flushall", &[b"flushall".to_vec()]); + assert_eq!( + error_text(&reply), + "ERR FLUSHALL is not supported in cluster mode yet" + ); + } + + #[test] + fn flush_commands_execute_when_gate_allows() { + let table = create_command_table_with_gates( + Arc::new(|| None), + CommandTableGates::from_flags(true, true, true), + ); + let reply = run_command(&table, "flushdb", &[b"flushdb".to_vec()]); + assert!( + matches!(reply, RespData::SimpleString(ref s) if s.as_ref() == b"OK"), + "flushdb should run when the gate allows it, got {reply:?}" + ); + } + + #[test] + fn flush_commands_execute_with_default_gates() { + // Default gates model standalone mode: nothing is blocked. + let table = create_command_table(Arc::new(|| None)); + let reply = run_command(&table, "flushdb", &[b"flushdb".to_vec()]); + assert!( + matches!(reply, RespData::SimpleString(ref s) if s.as_ref() == b"OK"), + "standalone flushdb should be unaffected, got {reply:?}" + ); + } + #[test] fn hello_bare_with_requirepass_returns_noauth() { let table = create_command_table(Arc::new(|| Some("secret".to_string()))); diff --git a/src/cmd/src/vector/mod.rs b/src/cmd/src/vector/mod.rs index 00c82a0c..a12aacbd 100644 --- a/src/cmd/src/vector/mod.rs +++ b/src/cmd/src/vector/mod.rs @@ -90,6 +90,10 @@ pub(crate) enum MissingError { pub(crate) fn storage_error_reply(error: Error, missing: MissingError) -> RespData { match &error { Error::RedisErr { message, .. } => error_reply(message.clone()), + // FLAT governance failures carry a client-ready "ERR ..." display. + Error::VectorFlatQueryTimeout { .. } + | Error::VectorFlatQueryCancelled { .. } + | Error::VectorFlatScanBudgetExceeded { .. } => error_reply(error.to_string()), Error::InvalidArgument { message, .. } if message.contains("dimension mismatch") => { error_reply(ERR_VECTOR_DIMENSION) } @@ -152,6 +156,7 @@ pub mod vadd; pub mod vcard; pub mod vdim; pub mod vemb; +pub mod vinfo; pub mod vismember; pub mod vrem; pub mod vsim; @@ -160,6 +165,7 @@ pub use vadd::VAddCmd; pub use vcard::VCardCmd; pub use vdim::VDimCmd; pub use vemb::VEmbCmd; +pub use vinfo::VInfoCmd; pub use vismember::VIsMemberCmd; pub use vrem::VRemCmd; pub use vsim::VSimCmd; @@ -178,6 +184,7 @@ mod tests { assert_eq!(VCardCmd::new().meta().arity, 2); assert_eq!(VDimCmd::new().meta().arity, 2); assert_eq!(VEmbCmd::new().meta().arity, -3); + assert_eq!(VInfoCmd::new().meta().arity, 2); assert_eq!(VIsMemberCmd::new().meta().arity, 3); } } diff --git a/src/cmd/src/vector/vinfo.rs b/src/cmd/src/vector/vinfo.rs new file mode 100644 index 00000000..87e64cdc --- /dev/null +++ b/src/cmd/src/vector/vinfo.rs @@ -0,0 +1,125 @@ +// Copyright (c) 2024-present, arana-db Community. All rights reserved. +// +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use bytes::Bytes; +use client::Client; +use resp::RespData; +use storage::{VectorInfo, storage::Storage}; + +use crate::{AclCategory, Cmd, CmdFlags, CmdMeta, impl_cmd_clone_box, impl_cmd_meta}; + +use super::{MissingError, integer_reply, storage_error_reply}; + +crate::define_vector_command!( + VInfoCmd, + "vinfo", + 2, // VINFO key + CmdFlags::READONLY | CmdFlags::FAST, + AclCategory::KEYSPACE | AclCategory::READ +); + +/// Phase 1 FLAT reply: nine Redis-compatible fields, no Kiwi-private extras +/// (private diagnostics go to INFO VECTOR). HNSW fields carry their FLAT +/// sentinel values; the encoder degrades the map to a flat array on RESP2. +fn vinfo_reply(info: &VectorInfo) -> RespData { + let field = |name: &'static str, value: RespData| { + ( + RespData::BulkString(Some(Bytes::from_static(name.as_bytes()))), + value, + ) + }; + let integer = |value: u64| integer_reply(value); + RespData::Map(vec![ + field( + "quant-type", + RespData::BulkString(Some(Bytes::from_static(b"fp32"))), + ), + field("hnsw-m", RespData::Integer(0)), + field("vector-dim", integer(u64::from(info.dimension))), + field("projection-input-dim", RespData::Integer(0)), + field("size", integer(info.size)), + field("max-level", RespData::Integer(0)), + field("attributes-count", RespData::Integer(0)), + field("vset-uid", integer(info.generation)), + field("hnsw-max-node-uid", RespData::Integer(0)), + ]) +} + +impl Cmd for VInfoCmd { + impl_cmd_meta!(); + impl_cmd_clone_box!(); + + fn do_initial(&self, client: &Client) -> bool { + super::set_command_key(client) + } + + fn do_cmd(&self, client: &Client, storage: Arc) { + let reply = match storage.vinfo(&client.key()) { + // Redis null array for a missing key (RESP2 *-1, RESP3 _). + Ok(Some(info)) => vinfo_reply(&info), + Ok(None) => RespData::Array(None), + Err(error) => storage_error_reply(error, MissingError::Key), + }; + client.set_reply(reply); + } +} + +#[allow(clippy::unwrap_used)] +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn vinfo_reply_has_nine_phase1_fields() { + let reply = vinfo_reply(&VectorInfo { + dimension: 3, + size: 42, + generation: 7, + }); + let RespData::Map(pairs) = reply else { + panic!("vinfo reply must be a map"); + }; + let fields: Vec<(String, RespData)> = pairs + .into_iter() + .map(|(name, value)| { + let RespData::BulkString(Some(name)) = name else { + panic!("field name must be a bulk string"); + }; + (String::from_utf8(name.to_vec()).expect("utf8 name"), value) + }) + .collect(); + assert_eq!( + fields, + vec![ + ( + "quant-type".to_string(), + RespData::BulkString(Some(Bytes::from_static(b"fp32"))) + ), + ("hnsw-m".to_string(), RespData::Integer(0)), + ("vector-dim".to_string(), RespData::Integer(3)), + ("projection-input-dim".to_string(), RespData::Integer(0)), + ("size".to_string(), RespData::Integer(42)), + ("max-level".to_string(), RespData::Integer(0)), + ("attributes-count".to_string(), RespData::Integer(0)), + ("vset-uid".to_string(), RespData::Integer(7)), + ("hnsw-max-node-uid".to_string(), RespData::Integer(0)), + ] + ); + } +} diff --git a/src/common/runtime/error.rs b/src/common/runtime/error.rs index 0e20b0c5..1a4ee269 100644 --- a/src/common/runtime/error.rs +++ b/src/common/runtime/error.rs @@ -182,6 +182,11 @@ impl DualRuntimeError { StorageError::Unknown { .. } => false, StorageError::OptionNone { .. } => false, StorageError::RedisErr { .. } => true, // Redis protocol errors are typically recoverable + // FLAT query governance: timeouts are transient (retry may + // succeed), cancellation and budget exhaustion are deterministic. + StorageError::VectorFlatQueryTimeout { .. } => true, + StorageError::VectorFlatQueryCancelled { .. } => false, + StorageError::VectorFlatScanBudgetExceeded { .. } => false, StorageError::LogIndex { message, .. } => { let msg = message.to_lowercase(); // Structural/programming errors are not retryable diff --git a/src/common/runtime/lib.rs b/src/common/runtime/lib.rs index 9c672f55..c2e16b3a 100644 --- a/src/common/runtime/lib.rs +++ b/src/common/runtime/lib.rs @@ -67,4 +67,5 @@ pub use storage_server::{ BackgroundTaskConfig, BackgroundTaskManager, BackgroundTaskStats, BatchConfig, BatchProcessor, BatchStats, RocksDbStats, StorageAccessGate, StorageAccessPermit, StorageServer, StorageServerConfig, StorageServerPauseController, initialize_storage_command_table, + initialize_storage_command_table_with_gates, }; diff --git a/src/common/runtime/storage_server.rs b/src/common/runtime/storage_server.rs index 2dbcd42f..52545df4 100644 --- a/src/common/runtime/storage_server.rs +++ b/src/common/runtime/storage_server.rs @@ -24,7 +24,7 @@ use std::time::Instant; use client::{Client, StreamTrait}; use cmd::auth::RequirepassProvider; -use cmd::table::{CmdTable, create_command_table}; +use cmd::table::{CmdTable, CommandTableGates, create_command_table_with_gates}; use log::{debug, error, info, warn}; use tokio::sync::mpsc; @@ -169,7 +169,18 @@ impl Drop for StorageAccessPermit { /// used by the network runtime, so AUTH validates against the configured /// `requirepass` regardless of which runtime executes the command. pub fn initialize_storage_command_table(requirepass_provider: RequirepassProvider) { - let _ = STORAGE_COMMAND_TABLE.get_or_init(|| create_command_table(requirepass_provider)); + initialize_storage_command_table_with_gates(requirepass_provider, CommandTableGates::default()); +} + +/// Same as [`initialize_storage_command_table`], additionally applying feature +/// gates (vector commands, cluster FLUSHDB/FLUSHALL) consistent with the +/// network-runtime command table. +pub fn initialize_storage_command_table_with_gates( + requirepass_provider: RequirepassProvider, + gates: CommandTableGates, +) { + let _ = STORAGE_COMMAND_TABLE + .get_or_init(|| create_command_table_with_gates(requirepass_provider, gates)); } struct RuntimeCommandStream; diff --git a/src/conf/src/config.rs b/src/conf/src/config.rs index 562a09d3..59c1119a 100644 --- a/src/conf/src/config.rs +++ b/src/conf/src/config.rs @@ -23,6 +23,7 @@ use validator::Validate; use crate::de_func::{parse_bool_from_string, parse_memory, parse_redis_config}; use crate::error::Error; use crate::runtime_config::RuntimeConfig; +use crate::vector_config::VectorConfig; /// Compression algorithm for RocksDB column families. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] @@ -127,6 +128,15 @@ pub struct Config { /// Dual-runtime configuration. #[serde(default)] pub runtime: RuntimeConfig, + + /// Vector Set feature configuration. + #[serde(default)] + pub vector: VectorConfig, + + /// Allow FLUSHDB/FLUSHALL in cluster mode. Defaults to false: both commands + /// are deterministically rejected while `raft` is configured. + #[serde(default)] + pub cluster_flush_enabled: bool, } impl std::fmt::Debug for Config { @@ -204,6 +214,8 @@ impl std::fmt::Debug for Config { ) .field("raft", &self.raft) .field("runtime", &self.runtime) + .field("vector", &self.vector) + .field("cluster_flush_enabled", &self.cluster_flush_enabled) .finish() } } @@ -269,6 +281,8 @@ impl Default for Config { requirepass: None, raft: None, runtime: RuntimeConfig::default(), + vector: VectorConfig::default(), + cluster_flush_enabled: false, } } } @@ -294,6 +308,7 @@ fn validate_loaded_config(config: &Config) -> Result<(), Error> { .validate() .map_err(|e| Error::ValidConfigFail { source: e })?; config.runtime.validate().map_err(invalid_config)?; + config.vector.validate().map_err(invalid_config)?; if let Some(raft) = config.raft.as_ref() { raft.validate() .map_err(|e| Error::ValidConfigFail { source: e })?; @@ -677,6 +692,50 @@ impl Config { "runtime-fault-injection-log-events" | "runtime-fault_injection-log_events" => { config.runtime.fault_injection.log_events = parse_bool_value(&key, &value)?; } + "vector-enabled" => { + config.vector.enabled = parse_bool_value(&key, &value)?; + } + "vector-cluster-enabled" => { + config.vector.cluster_enabled = parse_bool_value(&key, &value)?; + } + "vector-max-dimension" => { + config.vector.max_dimension = value + .parse() + .map_err(|e| invalid_config(format!("Invalid {}: {}", key, e)))?; + } + "vector-max-k" => { + config.vector.max_k = parse_usize_value(&key, &value)?; + } + "vector-max-element-bytes" => { + config.vector.max_element_bytes = parse_usize_value(&key, &value)?; + } + "vector-max-vector-bytes" => { + config.vector.max_vector_bytes = parse_usize_value(&key, &value)?; + } + "vector-max-concurrent-flat-queries" => { + config.vector.max_concurrent_flat_queries = parse_usize_value(&key, &value)?; + } + "vector-flat-query-timeout-ms" => { + config.vector.flat_query_timeout_ms = value + .parse() + .map_err(|e| invalid_config(format!("Invalid {}: {}", key, e)))?; + } + "vector-flat-cancel-check-interval" => { + config.vector.flat_cancel_check_interval = parse_usize_value(&key, &value)?; + } + "vector-flat-scan-max-entries" => { + config.vector.flat_scan_max_entries = value + .parse() + .map_err(|e| invalid_config(format!("Invalid {}: {}", key, e)))?; + } + "vector-flat-scan-max-bytes" => { + config.vector.flat_scan_max_bytes = value + .parse() + .map_err(|e| invalid_config(format!("Invalid {}: {}", key, e)))?; + } + "cluster-flush-enabled" => { + config.cluster_flush_enabled = parse_bool_value(&key, &value)?; + } _ => { log::warn!("unknown config key: {}", key); continue; diff --git a/src/conf/src/lib.rs b/src/conf/src/lib.rs index af556f22..6182732e 100644 --- a/src/conf/src/lib.rs +++ b/src/conf/src/lib.rs @@ -19,6 +19,7 @@ pub mod de_func; pub mod error; pub mod raft_type; pub mod runtime_config; +pub mod vector_config; #[allow(clippy::unwrap_used)] #[cfg(test)] @@ -117,6 +118,8 @@ mod tests { requirepass: None, raft: None, runtime: runtime_config::RuntimeConfig::default(), + vector: vector_config::VectorConfig::default(), + cluster_flush_enabled: false, }; assert!(invalid_config.validate().is_err()); @@ -348,4 +351,79 @@ mod tests { "raft-node-id 0 should be rejected by validation" ); } + + #[test] + fn test_vector_config_defaults() { + use std::io::Write; + + let mut config_file = tempfile::NamedTempFile::new().unwrap(); + writeln!(config_file, "port 7379").unwrap(); + + let loaded = Config::load(config_file.path().to_str().unwrap()).unwrap(); + assert!(loaded.vector.enabled); + assert_eq!(4096, loaded.vector.max_dimension); + assert_eq!(1000, loaded.vector.max_k); + assert_eq!(1048576, loaded.vector.max_element_bytes); + assert_eq!(16777216, loaded.vector.max_vector_bytes); + assert_eq!(4, loaded.vector.max_concurrent_flat_queries); + assert_eq!(5000, loaded.vector.flat_query_timeout_ms); + assert_eq!(256, loaded.vector.flat_cancel_check_interval); + assert_eq!(1000000, loaded.vector.flat_scan_max_entries); + assert_eq!(1073741824, loaded.vector.flat_scan_max_bytes); + assert!(!loaded.cluster_flush_enabled); + } + + #[test] + fn test_vector_config_parsing() { + use std::io::Write; + + let mut config_file = tempfile::NamedTempFile::new().unwrap(); + writeln!(config_file, "port 7379").unwrap(); + writeln!(config_file, "vector-enabled no").unwrap(); + writeln!(config_file, "vector-max-dimension 1024").unwrap(); + writeln!(config_file, "vector-max-k 100").unwrap(); + writeln!(config_file, "vector-max-element-bytes 4096").unwrap(); + writeln!(config_file, "vector-max-vector-bytes 65536").unwrap(); + writeln!(config_file, "vector-max-concurrent-flat-queries 8").unwrap(); + writeln!(config_file, "vector-flat-query-timeout-ms 1000").unwrap(); + writeln!(config_file, "vector-flat-cancel-check-interval 64").unwrap(); + writeln!(config_file, "vector-flat-scan-max-entries 500").unwrap(); + writeln!(config_file, "vector-flat-scan-max-bytes 1048576").unwrap(); + writeln!(config_file, "cluster-flush-enabled yes").unwrap(); + + let loaded = Config::load(config_file.path().to_str().unwrap()).unwrap(); + assert!(!loaded.vector.enabled); + assert_eq!(1024, loaded.vector.max_dimension); + assert_eq!(100, loaded.vector.max_k); + assert_eq!(4096, loaded.vector.max_element_bytes); + assert_eq!(65536, loaded.vector.max_vector_bytes); + assert_eq!(8, loaded.vector.max_concurrent_flat_queries); + assert_eq!(1000, loaded.vector.flat_query_timeout_ms); + assert_eq!(64, loaded.vector.flat_cancel_check_interval); + assert_eq!(500, loaded.vector.flat_scan_max_entries); + assert_eq!(1048576, loaded.vector.flat_scan_max_bytes); + assert!(loaded.cluster_flush_enabled); + } + + #[test] + fn test_vector_config_validation_is_applied() { + use std::io::Write; + + let mut config_file = tempfile::NamedTempFile::new().unwrap(); + writeln!(config_file, "port 7379").unwrap(); + writeln!(config_file, "vector-max-dimension 70000").unwrap(); + + let loaded = Config::load(config_file.path().to_str().unwrap()); + assert!( + loaded.is_err(), + "vector validation should reject max_dimension above the hard limit" + ); + + let mut config_file = tempfile::NamedTempFile::new().unwrap(); + writeln!(config_file, "port 7379").unwrap(); + writeln!(config_file, "vector-max-k 0").unwrap(); + + let loaded = Config::load(config_file.path().to_str().unwrap()); + assert!(loaded.is_err(), "vector validation should reject max_k 0"); + } } diff --git a/src/conf/src/vector_config.rs b/src/conf/src/vector_config.rs new file mode 100644 index 00000000..67f95d12 --- /dev/null +++ b/src/conf/src/vector_config.rs @@ -0,0 +1,159 @@ +// Copyright (c) 2024-present, arana-db Community. All rights reserved. +// +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use serde::{Deserialize, Serialize}; + +/// Hard upper bound for `max_dimension`. +pub const MAX_VECTOR_DIMENSION: u32 = 65536; + +/// Configuration for the Vector Set feature (VADD/VSIM/... commands). +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct VectorConfig { + /// Whether vector commands are enabled at all + pub enabled: bool, + /// Allow vector commands in cluster mode. Defaults to false: vector + /// commands are deterministically rejected while `raft` is configured, + /// until the Raft apply-correctness contract (PR0) lands. + pub cluster_enabled: bool, + /// Maximum accepted vector dimension + pub max_dimension: u32, + /// Maximum number of neighbors a query may request + pub max_k: usize, + /// Maximum byte size of a single element payload + pub max_element_bytes: usize, + /// Maximum byte size of a single vector blob + pub max_vector_bytes: usize, + /// Maximum number of flat (brute-force) queries running concurrently + pub max_concurrent_flat_queries: usize, + /// Timeout for a single flat query in milliseconds + pub flat_query_timeout_ms: u64, + /// How many scanned entries between cancellation checks + pub flat_cancel_check_interval: usize, + /// Maximum entries a flat scan may visit + pub flat_scan_max_entries: u64, + /// Maximum bytes a flat scan may read + pub flat_scan_max_bytes: u64, +} + +impl Default for VectorConfig { + fn default() -> Self { + Self { + enabled: true, + cluster_enabled: false, + max_dimension: 4096, + max_k: 1000, + max_element_bytes: 1048576, + max_vector_bytes: 16777216, + max_concurrent_flat_queries: 4, + flat_query_timeout_ms: 5000, + flat_cancel_check_interval: 256, + flat_scan_max_entries: 1000000, + flat_scan_max_bytes: 1073741824, + } + } +} + +impl VectorConfig { + /// Validate the configuration parameters + pub fn validate(&self) -> Result<(), String> { + if self.max_dimension == 0 { + return Err("vector.max_dimension must be greater than 0".to_string()); + } + + if self.max_dimension > MAX_VECTOR_DIMENSION { + return Err(format!( + "vector.max_dimension must be <= {MAX_VECTOR_DIMENSION}" + )); + } + + if self.max_k == 0 { + return Err("vector.max_k must be greater than 0".to_string()); + } + + if self.max_concurrent_flat_queries == 0 { + return Err("vector.max_concurrent_flat_queries must be greater than 0".to_string()); + } + + if self.flat_cancel_check_interval == 0 { + return Err("vector.flat_cancel_check_interval must be greater than 0".to_string()); + } + + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn default_config_is_valid() { + let config = VectorConfig::default(); + assert!(config.validate().is_ok()); + assert!(config.enabled); + assert_eq!(4096, config.max_dimension); + assert_eq!(1000, config.max_k); + assert_eq!(1048576, config.max_element_bytes); + assert_eq!(16777216, config.max_vector_bytes); + assert_eq!(4, config.max_concurrent_flat_queries); + assert_eq!(5000, config.flat_query_timeout_ms); + assert_eq!(256, config.flat_cancel_check_interval); + assert_eq!(1000000, config.flat_scan_max_entries); + assert_eq!(1073741824, config.flat_scan_max_bytes); + } + + #[test] + fn max_dimension_bounds_are_enforced() { + let mut config = VectorConfig::default(); + + config.max_dimension = 0; + assert!(config.validate().is_err()); + + config.max_dimension = MAX_VECTOR_DIMENSION + 1; + assert!(config.validate().is_err()); + + config.max_dimension = MAX_VECTOR_DIMENSION; + assert!(config.validate().is_ok()); + } + + #[test] + fn zero_max_k_is_rejected() { + let config = VectorConfig { + max_k: 0, + ..Default::default() + }; + assert!(config.validate().is_err()); + } + + #[test] + fn zero_concurrent_flat_queries_is_rejected() { + let config = VectorConfig { + max_concurrent_flat_queries: 0, + ..Default::default() + }; + assert!(config.validate().is_err()); + } + + #[test] + fn zero_cancel_check_interval_is_rejected() { + let config = VectorConfig { + flat_cancel_check_interval: 0, + ..Default::default() + }; + assert!(config.validate().is_err()); + } +} diff --git a/src/net/src/executor_ext.rs b/src/net/src/executor_ext.rs index f3ab5d46..47796f6c 100644 --- a/src/net/src/executor_ext.rs +++ b/src/net/src/executor_ext.rs @@ -34,6 +34,12 @@ use storage::storage::Storage; use crate::network_execution::NetworkCmdExecution; +/// RESP vector read commands that require a leader-linearizable read in +/// cluster mode. Phase 1 serves them only on the leader: followers redirect +/// exactly like writes, and the leader passes a linearizable-read barrier +/// before dispatch. +const VECTOR_READ_COMMANDS: &[&str] = &["vsim", "vcard", "vdim", "vemb", "vismember", "vinfo"]; + /// Extension trait for CmdExecutor to support network operations pub trait CmdExecutorNetworkExt { /// Execute a network command using StorageClient for dual runtime architecture @@ -66,17 +72,29 @@ impl CmdExecutorNetworkExt for CmdExecutor { && exec.cmd.has_flag(CmdFlags::WRITE) && !gate.is_leader() { - // Simplified redirect: Kiwi returns "MOVED " (no hash slot, - // unlike Redis Cluster's "MOVED "). Clients are - // expected to reconnect to the returned leader address directly. - let reply = match gate.leader_resp_addr() { - Some(addr) => format!("MOVED {addr}"), - None => "ERR not leader".to_string(), - }; - exec.client.set_reply(RespData::Error(reply.into())); + exec.client + .set_reply(RespData::Error(not_leader_reply(gate.as_ref()).into())); return Ok(()); } + // Cluster-mode vector reads are leader-linearizable: followers + // redirect like writes, and the leader must pass the read barrier + // before the command is dispatched to storage. + if let Some(gate) = exec.leader_gate.as_ref() + && VECTOR_READ_COMMANDS.contains(&cmd_name.as_str()) + { + if !gate.is_leader() { + exec.client + .set_reply(RespData::Error(not_leader_reply(gate.as_ref()).into())); + return Ok(()); + } + if let Err(message) = gate.ensure_linearizable_read().await { + exec.client + .set_reply(RespData::Error(format!("ERR {message}").into())); + return Ok(()); + } + } + // Execute do_initial if needed if !exec.cmd.do_initial(&exec.client) { debug!("Command initial check failed for: {}", cmd_name); @@ -136,6 +154,18 @@ async fn execute_generic_command(exec: &NetworkCmdExecution) -> Result<(), DualR Ok(()) } +/// Build the redirect reply for a command that must run on the leader. +/// +/// Simplified redirect: Kiwi returns "MOVED " (no hash slot, unlike +/// Redis Cluster's "MOVED "). Clients are expected to +/// reconnect to the returned leader address directly. +fn not_leader_reply(gate: &dyn raft::leader_gate::LeaderGate) -> String { + match gate.leader_resp_addr() { + Some(addr) => format!("MOVED {addr}"), + None => "ERR not leader".to_string(), + } +} + /// Format storage error for RESP response fn format_storage_error(command: &str, error: &DualRuntimeError) -> String { match error { diff --git a/src/net/src/lib.rs b/src/net/src/lib.rs index 68a66fbe..82baac25 100644 --- a/src/net/src/lib.rs +++ b/src/net/src/lib.rs @@ -39,7 +39,7 @@ use async_trait::async_trait; use crate::network_server::NetworkServer; use crate::storage_client::StorageClient; use crate::tcp::TcpServer; -use cmd::table::create_command_table; +use cmd::table::{CommandTableGates, create_command_table_with_gates}; use executor::CmdExecutorBuilder; use runtime::RuntimeManager; use std::sync::Arc; @@ -58,10 +58,17 @@ impl ServerFactory { runtime_manager: &RuntimeManager, requirepass: Option, leader_gate: Option>, + gates: CommandTableGates, ) -> Option> { match protocol.to_lowercase().as_str() { "tcp" => { - match Self::create_network_server(addr, runtime_manager, requirepass, leader_gate) { + match Self::create_network_server( + addr, + runtime_manager, + requirepass, + leader_gate, + gates, + ) { Ok(server) => Some(Box::new(server) as Box), Err(e) => { log::error!("Failed to create NetworkServer: {}", e); @@ -117,6 +124,7 @@ impl ServerFactory { runtime_manager: &RuntimeManager, requirepass: Option, leader_gate: Option>, + gates: CommandTableGates, ) -> Result> { // Get the storage client from RuntimeManager let runtime_storage_client = runtime_manager.storage_client().map_err(|e| { @@ -131,9 +139,10 @@ impl ServerFactory { // Create command table with requirepass provider let requirepass_for_provider = requirepass.clone(); - let cmd_table = Arc::new(create_command_table(Arc::new(move || { - requirepass_for_provider.clone() - }))); + let cmd_table = Arc::new(create_command_table_with_gates( + Arc::new(move || requirepass_for_provider.clone()), + gates, + )); let executor = Arc::new(CmdExecutorBuilder::new().build()); NetworkServer::new( diff --git a/src/net/tests/storage_command_e2e_tests.rs b/src/net/tests/storage_command_e2e_tests.rs index 371a3456..92a417a5 100644 --- a/src/net/tests/storage_command_e2e_tests.rs +++ b/src/net/tests/storage_command_e2e_tests.rs @@ -26,8 +26,10 @@ //! -> StorageServer::execute_storage_command //! ``` +use std::future::Future; use std::net::SocketAddr; use std::path::PathBuf; +use std::pin::Pin; use std::sync::Arc; use std::time::Duration; @@ -64,6 +66,28 @@ impl LeaderGate for FollowerGate { } } +/// Leader gate with a controllable linearizable-read barrier outcome. +struct LeaderBarrierGate { + barrier_result: Result<(), String>, +} + +impl LeaderGate for LeaderBarrierGate { + fn is_leader(&self) -> bool { + true + } + + fn leader_resp_addr(&self) -> Option { + None + } + + fn ensure_linearizable_read( + &self, + ) -> Pin> + Send + '_>> { + let result = self.barrier_result.clone(); + Box::pin(async move { result }) + } +} + impl TestServer { /// Start the network and storage runtimes, open a real storage DB, and bind /// a `NetworkServer` to an ephemeral port. @@ -699,3 +723,95 @@ async fn storage_command_e2e_rank_nulls_follow_negotiated_wire_protocol() { server.shutdown().await; } + +#[tokio::test] +async fn storage_command_e2e_vector_read_commands_redirect_on_follower() { + let server = TestServer::start_with_leader_gate(None, Some(Arc::new(FollowerGate))).await; + let mut stream = tokio::net::TcpStream::connect(server.addr) + .await + .expect("connect to server"); + let expected = RespData::Error(Bytes::from_static(b"MOVED 127.0.0.1:7380")); + + for args in [ + &["VCARD", "vectors"][..], + &["VDIM", "vectors"][..], + &["VEMB", "vectors", "member"][..], + &["VISMEMBER", "vectors", "member"][..], + &["VINFO", "vectors"][..], + &["VSIM", "vectors", "VALUES", "2", "1", "0"][..], + ] { + let reply = send_command(&mut stream, args).await; + assert_eq!(reply, expected, "vector read {args:?} must redirect"); + } + + // Non-vector reads are unaffected by the leader gate. + let reply = send_command(&mut stream, &["GET", "missing"]).await; + assert!( + !matches!(reply, RespData::Error(_)), + "GET must not be gated, got {reply:?}" + ); + + server.shutdown().await; +} + +#[tokio::test] +async fn storage_command_e2e_vector_read_barrier_failure_returns_err() { + let gate = LeaderBarrierGate { + barrier_result: Err("injected barrier failure".to_string()), + }; + let server = TestServer::start_with_leader_gate(None, Some(Arc::new(gate))).await; + let mut stream = tokio::net::TcpStream::connect(server.addr) + .await + .expect("connect to server"); + + let reply = send_command(&mut stream, &["VCARD", "vectors"]).await; + assert_eq!( + reply, + RespData::Error(Bytes::from_static(b"ERR injected barrier failure")) + ); + + // Non-vector reads skip the barrier entirely. + let reply = send_command(&mut stream, &["GET", "missing"]).await; + assert!( + !matches!(reply, RespData::Error(_)), + "GET must not hit the barrier, got {reply:?}" + ); + + server.shutdown().await; +} + +#[tokio::test] +async fn storage_command_e2e_vector_read_barrier_success_dispatches() { + let gate = LeaderBarrierGate { + barrier_result: Ok(()), + }; + let server = TestServer::start_with_leader_gate(None, Some(Arc::new(gate))).await; + let mut stream = tokio::net::TcpStream::connect(server.addr) + .await + .expect("connect to server"); + + assert_eq!( + send_command(&mut stream, &["VCARD", "vectors"]).await, + RespData::Integer(0) + ); + assert_eq!( + send_command( + &mut stream, + &[ + "VADD", "vectors", "VALUES", "2", "1", "0", "member", "NOQUANT" + ], + ) + .await, + RespData::Integer(1) + ); + assert_eq!( + send_command(&mut stream, &["VCARD", "vectors"]).await, + RespData::Integer(1) + ); + assert_eq!( + send_command(&mut stream, &["VISMEMBER", "vectors", "member"]).await, + RespData::Integer(1) + ); + + server.shutdown().await; +} diff --git a/src/raft/proto/admin.proto b/src/raft/proto/admin.proto index a42b4b13..5a1773c0 100644 --- a/src/raft/proto/admin.proto +++ b/src/raft/proto/admin.proto @@ -39,6 +39,9 @@ service RaftAdminService { // 移除节点 rpc RemoveNode (RemoveNodeRequest) returns (RemoveNodeResponse); + + // 查询节点 capability 清单（滚动升级通告） + rpc GetNodeCapabilities (GetNodeCapabilitiesRequest) returns (GetNodeCapabilitiesResponse); } // Initialize - 初始化集群 @@ -79,3 +82,13 @@ message RemoveNodeRequest { message RemoveNodeResponse { Response response = 1; } + +// GetNodeCapabilities - 查询本节点支持的 capability 清单 +message GetNodeCapabilitiesRequest { +} + +message GetNodeCapabilitiesResponse { + Response response = 1; + uint64 node_id = 2; + repeated string capabilities = 3; +} diff --git a/src/raft/src/capabilities.rs b/src/raft/src/capabilities.rs new file mode 100644 index 00000000..887eca08 --- /dev/null +++ b/src/raft/src/capabilities.rs @@ -0,0 +1,226 @@ +// Copyright (c) 2024-present, arana-db Community. All rights reserved. +// +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Node capability advertisement and cluster-wide capability checks for +//! rolling upgrades (Discussion #331 §19). +//! +//! Rolling upgrades deploy in two phases: first a binary that supports the new +//! vector-set storage / Raft mutation / snapshot schema paths runs with +//! `vector.enabled=false`; the feature is enabled only after every voting +//! member advertises all [`REQUIRED_VECTOR_SET_CAPABILITIES`]. Each node +//! reports its own list through the `GetNodeCapabilities` admin RPC; +//! [`check_cluster_capabilities`] queries a given member address list and +//! summarizes whether the whole set is ready. + +use std::fmt; + +use crate::raft_proto::GetNodeCapabilitiesRequest; +use crate::raft_proto::raft_admin_service_client::RaftAdminServiceClient; + +/// Vector set data can be stored in VectorDataCF with the v1 codec. +pub const CAP_VECTOR_SET_STORAGE_V1: &str = "vector_set_storage_v1"; +/// Vector set mutations can be applied through the Raft binlog path. +pub const CAP_VECTOR_SET_RAFT_MUTATION_V1: &str = "vector_set_raft_mutation_v1"; +/// Raft snapshot metadata carries the v2 storage schema description. +pub const CAP_SNAPSHOT_SCHEMA_V2: &str = "snapshot_schema_v2"; + +/// Capabilities every voting member must advertise before vector sets may be +/// enabled cluster-wide. +pub const REQUIRED_VECTOR_SET_CAPABILITIES: [&str; 3] = [ + CAP_VECTOR_SET_STORAGE_V1, + CAP_VECTOR_SET_RAFT_MUTATION_V1, + CAP_SNAPSHOT_SCHEMA_V2, +]; + +/// Capabilities supported by this binary. +pub fn node_capabilities() -> Vec { + REQUIRED_VECTOR_SET_CAPABILITIES + .iter() + .map(|capability| capability.to_string()) + .collect() +} + +/// Capabilities reported by one cluster member. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct MemberCapabilities { + pub addr: String, + pub capabilities: Vec, +} + +/// Reasons a cluster-wide capability check failed. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CapabilityCheckError { + pub failures: Vec, +} + +impl fmt::Display for CapabilityCheckError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "capability check failed: {}", self.failures.join("; ")) + } +} + +impl std::error::Error for CapabilityCheckError {} + +/// Required capabilities missing from one member's advertised list. +pub fn missing_capabilities(capabilities: &[String]) -> Vec<&'static str> { + REQUIRED_VECTOR_SET_CAPABILITIES + .iter() + .filter(|required| !capabilities.iter().any(|held| held == *required)) + .copied() + .collect() +} + +/// Summarize already-fetched member capabilities: ok only when every member +/// advertises every required capability. +pub fn summarize_capabilities(members: &[MemberCapabilities]) -> Result<(), CapabilityCheckError> { + let failures: Vec = members + .iter() + .filter_map(|member| { + let missing = missing_capabilities(&member.capabilities); + (!missing.is_empty()).then(|| { + format!( + "{} is missing capabilities: {}", + member.addr, + missing.join(", ") + ) + }) + }) + .collect(); + if failures.is_empty() { + Ok(()) + } else { + Err(CapabilityCheckError { failures }) + } +} + +/// Fetch one node's advertised capabilities via the admin gRPC service. +pub async fn fetch_node_capabilities(addr: &str) -> Result, String> { + let endpoint = tonic::transport::Endpoint::from_shared(format!("http://{addr}")) + .map_err(|error| format!("invalid admin address: {error}"))?; + let mut client = RaftAdminServiceClient::connect(endpoint) + .await + .map_err(|error| format!("failed to connect: {error}"))?; + let response = client + .get_node_capabilities(GetNodeCapabilitiesRequest {}) + .await + .map_err(|error| format!("GetNodeCapabilities failed: {error}"))? + .into_inner(); + Ok(response.capabilities) +} + +/// Query every given member address and verify that all of them advertise all +/// [`REQUIRED_VECTOR_SET_CAPABILITIES`]. Returns the fetched per-member lists +/// on success. +pub async fn check_cluster_capabilities( + addrs: &[String], +) -> Result, CapabilityCheckError> { + let mut members = Vec::with_capacity(addrs.len()); + let mut failures = Vec::new(); + for addr in addrs { + match fetch_node_capabilities(addr).await { + Ok(capabilities) => members.push(MemberCapabilities { + addr: addr.clone(), + capabilities, + }), + Err(error) => failures.push(format!("{addr}: {error}")), + } + } + if let Err(summary) = summarize_capabilities(&members) { + failures.extend(summary.failures); + } + if failures.is_empty() { + Ok(members) + } else { + Err(CapabilityCheckError { failures }) + } +} + +#[allow(clippy::unwrap_used)] +#[cfg(test)] +mod tests { + use super::*; + + fn member(addr: &str, capabilities: &[&str]) -> MemberCapabilities { + MemberCapabilities { + addr: addr.to_string(), + capabilities: capabilities.iter().map(|c| c.to_string()).collect(), + } + } + + #[test] + fn node_capabilities_cover_all_required() { + let capabilities = node_capabilities(); + assert!(missing_capabilities(&capabilities).is_empty()); + for required in REQUIRED_VECTOR_SET_CAPABILITIES { + assert!(capabilities.contains(&required.to_string())); + } + } + + #[test] + fn summarize_accepts_full_cluster() { + let all: Vec = node_capabilities(); + let members = vec![ + MemberCapabilities { + addr: "127.0.0.1:7401".to_string(), + capabilities: all.clone(), + }, + MemberCapabilities { + addr: "127.0.0.1:7402".to_string(), + capabilities: all, + }, + ]; + assert!(summarize_capabilities(&members).is_ok()); + } + + #[test] + fn summarize_rejects_member_missing_capability() { + let members = vec![ + MemberCapabilities { + addr: "127.0.0.1:7401".to_string(), + capabilities: node_capabilities(), + }, + member( + "127.0.0.1:7402", + &[CAP_VECTOR_SET_STORAGE_V1, CAP_VECTOR_SET_RAFT_MUTATION_V1], + ), + ]; + let error = summarize_capabilities(&members).unwrap_err(); + assert_eq!(error.failures.len(), 1); + assert!(error.failures[0].contains("127.0.0.1:7402")); + assert!(error.failures[0].contains(CAP_SNAPSHOT_SCHEMA_V2)); + assert!(error.to_string().contains("capability check failed")); + } + + #[test] + fn summarize_rejects_all_members_missing_capabilities() { + let members = vec![ + member("127.0.0.1:7401", &[CAP_VECTOR_SET_STORAGE_V1]), + member("127.0.0.1:7402", &[]), + ]; + let error = summarize_capabilities(&members).unwrap_err(); + assert_eq!(error.failures.len(), 2); + } + + #[test] + fn missing_capabilities_reports_exact_set() { + let missing = missing_capabilities(&[CAP_VECTOR_SET_STORAGE_V1.to_string()]); + assert_eq!( + missing, + vec![CAP_VECTOR_SET_RAFT_MUTATION_V1, CAP_SNAPSHOT_SCHEMA_V2] + ); + } +} diff --git a/src/raft/src/grpc/admin.rs b/src/raft/src/grpc/admin.rs index 30f4479b..dddd6c2c 100644 --- a/src/raft/src/grpc/admin.rs +++ b/src/raft/src/grpc/admin.rs @@ -26,8 +26,8 @@ use std::sync::Arc; // 导入 proto 生成的类型 use crate::raft_proto::{ AddLearnerRequest, AddLearnerResponse, ChangeMembershipRequest, ChangeMembershipResponse, - InitializeRequest, InitializeResponse, RemoveNodeRequest, RemoveNodeResponse, - Response as ProtoResponse, + GetNodeCapabilitiesRequest, GetNodeCapabilitiesResponse, InitializeRequest, InitializeResponse, + RemoveNodeRequest, RemoveNodeResponse, Response as ProtoResponse, raft_admin_service_server::{RaftAdminService, RaftAdminServiceServer}, }; use tonic::{Request, Response as TonicResponse, Status}; @@ -209,6 +209,18 @@ impl RaftAdminService for RaftAdminServiceImpl { } } + /// 查询本节点的 capability 清单（滚动升级通告） + async fn get_node_capabilities( + &self, + _request: Request, + ) -> Result, Status> { + Ok(TonicResponse::new(GetNodeCapabilitiesResponse { + response: Some(ok_response()), + node_id: self.app.node_id, + capabilities: crate::capabilities::node_capabilities(), + })) + } + /// 移除节点 async fn remove_node( &self, diff --git a/src/raft/src/leader_gate.rs b/src/raft/src/leader_gate.rs index 385e90d2..34f9c5d2 100644 --- a/src/raft/src/leader_gate.rs +++ b/src/raft/src/leader_gate.rs @@ -25,6 +25,21 @@ pub trait LeaderGate: Send + Sync { /// Current leader's RESP address for client redirect, if known. fn leader_resp_addr(&self) -> Option; + + /// Linearizable-read barrier for cluster-mode reads served by the leader. + /// + /// Resolves once this node has confirmed its leadership with a quorum and + /// its state machine has applied every entry up to the read index, so a + /// subsequent local read is linearizable. The error string is a + /// client-readable message (without the `ERR ` prefix). + /// + /// The default is a no-op so standalone deployments and test gates need + /// no Raft machinery. + fn ensure_linearizable_read( + &self, + ) -> std::pin::Pin> + Send + '_>> { + Box::pin(async { Ok(()) }) + } } #[allow(clippy::unwrap_used)] diff --git a/src/raft/src/lib.rs b/src/raft/src/lib.rs index dd334c9f..569cd053 100644 --- a/src/raft/src/lib.rs +++ b/src/raft/src/lib.rs @@ -19,6 +19,7 @@ //! //! Re-exports logindex types from storage::logindex to avoid code duplication. +pub mod capabilities; pub mod conversion; pub mod db_access; // Shim for backward compatibility with tests pub mod grpc; diff --git a/src/raft/src/node.rs b/src/raft/src/node.rs index d2d6a0d3..603273d1 100644 --- a/src/raft/src/node.rs +++ b/src/raft/src/node.rs @@ -99,6 +99,24 @@ impl crate::leader_gate::LeaderGate for RaftApp { fn leader_resp_addr(&self) -> Option { self.get_leader().map(|(_, node)| node.resp_addr) } + + fn ensure_linearizable_read( + &self, + ) -> std::pin::Pin> + Send + '_>> { + Box::pin(async move { + // openraft's `ensure_linearizable` confirms leadership with a + // quorum and then blocks until the state machine has applied up + // to the read log id (`wait().applied_index_at_least` on the + // raft-core metrics, which advance only after + // `KiwiStateMachine::apply` has written the entries to RocksDB + // and returned), so no extra applied-watch is needed here. + self.raft + .ensure_linearizable() + .await + .map(|_| ()) + .map_err(|e| format!("linearizable read failed: {e}")) + }) + } } pub struct RaftConfig { diff --git a/src/raft/src/snapshot_archive.rs b/src/raft/src/snapshot_archive.rs index 34fce4a3..dc25b237 100644 --- a/src/raft/src/snapshot_archive.rs +++ b/src/raft/src/snapshot_archive.rs @@ -131,7 +131,7 @@ mod tests { let root = unpacked_checkpoint_root(&unpack); assert!(root.join("0").join("marker").exists()); let m = RaftSnapshotMeta::read_from_dir(&root).unwrap(); - assert_eq!(m.version, 1); + assert_eq!(m.version, storage::CURRENT_SNAPSHOT_VERSION); assert_eq!(m.last_included_index, 7); assert_eq!(m.last_included_term, 3); } diff --git a/src/raft/src/state_machine.rs b/src/raft/src/state_machine.rs index 2ebd7643..2a58d029 100644 --- a/src/raft/src/state_machine.rs +++ b/src/raft/src/state_machine.rs @@ -73,6 +73,10 @@ const CURRENT_SNAPSHOT_DATA: &str = "current_snapshot.tar"; const CURRENT_SNAPSHOT_META: &str = "current_snapshot_meta.json"; const SNAPSHOT_INSTALL_MARKER_VERSION: u32 = 1; const SNAPSHOT_INSTALL_MARKER_SUFFIX: &str = ".snapshot-install-in-progress.json"; + +/// Number of vector metas / member entries decoded per instance when +/// validating restored snapshot data (sampling, not a full scan). +const RESTORED_VECTOR_SAMPLE_SIZE: usize = 64; const SNAPSHOT_INSTALL_CLEANUP_SUFFIX: &str = ".cleanup-pending"; #[cfg(test)] @@ -562,18 +566,26 @@ impl RaftStateMachine for KiwiStateMachine { )); } + // Loading the current Storage here is safe: schema validation and + // prepare only read and copy checkpoint input, and the owner remains + // live until the durable marker has been written after pause/drain. + let current_storage = self.storage_swap.load_full(); + let db_instance_num = current_storage.db_instance_num; + let db_id = current_storage.db_id; + + // Deterministically reject snapshots whose storage schema this binary + // cannot consume (version, instance count, column families, vector + // value format), before touching live storage. + file_meta + .validate_for_restore(db_instance_num) + .map_err(io_err_to_raft)?; + log::info!( "Snapshot metadata validated: index={}, term={}", file_meta.last_included_index, file_meta.last_included_term ); - // Loading the current Storage here is safe: prepare only reads and copies - // checkpoint input, and the owner remains live until the durable marker - // has been written after pause/drain. - let current_storage = self.storage_swap.load_full(); - let db_instance_num = current_storage.db_instance_num; - let db_id = current_storage.db_id; let prepared = prepare_checkpoint_restore(&checkpoint_root, &self.db_path, db_instance_num) .map_err(io_err_to_raft)?; @@ -635,6 +647,11 @@ impl RaftStateMachine for KiwiStateMachine { new_storage .open(options, &self.db_path) .map_err(|error| post_marker_error("opening the restored storage", &error))?; + // Sample-decode restored vector data; a corrupt sample rejects the + // install before the restored storage starts serving traffic. + new_storage + .validate_vector_data_sample(RESTORED_VECTOR_SAMPLE_SIZE) + .map_err(|error| post_marker_error("sampling restored vector data", &error))?; self.rearm_append_log_fn(&new_storage); self.storage_swap.swap(Arc::new(new_storage)); @@ -753,7 +770,8 @@ impl KiwiSnapshotBuilder { let collectors: Vec<_> = (0..self.storage.db_instance_num) .filter_map(|i| self.storage.get_logindex_collector(i)) .collect(); - let raft_meta = RaftSnapshotMeta::with_collector_states(last_idx, last_term, &collectors); + let raft_meta = + RaftSnapshotMeta::for_storage(last_idx, last_term, &collectors, &self.storage)?; self.storage.create_checkpoint(dir, &raft_meta)?; drop(snapshot_state_guard); diff --git a/src/raft/tests/snapshot_roundtrip_test.rs b/src/raft/tests/snapshot_roundtrip_test.rs index 1506df9a..c60f6069 100644 --- a/src/raft/tests/snapshot_roundtrip_test.rs +++ b/src/raft/tests/snapshot_roundtrip_test.rs @@ -1179,12 +1179,10 @@ async fn missing_checkpoint_instance_does_not_pause_or_replace_live_storage() -> Box::new(std::io::Cursor::new(snapshot_bytes)), ) .await - .expect_err("a checkpoint missing instance 1 must fail during prepare"); + .expect_err("a checkpoint built for 1 instance must fail schema validation against a 2-instance target"); assert!( - error - .to_string() - .contains("missing checkpoint instance directory"), - "unexpected prepare failure: {error}" + error.to_string().contains("db_instance_num"), + "unexpected validation failure: {error}" ); assert_eq!(pause_controller.pause_count.load(Ordering::SeqCst), 0); assert_eq!(pause_controller.resume_count.load(Ordering::SeqCst), 0); diff --git a/src/server/Cargo.toml b/src/server/Cargo.toml index 1456a2dd..67c087d4 100644 --- a/src/server/Cargo.toml +++ b/src/server/Cargo.toml @@ -15,6 +15,7 @@ path = "src/main.rs" net.workspace = true runtime.workspace = true storage.workspace = true +cmd.workspace = true tokio.workspace = true env_logger.workspace = true log.workspace = true diff --git a/src/server/src/main.rs b/src/server/src/main.rs index 698ea96e..d1df3d1f 100644 --- a/src/server/src/main.rs +++ b/src/server/src/main.rs @@ -38,6 +38,21 @@ struct PausePermitWrapper { _permit: RuntimeStorageAccessPermit, } +/// Build the command-table feature gates from the loaded configuration: +/// vector commands follow `vector-enabled`, are additionally rejected in +/// cluster mode unless `vector-cluster-enabled` is set, and FLUSHDB/FLUSHALL +/// are only allowed outside cluster mode unless `cluster-flush-enabled` is set. +fn command_table_gates(config: &Config) -> cmd::table::CommandTableGates { + let vector_enabled = config.vector.enabled; + let vector_cluster_allowed = config.raft.is_none() || config.vector.cluster_enabled; + let cluster_flush_allowed = config.raft.is_none() || config.cluster_flush_enabled; + cmd::table::CommandTableGates::from_flags( + vector_enabled, + vector_cluster_allowed, + cluster_flush_allowed, + ) +} + impl StorageAccessPermit for PausePermitWrapper {} impl PauseController for PauseControllerWrapper { @@ -166,6 +181,7 @@ fn main() -> std::io::Result<()> { let storage_for_server = storage.clone(); let requirepass_for_storage_server = config.requirepass.clone(); + let gates_for_storage_server = command_table_gates(&config); storage_handle.spawn(async move { info!("Initializing storage server..."); match initialize_storage_server( @@ -173,6 +189,7 @@ fn main() -> std::io::Result<()> { storage_for_server, pause_controller, requirepass_for_storage_server, + gates_for_storage_server, ) .await { @@ -295,12 +312,16 @@ async fn initialize_storage_server( global_storage: GlobalStorage, pause_controller: StorageServerPauseController, requirepass: Option, + gates: cmd::table::CommandTableGates, ) -> Result<(), DualRuntimeError> { info!("Initializing storage server..."); // Initialize the storage-runtime command table with the same password // provider used by the network runtime, so AUTH behaves consistently. - runtime::initialize_storage_command_table(Arc::new(move || requirepass.clone())); + runtime::initialize_storage_command_table_with_gates( + Arc::new(move || requirepass.clone()), + gates, + ); let storage_server = StorageServer::with_pause_controller(global_storage, request_receiver, pause_controller); @@ -449,6 +470,7 @@ async fn start_server( runtime_manager, config.requirepass.clone(), leader_gate, + command_table_gates(config), ) { Some(server) => { tokio::spawn(async move { diff --git a/src/storage/Cargo.toml b/src/storage/Cargo.toml index 365e306a..8afa9959 100644 --- a/src/storage/Cargo.toml +++ b/src/storage/Cargo.toml @@ -54,4 +54,7 @@ conf.workspace = true [target.'cfg(windows)'.dependencies] windows-sys.workspace = true +[dev-dependencies] +proptest.workspace = true + diff --git a/src/storage/src/batch.rs b/src/storage/src/batch.rs index 0bec4d9b..f32caf2b 100644 --- a/src/storage/src/batch.rs +++ b/src/storage/src/batch.rs @@ -45,6 +45,7 @@ use snafu::ResultExt; use crate::ColumnFamilyIndex; use crate::error::{BatchSnafu, InvalidFormatSnafu, Result, RocksSnafu}; +use crate::format_vector_member_key::ParsedVectorMemberDataKey; use crate::slot_indexer::key_to_slot_id; use crate::storage_define::{PREFIX_RESERVE_LENGTH, decode_user_key, seek_userkey_delim}; use bytes::BytesMut; @@ -256,6 +257,13 @@ impl BinlogBatch { .fail(); } + // Vector member keys use the V1 vector codec, not the shared + // MemberDataKey layout. + if cf_idx == ColumnFamilyIndex::VectorDataCF as u32 { + let parsed = ParsedVectorMemberDataKey::decode(encoded_key)?; + return Ok(parsed.key().to_vec()); + } + if encoded_key.len() <= PREFIX_RESERVE_LENGTH { return InvalidFormatSnafu { message: format!( diff --git a/src/storage/src/checkpoint.rs b/src/storage/src/checkpoint.rs index 32cf5d59..636e3a1c 100644 --- a/src/storage/src/checkpoint.rs +++ b/src/storage/src/checkpoint.rs @@ -35,7 +35,9 @@ use serde::{Deserialize, Serialize}; #[cfg(any(test, feature = "test-fault-injection"))] use parking_lot::Mutex; +use crate::format_vector::VECTOR_VALUE_FORMAT; use crate::logindex::LogIndexAndSequenceCollector; +use crate::redis::ColumnFamilyIndex; use crate::{sync_directory, sync_parent_directory}; /// File name for JSON metadata at the checkpoint root (not OpenRaft's `SnapshotMeta`). @@ -85,8 +87,23 @@ pub fn fail_next_restore_parent_sync_after_rename( RestoreParentSyncFailureGuard { target_db_path } } -/// Current snapshot format version -pub const CURRENT_SNAPSHOT_VERSION: u32 = 1; +/// Current snapshot format version. +/// +/// Version 1 was a development-phase format that was never released; only v2 +/// snapshots (which carry the storage schema description) are accepted. +pub const CURRENT_SNAPSHOT_VERSION: u32 = 2; + +/// Version of the storage on-disk schema understood by this binary. +pub const STORAGE_SCHEMA_VERSION: u32 = 1; + +/// Column families every instance of the checkpoint must contain, in +/// declaration order. +pub fn expected_column_families() -> Vec { + ColumnFamilyIndex::ALL + .iter() + .map(|cf| cf.name().to_string()) + .collect() +} /// Metadata persisted next to per-instance checkpoint directories. #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] @@ -97,6 +114,23 @@ pub struct RaftSnapshotMeta { pub last_included_index: u64, /// Last log term included in the snapshot pub last_included_term: u64, + /// Storage on-disk schema version understood by the snapshot writer. + #[serde(default)] + pub storage_schema_version: u32, + /// Storage incarnation of each RocksDB instance, ordered by instance id. + /// A restored database adopts these via the per-instance manifest files + /// carried inside the checkpoint; they are validated for structure only. + #[serde(default)] + pub storage_incarnations: Vec, + /// Number of RocksDB instances in the checkpoint. + #[serde(default)] + pub db_instance_num: u32, + /// Column families each checkpoint instance must contain, in declaration order. + #[serde(default)] + pub column_families: Vec, + /// Highest vector value format byte the snapshot writer can emit. + #[serde(default)] + pub vector_value_format_max: u8, /// LogIndex collector states, one entry per Storage instance. /// Outer index is the instance id; inner Vec holds `"log_index:seqno"` pairs. #[serde(default)] @@ -110,6 +144,11 @@ impl RaftSnapshotMeta { version: CURRENT_SNAPSHOT_VERSION, last_included_index, last_included_term, + storage_schema_version: STORAGE_SCHEMA_VERSION, + storage_incarnations: Vec::new(), + db_instance_num: 0, + column_families: expected_column_families(), + vector_value_format_max: VECTOR_VALUE_FORMAT, logindex_collector_states: Vec::new(), } } @@ -121,13 +160,32 @@ impl RaftSnapshotMeta { collectors: &[Arc], ) -> Self { Self { - version: CURRENT_SNAPSHOT_VERSION, - last_included_index, - last_included_term, + db_instance_num: collectors.len() as u32, logindex_collector_states: collectors.iter().map(|c| c.export_state()).collect(), + ..Self::new(last_included_index, last_included_term) } } + /// Create snapshot meta describing the given live Storage: per-instance + /// storage incarnations, the instance count, and the column-family list. + pub fn for_storage( + last_included_index: u64, + last_included_term: u64, + collectors: &[Arc], + storage: &crate::storage::Storage, + ) -> crate::error::Result { + let storage_incarnations = storage + .insts + .iter() + .map(|inst| inst.storage_incarnation()) + .collect::>>()?; + Ok(Self { + storage_incarnations, + db_instance_num: storage.db_instance_num as u32, + ..Self::with_collector_states(last_included_index, last_included_term, collectors) + }) + } + /// Restore collector states for each Storage instance from snapshot metadata. /// /// `collectors[i]` receives the entries originally exported from instance `i`. Extra @@ -193,13 +251,14 @@ impl RaftSnapshotMeta { let meta: Self = serde_json::from_slice(&bytes) .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; - // Validate version: reject old unsupported versions for safety, - // but allow higher versions for forward compatibility during rolling upgrades. - if meta.version < CURRENT_SNAPSHOT_VERSION { + // Only the exact current version is accepted. v1 was a development-phase + // format that never shipped, and a higher version comes from a newer + // binary whose schema this node cannot safely consume. + if meta.version != CURRENT_SNAPSHOT_VERSION { return Err(io::Error::new( io::ErrorKind::InvalidData, format!( - "unsupported snapshot version: {}, expected >= {}", + "unsupported snapshot version: {}, expected {}", meta.version, CURRENT_SNAPSHOT_VERSION ), )); @@ -207,6 +266,55 @@ impl RaftSnapshotMeta { Ok(meta) } + + /// Validate the storage schema description against this binary and the + /// local configuration. Restore must deterministically reject snapshots it + /// cannot consume. The storage incarnations are structural metadata only: + /// a restored database adopts the snapshot writer's incarnations via the + /// per-instance manifest files, so differing values are never an error. + pub fn validate_for_restore(&self, expected_db_instance_num: usize) -> io::Result<()> { + let invalid = |message: String| io::Error::new(io::ErrorKind::InvalidData, message); + + if self.version != CURRENT_SNAPSHOT_VERSION { + return Err(invalid(format!( + "unsupported snapshot version: {}, expected {}", + self.version, CURRENT_SNAPSHOT_VERSION + ))); + } + if self.storage_schema_version > STORAGE_SCHEMA_VERSION { + return Err(invalid(format!( + "unsupported storage schema version: {} > {}", + self.storage_schema_version, STORAGE_SCHEMA_VERSION + ))); + } + if self.db_instance_num as usize != expected_db_instance_num { + return Err(invalid(format!( + "snapshot db_instance_num {} does not match local configuration {}", + self.db_instance_num, expected_db_instance_num + ))); + } + if self.storage_incarnations.len() != self.db_instance_num as usize { + return Err(invalid(format!( + "snapshot carries {} storage incarnations for {} instances", + self.storage_incarnations.len(), + self.db_instance_num + ))); + } + let expected_column_families = expected_column_families(); + if self.column_families != expected_column_families { + return Err(invalid(format!( + "snapshot column families {:?} do not match expected {:?}", + self.column_families, expected_column_families + ))); + } + if self.vector_value_format_max > VECTOR_VALUE_FORMAT { + return Err(invalid(format!( + "snapshot vector value format {} exceeds supported format {}", + self.vector_value_format_max, VECTOR_VALUE_FORMAT + ))); + } + Ok(()) + } } pub fn copy_dir_all(src: &Path, dst: &Path) -> io::Result<()> { diff --git a/src/storage/src/data_compaction_filter.rs b/src/storage/src/data_compaction_filter.rs index b087c423..2987352d 100644 --- a/src/storage/src/data_compaction_filter.rs +++ b/src/storage/src/data_compaction_filter.rs @@ -32,8 +32,10 @@ use rocksdb::{ use crate::{ DataType, coding::decode_fixed, + format_base_key::BaseMetaKey, format_base_meta_value::ParsedBaseMetaValue, format_list_meta_value::ParsedListsMetaValue, + format_vector_member_key::ParsedVectorMemberDataKey, storage_define::{ ENCODED_KEY_DELIM_SIZE, NEED_TRANSFORM_CHARACTER, PREFIX_RESERVE_LENGTH, SUFFIX_RESERVE_LENGTH, VERSION_LENGTH, seek_userkey_delim, @@ -171,6 +173,7 @@ enum MetaLookup { pub struct DataCompactionFilter { db: Option>, data_type: DataType, + storage_incarnation: Arc>, cur_key: BytesMut, meta_not_found: bool, cur_meta_version: u64, @@ -178,10 +181,15 @@ pub struct DataCompactionFilter { } impl DataCompactionFilter { - pub fn new(db: Option>, data_type: DataType) -> Self { + pub fn new( + db: Option>, + data_type: DataType, + storage_incarnation: Arc>, + ) -> Self { Self { db, data_type, + storage_incarnation, cur_key: BytesMut::new(), meta_not_found: false, cur_meta_version: 0, @@ -313,6 +321,50 @@ impl DataCompactionFilter { MetaLookup::Valid } } + + /// Filter a vector set member key (V1 codec, see + /// `format_vector_member_key.rs`). + /// + /// A member is removed when it is provably stale: its meta is gone or + /// expired, it was written by a different storage incarnation, or its + /// generation sequence differs from the live generation in the meta + /// (`meta.version`). Anything that cannot be parsed or proven stale is + /// kept. + fn filter_vector_member(&mut self, key: &[u8]) -> CompactionDecision { + let Ok(parsed) = ParsedVectorMemberDataKey::decode(key) else { + return CompactionDecision::Keep; + }; + let Ok(meta_key) = BaseMetaKey::new(parsed.key()).encode() else { + return CompactionDecision::Keep; + }; + + match self.ensure_meta_state(&meta_key) { + MetaLookup::Unavailable => CompactionDecision::Keep, + MetaLookup::NotFound => CompactionDecision::Remove, + MetaLookup::Valid => { + let cur_time = Utc::now().timestamp_micros() as u64; + if self.cur_meta_etime != 0 && self.cur_meta_etime < cur_time { + return CompactionDecision::Remove; + } + + // Members written by another storage (e.g. data written before + // the instance identity was rebuilt) are always stale here. + let Some(current_incarnation) = self.storage_incarnation.get() else { + return CompactionDecision::Keep; + }; + if parsed.storage_incarnation() != *current_incarnation { + return CompactionDecision::Remove; + } + + // The meta version is the live generation sequence. + if parsed.generation_sequence() != self.cur_meta_version { + return CompactionDecision::Remove; + } + + CompactionDecision::Keep + } + } + } } impl CompactionFilter for DataCompactionFilter { @@ -324,6 +376,10 @@ impl CompactionFilter for DataCompactionFilter { #[cfg(test)] block_once_for_compaction_filter_test(key); + if self.data_type == DataType::VectorSet { + return self.filter_vector_member(key); + } + let Some(meta_key) = Self::build_meta_key(key) else { return CompactionDecision::Keep; }; @@ -349,11 +405,20 @@ impl CompactionFilter for DataCompactionFilter { pub struct DataCompactionFilterFactory { db: Arc>>, data_type: DataType, + storage_incarnation: Arc>, } impl DataCompactionFilterFactory { - pub fn new(db: Arc>>, data_type: DataType) -> Self { - Self { db, data_type } + pub fn new( + db: Arc>>, + data_type: DataType, + storage_incarnation: Arc>, + ) -> Self { + Self { + db, + data_type, + storage_incarnation, + } } } @@ -365,7 +430,7 @@ impl CompactionFilterFactory for DataCompactionFilterFactory { _context: rocksdb::compaction_filter_factory::CompactionFilterContext, ) -> Self::Filter { let db = self.db.get().and_then(Weak::upgrade); - DataCompactionFilter::new(db, self.data_type) + DataCompactionFilter::new(db, self.data_type, Arc::clone(&self.storage_incarnation)) } fn name(&self) -> &std::ffi::CStr { @@ -380,6 +445,7 @@ mod tests { use crate::format_base_key::BaseKey; use crate::format_base_meta_value::BaseMetaValue; use crate::format_list_meta_value::ListsMetaValue; + use crate::format_vector_member_key::VectorMemberDataKey; use crate::storage_define::SUFFIX_RESERVE_LENGTH; use crate::unique_test_db_path; use bytes::BufMut; @@ -387,6 +453,15 @@ mod tests { ColumnFamilyDescriptor, Options, compaction_filter_factory::CompactionFilterContext, }; + const TEST_INCARNATION: u64 = 7; + + fn test_incarnation_cell(value: u64) -> Arc> { + let cell = Arc::new(OnceCell::new()); + cell.set(value) + .expect("incarnation cell should be set once"); + cell + } + fn setup_db_for_filter_test(path: &std::path::Path) -> (Arc>>, Arc) { let mut db_opts = Options::default(); db_opts.create_if_missing(true); @@ -419,6 +494,18 @@ mod tests { encoded.to_vec() } + /// Helper to create a V1 vector member key for testing purposes. + fn encode_vector_member_key(key: &[u8], incarnation: u64, generation: u64) -> Vec { + VectorMemberDataKey { + key, + storage_incarnation: incarnation, + generation_sequence: generation, + element: b"member", + } + .encode_full() + .expect("encode vector member key") + } + fn put_meta(db: &Arc, user_key: &[u8], data_type: DataType, version: u64, etime: u64) { let meta_key = BaseKey::new(user_key).encode().unwrap(); match data_type { @@ -446,7 +533,11 @@ mod tests { let path = unique_test_db_path(); let (db_cell, _db) = setup_db_for_filter_test(&path); - let mut factory = DataCompactionFilterFactory::new(db_cell, DataType::Hash); + let mut factory = DataCompactionFilterFactory::new( + db_cell, + DataType::Hash, + test_incarnation_cell(TEST_INCARNATION), + ); let context = CompactionFilterContext { is_full_compaction: false, is_manual_compaction: false, @@ -473,7 +564,11 @@ mod tests { .expect("DB cell should be set once"); } - let mut factory = DataCompactionFilterFactory::new(db_cell, DataType::Hash); + let mut factory = DataCompactionFilterFactory::new( + db_cell, + DataType::Hash, + test_incarnation_cell(TEST_INCARNATION), + ); let context = CompactionFilterContext { is_full_compaction: false, is_manual_compaction: false, @@ -497,7 +592,11 @@ mod tests { // Allow some time for the snapshot to see the write std::thread::sleep(std::time::Duration::from_millis(10)); - let mut factory = DataCompactionFilterFactory::new(db_cell, DataType::Hash); + let mut factory = DataCompactionFilterFactory::new( + db_cell, + DataType::Hash, + test_incarnation_cell(TEST_INCARNATION), + ); let context = CompactionFilterContext { is_full_compaction: false, is_manual_compaction: false, @@ -519,7 +618,11 @@ mod tests { std::thread::sleep(std::time::Duration::from_millis(10)); - let mut factory = DataCompactionFilterFactory::new(db_cell, DataType::Set); + let mut factory = DataCompactionFilterFactory::new( + db_cell, + DataType::Set, + test_incarnation_cell(TEST_INCARNATION), + ); let context = CompactionFilterContext { is_full_compaction: false, is_manual_compaction: false, @@ -541,7 +644,11 @@ mod tests { std::thread::sleep(std::time::Duration::from_millis(10)); - let mut factory = DataCompactionFilterFactory::new(db_cell, DataType::ZSet); + let mut factory = DataCompactionFilterFactory::new( + db_cell, + DataType::ZSet, + test_incarnation_cell(TEST_INCARNATION), + ); let context = CompactionFilterContext { is_full_compaction: false, is_manual_compaction: false, @@ -559,23 +666,154 @@ mod tests { let path = unique_test_db_path(); let (db_cell, db) = setup_db_for_filter_test(&path); - put_meta(&db, b"vector_key", DataType::VectorSet, 1, 0); + // Meta version holds the live generation sequence. + put_meta(&db, b"vector_key", DataType::VectorSet, 3, 0); std::thread::sleep(std::time::Duration::from_millis(10)); - let mut factory = DataCompactionFilterFactory::new(db_cell, DataType::VectorSet); + let mut factory = DataCompactionFilterFactory::new( + db_cell, + DataType::VectorSet, + test_incarnation_cell(TEST_INCARNATION), + ); let context = CompactionFilterContext { is_full_compaction: false, is_manual_compaction: false, }; let mut filter = factory.create(context); - let data_key = encode_data_key(b"vector_key", 1); + let data_key = encode_vector_member_key(b"vector_key", TEST_INCARNATION, 3); let decision = filter.filter(0, &data_key, b""); assert!(matches!(decision, CompactionDecision::Keep)); } + #[test] + fn test_removes_vector_member_on_incarnation_mismatch() { + let path = unique_test_db_path(); + let (db_cell, db) = setup_db_for_filter_test(&path); + + put_meta(&db, b"vector_key", DataType::VectorSet, 3, 0); + + std::thread::sleep(std::time::Duration::from_millis(10)); + + let mut factory = DataCompactionFilterFactory::new( + db_cell, + DataType::VectorSet, + test_incarnation_cell(TEST_INCARNATION), + ); + let context = CompactionFilterContext { + is_full_compaction: false, + is_manual_compaction: false, + }; + let mut filter = factory.create(context); + + // Written by another storage incarnation: stale even though the + // generation matches. + let data_key = encode_vector_member_key(b"vector_key", TEST_INCARNATION + 1, 3); + let decision = filter.filter(0, &data_key, b""); + + assert!(matches!(decision, CompactionDecision::Remove)); + } + + #[test] + fn test_removes_vector_member_on_generation_mismatch() { + let path = unique_test_db_path(); + let (db_cell, db) = setup_db_for_filter_test(&path); + + put_meta(&db, b"vector_key", DataType::VectorSet, 3, 0); + + std::thread::sleep(std::time::Duration::from_millis(10)); + + let mut factory = DataCompactionFilterFactory::new( + db_cell, + DataType::VectorSet, + test_incarnation_cell(TEST_INCARNATION), + ); + let context = CompactionFilterContext { + is_full_compaction: false, + is_manual_compaction: false, + }; + let mut filter = factory.create(context); + + // Member of a previous lifecycle of the same key. + let data_key = encode_vector_member_key(b"vector_key", TEST_INCARNATION, 2); + let decision = filter.filter(0, &data_key, b""); + + assert!(matches!(decision, CompactionDecision::Remove)); + } + + #[test] + fn test_keeps_vector_member_when_meta_is_unreadable() { + // Without a live DB handle the meta cannot be read, so nothing can be + // proven stale. + let db_cell: Arc>> = Arc::new(OnceCell::new()); + let mut factory = DataCompactionFilterFactory::new( + db_cell, + DataType::VectorSet, + test_incarnation_cell(TEST_INCARNATION), + ); + let context = CompactionFilterContext { + is_full_compaction: false, + is_manual_compaction: false, + }; + let mut filter = factory.create(context); + + let data_key = encode_vector_member_key(b"vector_key", TEST_INCARNATION, 1); + assert!(matches!( + filter.filter(0, &data_key, b""), + CompactionDecision::Keep + )); + } + + #[test] + fn test_keeps_vector_member_when_key_is_malformed() { + let path = unique_test_db_path(); + let (db_cell, _db) = setup_db_for_filter_test(&path); + + let mut factory = DataCompactionFilterFactory::new( + db_cell, + DataType::VectorSet, + test_incarnation_cell(TEST_INCARNATION), + ); + let context = CompactionFilterContext { + is_full_compaction: false, + is_manual_compaction: false, + }; + let mut filter = factory.create(context); + + // Truncated inside the fixed header: cannot be parsed, must be kept. + let mut data_key = encode_vector_member_key(b"vector_key", TEST_INCARNATION, 1); + data_key.truncate(10); + assert!(matches!( + filter.filter(0, &data_key, b""), + CompactionDecision::Keep + )); + } + + #[test] + fn test_removes_vector_member_if_meta_is_missing() { + let path = unique_test_db_path(); + let (db_cell, _db) = setup_db_for_filter_test(&path); + + let mut factory = DataCompactionFilterFactory::new( + db_cell, + DataType::VectorSet, + test_incarnation_cell(TEST_INCARNATION), + ); + let context = CompactionFilterContext { + is_full_compaction: false, + is_manual_compaction: false, + }; + let mut filter = factory.create(context); + + let data_key = encode_vector_member_key(b"vector_key", TEST_INCARNATION, 1); + assert!(matches!( + filter.filter(0, &data_key, b""), + CompactionDecision::Remove + )); + } + #[test] fn test_keeps_data_if_meta_is_valid_for_list() { let path = unique_test_db_path(); @@ -585,7 +823,11 @@ mod tests { std::thread::sleep(std::time::Duration::from_millis(10)); - let mut factory = DataCompactionFilterFactory::new(db_cell, DataType::List); + let mut factory = DataCompactionFilterFactory::new( + db_cell, + DataType::List, + test_incarnation_cell(TEST_INCARNATION), + ); let context = CompactionFilterContext { is_full_compaction: false, is_manual_compaction: false, @@ -613,7 +855,11 @@ mod tests { std::thread::sleep(std::time::Duration::from_millis(10)); - let mut factory = DataCompactionFilterFactory::new(db_cell, DataType::Hash); + let mut factory = DataCompactionFilterFactory::new( + db_cell, + DataType::Hash, + test_incarnation_cell(TEST_INCARNATION), + ); let context = CompactionFilterContext { is_full_compaction: false, is_manual_compaction: false, @@ -640,7 +886,11 @@ mod tests { std::thread::sleep(std::time::Duration::from_millis(10)); - let mut factory = DataCompactionFilterFactory::new(db_cell, DataType::Hash); + let mut factory = DataCompactionFilterFactory::new( + db_cell, + DataType::Hash, + test_incarnation_cell(TEST_INCARNATION), + ); let context = CompactionFilterContext { is_full_compaction: false, is_manual_compaction: false, @@ -663,7 +913,11 @@ mod tests { std::thread::sleep(std::time::Duration::from_millis(10)); // Filter is for Hashes - let mut factory = DataCompactionFilterFactory::new(db_cell, DataType::Hash); + let mut factory = DataCompactionFilterFactory::new( + db_cell, + DataType::Hash, + test_incarnation_cell(TEST_INCARNATION), + ); let context = CompactionFilterContext { is_full_compaction: false, is_manual_compaction: false, @@ -681,7 +935,11 @@ mod tests { let path = unique_test_db_path(); let (db_cell, _) = setup_db_for_filter_test(&path); - let mut factory = DataCompactionFilterFactory::new(db_cell, DataType::Hash); + let mut factory = DataCompactionFilterFactory::new( + db_cell, + DataType::Hash, + test_incarnation_cell(TEST_INCARNATION), + ); let context = CompactionFilterContext { is_full_compaction: false, is_manual_compaction: false, diff --git a/src/storage/src/error.rs b/src/storage/src/error.rs index 39dfc780..07a8cffc 100644 --- a/src/storage/src/error.rs +++ b/src/storage/src/error.rs @@ -147,6 +147,26 @@ pub enum Error { location: Location, }, + // FLAT vector query governance failures; kept as distinct variants so + // metrics can tell timeout / cancellation / budget exhaustion apart. + #[snafu(display("ERR vector flat query timeout"))] + VectorFlatQueryTimeout { + #[snafu(implicit)] + location: Location, + }, + + #[snafu(display("ERR vector flat query cancelled"))] + VectorFlatQueryCancelled { + #[snafu(implicit)] + location: Location, + }, + + #[snafu(display("ERR vector flat query scan budget exceeded"))] + VectorFlatScanBudgetExceeded { + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("LogIndex error: {}", message))] LogIndex { message: String, diff --git a/src/storage/src/format_member_data_key.rs b/src/storage/src/format_member_data_key.rs index add8768e..3b54b1b9 100644 --- a/src/storage/src/format_member_data_key.rs +++ b/src/storage/src/format_member_data_key.rs @@ -16,10 +16,11 @@ // limitations under the License. use bytes::{BufMut, Bytes, BytesMut}; +use snafu::ensure; use crate::storage_define::seek_userkey_delim; use crate::{ - error::Result, + error::{InvalidFormatSnafu, Result}, storage_define::{ ENCODED_KEY_DELIM_SIZE, PREFIX_RESERVE_LENGTH, SUFFIX_RESERVE_LENGTH, decode_user_key, encode_user_key, @@ -29,8 +30,7 @@ use crate::{ // | reserve1 | key | version | data | reserve2 | // | 8B | | 8B | | 16B | // -// For vector sets, `data` is the element name and the value stored at this key -// is a VectorDataValue (see format_vector.rs). +// Vector sets use their own V1 codec (see format_vector_member_key.rs). #[derive(Debug, Clone)] pub struct MemberDataKey { @@ -101,6 +101,21 @@ impl ParsedMemberDataKey { pub fn new(encoded_key: &[u8]) -> Result { let mut key_str = BytesMut::new(); + // minimum: reserve1 + encoded empty key (delim only) + version + reserve2 + const MIN_LENGTH: usize = PREFIX_RESERVE_LENGTH + + ENCODED_KEY_DELIM_SIZE + + size_of::() + + SUFFIX_RESERVE_LENGTH; + ensure!( + encoded_key.len() >= MIN_LENGTH, + InvalidFormatSnafu { + message: format!( + "member data key too short: {} < {MIN_LENGTH}", + encoded_key.len() + ) + } + ); + let start_idx = PREFIX_RESERVE_LENGTH; let end_idx = encoded_key.len() - SUFFIX_RESERVE_LENGTH; @@ -110,13 +125,24 @@ impl ParsedMemberDataKey { reserve1.copy_from_slice(reserve_slice); // key - let key_end_idx = start_idx + seek_userkey_delim(&encoded_key[start_idx..]); + let key_end_idx = start_idx + seek_userkey_delim(&encoded_key[start_idx..end_idx]); decode_user_key(&encoded_key[start_idx..key_end_idx], &mut key_str)?; // version let version_end_idx = key_end_idx + size_of::(); + ensure!( + version_end_idx <= end_idx, + InvalidFormatSnafu { + message: "member data key version field out of bounds".to_string() + } + ); let version_slice = &encoded_key[key_end_idx..version_end_idx]; - let version = u64::from_le_bytes(version_slice.try_into().expect("slice length mismatch")); + let version = u64::from_le_bytes(version_slice.try_into().map_err(|_| { + InvalidFormatSnafu { + message: "member data key version field length mismatch".to_string(), + } + .build() + })?); // data let data_slice = &encoded_key[version_end_idx..end_idx]; @@ -209,6 +235,21 @@ mod tests { assert_eq!(&encoded[data_begin..data_end], test_data); } + #[test] + fn mv_test_member_data_key_rejects_malformed_keys() { + // shorter than reserve1 + delim + version + reserve2 + let short = vec![0u8; 10]; + assert!(ParsedMemberDataKey::new(&short).is_err()); + + // long enough, but the version field would run into reserve2 + let mut no_room_for_version = Vec::new(); + no_room_for_version.extend_from_slice(&[0u8; PREFIX_RESERVE_LENGTH]); + no_room_for_version.extend_from_slice(b"\x00\x00"); // empty encoded key + no_room_for_version.extend_from_slice(&[0u8; 4]); // only 4 bytes before reserve2 + no_room_for_version.extend_from_slice(&[0u8; SUFFIX_RESERVE_LENGTH]); + assert!(ParsedMemberDataKey::new(&no_room_for_version).is_err()); + } + #[test] fn mv_test_member_data_key_with_empty_key_and_data() { let test_key = b""; diff --git a/src/storage/src/format_vector.rs b/src/storage/src/format_vector.rs index c72fd0c2..86665c49 100644 --- a/src/storage/src/format_vector.rs +++ b/src/storage/src/format_vector.rs @@ -37,8 +37,7 @@ pub const VECTOR_METRIC_COSINE: u8 = 1; pub const VECTOR_VALUE_MAGIC: u8 = 0x56; pub const VECTOR_VALUE_FORMAT: u8 = 1; -const VECTOR_META_ZERO_RESERVE_LENGTH: usize = 8; -const VECTOR_VALUE_HEADER_LENGTH: usize = 12; +const VECTOR_VALUE_HEADER_LENGTH: usize = 16; const VECTOR_Q8_PARAMS_LENGTH: usize = 8; /// Similarity metric used to compare vectors in a vector set. @@ -98,11 +97,15 @@ impl SimilarityMetric { // Vector set meta value layout stored in MetaCF: // -// | data_type | count | version | format | quant | metric | flags | dimension | zero_reserve | ctime | etime | -// | 1B | 8B | 8B | 1B | 1B | 1B | 1B | 4B | 8B | 8B | 8B | +// | data_type | count | version | format | quant | metric | flags | dimension | data_revision | ctime | etime | +// | 1B | 8B | 8B | 1B | 1B | 1B | 1B | 4B | 8B | 8B | 8B | // // `data_type` is DataType::VectorSet, `quant` is the quantization type (NOQUANT/BIN/Q8), // and `metric` is the similarity metric used for VSIM (e.g. cosine). +// `version` holds the set's generation sequence: the identifier of one +// lifecycle of the key, assigned by the storage generation generator (or the +// creating Raft log index in cluster mode). `data_revision` starts at 1 when +// the set is created and is incremented by every successful VADD/VREM. #[derive(Debug, Clone, PartialEq, Eq)] pub struct VectorMeta { count: u64, @@ -110,28 +113,26 @@ pub struct VectorMeta { dimension: u32, quantization: QuantizationType, metric: SimilarityMetric, + data_revision: u64, ctime: u64, etime: u64, } impl VectorMeta { - pub(crate) fn new_after( + pub(crate) fn new( count: u64, dimension: u32, quantization: QuantizationType, - previous_version: u64, + generation: u64, ) -> Self { let now = Utc::now().timestamp_micros() as u64; - let version = match previous_version >= now { - true => previous_version + 1, - false => now, - }; Self { count, - version, + version: generation, dimension, quantization, metric: SimilarityMetric::Cosine, + data_revision: 1, ctime: now, etime: 0, } @@ -155,7 +156,7 @@ impl VectorMeta { output.put_u8(self.metric.to_u8()); output.put_u8(0); output.put_u32_le(self.dimension); - output.put_bytes(0, VECTOR_META_ZERO_RESERVE_LENGTH); + output.put_u64_le(self.data_revision); output.put_u64_le(self.ctime); output.put_u64_le(self.etime); output @@ -189,8 +190,7 @@ impl VectorMeta { let metric = reader.get_u8(); let flags = reader.get_u8(); let dimension = reader.get_u32_le(); - let zero_reserve = &reader[..VECTOR_META_ZERO_RESERVE_LENGTH]; - reader.advance(VECTOR_META_ZERO_RESERVE_LENGTH); + let data_revision = reader.get_u64_le(); let ctime = reader.get_u64_le(); let etime = reader.get_u64_le(); @@ -203,9 +203,9 @@ impl VectorMeta { let quantization = QuantizationType::from_u8(quant)?; let metric = SimilarityMetric::from_u8(metric)?; ensure!( - flags == 0 && zero_reserve.iter().all(|byte| *byte == 0), + flags == 0, InvalidFormatSnafu { - message: "invalid non-zero vector meta reserve".to_string() + message: "invalid non-zero vector meta flags".to_string() } ); ensure!( @@ -221,6 +221,7 @@ impl VectorMeta { dimension, quantization, metric, + data_revision, ctime, etime, }) @@ -246,6 +247,14 @@ impl VectorMeta { self.dimension } + pub fn data_revision(&self) -> u64 { + self.data_revision + } + + pub(crate) fn bump_data_revision(&mut self) { + self.data_revision = self.data_revision.saturating_add(1); + } + pub(crate) fn is_stale(&self) -> bool { self.etime != 0 && self.etime < Utc::now().timestamp_micros() as u64 } @@ -257,13 +266,14 @@ impl VectorMeta { // Vector member data value layout stored in VectorDataCF: // -// | magic | format | quant | flags | dimension | original_l2 | [quant_params] | payload | -// | 1B | 1B | 1B | 1B | 4B | 4B | 0B or 8B | varies | +// | magic | format | quant | flags | dimension | original_l2 | payload_len | [quant_params] | payload | +// | 1B | 1B | 1B | 1B | 4B | 4B | 4B | 0B or 8B | varies | // // `magic` is VECTOR_VALUE_MAGIC and `original_l2` preserves the pre-normalization // L2 norm so VEMB can reconstruct the original FP32 vector. `flags` is reserved // for optional sections (e.g. bit 0 = trailing SETATTR attributes JSON) and -// must be zero until such a section is implemented. The payload layout +// must be zero until such a section is implemented. `payload_len` is the byte +// length of the payload section (excluding `quant_params`). The payload layout // depends on `quant`: // NOQUANT: 4B * dimension FP32 components, no quant_params. // BIN: ceil(dimension / 8) bitmap bytes, no quant_params. @@ -282,18 +292,26 @@ impl VectorDataValue { pub(crate) fn encode(&self) -> BytesMut { let canonical = &self.canonical; + let params_length = match canonical.data() { + VectorData::Int8 { .. } => VECTOR_Q8_PARAMS_LENGTH, + _ => 0, + }; let payload_length = match canonical.data() { VectorData::Fp32(values) => values.len() * size_of::(), VectorData::Binary(bits) => bits.len(), - VectorData::Int8 { values, .. } => VECTOR_Q8_PARAMS_LENGTH + values.len(), + VectorData::Int8 { values, .. } => values.len(), }; - let mut output = BytesMut::with_capacity(VECTOR_VALUE_HEADER_LENGTH + payload_length); + let payload_len = u32::try_from(payload_length) + .expect("vector payload length always fits in u32 for a valid canonical vector"); + let mut output = + BytesMut::with_capacity(VECTOR_VALUE_HEADER_LENGTH + params_length + payload_length); output.put_u8(VECTOR_VALUE_MAGIC); output.put_u8(VECTOR_VALUE_FORMAT); output.put_u8(canonical.quantization().to_u8()); output.put_u8(0); // flags: reserved, no optional sections yet output.put_u32_le(canonical.dimension()); output.put_f32_le(canonical.original_l2()); + output.put_u32_le(payload_len); match canonical.data() { VectorData::Fp32(values) => { for component in values { @@ -333,6 +351,7 @@ impl VectorDataValue { let flags = reader.get_u8(); let dimension = reader.get_u32_le(); let original_l2 = reader.get_f32_le(); + let payload_len = reader.get_u32_le() as usize; ensure!( magic == VECTOR_VALUE_MAGIC, @@ -368,11 +387,12 @@ impl VectorDataValue { let dimension = dimension as usize; let data = match quantization { QuantizationType::None => { + let expected_payload = dimension * size_of::(); ensure!( - reader.len() == dimension * size_of::(), + payload_len == expected_payload && reader.len() == expected_payload, InvalidFormatSnafu { message: format!( - "invalid vector payload length: {} for dimension {}", + "invalid vector payload length: payload_len {payload_len}, actual {} for dimension {}", reader.len(), dimension ) @@ -391,11 +411,12 @@ impl VectorDataValue { VectorData::Fp32(normalized) } QuantizationType::Binary => { + let expected_payload = dimension.div_ceil(8); ensure!( - reader.len() == dimension.div_ceil(8), + payload_len == expected_payload && reader.len() == expected_payload, InvalidFormatSnafu { message: format!( - "invalid binary vector payload length: {} for dimension {}", + "invalid binary vector payload length: payload_len {payload_len}, actual {} for dimension {}", reader.len(), dimension ) @@ -405,10 +426,10 @@ impl VectorDataValue { } QuantizationType::Int8 => { ensure!( - reader.len() == VECTOR_Q8_PARAMS_LENGTH + dimension, + payload_len == dimension && reader.len() == VECTOR_Q8_PARAMS_LENGTH + dimension, InvalidFormatSnafu { message: format!( - "invalid q8 vector payload length: {} for dimension {}", + "invalid q8 vector payload length: payload_len {payload_len}, actual {} for dimension {}", reader.len(), dimension ) @@ -448,25 +469,10 @@ impl VectorDataValue { #[cfg(test)] mod tests { - use crate::{ - format_member_data_key::{MemberDataKey, ParsedMemberDataKey}, - vector::CanonicalVector, - }; + use crate::vector::CanonicalVector; use super::*; - #[test] - fn member_data_key_round_trips_empty_binary_element() { - let encoded = MemberDataKey::new(b"vectors\0key", 42, b"") - .encode() - .expect("encode member key"); - let decoded = ParsedMemberDataKey::new(&encoded).expect("decode member key"); - - assert_eq!(decoded.key(), b"vectors\0key"); - assert_eq!(decoded.version(), 42); - assert_eq!(decoded.data(), b""); - } - #[test] fn vector_data_value_round_trips() { let canonical = CanonicalVector::from_values(&[3.0, 4.0]).expect("valid vector"); @@ -491,14 +497,15 @@ mod tests { #[test] fn vector_meta_round_trips() { - let mut meta = VectorMeta::new_after(2, 2, QuantizationType::None, 0); - meta.version = 42; + let mut meta = VectorMeta::new(2, 2, QuantizationType::None, 42); + meta.bump_data_revision(); let encoded = meta.encode(); let decoded = VectorMeta::decode(&encoded).expect("decode vector meta"); assert_eq!(decoded.count(), 2); assert_eq!(decoded.version(), 42); assert_eq!(decoded.dimension(), 2); + assert_eq!(decoded.data_revision(), 2); assert!(!decoded.is_stale()); } @@ -523,11 +530,16 @@ mod tests { zero_dimension[4..8].copy_from_slice(&0_u32.to_le_bytes()); assert!(VectorDataValue::decode(&zero_dimension).is_err()); + // payload_len that disagrees with the actual payload size. + let mut bad_payload_len = encoded_value.clone(); + bad_payload_len[12..16].copy_from_slice(&4_u32.to_le_bytes()); + assert!(VectorDataValue::decode(&bad_payload_len).is_err()); + let mut non_finite_payload = encoded_value; - non_finite_payload[12..16].copy_from_slice(&f32::NAN.to_le_bytes()); + non_finite_payload[16..20].copy_from_slice(&f32::NAN.to_le_bytes()); assert!(VectorDataValue::decode(&non_finite_payload).is_err()); - let encoded_meta = VectorMeta::new_after(2, 2, QuantizationType::None, 0).encode(); + let encoded_meta = VectorMeta::new(2, 2, QuantizationType::None, 1).encode(); assert!(VectorMeta::decode(&encoded_meta[..encoded_meta.len() - 1]).is_err()); let mut bad_meta_format = encoded_meta; @@ -539,4 +551,95 @@ mod tests { bad_metric[19] = 0xFF; assert!(VectorMeta::decode(&bad_metric).is_err()); } + + #[test] + fn canonical_vector_rejects_infinite_components() { + // NaN is covered in vector.rs; the same finiteness check must also + // reject both infinities in any component position. + assert!(CanonicalVector::from_values(&[f32::INFINITY, 1.0]).is_err()); + assert!(CanonicalVector::from_values(&[f32::NEG_INFINITY, 1.0]).is_err()); + assert!(CanonicalVector::from_values(&[1.0, f32::INFINITY]).is_err()); + assert!(CanonicalVector::from_values(&[1.0, f32::NEG_INFINITY]).is_err()); + let mut blob = Vec::new(); + blob.extend_from_slice(&f32::INFINITY.to_le_bytes()); + blob.extend_from_slice(&1.0_f32.to_le_bytes()); + assert!(CanonicalVector::from_fp32_le(&blob).is_err()); + } + + mod props { + use proptest::prelude::*; + + use super::*; + + proptest! { + #![proptest_config(ProptestConfig::with_cases(64))] + + /// Any combination of valid meta fields must round trip through + /// the frozen byte layout unchanged. The encoded bytes are crafted + /// directly so count/version/dimension/data_revision and the + /// timestamps can take arbitrary values (the struct setters cannot + /// reach every field). + #[test] + fn vector_meta_round_trips_arbitrary_valid_fields( + count in any::(), + version in any::(), + dimension in 1u32..=u32::MAX, + quantization in prop_oneof![ + Just(QuantizationType::None), + Just(QuantizationType::Binary), + Just(QuantizationType::Int8), + ], + data_revision in any::(), + ctime in any::(), + etime in any::(), + ) { + let mut raw = BytesMut::new(); + raw.put_u8(DataType::VectorSet as u8); + raw.put_u64_le(count); + raw.put_u64_le(version); + raw.put_u8(VECTOR_META_FORMAT); + raw.put_u8(quantization.to_u8()); + raw.put_u8(SimilarityMetric::Cosine.to_u8()); + raw.put_u8(0); // flags + raw.put_u32_le(dimension); + raw.put_u64_le(data_revision); + raw.put_u64_le(ctime); + raw.put_u64_le(etime); + + let decoded = VectorMeta::decode(&raw).expect("decode valid meta"); + prop_assert_eq!(decoded.count(), count); + prop_assert_eq!(decoded.version(), version); + prop_assert_eq!(decoded.dimension(), dimension); + prop_assert_eq!(decoded.quantization(), quantization); + prop_assert_eq!(decoded.metric(), SimilarityMetric::Cosine); + prop_assert_eq!(decoded.data_revision(), data_revision); + prop_assert_eq!(&decoded.encode()[..], &raw[..]); + } + + /// Any finite FP32 vector (dimension 1..=128, positive finite L2 + /// norm) must survive a NOQUANT encode/decode round trip exactly. + #[test] + fn vector_data_value_round_trips_finite_fp32( + values in prop::collection::vec(-1.0e6f32..1.0e6f32, 1..=128usize), + ) { + // All-zero vectors are rejected by design; skip those cases. + prop_assume!(CanonicalVector::from_values(&values).is_ok()); + let canonical = CanonicalVector::from_values(&values).expect("valid vector"); + let encoded = VectorDataValue::from_canonical(&canonical).encode(); + let decoded = VectorDataValue::decode(&encoded).expect("decode vector value"); + prop_assert_eq!(decoded.canonical(), &canonical); + } + + /// Arbitrary bytes fed to either decoder must only produce Ok or + /// Err: never a panic, an out-of-bounds read, or an allocation + /// sized by hostile header fields. + #[test] + fn decoders_never_panic_on_arbitrary_bytes( + data in prop::collection::vec(any::(), 0..=300), + ) { + let _ = VectorMeta::decode(&data); + let _ = VectorDataValue::decode(&data); + } + } + } } diff --git a/src/storage/src/format_vector_member_key.rs b/src/storage/src/format_vector_member_key.rs new file mode 100644 index 00000000..a261af6a --- /dev/null +++ b/src/storage/src/format_vector_member_key.rs @@ -0,0 +1,508 @@ +// Copyright (c) 2024-present, arana-db Community. All rights reserved. +// +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Vector set member data key codec (V1, frozen layout). +//! +//! This codec is vector-set specific. Hash/Set/ZSet keep using the shared +//! `MemberDataKey` in `format_member_data_key.rs`; do not unify them. +//! +//! Layout stored in VectorDataCF: +//! +//! | codec_version | key_len | user_key | storage_incarnation | generation_sequence | element | +//! | 1B | 4B BE | key_len | 8B BE | 8B BE | remainder | +//! +//! - `codec_version` is `VECTOR_MEMBER_KEY_CODEC_VERSION` (1). +//! - `key_len` is a u32 big-endian length of `user_key`; keys that do not fit +//! in u32 are rejected at encode time. +//! - `storage_incarnation` identifies the RocksDB instance that wrote the +//! member (see `storage_manifest.rs`). +//! - `generation_sequence` identifies one lifecycle of the vector set; it is +//! stored in `VectorMeta::version`. +//! - `element` occupies all remaining bytes and may be empty. +//! +//! All fixed-width integers are big-endian so members of one +//! (key, incarnation, generation) prefix stay contiguous under the default +//! RocksDB bytewise comparator. + +use bytes::{Buf, BufMut}; +use snafu::ensure; + +use crate::error::{InvalidArgumentSnafu, InvalidFormatSnafu, Result}; + +pub const VECTOR_MEMBER_KEY_CODEC_VERSION: u8 = 1; + +const CODEC_VERSION_LENGTH: usize = 1; +const KEY_LEN_LENGTH: usize = 4; +const INCARNATION_LENGTH: usize = 8; +const GENERATION_LENGTH: usize = 8; +const HEADER_LENGTH: usize = CODEC_VERSION_LENGTH + KEY_LEN_LENGTH; +const PREFIX_TRAILER_LENGTH: usize = INCARNATION_LENGTH + GENERATION_LENGTH; +const MIN_KEY_LENGTH: usize = HEADER_LENGTH + PREFIX_TRAILER_LENGTH; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct VectorMemberDataKey<'a> { + pub key: &'a [u8], + pub storage_incarnation: u64, + pub generation_sequence: u64, + pub element: &'a [u8], +} + +impl<'a> VectorMemberDataKey<'a> { + /// Encode the full member key including the element. + pub fn encode_full(&self) -> Result> { + let mut output = Self::encode_key_prefix(self.key)?; + output.put_u64(self.storage_incarnation); + output.put_u64(self.generation_sequence); + output.put_slice(self.element); + Ok(output) + } + + /// Encode the iteration prefix ending exactly after `generation_sequence`. + pub fn encode_prefix( + key: &[u8], + storage_incarnation: u64, + generation_sequence: u64, + ) -> Result> { + let mut output = Self::encode_key_prefix(key)?; + output.put_u64(storage_incarnation); + output.put_u64(generation_sequence); + Ok(output) + } + + /// Encode the prefix covering all incarnations and generations of `key` + /// (ends exactly after `user_key`). + pub fn encode_key_prefix(key: &[u8]) -> Result> { + let key_len = u32::try_from(key.len()).map_err(|_| { + InvalidArgumentSnafu { + message: format!("vector member key too long: {} bytes", key.len()), + } + .build() + })?; + let mut output = Vec::with_capacity(HEADER_LENGTH + key.len()); + output.put_u8(VECTOR_MEMBER_KEY_CODEC_VERSION); + output.put_u32(key_len); + output.put_slice(key); + Ok(output) + } + + /// Compute the exclusive upper bound for a prefix: the lexicographic + /// successor with trailing 0xFF bytes folded by incrementing the last + /// non-0xFF byte. Returns None when the prefix is empty or all 0xFF. + pub fn prefix_upper_bound(prefix: &[u8]) -> Option> { + let mut bound = prefix.to_vec(); + while let Some(last) = bound.last_mut() { + if *last == 0xFF { + bound.pop(); + } else { + *last += 1; + return Some(bound); + } + } + None + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ParsedVectorMemberDataKey<'a> { + key: &'a [u8], + storage_incarnation: u64, + generation_sequence: u64, + element: &'a [u8], +} + +impl<'a> ParsedVectorMemberDataKey<'a> { + pub fn decode(encoded_key: &'a [u8]) -> Result { + ensure!( + encoded_key.len() >= MIN_KEY_LENGTH, + InvalidFormatSnafu { + message: format!( + "vector member key too short: {} < {MIN_KEY_LENGTH}", + encoded_key.len() + ) + } + ); + + let mut reader = encoded_key; + let codec_version = reader.get_u8(); + ensure!( + codec_version == VECTOR_MEMBER_KEY_CODEC_VERSION, + InvalidFormatSnafu { + message: format!("unsupported vector member key codec version: {codec_version}") + } + ); + + let key_len = reader.get_u32() as usize; + // Bound-check against the remaining bytes before slicing so a hostile + // key_len can never cause an out-of-bounds read or a large allocation. + ensure!( + key_len <= reader.len().saturating_sub(PREFIX_TRAILER_LENGTH), + InvalidFormatSnafu { + message: format!( + "vector member key_len {} exceeds remaining {} bytes", + key_len, + reader.len() + ) + } + ); + + let key = &reader[..key_len]; + reader.advance(key_len); + let storage_incarnation = reader.get_u64(); + let generation_sequence = reader.get_u64(); + let element = reader; + + Ok(Self { + key, + storage_incarnation, + generation_sequence, + element, + }) + } + + pub fn key(&self) -> &'a [u8] { + self.key + } + + pub fn storage_incarnation(&self) -> u64 { + self.storage_incarnation + } + + pub fn generation_sequence(&self) -> u64 { + self.generation_sequence + } + + pub fn element(&self) -> &'a [u8] { + self.element + } +} + +#[allow(clippy::unwrap_used)] +#[cfg(test)] +mod tests { + use super::*; + + fn member<'a>(key: &'a [u8], element: &'a [u8]) -> VectorMemberDataKey<'a> { + VectorMemberDataKey { + key, + storage_incarnation: 0x0102_0304_0506_0708, + generation_sequence: 0x1112_1314_1516_1718, + element, + } + } + + #[test] + fn encode_full_matches_frozen_golden_layout() { + let encoded = member(b"vec", b"e1").encode_full().expect("encode"); + let mut expected = Vec::new(); + expected.push(1u8); // codec_version + expected.extend_from_slice(&3u32.to_be_bytes()); // key_len + expected.extend_from_slice(b"vec"); // user_key + expected.extend_from_slice(&0x0102_0304_0506_0708u64.to_be_bytes()); + expected.extend_from_slice(&0x1112_1314_1516_1718u64.to_be_bytes()); + expected.extend_from_slice(b"e1"); // element + assert_eq!(encoded, expected); + } + + #[test] + fn encode_decode_round_trip() { + for (key, element) in [ + (&b"vectors"[..], &b"member"[..]), + (b"vectors", b""), + (b"", b"member"), + (b"", b""), + (b"key\x00with\x00zeros", b"el\x00ement"), + (b"vec", &[0x00, 0xFF, 0x00][..]), + ] { + let encoded = member(key, element).encode_full().expect("encode"); + let parsed = ParsedVectorMemberDataKey::decode(&encoded).expect("decode"); + assert_eq!(parsed.key(), key); + assert_eq!(parsed.element(), element); + assert_eq!(parsed.storage_incarnation(), 0x0102_0304_0506_0708); + assert_eq!(parsed.generation_sequence(), 0x1112_1314_1516_1718); + } + } + + #[test] + fn encode_prefix_is_a_prefix_of_the_full_key() { + let full = member(b"vec\0key", b"elem") + .encode_full() + .expect("encode full"); + let prefix = VectorMemberDataKey::encode_prefix( + b"vec\0key", + 0x0102_0304_0506_0708, + 0x1112_1314_1516_1718, + ) + .expect("encode prefix"); + + assert!(full.starts_with(&prefix)); + assert_eq!( + prefix.len(), + HEADER_LENGTH + b"vec\0key".len() + PREFIX_TRAILER_LENGTH + ); + let parsed = ParsedVectorMemberDataKey::decode(&full).expect("decode"); + assert_eq!(&full[prefix.len()..], parsed.element()); + } + + #[test] + fn keys_with_prefix_relationship_do_not_collide() { + // "vec" is a prefix of "vectors"; key_len keeps their encodings apart. + let short = VectorMemberDataKey::encode_prefix(b"vec", 1, 1).expect("short"); + let long = VectorMemberDataKey::encode_prefix(b"vectors", 1, 1).expect("long"); + assert!(!long.starts_with(&short)); + assert!(!short.starts_with(&long)); + } + + #[test] + fn same_prefix_members_sort_contiguously() { + let generation = 7u64; + let make = |element: &[u8]| { + VectorMemberDataKey { + key: b"vec", + storage_incarnation: 9, + generation_sequence: generation, + element, + } + .encode_full() + .expect("encode") + }; + let mut keys = [make(b"b"), make(b""), make(b"a")]; + keys.sort(); + let prefix = VectorMemberDataKey::encode_prefix(b"vec", 9, generation).expect("prefix"); + assert!(keys.iter().all(|key| key.starts_with(&prefix))); + assert!(keys.windows(2).all(|pair| pair[0] < pair[1])); + + // A different generation of the same key sorts outside the prefix range. + let other = VectorMemberDataKey { + key: b"vec", + storage_incarnation: 9, + generation_sequence: generation + 1, + element: b"", + } + .encode_full() + .expect("encode"); + assert!(!other.starts_with(&prefix)); + assert!(other > *keys.last().expect("non-empty")); + } + + #[test] + fn decode_rejects_malformed_keys_without_panicking() { + let valid = member(b"vec", b"e").encode_full().expect("encode"); + + // Truncated at every boundary, including mid-header. Trimming the + // single element byte would still be a valid key (empty element), so + // the loop stops one byte earlier. + for len in 0..valid.len() - 1 { + assert!( + ParsedVectorMemberDataKey::decode(&valid[..len]).is_err(), + "truncated length {len} must be rejected" + ); + } + + // Bad codec version. + let mut bad_version = valid.clone(); + bad_version[0] = 2; + assert!(ParsedVectorMemberDataKey::decode(&bad_version).is_err()); + + // key_len larger than the remaining bytes (no huge allocation). + let mut bad_len = valid.clone(); + bad_len[1..5].copy_from_slice(&u32::MAX.to_be_bytes()); + assert!(ParsedVectorMemberDataKey::decode(&bad_len).is_err()); + + // key_len consuming the incarnation/generation trailer. + let mut greedy_len = valid.clone(); + greedy_len[1..5].copy_from_slice(&10u32.to_be_bytes()); + assert!(ParsedVectorMemberDataKey::decode(&greedy_len).is_err()); + } + + #[test] + fn prefix_upper_bound_increments_and_carries() { + assert_eq!( + VectorMemberDataKey::prefix_upper_bound(&[0x01, 0x00, 0x01]), + Some(vec![0x01, 0x00, 0x02]) + ); + // 0xFF carry folds into the preceding byte. + assert_eq!( + VectorMemberDataKey::prefix_upper_bound(&[0x01, 0xFF, 0xFF]), + Some(vec![0x02]) + ); + assert_eq!(VectorMemberDataKey::prefix_upper_bound(&[0xFF, 0xFF]), None); + assert_eq!(VectorMemberDataKey::prefix_upper_bound(&[]), None); + + // A real prefix upper bound excludes every member of the prefix. + let prefix = + VectorMemberDataKey::encode_prefix(b"vec", u64::MAX, u64::MAX).expect("prefix"); + let upper = VectorMemberDataKey::prefix_upper_bound(&prefix).expect("bound"); + let member_key = VectorMemberDataKey { + key: b"vec", + storage_incarnation: u64::MAX, + generation_sequence: u64::MAX, + element: b"x", + } + .encode_full() + .expect("encode"); + assert!(member_key >= prefix); + assert!(member_key < upper); + } + + mod props { + use std::cmp::Ordering; + + use proptest::prelude::*; + + use super::*; + + fn arb_bytes(max_len: usize) -> impl Strategy> { + prop::collection::vec(any::(), 0..=max_len) + } + + proptest! { + #![proptest_config(ProptestConfig::with_cases(64))] + + /// Any (key, incarnation, generation, element) tuple must survive + /// an encode/decode round trip with every field intact, including + /// empty and 0x00-containing keys and elements. + #[test] + fn encode_decode_round_trip_preserves_all_fields( + key in arb_bytes(1024), + incarnation in any::(), + generation in any::(), + element in arb_bytes(512), + ) { + let member = VectorMemberDataKey { + key: &key, + storage_incarnation: incarnation, + generation_sequence: generation, + element: &element, + }; + let encoded = member.encode_full().expect("encode full key"); + let parsed = + ParsedVectorMemberDataKey::decode(&encoded).expect("decode encoded key"); + prop_assert_eq!(parsed.key(), &key[..]); + prop_assert_eq!(parsed.storage_incarnation(), incarnation); + prop_assert_eq!(parsed.generation_sequence(), generation); + prop_assert_eq!(parsed.element(), &element[..]); + } + + /// `encode_key_prefix` ends after the user key, `encode_prefix` + /// after the generation, and `encode_full` after the element, so + /// each must be a prefix of the next. + #[test] + fn prefixes_are_strictly_nested( + key in arb_bytes(256), + incarnation in any::(), + generation in any::(), + element in arb_bytes(128), + ) { + let key_prefix = + VectorMemberDataKey::encode_key_prefix(&key).expect("encode key prefix"); + let prefix = VectorMemberDataKey::encode_prefix(&key, incarnation, generation) + .expect("encode prefix"); + let full = VectorMemberDataKey { + key: &key, + storage_incarnation: incarnation, + generation_sequence: generation, + element: &element, + } + .encode_full() + .expect("encode full key"); + + // The trailer always makes the key prefix strictly shorter. + prop_assert!(prefix.starts_with(&key_prefix)); + prop_assert!(key_prefix.len() < prefix.len()); + + prop_assert!(full.starts_with(&prefix)); + prop_assert!(prefix.len() <= full.len()); + // Strictly a prefix whenever the element is non-empty. + prop_assert_eq!(prefix.len() < full.len(), !element.is_empty()); + } + + /// Every member of one (key, incarnation, generation) prefix must + /// sort at or after the prefix and, when a lexicographic successor + /// exists, strictly below the exclusive upper bound. + #[test] + fn member_keys_stay_inside_prefix_range( + key in arb_bytes(256), + incarnation in any::(), + generation in any::(), + element in arb_bytes(128), + ) { + let prefix = VectorMemberDataKey::encode_prefix(&key, incarnation, generation) + .expect("encode prefix"); + let full = VectorMemberDataKey { + key: &key, + storage_incarnation: incarnation, + generation_sequence: generation, + element: &element, + } + .encode_full() + .expect("encode full key"); + + prop_assert!(full >= prefix); + if let Some(upper) = VectorMemberDataKey::prefix_upper_bound(&prefix) { + prop_assert!(full < upper); + } + } + + /// Prefixes of the same key order by (incarnation, generation) in + /// big-endian byte order, and members of an earlier prefix always + /// sort before members of a later one under the bytewise + /// comparator, regardless of the element bytes. + #[test] + fn prefixes_order_by_incarnation_then_generation( + key in arb_bytes(64), + left in (any::(), any::()), + right in (any::(), any::()), + left_element in arb_bytes(32), + right_element in arb_bytes(32), + ) { + let left_prefix = + VectorMemberDataKey::encode_prefix(&key, left.0, left.1).expect("left prefix"); + let right_prefix = VectorMemberDataKey::encode_prefix(&key, right.0, right.1) + .expect("right prefix"); + prop_assert_eq!(left_prefix.cmp(&right_prefix), left.cmp(&right)); + + if left.cmp(&right) == Ordering::Less { + let left_full = VectorMemberDataKey { + key: &key, + storage_incarnation: left.0, + generation_sequence: left.1, + element: &left_element, + } + .encode_full() + .expect("left full"); + let right_full = VectorMemberDataKey { + key: &key, + storage_incarnation: right.0, + generation_sequence: right.1, + element: &right_element, + } + .encode_full() + .expect("right full"); + prop_assert!(left_full < right_full); + } + } + + /// Feeding arbitrary bytes to decode must only ever produce Ok or + /// Err: never a panic, an out-of-bounds read, or a huge allocation + /// driven by a hostile key_len. + #[test] + fn decode_arbitrary_bytes_never_panics(data in arb_bytes(300)) { + let _ = ParsedVectorMemberDataKey::decode(&data); + } + } + } +} diff --git a/src/storage/src/lib.rs b/src/storage/src/lib.rs index b146db29..3e6bbb0e 100644 --- a/src/storage/src/lib.rs +++ b/src/storage/src/lib.rs @@ -20,7 +20,13 @@ extern crate core; mod format_base_data_value; mod format_member_data_key; pub mod format_vector; +pub mod format_vector_member_key; +mod storage_manifest; pub mod vector; +pub mod vector_fault; +mod vector_flat; +pub mod vector_metrics; +pub mod vector_mutation; mod data_compaction_filter; mod durable_fs; @@ -68,8 +74,8 @@ pub mod logindex; pub use batch::{AppendLogFn, Batch, BinlogBatch, RocksBatch}; pub use checkpoint::{ - PreparedCheckpointRestore, RAFT_SNAPSHOT_META_FILE, RaftSnapshotMeta, - prepare_checkpoint_restore, restore_checkpoint_layout, + CURRENT_SNAPSHOT_VERSION, PreparedCheckpointRestore, RAFT_SNAPSHOT_META_FILE, RaftSnapshotMeta, + STORAGE_SCHEMA_VERSION, prepare_checkpoint_restore, restore_checkpoint_layout, }; pub use durable_fs::{sync_directory, sync_parent_directory}; pub use error::Result; @@ -78,12 +84,20 @@ pub use format_base_key::BaseMetaKey; pub use format_base_value::*; pub use format_zset_score_key::{ScoreMember, ZsetScoreMember}; pub use options::StorageOptions; -pub use redis::{ColumnFamilyIndex, Redis, TypeCheckState}; +pub use redis::{ColumnFamilyIndex, GenerationProvider, Redis, TypeCheckState}; +pub use redis_vectors::VectorDataSample; pub use statistics::KeyStatistics; pub use storage::{BgTask, BgTaskHandler}; pub use storage_impl::BeforeOrAfter; +pub use storage_manifest::STORAGE_MANIFEST_FILE; pub use util::{safe_cleanup_test_db, unique_test_db_path}; pub use vector::{ - CanonicalVector, QuantizationType, VectorHit, VectorQuery, VectorSearchEngine, + CanonicalVector, QuantizationType, VectorHit, VectorInfo, VectorQuery, VectorSearchEngine, VectorSearchMode, VectorSearchOptions, }; +pub use vector_fault::VectorFaultHooks; +pub use vector_flat::{FlatQueryCancel, FlatQueryGate, FlatScanGuard}; +pub use vector_metrics::{VectorMetrics, VectorMetricsSnapshot}; +pub use vector_mutation::{ + VectorSetApplyError, VectorSetApplyResult, VectorSetBusinessError, VectorSetMutationV1, +}; diff --git a/src/storage/src/options.rs b/src/storage/src/options.rs index 650809e9..6a1d5acd 100644 --- a/src/storage/src/options.rs +++ b/src/storage/src/options.rs @@ -74,6 +74,8 @@ pub struct StorageOptions { pub max_gap: i64, /// Memory manager size pub mem_manager_size: usize, + /// Vector Set feature configuration + pub vector: conf::vector_config::VectorConfig, } impl Default for StorageOptions { @@ -100,6 +102,7 @@ impl Default for StorageOptions { raft_timeout_s: u32::MAX, max_gap: 1000, mem_manager_size: 100_000_000, + vector: conf::vector_config::VectorConfig::default(), } } } @@ -119,6 +122,7 @@ impl StorageOptions { small_compaction_threshold: config.small_compaction_threshold, small_compaction_duration_threshold: config.small_compaction_duration_threshold, db_instance_num: config.db_instance_num, + vector: config.vector.clone(), ..Self::default() } } diff --git a/src/storage/src/redis.rs b/src/storage/src/redis.rs index e3af0459..b2cd5f40 100644 --- a/src/storage/src/redis.rs +++ b/src/storage/src/redis.rs @@ -51,6 +51,14 @@ use crate::options::{OptionType, StorageOptions}; use crate::statistics::KeyStatistics; use crate::storage::BgTaskHandler; use crate::storage_define::TYPE_LENGTH; +use crate::storage_manifest::StorageManifest; + +/// Injection point for vector set generation sequences. +/// +/// Standalone mode leaves this unset and falls back to the persistent +/// manifest generator. Cluster mode installs a provider returning the Raft +/// log index that created the key (wired up by the raft layer later). +pub type GenerationProvider = Arc Result + Send + Sync>; // Import logindex types for use in Storage @@ -71,6 +79,17 @@ impl ColumnFamilyIndex { /// This constant is used by batch.rs for validation. pub const COUNT: usize = 7; + /// All column families in declaration order (by discriminant). + pub const ALL: [ColumnFamilyIndex; Self::COUNT] = [ + ColumnFamilyIndex::MetaCF, + ColumnFamilyIndex::HashesDataCF, + ColumnFamilyIndex::SetsDataCF, + ColumnFamilyIndex::ListsDataCF, + ColumnFamilyIndex::ZsetsDataCF, + ColumnFamilyIndex::ZsetsScoreCF, + ColumnFamilyIndex::VectorDataCF, + ]; + pub fn name(&self) -> &'static str { match self { ColumnFamilyIndex::MetaCF => "default", @@ -284,6 +303,24 @@ pub struct Redis { // For cluster mode: when set, create_batch returns a BinlogBatch. pub append_log_fn: OnceLock, + + // Per-instance storage identity (incarnation) and the standalone-mode + // generation sequence generator, loaded in open(). + pub(crate) manifest: Option, + + // Cluster-mode injection point for vector generation sequences. + pub generation_provider: OnceLock, + + // Concurrency gate for FLAT vector queries, sized from + // `StorageOptions.vector.max_concurrent_flat_queries`. + pub flat_query_gate: crate::vector_flat::FlatQueryGate, + + // Counters for FLAT vector query execution, surfaced via INFO VECTOR. + pub vector_metrics: crate::vector_metrics::VectorMetrics, + + // Fault injection hooks for the vector set storage path; armed by + // tests only, all flags off in production. + pub vector_fault_hooks: crate::vector_fault::VectorFaultHooks, } impl Redis { @@ -300,6 +337,9 @@ impl Redis { let statistics_store: Cache = CacheBuilder::new(storage.statistics_max_size).build(); + let flat_query_gate = + crate::vector_flat::FlatQueryGate::new(storage.vector.max_concurrent_flat_queries); + Self { index, is_starting: AtomicBool::new(true), @@ -312,6 +352,9 @@ impl Redis { write_options: WriteOptions::default(), read_options: ReadOptions::default(), compact_options, + flat_query_gate, + vector_metrics: crate::vector_metrics::VectorMetrics::default(), + vector_fault_hooks: crate::vector_fault::VectorFaultHooks::default(), statistics_store: Arc::new(statistics_store), scan_cursors_store: Mutex::new(CacheBuilder::new(5000).build()), @@ -325,6 +368,8 @@ impl Redis { logindex_cf_tracker: None, append_log_fn: OnceLock::new(), + manifest: None, + generation_provider: OnceLock::new(), } } @@ -342,6 +387,12 @@ impl Redis { // background threads after open) can reach the live DB handle. let db_once_cell: Arc>> = Arc::new(OnceCell::new()); + // Shared with the data compaction filters so the VectorSet branch can + // compare member-key incarnations against this instance's identity. + // Populated once the storage manifest is loaded below; filters treat + // an unset cell as "cannot prove stale" and keep the data. + let incarnation_cell: Arc> = Arc::new(OnceCell::new()); + // Snapshot trigger: emitted by the LogIndex purger every N flushes to suggest // that the Raft layer build a snapshot. Wiring this into Raft requires a // back-channel from storage -> raft that does not currently exist (the raft @@ -411,6 +462,7 @@ impl Redis { *use_bloom, *block_size, Some(&db_once_cell), + &incarnation_cell, &collector, ) }) @@ -440,6 +492,23 @@ impl Redis { self.logindex_cf_tracker = Some(cf_tracker); self.db = Some(db); + // Load (or create) the per-instance storage manifest. A manifest that + // is missing while the database holds data means the data predates + // the incarnation mechanism; refuse to open rather than reinterpret. + let has_entries = self.handles.iter().any(|name| { + let Some(db) = self.db.as_ref() else { + return false; + }; + db.cf_handle(name).is_some_and(|cf| { + db.iterator_cf(&cf, rocksdb::IteratorMode::Start) + .next() + .is_some() + }) + }); + let manifest = StorageManifest::open(Path::new(db_path), has_entries)?; + let _ = incarnation_cell.set(manifest.storage_incarnation()); + self.manifest = Some(manifest); + Ok(()) } @@ -450,6 +519,7 @@ impl Redis { use_bloom_filter: bool, block_size: Option, db_once_cell: Option<&Arc>>>, + incarnation_cell: &Arc>, collector: &Arc, ) -> ColumnFamilyDescriptor { let mut cf_opts = storage_options.options.clone(); @@ -507,7 +577,11 @@ impl Redis { .find(|cf| cf.name() == cf_name) .and_then(|cf| cf.data_type()) { - let factory = DataCompactionFilterFactory::new(Arc::clone(db_once_cell), data_type); + let factory = DataCompactionFilterFactory::new( + Arc::clone(db_once_cell), + data_type, + Arc::clone(incarnation_cell), + ); cf_opts.set_compaction_filter_factory(factory); } @@ -530,7 +604,14 @@ impl Redis { message: "Database is not initialized".to_string(), })?; let checkpoint = rocksdb::checkpoint::Checkpoint::new(db.as_ref()).context(RocksSnafu)?; - checkpoint.create_checkpoint(path).context(RocksSnafu) + checkpoint.create_checkpoint(path).context(RocksSnafu)?; + // RocksDB checkpoints only cover database files; copy the storage + // manifest alongside so a restored snapshot keeps the storage + // incarnation and generation sequence that wrote the data. + if let Some(manifest) = self.manifest.as_ref() { + manifest.copy_to(path)?; + } + Ok(()) } /// Compact database range @@ -586,6 +667,39 @@ impl Redis { let _ = self.append_log_fn.set(f); } + /// The stable identity of this storage instance, embedded in every vector + /// member key written by it. + pub fn storage_incarnation(&self) -> Result { + self.manifest + .as_ref() + .map(StorageManifest::storage_incarnation) + .context(OptionNoneSnafu { + message: "storage manifest is not initialized".to_string(), + }) + } + + /// Allocate the generation sequence for a newly created vector set. + /// + /// Uses the injected provider when present (cluster mode: the creating + /// Raft log index), otherwise the persistent manifest generator. + pub(crate) fn allocate_vector_generation(&self) -> Result { + if let Some(provider) = self.generation_provider.get() { + return provider(); + } + self.manifest + .as_ref() + .context(OptionNoneSnafu { + message: "storage manifest is not initialized".to_string(), + })? + .allocate_generation() + } + + /// Inject the cluster-mode generation provider. Idempotent: subsequent + /// calls are ignored (OnceLock semantics). + pub fn set_generation_provider(&self, provider: GenerationProvider) { + let _ = self.generation_provider.set(provider); + } + /// Create a new batch for atomic write operations. /// /// This method creates a batch appropriate for the current deployment mode: diff --git a/src/storage/src/redis_strings.rs b/src/storage/src/redis_strings.rs index 1a7515c4..7dbd48de 100644 --- a/src/storage/src/redis_strings.rs +++ b/src/storage/src/redis_strings.rs @@ -33,6 +33,7 @@ use crate::{ format_base_meta_value::ParsedBaseMetaValue, format_list_meta_value::ParsedListsMetaValue, format_strings_value::{ParsedStringsValue, StringValue}, + format_vector_member_key::VectorMemberDataKey, redis_sets::glob_match_bytes, }; @@ -2020,9 +2021,10 @@ impl Redis { let encoded = string_key.encode()?; // Build correct prefix for data CF scanning: - // Data keys format: | reserve1 (8B) | encoded_user_key | version (8B) | data | reserve2 | + // Shared data key format: | reserve1 (8B) | encoded_user_key | version (8B) | data | reserve2 | // We need prefix: | reserve1 (8B) | encoded_user_key (with \x00\x00 delimiter) | // Note: BaseKey.encode() includes reserve2, which would not match data keys. + // (VectorDataCF uses the V1 vector codec; its prefix is built per-CF below.) let mut data_key_prefix = bytes::BytesMut::with_capacity(PREFIX_RESERVE_LENGTH + key.len() * 2 + 2); data_key_prefix.put_slice(&[0u8; PREFIX_RESERVE_LENGTH]); @@ -2050,14 +2052,21 @@ impl Redis { ] { let cf_handle = self.get_cf_handle(cf_index); if let Some(cf) = cf_handle { + // Vector member keys use the V1 vector codec; all other + // data CFs share the reserve1 | encoded_user_key layout. + let prefix = if cf_index == ColumnFamilyIndex::VectorDataCF { + VectorMemberDataKey::encode_key_prefix(key)? + } else { + data_key_prefix.to_vec() + }; // Prefix-scan data CF and delete all derived keys let iter = db.iterator_cf( &cf, - rocksdb::IteratorMode::From(&data_key_prefix, rocksdb::Direction::Forward), + rocksdb::IteratorMode::From(&prefix, rocksdb::Direction::Forward), ); for item in iter { let (k, _) = item.context(RocksSnafu)?; - if !k.starts_with(&data_key_prefix) { + if !k.starts_with(&prefix) { break; } keys_to_delete.push((cf_index, k.to_vec())); diff --git a/src/storage/src/redis_vectors.rs b/src/storage/src/redis_vectors.rs index 1f407b98..ccc6a8ba 100644 --- a/src/storage/src/redis_vectors.rs +++ b/src/storage/src/redis_vectors.rs @@ -15,23 +15,45 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::time::{Duration, Instant}; + use kstd::lock_mgr::ScopeRecordLock; -use rocksdb::{Direction, IteratorMode, ReadOptions}; +use rocksdb::{IteratorMode, ReadOptions}; use snafu::{OptionExt, ResultExt}; use crate::{ CanonicalVector, ColumnFamilyIndex, DataType, Redis, Result, TypeCheckState, VectorHit, - VectorQuery, VectorSearchEngine, VectorSearchOptions, + VectorQuery, VectorSearchEngine, VectorSearchMode, VectorSearchOptions, error::{ - InvalidArgumentSnafu, InvalidFormatSnafu, KeyNotFoundSnafu, OptionNoneSnafu, RocksSnafu, + BatchSnafu, InvalidArgumentSnafu, InvalidFormatSnafu, KeyNotFoundSnafu, OptionNoneSnafu, + RedisErrSnafu, RocksSnafu, SystemSnafu, VectorFlatQueryTimeoutSnafu, }, format_base_key::BaseMetaKey, - format_member_data_key::MemberDataKey, format_vector::{VectorDataValue, VectorMeta}, - storage_define::SUFFIX_RESERVE_LENGTH, + format_vector_member_key::{ParsedVectorMemberDataKey, VectorMemberDataKey}, + vector_flat::{FlatQueryCancel, FlatScanGuard}, + vector_mutation::{ + VectorSetApplyError, VectorSetApplyResult, VectorSetBusinessError, VectorSetMutationV1, + }, }; +/// Clamp for the configured FLAT query timeout so the `Instant` deadline +/// computation cannot overflow even for pathological config values. +const MAX_FLAT_QUERY_TIMEOUT: Duration = Duration::from_secs(3600); + impl Redis { + /// Build the V1 member key for `element` under the set's current + /// generation (stored in `meta.version`). + fn vector_member_key(&self, key: &[u8], generation: u64, element: &[u8]) -> Result> { + VectorMemberDataKey { + key, + storage_incarnation: self.storage_incarnation()?, + generation_sequence: generation, + element, + } + .encode_full() + } + /// Decode raw meta bytes into a `VectorMeta` without liveness filtering: /// stale or emptied sets are still returned so callers can inspect their /// version. Returns an error when the key holds another live data type. @@ -59,14 +81,89 @@ impl Redis { } pub fn vadd(&self, key: &[u8], element: &[u8], vector: &CanonicalVector) -> Result { + // Resource limits apply before any write so standalone and cluster + // modes reject oversized input identically. + let vector_config = &self.storage.vector; + if vector.dimension() > vector_config.max_dimension { + return RedisErrSnafu { + message: "ERR vector dimension exceeds max_dimension".to_string(), + } + .fail(); + } + if element.len() > vector_config.max_element_bytes { + return RedisErrSnafu { + message: "ERR vector element exceeds max_element_bytes".to_string(), + } + .fail(); + } + // The client supplies either an FP32 blob or VALUES floats; both are + // `dimension * 4` bytes on the wire. + if u64::from(vector.dimension()) * size_of::() as u64 + > vector_config.max_vector_bytes as u64 + { + return RedisErrSnafu { + message: "ERR vector exceeds max_vector_bytes".to_string(), + } + .fail(); + } + + let mutation = VectorSetMutationV1::add_from_canonical(element, vector)?; + let result = self.apply_vector_set_mutation(key, &mutation, None)?; + Ok(matches!(result, VectorSetApplyResult::Added)) + } + + pub fn vrem(&self, key: &[u8], element: &[u8]) -> Result { + let mutation = VectorSetMutationV1::Remove { + element: element.to_vec(), + }; + let result = self.apply_vector_set_mutation(key, &mutation, None)?; + Ok(matches!(result, VectorSetApplyResult::Removed)) + } + + /// Single decision point for vector-set mutations. + /// + /// Applies a logical mutation to `key`: reads the current meta, decides + /// the outcome (create / add / update / remove / miss), and commits the + /// derived state (member record, count, data_revision, generation on + /// create) in one atomic batch. Standalone VADD/VREM are thin wrappers + /// around this entry; the Raft state machine will replay logical + /// mutations through it in log order so every replica derives identical + /// state. + /// + /// `create_generation` is the generation assigned when the mutation + /// creates the set (cluster mode: the creating Raft log index); `None` + /// falls back to `allocate_vector_generation`. + /// + /// Business rejections (WRONGTYPE, dimension mismatch) are returned as + /// deterministic `VectorSetApplyError::Business` errors; everything else + /// is a fatal storage error. + pub fn apply_vector_set_mutation( + &self, + key: &[u8], + mutation: &VectorSetMutationV1, + create_generation: Option, + ) -> std::result::Result { + match mutation { + VectorSetMutationV1::Add { element, .. } => { + let vector = mutation.canonical_vector()?.context(OptionNoneSnafu { + message: "add mutation carries no vector".to_string(), + })?; + self.apply_vector_add(key, element, vector, create_generation) + } + VectorSetMutationV1::Remove { element } => self.apply_vector_remove(key, element), + } + } + + fn apply_vector_add( + &self, + key: &[u8], + element: &[u8], + vector: CanonicalVector, + create_generation: Option, + ) -> std::result::Result { let db = self.db.as_ref().context(OptionNoneSnafu { message: "db is not initialized".to_string(), })?; - let meta_cf = self - .get_cf_handle(ColumnFamilyIndex::MetaCF) - .context(OptionNoneSnafu { - message: "MetaCF is not initialized".to_string(), - })?; let vector_cf = self .get_cf_handle(ColumnFamilyIndex::VectorDataCF) .context(OptionNoneSnafu { @@ -76,40 +173,34 @@ impl Redis { let lock_key = String::from_utf8_lossy(key); let _lock = ScopeRecordLock::new(self.lock_mgr.as_ref(), &lock_key); let meta_key = BaseMetaKey::new(key).encode()?; - let stored_raw = db.get_cf(&meta_cf, &meta_key).context(RocksSnafu)?; - let stored_meta = match stored_raw.as_deref() { - Some(value) => self.decode_vector_meta(value)?, - None => None, - }; - // Keep the previous version even for stale or emptied sets so a - // recreated set always gets a fresh, monotonically increasing generation. - let previous_generation = stored_meta.as_ref().map_or(0, VectorMeta::version); - let live_meta = stored_meta.filter(|meta| !meta.is_stale() && meta.count() != 0); + let live_meta = self.read_live_vector_meta_for_apply(&meta_key)?; let is_new_set = live_meta.is_none(); let mut meta = match live_meta { Some(meta) => { if meta.dimension() != vector.dimension() { - return InvalidArgumentSnafu { - message: format!( - "vector dimension mismatch: expected {}, got {}", - meta.dimension(), - vector.dimension() - ), - } - .fail(); + return Err(VectorSetApplyError::Business( + VectorSetBusinessError::DimensionMismatch { + expected: meta.dimension(), + got: vector.dimension(), + }, + )); } meta } - None => VectorMeta::new_after( - 1, - vector.dimension(), - vector.quantization(), - previous_generation, - ), + // Creating (or recreating after expiry/deletion) allocates a fresh + // generation sequence from the persistent generator, so stale + // members of a previous lifecycle never collide with the new set. + None => { + let generation = match create_generation { + Some(generation) => generation, + None => self.allocate_vector_generation()?, + }; + VectorMeta::new(1, vector.dimension(), vector.quantization(), generation) + } }; - let member_key = MemberDataKey::new(key, meta.version(), element).encode()?; + let member_key = self.vector_member_key(key, meta.version(), element)?; let inserted = if is_new_set { true } else { @@ -118,12 +209,13 @@ impl Redis { .is_none() }; if inserted && !is_new_set { - let count = meta.count().checked_add(1); - let Some(count) = count else { - return InvalidArgumentSnafu { - message: "vector set size overflow".to_string(), - } - .fail(); + let Some(count) = meta.count().checked_add(1) else { + return Err(VectorSetApplyError::Storage( + InvalidArgumentSnafu { + message: "vector set size overflow".to_string(), + } + .build(), + )); }; meta.set_count(count); } @@ -132,15 +224,37 @@ impl Redis { // quantization regardless of the form the client supplied. let vector = vector.to_quantized(meta.quantization())?; let member_value = VectorDataValue::from_canonical(&vector).encode(); + if !is_new_set { + meta.bump_data_revision(); + } let meta_value = meta.encode(); let mut batch = self.create_batch()?; batch.put(ColumnFamilyIndex::VectorDataCF, &member_key, &member_value)?; batch.put(ColumnFamilyIndex::MetaCF, &meta_key, &meta_value)?; + if self.vector_fault_hooks.fail_batch_commit() { + // Dropping the uncommitted batch leaves meta and member + // untouched, exactly like a failed commit: the mutation is + // all-or-nothing. + return Err(VectorSetApplyError::Storage( + BatchSnafu { + message: "injected fault: vector batch commit failed".to_string(), + } + .build(), + )); + } batch.commit()?; - Ok(inserted) + Ok(if inserted { + VectorSetApplyResult::Added + } else { + VectorSetApplyResult::Updated + }) } - pub fn vrem(&self, key: &[u8], element: &[u8]) -> Result { + fn apply_vector_remove( + &self, + key: &[u8], + element: &[u8], + ) -> std::result::Result { let db = self.db.as_ref().context(OptionNoneSnafu { message: "db is not initialized".to_string(), })?; @@ -152,31 +266,90 @@ impl Redis { let lock_key = String::from_utf8_lossy(key); let _lock = ScopeRecordLock::new(self.lock_mgr.as_ref(), &lock_key); - let Some(mut meta) = self.read_vector_meta(key)? else { - return Ok(false); + let meta_key = BaseMetaKey::new(key).encode()?; + let Some(mut meta) = self.read_live_vector_meta_for_apply(&meta_key)? else { + return Ok(VectorSetApplyResult::RemoveMissed); }; - let member_key = MemberDataKey::new(key, meta.version(), element).encode()?; + let member_key = self.vector_member_key(key, meta.version(), element)?; if db .get_cf(&vector_cf, &member_key) .context(RocksSnafu)? .is_none() { - return Ok(false); + return Ok(VectorSetApplyResult::RemoveMissed); } - let meta_key = BaseMetaKey::new(key).encode()?; let mut batch = self.create_batch()?; batch.delete(ColumnFamilyIndex::VectorDataCF, &member_key)?; if meta.count() > 1 { meta.set_count(meta.count() - 1); + meta.bump_data_revision(); let meta_value = meta.encode(); batch.put(ColumnFamilyIndex::MetaCF, &meta_key, &meta_value)?; } else { batch.delete(ColumnFamilyIndex::MetaCF, &meta_key)?; } + if self.vector_fault_hooks.fail_batch_commit() { + // Dropping the uncommitted batch leaves meta and member + // untouched, exactly like a failed commit: the mutation is + // all-or-nothing. + return Err(VectorSetApplyError::Storage( + BatchSnafu { + message: "injected fault: vector batch commit failed".to_string(), + } + .build(), + )); + } batch.commit()?; - Ok(true) + Ok(VectorSetApplyResult::Removed) + } + + /// Read the live vector meta for `meta_key` during mutation apply. + /// + /// A live non-vector value is classified as the WRONGTYPE business error + /// so a Raft apply loop can treat it as deterministic; decode failures of + /// vector metadata indicate corruption and stay fatal storage errors. + fn read_live_vector_meta_for_apply( + &self, + meta_key: &[u8], + ) -> std::result::Result, VectorSetApplyError> { + let db = self.db.as_ref().context(OptionNoneSnafu { + message: "db is not initialized".to_string(), + })?; + let meta_cf = self + .get_cf_handle(ColumnFamilyIndex::MetaCF) + .context(OptionNoneSnafu { + message: "MetaCF is not initialized".to_string(), + })?; + if self.vector_fault_hooks.fail_meta_read() { + return Err(VectorSetApplyError::Storage( + SystemSnafu { + message: "injected fault: vector meta read failed".to_string(), + } + .build(), + )); + } + let stored_raw = db.get_cf(&meta_cf, meta_key).context(RocksSnafu)?; + let Some(value) = stored_raw.as_deref() else { + return Ok(None); + }; + if value.is_empty() { + return Ok(None); + } + if value[0] != DataType::VectorSet as u8 { + // Mirrors `decode_vector_meta`: stale non-vector metadata counts + // as absent, a live non-vector value is WRONGTYPE. + return match self.is_stale(value) { + Ok(true) => Ok(None), + Ok(false) => Err(VectorSetApplyError::Business( + VectorSetBusinessError::WrongType, + )), + Err(error) => Err(VectorSetApplyError::Storage(error)), + }; + } + let meta = VectorMeta::decode(value)?; + Ok((!meta.is_stale() && meta.count() != 0).then_some(meta)) } pub fn vcard(&self, key: &[u8]) -> Result { @@ -193,6 +366,16 @@ impl Redis { } } + /// O(1) per-set metadata for VINFO; `None` when the key is missing (or + /// stale/emptied), WRONGTYPE when it holds another live data type. + pub fn vinfo(&self, key: &[u8]) -> Result> { + Ok(self.read_vector_meta(key)?.map(|meta| crate::VectorInfo { + dimension: meta.dimension(), + size: meta.count(), + generation: meta.version(), + })) + } + pub fn vemb(&self, key: &[u8], element: &[u8]) -> Result>> { let db = self.db.as_ref().context(OptionNoneSnafu { message: "db is not initialized".to_string(), @@ -209,7 +392,7 @@ impl Redis { let Some(meta) = self.read_vector_meta_opt(key, Some(&read_options))? else { return Ok(None); }; - let member_key = MemberDataKey::new(key, meta.version(), element).encode()?; + let member_key = self.vector_member_key(key, meta.version(), element)?; let Some(value_raw) = db .get_cf_opt(&vector_cf, &member_key, &read_options) .context(RocksSnafu)? @@ -246,7 +429,7 @@ impl Redis { let Some(meta) = self.read_vector_meta_opt(key, Some(&read_options))? else { return Ok(false); }; - let member_key = MemberDataKey::new(key, meta.version(), element).encode()?; + let member_key = self.vector_member_key(key, meta.version(), element)?; Ok(db .get_cf_opt(&vector_cf, &member_key, &read_options) .context(RocksSnafu)? @@ -259,12 +442,63 @@ impl Redis { query: VectorQuery, options: VectorSearchOptions, ) -> Result> { + self.vsim_with_cancel(key, query, options, &FlatQueryCancel::default()) + } + + /// `vsim` with an explicit cancellation token. The token is checked + /// cooperatively during the scan; no trigger source is wired into the + /// command path yet (client disconnects are follow-up work). + pub fn vsim_with_cancel( + &self, + key: &[u8], + query: VectorQuery, + options: VectorSearchOptions, + cancel: &FlatQueryCancel, + ) -> Result> { + let vector_config = &self.storage.vector; if options.count == 0 { return InvalidArgumentSnafu { message: "vector search count must be greater than zero".to_string(), } .fail(); } + if options.count > vector_config.max_k { + return RedisErrSnafu { + message: "ERR vector search count exceeds max_k".to_string(), + } + .fail(); + } + + // The deadline starts before gate acquisition so queue wait counts + // against the configured timeout. + let timeout = + Duration::from_millis(vector_config.flat_query_timeout_ms).min(MAX_FLAT_QUERY_TIMEOUT); + let deadline = Instant::now() + timeout; + let Some(_permit) = self.flat_query_gate.acquire(deadline) else { + self.vector_metrics.record_capacity_rejected(); + return VectorFlatQueryTimeoutSnafu.fail(); + }; + self.vector_metrics.record_query_started(); + let started = Instant::now(); + let result = self.vsim_scan(key, query, options, cancel, deadline); + self.vector_metrics + .record_query_finished(started.elapsed(), result.as_ref().err()); + result + } + + /// The scan body of `vsim_with_cancel`, run while holding a gate permit. + /// Metrics are recorded by the caller around this function. + fn vsim_scan( + &self, + key: &[u8], + query: VectorQuery, + options: VectorSearchOptions, + cancel: &FlatQueryCancel, + deadline: Instant, + ) -> Result> { + let vector_config = &self.storage.vector; + let mut scan_guard = FlatScanGuard::new(vector_config, deadline, cancel); + scan_guard.check_signals()?; let db = self.db.as_ref().context(OptionNoneSnafu { message: "db is not initialized".to_string(), @@ -284,7 +518,7 @@ impl Redis { let query_vector = match query { VectorQuery::Element(element) => { - let query_key = MemberDataKey::new(key, meta.version(), &element).encode()?; + let query_key = self.vector_member_key(key, meta.version(), &element)?; let Some(query_raw) = db .get_cf_opt(&vector_cf, &query_key, &point_read_options) .context(RocksSnafu)? @@ -312,15 +546,26 @@ impl Redis { // the stored members. let query_vector = query_vector.to_quantized(meta.quantization())?; - let prefix = MemberDataKey::new(key, meta.version(), b"").encode_seek_key()?; + // Scan exactly this set's generation range: an inclusive lower bound + // at the (key, incarnation, generation) prefix and its exclusive + // successor as the upper bound. The starts_with check stays as a + // defensive guard. + let prefix = + VectorMemberDataKey::encode_prefix(key, self.storage_incarnation()?, meta.version())?; let mut scan_options = ReadOptions::default(); scan_options.set_snapshot(&snapshot); - let iterator = db.iterator_cf_opt( - &vector_cf, - scan_options, - IteratorMode::From(&prefix, Direction::Forward), - ); - let engine = VectorSearchEngine::Flat; + scan_options.set_iterate_lower_bound(prefix.clone()); + if let Some(upper_bound) = VectorMemberDataKey::prefix_upper_bound(&prefix) { + scan_options.set_iterate_upper_bound(upper_bound); + } + let iterator = db.iterator_cf_opt(&vector_cf, scan_options, IteratorMode::Start); + // Phase 1: both modes run the exhaustive FLAT scan. TRUTH is wired + // through explicitly so the mode selects the engine once an + // approximate index exists. + let engine = match options.mode { + VectorSearchMode::Approximate => VectorSearchEngine::Flat, + VectorSearchMode::Truth => VectorSearchEngine::Flat, + }; let candidates = iterator .take_while(|result| match result { Ok((encoded_key, _)) => encoded_key.starts_with(&prefix), @@ -328,16 +573,16 @@ impl Redis { }) .map(|entry| { let (encoded_key, encoded_value) = entry.context(RocksSnafu)?; - if encoded_key.len() < prefix.len() + SUFFIX_RESERVE_LENGTH { - return InvalidFormatSnafu { - message: "vector member key is shorter than its generation prefix" - .to_string(), + if self.vector_fault_hooks.fail_member_read() { + return SystemSnafu { + message: "injected fault: vector member read failed".to_string(), } .fail(); } - - let element_end = encoded_key.len() - SUFFIX_RESERVE_LENGTH; - let element = encoded_key[prefix.len()..element_end].to_vec(); + scan_guard.record(encoded_key.len(), encoded_value.len())?; + let element = ParsedVectorMemberDataKey::decode(&encoded_key)? + .element() + .to_vec(); let value = VectorDataValue::decode(&encoded_value)?; if value.dimension() != meta.dimension() { return InvalidFormatSnafu { @@ -369,6 +614,12 @@ impl Redis { message: "MetaCF is not initialized".to_string(), })?; let meta_key = BaseMetaKey::new(key).encode()?; + if self.vector_fault_hooks.fail_meta_read() { + return SystemSnafu { + message: "injected fault: vector meta read failed".to_string(), + } + .fail(); + } let value = match read_options { Some(opts) => db.get_cf_opt(&meta_cf, &meta_key, opts), None => db.get_cf(&meta_cf, &meta_key), @@ -383,4 +634,57 @@ impl Redis { fn read_vector_meta(&self, key: &[u8]) -> Result> { self.read_vector_meta_opt(key, None) } + + /// Sample-decode vector set metas (MetaCF) and member entries + /// (VectorDataCF) to verify the codec can parse this instance's data. + /// + /// Used to validate restored snapshot data: at most `sample_size` members + /// and `sample_size` metas are decoded (sampling, not a full scan); any + /// decode failure rejects the data. + pub fn validate_vector_data_sample(&self, sample_size: usize) -> Result { + let db = self.db.as_ref().context(OptionNoneSnafu { + message: "db is not initialized".to_string(), + })?; + let vector_cf = self + .get_cf_handle(ColumnFamilyIndex::VectorDataCF) + .context(OptionNoneSnafu { + message: "VectorDataCF is not initialized".to_string(), + })?; + let meta_cf = self + .get_cf_handle(ColumnFamilyIndex::MetaCF) + .context(OptionNoneSnafu { + message: "MetaCF is not initialized".to_string(), + })?; + + let mut sample = VectorDataSample::default(); + for entry in db + .iterator_cf(&vector_cf, IteratorMode::Start) + .take(sample_size) + { + let (encoded_key, encoded_value) = entry.context(RocksSnafu)?; + ParsedVectorMemberDataKey::decode(&encoded_key)?; + VectorDataValue::decode(&encoded_value)?; + sample.members += 1; + } + + for entry in db.iterator_cf(&meta_cf, IteratorMode::Start) { + if sample.metas >= sample_size { + break; + } + let (_encoded_key, encoded_value) = entry.context(RocksSnafu)?; + if encoded_value.first() == Some(&(DataType::VectorSet as u8)) { + VectorMeta::decode(&encoded_value)?; + sample.metas += 1; + } + } + + Ok(sample) + } +} + +/// Counts of vector entries decoded during a restore validation sample. +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] +pub struct VectorDataSample { + pub metas: usize, + pub members: usize, } diff --git a/src/storage/src/storage.rs b/src/storage/src/storage.rs index 5753bd6a..a23e329e 100644 --- a/src/storage/src/storage.rs +++ b/src/storage/src/storage.rs @@ -444,6 +444,22 @@ impl Storage { Ok(()) } + /// Sample-decode vector set metas and member entries on every instance to + /// verify the codec can parse the data (used to validate restored + /// snapshot data before it serves traffic). + pub fn validate_vector_data_sample( + &self, + sample_size: usize, + ) -> Result { + let mut total = crate::VectorDataSample::default(); + for inst in &self.insts { + let sample = inst.validate_vector_data_sample(sample_size)?; + total.metas += sample.metas; + total.members += sample.members; + } + Ok(total) + } + pub fn load_cursor_start_key(&self, dtype: DataType, cursor: i64) -> Result<(char, String)> { let index_key = format!("{}{}", data_type_to_tag(dtype), cursor); match self.cursors_store.get(&index_key) { diff --git a/src/storage/src/storage_impl.rs b/src/storage/src/storage_impl.rs index 9f2e1d3a..b2b2dd82 100644 --- a/src/storage/src/storage_impl.rs +++ b/src/storage/src/storage_impl.rs @@ -802,6 +802,21 @@ impl Storage { self.insts[instance_id].vismember(key, element) } + pub fn vinfo(&self, key: &[u8]) -> Result> { + let slot_id = key_to_slot_id(key); + let instance_id = self.slot_indexer.get_instance_id(slot_id); + self.insts[instance_id].vinfo(key) + } + + /// Aggregated FLAT vector query counters across all instances. + pub fn vector_metrics(&self) -> crate::vector_metrics::VectorMetricsSnapshot { + let mut total = crate::vector_metrics::VectorMetricsSnapshot::default(); + for inst in &self.insts { + total.add(&inst.vector_metrics.snapshot()); + } + total + } + // Sets Commands Implementation // Add the specified members to the set stored at key. Specified members that diff --git a/src/storage/src/storage_manifest.rs b/src/storage/src/storage_manifest.rs new file mode 100644 index 00000000..32ba6cb6 --- /dev/null +++ b/src/storage/src/storage_manifest.rs @@ -0,0 +1,243 @@ +// Copyright (c) 2024-present, arana-db Community. All rights reserved. +// +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Per-instance storage manifest: the durable home of `storage_incarnation` +//! and the standalone-mode generation sequence generator. +//! +//! Each `Redis` (RocksDB) instance owns one manifest file living next to the +//! database files in the instance's data directory. A plain file (rather than +//! a reserved MetaCF key) is used so the manifest never interferes with +//! MetaCF scans, compaction filters, or FLUSHDB, and so it can be carried +//! into Raft snapshot checkpoints with a simple copy. +//! +//! - `storage_incarnation` is generated once when an empty database directory +//! is first opened and never changes afterwards. It distinguishes data +//! written by different storages (or by a rebuilt data directory) inside +//! vector member keys. +//! - `next_generation` backs the monotonic generation sequence generator. The +//! incremented value is persisted *before* a generation is handed out, so a +//! restarted instance never reuses an allocated generation. +//! +//! A missing manifest on a non-empty database means the data predates this +//! mechanism (or the file was lost); opening then fails instead of silently +//! reinterpreting existing data. + +use std::fs; +use std::io::Write; +use std::path::{Path, PathBuf}; +use std::sync::Mutex; + +use rand::Rng; +use serde::{Deserialize, Serialize}; +use snafu::{OptionExt, ResultExt, ensure}; + +use crate::durable_fs::sync_parent_directory; +use crate::error::{InvalidFormatSnafu, IoSnafu, Result}; + +pub const STORAGE_MANIFEST_FILE: &str = "__kiwi_storage_manifest"; +const STORAGE_MANIFEST_VERSION: u32 = 1; +const FIRST_GENERATION: u64 = 1; + +#[derive(Debug, Serialize, Deserialize)] +struct ManifestFile { + version: u32, + storage_incarnation: u64, + next_generation: u64, +} + +pub(crate) struct StorageManifest { + path: PathBuf, + storage_incarnation: u64, + next_generation: Mutex, +} + +impl StorageManifest { + /// Load the manifest for the instance stored in `db_dir`, creating one + /// when the database is empty. `db_has_entries` must report whether any + /// column family of the already-open database contains at least one key. + pub(crate) fn open(db_dir: &Path, db_has_entries: bool) -> Result { + let path = db_dir.join(STORAGE_MANIFEST_FILE); + if path.exists() { + return Self::read(&path); + } + + ensure!( + !db_has_entries, + InvalidFormatSnafu { + message: format!( + "storage manifest {} is missing but the database is not empty; \ + refusing to reinterpret existing data", + path.display() + ) + } + ); + + let storage_incarnation = rand::thread_rng().r#gen::().max(1); + let manifest = Self { + path, + storage_incarnation, + next_generation: Mutex::new(FIRST_GENERATION), + }; + manifest.persist(FIRST_GENERATION)?; + Ok(manifest) + } + + pub(crate) fn storage_incarnation(&self) -> u64 { + self.storage_incarnation + } + + /// Allocate the next generation sequence. The incremented counter is + /// persisted before the generation is returned, so allocations survive + /// restarts and are never reused. + pub(crate) fn allocate_generation(&self) -> Result { + let mut next_generation = self + .next_generation + .lock() + .expect("storage manifest mutex should not be poisoned"); + let generation = *next_generation; + let successor = generation.checked_add(1).context(InvalidFormatSnafu { + message: "generation sequence exhausted".to_string(), + })?; + self.persist(successor)?; + *next_generation = successor; + Ok(generation) + } + + /// Copy the manifest file into `dir` (used when exporting a checkpoint so + /// the snapshot carries the storage identity with the data). + pub(crate) fn copy_to(&self, dir: &Path) -> Result<()> { + let target = dir.join(STORAGE_MANIFEST_FILE); + fs::copy(&self.path, &target).context(IoSnafu)?; + fs::File::open(&target) + .context(IoSnafu)? + .sync_all() + .context(IoSnafu)?; + Ok(()) + } + + fn read(path: &Path) -> Result { + let bytes = fs::read(path).context(IoSnafu)?; + let file: ManifestFile = serde_json::from_slice(&bytes).map_err(|error| { + InvalidFormatSnafu { + message: format!("invalid storage manifest {}: {error}", path.display()), + } + .build() + })?; + ensure!( + file.version == STORAGE_MANIFEST_VERSION, + InvalidFormatSnafu { + message: format!( + "unsupported storage manifest version {} in {}", + file.version, + path.display() + ) + } + ); + ensure!( + file.storage_incarnation != 0 && file.next_generation >= FIRST_GENERATION, + InvalidFormatSnafu { + message: format!( + "corrupt storage manifest {}: incarnation {}, next generation {}", + path.display(), + file.storage_incarnation, + file.next_generation + ) + } + ); + Ok(Self { + path: path.to_path_buf(), + storage_incarnation: file.storage_incarnation, + next_generation: Mutex::new(file.next_generation), + }) + } + + /// Atomically persist `next_generation` via write-temp-sync-rename, then + /// sync the directory so the rename reaches stable storage. + fn persist(&self, next_generation: u64) -> Result<()> { + let file = ManifestFile { + version: STORAGE_MANIFEST_VERSION, + storage_incarnation: self.storage_incarnation, + next_generation, + }; + let json = serde_json::to_vec(&file).map_err(|error| { + InvalidFormatSnafu { + message: format!("failed to serialize storage manifest: {error}"), + } + .build() + })?; + + let temp_path = self.path.with_extension("tmp"); + { + let mut temp = fs::File::create(&temp_path).context(IoSnafu)?; + temp.write_all(&json).context(IoSnafu)?; + temp.sync_all().context(IoSnafu)?; + } + fs::rename(&temp_path, &self.path).context(IoSnafu)?; + sync_parent_directory(&self.path).context(IoSnafu)?; + Ok(()) + } +} + +#[allow(clippy::unwrap_used)] +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn manifest_is_created_for_empty_db_and_survives_reopen() { + let dir = tempfile::tempdir().expect("temp dir"); + + let created = StorageManifest::open(dir.path(), false).expect("create manifest"); + let incarnation = created.storage_incarnation(); + assert_ne!(incarnation, 0); + assert!(dir.path().join(STORAGE_MANIFEST_FILE).exists()); + + let reopened = StorageManifest::open(dir.path(), true).expect("reopen manifest"); + assert_eq!(reopened.storage_incarnation(), incarnation); + } + + #[test] + fn allocations_are_monotonic_and_persisted() { + let dir = tempfile::tempdir().expect("temp dir"); + + let manifest = StorageManifest::open(dir.path(), false).expect("create manifest"); + let first = manifest.allocate_generation().expect("allocate first"); + let second = manifest.allocate_generation().expect("allocate second"); + assert_eq!(first, FIRST_GENERATION); + assert_eq!(second, FIRST_GENERATION + 1); + + // A reopen must never reuse an allocated generation. + let reopened = StorageManifest::open(dir.path(), false).expect("reopen manifest"); + let third = reopened + .allocate_generation() + .expect("allocate after reopen"); + assert!(third > second); + } + + #[test] + fn missing_manifest_on_non_empty_db_is_rejected() { + let dir = tempfile::tempdir().expect("temp dir"); + assert!(StorageManifest::open(dir.path(), true).is_err()); + } + + #[test] + fn corrupt_manifest_is_rejected() { + let dir = tempfile::tempdir().expect("temp dir"); + fs::write(dir.path().join(STORAGE_MANIFEST_FILE), b"not json").expect("write"); + assert!(StorageManifest::open(dir.path(), false).is_err()); + } +} diff --git a/src/storage/src/vector.rs b/src/storage/src/vector.rs index 3b0ced99..8d5d5ab4 100644 --- a/src/storage/src/vector.rs +++ b/src/storage/src/vector.rs @@ -356,6 +356,15 @@ pub struct VectorHit { pub score: f64, } +/// Per-set metadata reported by VINFO. Phase 1 only exposes what the stored +/// `VectorMeta` can answer in O(1); FLAT sentinels live in the command layer. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct VectorInfo { + pub dimension: u32, + pub size: u64, + pub generation: u64, +} + #[derive(Debug, Clone, PartialEq)] struct ScoredCandidate { element: Vec, diff --git a/src/storage/src/vector_fault.rs b/src/storage/src/vector_fault.rs new file mode 100644 index 00000000..5d606521 --- /dev/null +++ b/src/storage/src/vector_fault.rs @@ -0,0 +1,81 @@ +// Copyright (c) 2024-present, arana-db Community. All rights reserved. +// +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Fault injection hooks for the vector set storage path. +//! +//! The runtime-level `FaultInjectionConfig` only describes network delay and +//! message drop semantics for the dual-runtime message channel, so it cannot +//! express RocksDB-layer failures inside storage. These hooks close that gap +//! for the vector set path: meta reads, member reads during a FLAT scan, and +//! the atomic meta+member WriteBatch commit. +//! +//! The hooks are compiled in unconditionally (instead of behind a cargo +//! feature) because the storage crate's own integration tests under +//! `src/storage/tests/` link against the library without `cfg(test)` and +//! without self-enabled features. The cost in production is three relaxed +//! atomic loads per vector operation with all flags off. Flags are meant to +//! be flipped by tests only; nothing in the server startup path sets them. + +use std::sync::atomic::{AtomicBool, Ordering}; + +/// Switchboard for vector storage fault injection. One instance lives on +/// each `Redis`; tests arm and disarm flags between operations. +#[derive(Debug, Default)] +pub struct VectorFaultHooks { + fail_meta_read: AtomicBool, + fail_member_read: AtomicBool, + fail_batch_commit: AtomicBool, +} + +impl VectorFaultHooks { + /// Whether vector meta reads (MetaCF point lookups) should fail. + pub fn fail_meta_read(&self) -> bool { + self.fail_meta_read.load(Ordering::Relaxed) + } + + /// Whether member reads during a FLAT scan should fail. + pub fn fail_member_read(&self) -> bool { + self.fail_member_read.load(Ordering::Relaxed) + } + + /// Whether the meta+member WriteBatch commit should fail. + pub fn fail_batch_commit(&self) -> bool { + self.fail_batch_commit.load(Ordering::Relaxed) + } + + /// Arm/disarm meta read failure (test-only). + pub fn set_fail_meta_read(&self, enabled: bool) { + self.fail_meta_read.store(enabled, Ordering::Relaxed); + } + + /// Arm/disarm member read failure (test-only). + pub fn set_fail_member_read(&self, enabled: bool) { + self.fail_member_read.store(enabled, Ordering::Relaxed); + } + + /// Arm/disarm batch commit failure (test-only). + pub fn set_fail_batch_commit(&self, enabled: bool) { + self.fail_batch_commit.store(enabled, Ordering::Relaxed); + } + + /// Disarm every flag (test-only). + pub fn reset(&self) { + self.set_fail_meta_read(false); + self.set_fail_member_read(false); + self.set_fail_batch_commit(false); + } +} diff --git a/src/storage/src/vector_flat.rs b/src/storage/src/vector_flat.rs new file mode 100644 index 00000000..b6ad73e4 --- /dev/null +++ b/src/storage/src/vector_flat.rs @@ -0,0 +1,190 @@ +// Copyright (c) 2024-present, arana-db Community. All rights reserved. +// +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Runtime governance for FLAT (brute-force) vector similarity queries: +//! a concurrency gate, a cooperative cancellation token, and a scan guard +//! enforcing the deadline and scan budget. +//! +//! The pieces are synchronous by design: storage commands execute as +//! blocking code inside storage-runtime tasks and `Redis::vsim` is a plain +//! synchronous function, so a blocking semaphore matches every call site +//! without requiring an async context. + +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Condvar, Mutex}; +use std::time::Instant; + +use conf::vector_config::VectorConfig; + +use crate::error::{ + Result, VectorFlatQueryCancelledSnafu, VectorFlatQueryTimeoutSnafu, + VectorFlatScanBudgetExceededSnafu, +}; + +/// Cooperative cancellation signal for an in-flight FLAT query. +/// +/// Cheap to clone; the default token is never cancelled. No trigger source +/// is wired into the call chain yet (client-disconnect propagation is +/// follow-up work); callers that have one share a token and trip it via +/// [`FlatQueryCancel::cancel`]. +#[derive(Debug, Clone, Default)] +pub struct FlatQueryCancel { + flag: Arc, +} + +impl FlatQueryCancel { + /// Trip the signal; the owning query aborts at its next check. + pub fn cancel(&self) { + self.flag.store(true, Ordering::Relaxed); + } + + pub fn is_cancelled(&self) -> bool { + self.flag.load(Ordering::Relaxed) + } +} + +#[derive(Debug)] +struct FlatQueryGateInner { + available: Mutex, + released: Condvar, +} + +/// Counting semaphore bounding the number of concurrent FLAT scans. +#[derive(Debug)] +pub struct FlatQueryGate { + inner: Arc, +} + +/// RAII slot in a [`FlatQueryGate`]; dropping it returns the slot. +#[derive(Debug)] +pub struct FlatQueryPermit { + inner: Arc, +} + +impl Drop for FlatQueryPermit { + fn drop(&mut self) { + let mut available = self + .inner + .available + .lock() + .expect("flat query gate mutex poisoned"); + *available += 1; + drop(available); + self.inner.released.notify_one(); + } +} + +impl FlatQueryGate { + pub fn new(capacity: usize) -> Self { + Self { + inner: Arc::new(FlatQueryGateInner { + available: Mutex::new(capacity.max(1)), + released: Condvar::new(), + }), + } + } + + /// Block until a slot is free or `deadline` passes; `None` on timeout, + /// so queue wait counts against the caller's deadline. + pub fn acquire(&self, deadline: Instant) -> Option { + let mut available = self + .inner + .available + .lock() + .expect("flat query gate mutex poisoned"); + loop { + if *available > 0 { + *available -= 1; + return Some(FlatQueryPermit { + inner: Arc::clone(&self.inner), + }); + } + let now = Instant::now(); + if now >= deadline { + return None; + } + let (guard, _timeout) = self + .inner + .released + .wait_timeout(available, deadline - now) + .expect("flat query gate mutex poisoned"); + available = guard; + } + } + + /// Number of currently free slots. + pub fn available_permits(&self) -> usize { + *self + .inner + .available + .lock() + .expect("flat query gate mutex poisoned") + } +} + +/// Cooperative guard consulted while a FLAT scan streams candidates: the +/// scan budget is enforced on every entry, the deadline and cancellation +/// signal every `flat_cancel_check_interval` entries. +pub struct FlatScanGuard<'a> { + deadline: Instant, + cancel: &'a FlatQueryCancel, + check_interval: u64, + max_entries: u64, + max_bytes: u64, + entries: u64, + bytes: u64, +} + +impl<'a> FlatScanGuard<'a> { + pub fn new(config: &VectorConfig, deadline: Instant, cancel: &'a FlatQueryCancel) -> Self { + Self { + deadline, + cancel, + check_interval: (config.flat_cancel_check_interval as u64).max(1), + max_entries: config.flat_scan_max_entries, + max_bytes: config.flat_scan_max_bytes, + entries: 0, + bytes: 0, + } + } + + /// Account one scanned entry (raw key/value bytes) and enforce the + /// budget, deadline and cancellation policy. + pub fn record(&mut self, key_bytes: usize, value_bytes: usize) -> Result<()> { + self.entries += 1; + self.bytes += (key_bytes + value_bytes) as u64; + if self.entries > self.max_entries || self.bytes > self.max_bytes { + return VectorFlatScanBudgetExceededSnafu.fail(); + } + if self.entries.is_multiple_of(self.check_interval) { + self.check_signals()?; + } + Ok(()) + } + + /// Check the deadline and cancellation signal without accounting an + /// entry (used before the scan starts, right after gate acquisition). + pub fn check_signals(&self) -> Result<()> { + if self.cancel.is_cancelled() { + return VectorFlatQueryCancelledSnafu.fail(); + } + if Instant::now() >= self.deadline { + return VectorFlatQueryTimeoutSnafu.fail(); + } + Ok(()) + } +} diff --git a/src/storage/src/vector_metrics.rs b/src/storage/src/vector_metrics.rs new file mode 100644 index 00000000..46516c61 --- /dev/null +++ b/src/storage/src/vector_metrics.rs @@ -0,0 +1,104 @@ +// Copyright (c) 2024-present, arana-db Community. All rights reserved. +// +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Counters for FLAT vector query execution, surfaced through INFO VECTOR. +//! +//! One `VectorMetrics` lives on each `Redis` instance next to its FLAT query +//! gate; `Storage::vector_metrics` aggregates the per-instance snapshots. + +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::Duration; + +use crate::error::Error; + +#[derive(Debug, Default)] +pub struct VectorMetrics { + flat_queries_total: AtomicU64, + flat_query_timeouts_total: AtomicU64, + flat_query_errors_total: AtomicU64, + capacity_rejected_total: AtomicU64, + flat_query_duration_micros_total: AtomicU64, + flat_query_duration_count: AtomicU64, +} + +/// Point-in-time copy of the vector counters. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct VectorMetricsSnapshot { + pub flat_queries_total: u64, + pub flat_query_timeouts_total: u64, + pub flat_query_errors_total: u64, + pub capacity_rejected_total: u64, + pub flat_query_duration_micros_total: u64, + pub flat_query_duration_count: u64, +} + +impl VectorMetricsSnapshot { + pub fn add(&mut self, other: &Self) { + self.flat_queries_total += other.flat_queries_total; + self.flat_query_timeouts_total += other.flat_query_timeouts_total; + self.flat_query_errors_total += other.flat_query_errors_total; + self.capacity_rejected_total += other.capacity_rejected_total; + self.flat_query_duration_micros_total += other.flat_query_duration_micros_total; + self.flat_query_duration_count += other.flat_query_duration_count; + } +} + +impl VectorMetrics { + /// A query passed the gate and entered the scan path. + pub fn record_query_started(&self) { + self.flat_queries_total.fetch_add(1, Ordering::Relaxed); + } + + /// The gate could not hand out a permit before the caller's deadline. + pub fn record_capacity_rejected(&self) { + self.capacity_rejected_total.fetch_add(1, Ordering::Relaxed); + } + + /// A started query finished (successfully or not); `error` classifies + /// the abort path: scan-deadline timeouts are counted separately from + /// budget, cancellation and other storage errors. + pub fn record_query_finished(&self, elapsed: Duration, error: Option<&Error>) { + let micros = u64::try_from(elapsed.as_micros()).unwrap_or(u64::MAX); + self.flat_query_duration_micros_total + .fetch_add(micros, Ordering::Relaxed); + self.flat_query_duration_count + .fetch_add(1, Ordering::Relaxed); + match error { + None => {} + Some(Error::VectorFlatQueryTimeout { .. }) => { + self.flat_query_timeouts_total + .fetch_add(1, Ordering::Relaxed); + } + Some(_) => { + self.flat_query_errors_total.fetch_add(1, Ordering::Relaxed); + } + } + } + + pub fn snapshot(&self) -> VectorMetricsSnapshot { + VectorMetricsSnapshot { + flat_queries_total: self.flat_queries_total.load(Ordering::Relaxed), + flat_query_timeouts_total: self.flat_query_timeouts_total.load(Ordering::Relaxed), + flat_query_errors_total: self.flat_query_errors_total.load(Ordering::Relaxed), + capacity_rejected_total: self.capacity_rejected_total.load(Ordering::Relaxed), + flat_query_duration_micros_total: self + .flat_query_duration_micros_total + .load(Ordering::Relaxed), + flat_query_duration_count: self.flat_query_duration_count.load(Ordering::Relaxed), + } + } +} diff --git a/src/storage/src/vector_mutation.rs b/src/storage/src/vector_mutation.rs new file mode 100644 index 00000000..34183fd2 --- /dev/null +++ b/src/storage/src/vector_mutation.rs @@ -0,0 +1,525 @@ +// Copyright (c) 2024-present, arana-db Community. All rights reserved. +// +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Logical vector-set mutation contract. +//! +//! `VectorSetMutationV1` is the deterministic, logical description of a +//! VADD/VREM: it carries the client intent (which element, which canonical +//! vector) but none of the derived state (count, data_revision, generation). +//! The single decision point is `Redis::apply_vector_set_mutation`, which +//! reads the current meta, decides the outcome, and commits it atomically. +//! Standalone VADD/VREM already go through that entry; Raft log replication +//! of logical mutations (PR0) will reuse the same entry so the state machine +//! decides in log order. +//! +//! Wire format (all integers little-endian): +//! +//!
+//! | version(1B) | tag(1B) | payload                                        |
+//!
+//! tag 1 (Add):    | element_len(u32) | element | dimension(u32) |
+//!                 | original_l2(f32) | vector_len(u32) | canonical_vector |
+//! tag 2 (Remove): | element_len(u32) | element |
+//!
+//! +//! `canonical_vector` is the L2-normalized FP32 components as `dimension * 4` +//! little-endian bytes (NOQUANT form); the apply entry converts it to the +//! set's quantization. Length-prefixed fields let the decoder distinguish an +//! empty element from a missing one; malformed input yields an error, never +//! a panic. + +use snafu::ensure; + +use crate::error::{InvalidArgumentSnafu, InvalidFormatSnafu, RedisErrSnafu, Result}; +use crate::vector::{CanonicalVector, QuantizationType, VectorData}; + +/// Format version byte; bumped when the wire layout changes. +pub const VECTOR_SET_MUTATION_VERSION_V1: u8 = 1; + +const TAG_ADD: u8 = 1; +const TAG_REMOVE: u8 = 2; + +/// A logical vector-set mutation, version 1. +#[derive(Debug, Clone, PartialEq)] +pub enum VectorSetMutationV1 { + /// Add (or overwrite) one member. + Add { + element: Vec, + dimension: u32, + /// Pre-normalization L2 norm, preserved so VEMB can restore the + /// original vector. + original_l2: f32, + /// L2-normalized FP32 components, `dimension * 4` LE bytes. + canonical_vector: Vec, + }, + /// Remove one member. + Remove { element: Vec }, +} + +impl VectorSetMutationV1 { + /// Build an `Add` mutation from a canonical vector, storing the normalized + /// FP32 form so the apply entry can re-quantize to the set's quantization. + pub fn add_from_canonical(element: &[u8], vector: &CanonicalVector) -> Result { + let components = vector.as_fp32()?; + let mut canonical_vector = Vec::with_capacity(components.len() * size_of::()); + for component in components { + canonical_vector.extend_from_slice(&component.to_le_bytes()); + } + Ok(Self::Add { + element: element.to_vec(), + dimension: vector.dimension(), + original_l2: vector.original_l2(), + canonical_vector, + }) + } + + /// Reconstruct the canonical vector carried by an `Add` mutation. + /// Returns `None` for `Remove`. + pub(crate) fn canonical_vector(&self) -> Result> { + let Self::Add { + dimension, + original_l2, + canonical_vector, + .. + } = self + else { + return Ok(None); + }; + Ok(Some(Self::decode_canonical_vector( + *dimension, + *original_l2, + canonical_vector, + )?)) + } + + fn decode_canonical_vector( + dimension: u32, + original_l2: f32, + canonical_vector: &[u8], + ) -> Result { + ensure!( + dimension != 0, + InvalidFormatSnafu { + message: "vector mutation dimension must not be zero".to_string() + } + ); + ensure!( + original_l2.is_finite() && original_l2 > 0.0, + InvalidFormatSnafu { + message: "vector mutation L2 norm must be finite and positive".to_string() + } + ); + let expected_len = dimension as usize * size_of::(); + ensure!( + canonical_vector.len() == expected_len, + InvalidFormatSnafu { + message: format!( + "vector mutation payload length {} does not match dimension {}", + canonical_vector.len(), + dimension + ) + } + ); + let components = canonical_vector + .chunks_exact(size_of::()) + .map(|chunk| f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]])) + .collect::>(); + ensure!( + components.iter().all(|component| component.is_finite()), + InvalidFormatSnafu { + message: "vector mutation components must be finite".to_string() + } + ); + Ok(CanonicalVector::from_parts( + dimension, + original_l2, + QuantizationType::None, + VectorData::Fp32(components), + )) + } + + pub fn encode(&self) -> Vec { + let mut output = vec![VECTOR_SET_MUTATION_VERSION_V1]; + match self { + Self::Add { + element, + dimension, + original_l2, + canonical_vector, + } => { + output.push(TAG_ADD); + put_length_prefixed(&mut output, element); + output.extend_from_slice(&dimension.to_le_bytes()); + output.extend_from_slice(&original_l2.to_le_bytes()); + put_length_prefixed(&mut output, canonical_vector); + } + Self::Remove { element } => { + output.push(TAG_REMOVE); + put_length_prefixed(&mut output, element); + } + } + output + } + + pub fn decode(bytes: &[u8]) -> Result { + let mut reader = MutationReader::new(bytes); + let version = reader.take_u8()?; + ensure!( + version == VECTOR_SET_MUTATION_VERSION_V1, + InvalidFormatSnafu { + message: format!("unsupported vector mutation version: {version}") + } + ); + let tag = reader.take_u8()?; + let mutation = match tag { + TAG_ADD => { + let element = reader.take_length_prefixed()?; + let dimension = reader.take_u32_le()?; + let original_l2 = reader.take_f32_le()?; + let canonical_vector = reader.take_length_prefixed()?; + let mutation = Self::Add { + element, + dimension, + original_l2, + canonical_vector, + }; + // Structural validation of the vector payload happens here so + // a malformed mutation never reaches the apply entry. + mutation.canonical_vector()?; + mutation + } + TAG_REMOVE => Self::Remove { + element: reader.take_length_prefixed()?, + }, + _ => { + return InvalidFormatSnafu { + message: format!("unsupported vector mutation tag: {tag}"), + } + .fail(); + } + }; + ensure!( + reader.is_exhausted(), + InvalidFormatSnafu { + message: format!("vector mutation has {} trailing bytes", reader.remaining()) + } + ); + Ok(mutation) + } +} + +fn put_length_prefixed(output: &mut Vec, bytes: &[u8]) { + let len = u32::try_from(bytes.len()).expect("mutation field length always fits in u32"); + output.extend_from_slice(&len.to_le_bytes()); + output.extend_from_slice(bytes); +} + +/// Bounds-checked cursor over the encoded mutation; every read is fallible so +/// malformed input produces an error instead of a panic. +struct MutationReader<'a> { + bytes: &'a [u8], + position: usize, +} + +impl<'a> MutationReader<'a> { + fn new(bytes: &'a [u8]) -> Self { + Self { bytes, position: 0 } + } + + fn remaining(&self) -> usize { + self.bytes.len() - self.position + } + + fn is_exhausted(&self) -> bool { + self.remaining() == 0 + } + + fn take(&mut self, len: usize) -> Result<&'a [u8]> { + ensure!( + self.remaining() >= len, + InvalidFormatSnafu { + message: format!( + "truncated vector mutation: need {len} bytes, have {}", + self.remaining() + ) + } + ); + let start = self.position; + self.position += len; + Ok(&self.bytes[start..self.position]) + } + + fn take_u8(&mut self) -> Result { + Ok(self.take(1)?[0]) + } + + fn take_u32_le(&mut self) -> Result { + let bytes = self.take(size_of::())?; + Ok(u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]])) + } + + fn take_f32_le(&mut self) -> Result { + let bytes = self.take(size_of::())?; + Ok(f32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]])) + } + + fn take_length_prefixed(&mut self) -> Result> { + let len = self.take_u32_le()? as usize; + Ok(self.take(len)?.to_vec()) + } +} + +/// Outcome of applying a vector-set mutation. VADD maps `Added`/`Updated` to +/// 1/0; VREM maps `Removed`/`RemoveMissed` to 1/0. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum VectorSetApplyResult { + /// A new member was inserted (set created or member absent). + Added, + /// An existing member's vector was overwritten. + Updated, + /// The member existed and was removed. + Removed, + /// The member (or the set) did not exist; nothing changed. + RemoveMissed, +} + +/// Deterministic business rejections. A Raft apply loop must not treat these +/// as fatal: it records the deterministic error for the client and continues +/// with the next log entry. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum VectorSetBusinessError { + /// The key holds a live value of another data type. + WrongType, + /// The mutation's dimension differs from the existing set's dimension. + DimensionMismatch { expected: u32, got: u32 }, +} + +impl std::fmt::Display for VectorSetBusinessError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::WrongType => { + write!( + f, + "WRONGTYPE Operation against a key holding the wrong kind of value" + ) + } + Self::DimensionMismatch { expected, got } => { + write!( + f, + "vector dimension mismatch: expected {expected}, got {got}" + ) + } + } + } +} + +/// Error from `Redis::apply_vector_set_mutation`, split by how a Raft apply +/// loop should react: business errors are deterministic and skippable, storage +/// errors are fatal. +#[derive(Debug)] +pub enum VectorSetApplyError { + Business(VectorSetBusinessError), + Storage(crate::error::Error), +} + +impl std::fmt::Display for VectorSetApplyError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Business(business) => write!(f, "{business}"), + Self::Storage(error) => write!(f, "{error}"), + } + } +} + +impl std::error::Error for VectorSetApplyError {} + +impl From for VectorSetApplyError { + /// Anything not explicitly classified as a business error is fatal. + fn from(error: crate::error::Error) -> Self { + Self::Storage(error) + } +} + +impl From for crate::error::Error { + /// Map back to the storage error type with the exact variants and messages + /// VADD/VREM produced before the apply entry existed, so standalone + /// behavior is unchanged. + fn from(error: VectorSetApplyError) -> Self { + match error { + VectorSetApplyError::Business(VectorSetBusinessError::WrongType) => RedisErrSnafu { + message: "WRONGTYPE Operation against a key holding the wrong kind of value" + .to_string(), + } + .build(), + VectorSetApplyError::Business(VectorSetBusinessError::DimensionMismatch { + expected, + got, + }) => InvalidArgumentSnafu { + message: format!("vector dimension mismatch: expected {expected}, got {got}"), + } + .build(), + VectorSetApplyError::Storage(error) => error, + } + } +} + +#[allow(clippy::unwrap_used)] +#[cfg(test)] +mod tests { + use super::*; + + fn add_mutation(element: &[u8]) -> VectorSetMutationV1 { + let vector = CanonicalVector::from_values(&[3.0, 4.0]).expect("valid vector"); + VectorSetMutationV1::add_from_canonical(element, &vector).expect("mutation from vector") + } + + #[test] + fn add_mutation_round_trips() { + let mutation = add_mutation(b"member"); + let decoded = VectorSetMutationV1::decode(&mutation.encode()).expect("decode add"); + assert_eq!(decoded, mutation); + + let VectorSetMutationV1::Add { + dimension, + original_l2, + canonical_vector, + .. + } = &mutation + else { + panic!("expected add mutation"); + }; + assert_eq!(*dimension, 2); + assert!((*original_l2 - 5.0).abs() < 1e-6); + assert_eq!(canonical_vector.len(), 2 * size_of::()); + + let vector = mutation + .canonical_vector() + .expect("canonical vector") + .expect("add carries a vector"); + let restored = vector.restore(); + assert!((restored[0] - 3.0).abs() < 1e-6); + assert!((restored[1] - 4.0).abs() < 1e-6); + } + + #[test] + fn remove_mutation_round_trips() { + let mutation = VectorSetMutationV1::Remove { + element: b"member".to_vec(), + }; + let decoded = VectorSetMutationV1::decode(&mutation.encode()).expect("decode remove"); + assert_eq!(decoded, mutation); + assert!(decoded.canonical_vector().expect("no vector").is_none()); + } + + #[test] + fn empty_element_round_trips() { + for mutation in [ + add_mutation(b""), + VectorSetMutationV1::Remove { + element: Vec::new(), + }, + ] { + let decoded = VectorSetMutationV1::decode(&mutation.encode()).expect("decode"); + assert_eq!(decoded, mutation); + } + } + + #[test] + fn large_element_round_trips() { + let element = vec![0xABu8; 1 << 20]; + let mutation = add_mutation(&element); + let decoded = VectorSetMutationV1::decode(&mutation.encode()).expect("decode"); + assert_eq!(decoded, mutation); + } + + #[test] + fn decode_rejects_malformed_bytes_without_panicking() { + assert!(VectorSetMutationV1::decode(&[]).is_err()); + assert!(VectorSetMutationV1::decode(&[0, TAG_ADD]).is_err()); + assert!(VectorSetMutationV1::decode(&[VECTOR_SET_MUTATION_VERSION_V1, 0xFF]).is_err()); + + // Every strict prefix of a valid encoding is a truncation error. + let encoded = add_mutation(b"member").encode(); + for len in 0..encoded.len() { + assert!( + VectorSetMutationV1::decode(&encoded[..len]).is_err(), + "prefix of length {len} must fail to decode" + ); + } + + // Trailing garbage is rejected. + let mut trailing = encoded.clone(); + trailing.push(0); + assert!(VectorSetMutationV1::decode(&trailing).is_err()); + } + + #[test] + fn decode_rejects_invalid_vector_payload() { + let valid = add_mutation(b"member").encode(); + + // dimension = 0 + let mut zero_dimension = valid.clone(); + // layout: version(1) tag(1) elem_len(4) elem(6) dimension(4) + zero_dimension[12..16].copy_from_slice(&0_u32.to_le_bytes()); + assert!(VectorSetMutationV1::decode(&zero_dimension).is_err()); + + // canonical_vector length disagrees with dimension + let mut bad_len = valid.clone(); + bad_len[20..24].copy_from_slice(&4_u32.to_le_bytes()); + assert!(VectorSetMutationV1::decode(&bad_len).is_err()); + + // NaN original_l2 + let mut nan_l2 = valid.clone(); + nan_l2[16..20].copy_from_slice(&f32::NAN.to_le_bytes()); + assert!(VectorSetMutationV1::decode(&nan_l2).is_err()); + + // NaN component + let mut nan_component = valid.clone(); + let component_offset = 24; + nan_component[component_offset..component_offset + 4] + .copy_from_slice(&f32::NAN.to_le_bytes()); + assert!(VectorSetMutationV1::decode(&nan_component).is_err()); + } + + #[test] + fn apply_error_maps_back_to_storage_errors() { + let wrong_type: crate::error::Error = + VectorSetApplyError::Business(VectorSetBusinessError::WrongType).into(); + assert_eq!( + wrong_type.to_string(), + "WRONGTYPE Operation against a key holding the wrong kind of value" + ); + + let mismatch: crate::error::Error = + VectorSetApplyError::Business(VectorSetBusinessError::DimensionMismatch { + expected: 2, + got: 3, + }) + .into(); + assert_eq!( + mismatch.to_string(), + "Invalid argument: vector dimension mismatch: expected 2, got 3" + ); + + let storage = crate::error::Error::OptionNone { + message: "db is not initialized".to_string(), + location: Default::default(), + }; + let fatal = VectorSetApplyError::from(storage); + assert!(matches!(fatal, VectorSetApplyError::Storage(_))); + } +} diff --git a/src/storage/tests/checkpoint_test.rs b/src/storage/tests/checkpoint_test.rs index c7b14fae..08245c57 100644 --- a/src/storage/tests/checkpoint_test.rs +++ b/src/storage/tests/checkpoint_test.rs @@ -156,10 +156,10 @@ fn test_snapshot_meta_version() { let meta = RaftSnapshotMeta::new(100, 5); let json = serde_json::to_string(&meta).unwrap(); - assert!(json.contains("\"version\":1")); + assert!(json.contains("\"version\":2")); let deserialized: RaftSnapshotMeta = serde_json::from_str(&json).unwrap(); - assert_eq!(deserialized.version, 1); + assert_eq!(deserialized.version, 2); assert_eq!(deserialized.last_included_index, 100); assert_eq!(deserialized.last_included_term, 5); } @@ -191,13 +191,14 @@ fn test_snapshot_meta_rejects_unsupported_version() { } #[test] -fn test_snapshot_meta_accepts_future_version() { +fn test_snapshot_meta_rejects_future_version() { use std::fs; let tmp_dir = tempfile::tempdir().unwrap(); let meta_path = tmp_dir.path().join("__raft_snapshot_meta"); - // Write a future version (9999) to test forward compatibility for rolling upgrades + // A higher version comes from a newer binary whose schema this node + // cannot safely consume; it must be rejected deterministically. let json = r#"{ "version": 9999, "last_included_index": 42, @@ -206,9 +207,41 @@ fn test_snapshot_meta_accepts_future_version() { fs::write(&meta_path, json).unwrap(); let result = RaftSnapshotMeta::read_from_dir(tmp_dir.path()); + assert!(result.is_err(), "Higher versions must be rejected"); assert!( - result.is_ok(), - "Higher versions should be accepted for forward compatibility during rolling upgrades" + result + .unwrap_err() + .to_string() + .contains("unsupported snapshot version"), + "Error should mention unsupported version" + ); +} + +#[test] +fn test_snapshot_meta_rejects_v1() { + use std::fs; + + let tmp_dir = tempfile::tempdir().unwrap(); + let meta_path = tmp_dir.path().join("__raft_snapshot_meta"); + + // v1 was a development-phase format that never shipped; it lacks the + // storage schema description and is rejected outright. + let json = r#"{ + "version": 1, + "last_included_index": 42, + "last_included_term": 7, + "logindex_collector_states": [] + }"#; + fs::write(&meta_path, json).unwrap(); + + let result = RaftSnapshotMeta::read_from_dir(tmp_dir.path()); + assert!(result.is_err(), "v1 snapshot meta must be rejected"); + assert!( + result + .unwrap_err() + .to_string() + .contains("unsupported snapshot version"), + "Error should mention unsupported version" ); } @@ -216,9 +249,7 @@ fn test_snapshot_meta_accepts_future_version() { fn test_snapshot_meta_max_version() { let meta = RaftSnapshotMeta { version: u32::MAX, - last_included_index: 42, - last_included_term: 7, - logindex_collector_states: Vec::new(), + ..RaftSnapshotMeta::new(42, 7) }; let json = serde_json::to_string(&meta).unwrap(); @@ -230,13 +261,11 @@ fn test_snapshot_meta_max_version() { #[test] fn test_raft_snapshot_meta_with_collector_states() { let meta = RaftSnapshotMeta { - version: 1, - last_included_index: 300, - last_included_term: 1, logindex_collector_states: vec![ vec!["100:1000".to_string(), "200:2000".to_string()], vec!["150:1500".to_string()], ], + ..RaftSnapshotMeta::new(300, 1) }; let json = serde_json::to_string_pretty(&meta).unwrap(); @@ -254,7 +283,7 @@ fn test_raft_snapshot_meta_with_collector_states() { #[test] fn test_raft_snapshot_meta_defaults_empty_states() { let json = r#"{ - "version": 1, + "version": 2, "last_included_index": 100, "last_included_term": 5 }"#; @@ -300,3 +329,104 @@ fn test_collector_states_roundtrip() { assert_eq!(new_inst0.find_applied_log_index(3000), 300); assert_eq!(new_inst1.find_applied_log_index(1500), 150); } + +/// v2 meta built from a live Storage round trips through the checkpoint +/// directory and passes restore validation against the same instance count. +#[tokio::test] +async fn test_v2_meta_for_storage_roundtrip_and_validate() { + let db_path = unique_test_db_path(); + let cp_root = unique_test_db_path(); + + let mut storage = Storage::new(2, 0); + let options = Arc::new(StorageOptions::default()); + let _rx = storage.open(options, &db_path).unwrap(); + + let collectors: Vec<_> = (0..storage.db_instance_num) + .filter_map(|i| storage.get_logindex_collector(i)) + .collect(); + let meta = RaftSnapshotMeta::for_storage(42, 7, &collectors, &storage).unwrap(); + + assert_eq!(meta.version, 2); + assert_eq!(meta.db_instance_num, 2); + assert_eq!(meta.storage_incarnations.len(), 2); + for (i, inst) in storage.insts.iter().enumerate() { + assert_eq!( + meta.storage_incarnations[i], + inst.storage_incarnation().unwrap() + ); + } + assert_eq!( + meta.column_families, + storage::checkpoint::expected_column_families() + ); + assert_eq!( + meta.vector_value_format_max, + storage::format_vector::VECTOR_VALUE_FORMAT + ); + + storage.create_checkpoint(&cp_root, &meta).unwrap(); + let read_back = RaftSnapshotMeta::read_from_dir(&cp_root).unwrap(); + assert_eq!(read_back, meta); + read_back.validate_for_restore(2).unwrap(); +} + +#[test] +fn test_validate_for_restore_rejects_bad_schema() { + use storage::format_vector::VECTOR_VALUE_FORMAT; + + let valid = || RaftSnapshotMeta { + db_instance_num: 2, + storage_incarnations: vec![11, 22], + ..RaftSnapshotMeta::new(42, 7) + }; + + // Instance count does not match local configuration. + let err = valid().validate_for_restore(3).unwrap_err(); + assert!( + err.to_string().contains("db_instance_num"), + "unexpected error: {err}" + ); + + // Incarnation list length must match the instance count. Values + // themselves are never compared: the restore adopts the snapshot's + // incarnations via the per-instance manifest files. + let mut meta = valid(); + meta.storage_incarnations = vec![11]; + let err = meta.validate_for_restore(2).unwrap_err(); + assert!( + err.to_string().contains("storage incarnations"), + "unexpected error: {err}" + ); + + // Differing incarnation values are accepted. + let mut meta = valid(); + meta.storage_incarnations = vec![999, 888]; + meta.validate_for_restore(2).unwrap(); + + // Missing / mismatched column families. + let mut meta = valid(); + meta.column_families.pop(); + let err = meta.validate_for_restore(2).unwrap_err(); + assert!( + err.to_string().contains("column families"), + "unexpected error: {err}" + ); + + // Unknown (newer) vector value format. + let mut meta = valid(); + meta.vector_value_format_max = VECTOR_VALUE_FORMAT + 1; + let err = meta.validate_for_restore(2).unwrap_err(); + assert!( + err.to_string().contains("vector value format"), + "unexpected error: {err}" + ); + + // Unknown (newer) storage schema version. + let mut meta = valid(); + meta.storage_schema_version = storage::STORAGE_SCHEMA_VERSION + 1; + let err = meta.validate_for_restore(2).unwrap_err(); + assert!( + err.to_string().contains("storage schema version"), + "unexpected error: {err}" + ); +} diff --git a/src/storage/tests/redis_vector_test.rs b/src/storage/tests/redis_vector_test.rs index 830a5a62..3ade8be3 100644 --- a/src/storage/tests/redis_vector_test.rs +++ b/src/storage/tests/redis_vector_test.rs @@ -17,19 +17,20 @@ #![allow(clippy::unwrap_used)] -use std::{path::PathBuf, sync::Arc, sync::Mutex}; +use std::{path::PathBuf, sync::Arc, sync::Mutex, time::Duration, time::Instant}; +use conf::vector_config::VectorConfig; use kstd::lock_mgr::LockMgr; use rocksdb::{IteratorMode, ReadOptions}; use storage::{ - BaseMetaKey, BgTaskHandler, CanonicalVector, ColumnFamilyIndex, Redis, StorageOptions, - VectorQuery, VectorSearchMode, VectorSearchOptions, format_vector::VectorMeta, - safe_cleanup_test_db, unique_test_db_path, + BaseMetaKey, BgTaskHandler, CanonicalVector, ColumnFamilyIndex, FlatQueryCancel, Redis, + STORAGE_MANIFEST_FILE, StorageOptions, VectorQuery, VectorSearchMode, VectorSearchOptions, + VectorSetApplyError, VectorSetApplyResult, VectorSetBusinessError, VectorSetMutationV1, + error::Error, format_vector::VectorMeta, safe_cleanup_test_db, unique_test_db_path, }; use storage::{slot_indexer::key_to_slot_id, storage::Storage}; -fn open_redis(path: &PathBuf) -> Redis { - let storage_options = Arc::new(StorageOptions::default()); +fn open_redis_with_options(path: &PathBuf, storage_options: Arc) -> Redis { let (bg_task_handler, _) = BgTaskHandler::new(); let lock_mgr = Arc::new(LockMgr::new(1000)); let mut redis = Redis::new(storage_options, 1, Arc::new(bg_task_handler), lock_mgr); @@ -39,6 +40,10 @@ fn open_redis(path: &PathBuf) -> Redis { redis } +fn open_redis(path: &PathBuf) -> Redis { + open_redis_with_options(path, Arc::new(StorageOptions::default())) +} + fn with_redis(test: impl FnOnce(&Redis)) { let path = unique_test_db_path(); safe_cleanup_test_db(&path); @@ -50,6 +55,21 @@ fn with_redis(test: impl FnOnce(&Redis)) { safe_cleanup_test_db(&path); } +fn with_redis_vector_config(edit: impl FnOnce(&mut VectorConfig), test: impl FnOnce(&Redis)) { + let path = unique_test_db_path(); + safe_cleanup_test_db(&path); + let mut vector = VectorConfig::default(); + edit(&mut vector); + let mut storage_options = StorageOptions::default(); + storage_options.vector = vector; + let redis = open_redis_with_options(&path, Arc::new(storage_options)); + + test(&redis); + + drop(redis); + safe_cleanup_test_db(&path); +} + fn populate_search_vectors(redis: &Redis) -> CanonicalVector { let x = CanonicalVector::from_values(&[1.0, 0.0]).expect("x"); let y = CanonicalVector::from_values(&[0.0, 1.0]).expect("y"); @@ -73,6 +93,20 @@ fn count_cf_entries(redis: &Redis, cf_index: ColumnFamilyIndex) -> usize { .count() } +fn read_stored_vector_meta(redis: &Redis, key: &[u8]) -> VectorMeta { + let db = redis.db().expect("db is initialized"); + let meta_cf = redis + .get_cf_handle(ColumnFamilyIndex::MetaCF) + .expect("MetaCF exists"); + let meta_key = BaseMetaKey::new(key).encode().expect("meta key"); + VectorMeta::decode( + &db.get_cf(&meta_cf, &meta_key) + .expect("read vector meta") + .expect("vector meta exists"), + ) + .expect("decode vector meta") +} + #[test] fn test_vadd_create_update_and_dimension_guard() { with_redis(|redis| { @@ -445,25 +479,23 @@ fn test_type_returns_vectorset() { } #[test] -fn test_vadd_rebuilds_expired_vectorset_with_newer_generation() { +fn test_vadd_rebuilds_expired_vectorset_with_new_generation() { with_redis(|redis| { let key = b"expiring-vectors"; let vector = CanonicalVector::from_values(&[1.0, 0.0]).expect("vector"); redis.vadd(key, b"old", &vector).expect("insert old member"); + let previous_meta = read_stored_vector_meta(redis, key); + let previous_generation = previous_meta.version(); + assert_eq!(previous_meta.data_revision(), 1); + + // Expire the set in place, keeping the previous generation on disk. let db = redis.db().expect("db is initialized"); let meta_cf = redis .get_cf_handle(ColumnFamilyIndex::MetaCF) .expect("MetaCF exists"); let meta_key = BaseMetaKey::new(key).encode().expect("meta key"); - let mut meta = VectorMeta::decode( - &db.get_cf(&meta_cf, &meta_key) - .expect("read vector meta") - .expect("vector meta exists"), - ) - .expect("decode vector meta"); - let previous_generation = u64::MAX - 1; - meta.set_version(previous_generation); + let mut meta = previous_meta; meta.set_etime(1); db.put_cf(&meta_cf, &meta_key, &meta.encode()) .expect("store expired vector meta"); @@ -478,16 +510,107 @@ fn test_vadd_rebuilds_expired_vectorset_with_newer_generation() { assert_eq!(redis.vcard(key).expect("rebuilt card"), 1); assert!(!redis.vismember(key, b"old").expect("old membership")); assert!(redis.vismember(key, b"new").expect("new membership")); - let rebuilt_meta = VectorMeta::decode( - &db.get_cf(&meta_cf, &meta_key) - .expect("read rebuilt vector meta") - .expect("rebuilt vector meta exists"), - ) - .expect("decode rebuilt vector meta"); - assert!(rebuilt_meta.version() > previous_generation); + let rebuilt_meta = read_stored_vector_meta(redis, key); + assert!( + rebuilt_meta.version() > previous_generation, + "rebuilt set must get a fresh generation: {} <= {}", + rebuilt_meta.version(), + previous_generation + ); + assert_eq!(rebuilt_meta.data_revision(), 1); }); } +#[test] +fn test_data_revision_increments_on_successful_vadd_and_vrem() { + with_redis(|redis| { + let key = b"revision-vectors"; + let vector = CanonicalVector::from_values(&[1.0, 0.0]).expect("vector"); + + redis.vadd(key, b"a", &vector).expect("insert a"); + assert_eq!(read_stored_vector_meta(redis, key).data_revision(), 1); + + redis.vadd(key, b"b", &vector).expect("insert b"); + assert_eq!(read_stored_vector_meta(redis, key).data_revision(), 2); + + // Updating an existing member also bumps the revision. + assert!(!redis.vadd(key, b"a", &vector).expect("update a")); + assert_eq!(read_stored_vector_meta(redis, key).data_revision(), 3); + + redis.vrem(key, b"a").expect("remove a"); + assert_eq!(read_stored_vector_meta(redis, key).data_revision(), 4); + + // A failed removal leaves the revision untouched. + assert!(!redis.vrem(key, b"absent").expect("remove absent")); + assert_eq!(read_stored_vector_meta(redis, key).data_revision(), 4); + }); +} + +#[test] +fn test_generation_is_not_reused_across_restart() { + let path = unique_test_db_path(); + safe_cleanup_test_db(&path); + + let incarnation; + let last_generation; + { + let redis = open_redis(&path); + let vector = CanonicalVector::from_values(&[1.0, 0.0]).expect("vector"); + redis.vadd(b"first", b"a", &vector).expect("insert first"); + redis.vadd(b"second", b"a", &vector).expect("insert second"); + let first_generation = read_stored_vector_meta(&redis, b"first").version(); + last_generation = read_stored_vector_meta(&redis, b"second").version(); + assert!(last_generation > first_generation); + incarnation = redis.storage_incarnation().expect("incarnation"); + } + + let reopened = open_redis(&path); + assert_eq!( + reopened.storage_incarnation().expect("incarnation"), + incarnation, + "incarnation must be stable across restarts" + ); + let vector = CanonicalVector::from_values(&[1.0, 0.0]).expect("vector"); + reopened + .vadd(b"third", b"a", &vector) + .expect("insert third after reopen"); + let generation = read_stored_vector_meta(&reopened, b"third").version(); + assert!( + generation > last_generation, + "restarted instance must not reuse generations: {generation} <= {last_generation}" + ); + + drop(reopened); + safe_cleanup_test_db(&path); +} + +#[test] +fn test_open_refuses_non_empty_db_without_manifest() { + let path = unique_test_db_path(); + safe_cleanup_test_db(&path); + + { + let redis = open_redis(&path); + let vector = CanonicalVector::from_values(&[1.0, 0.0]).expect("vector"); + redis.vadd(b"vectors", b"a", &vector).expect("insert"); + } + + std::fs::remove_file(path.join(STORAGE_MANIFEST_FILE)).expect("remove manifest"); + + let storage_options = Arc::new(StorageOptions::default()); + let (bg_task_handler, _) = BgTaskHandler::new(); + let lock_mgr = Arc::new(LockMgr::new(1000)); + let mut redis = Redis::new(storage_options, 1, Arc::new(bg_task_handler), lock_mgr); + let result = redis.open(path.to_str().expect("test path must be UTF-8")); + assert!( + result.is_err(), + "opening a non-empty database without a manifest must fail" + ); + + drop(redis); + safe_cleanup_test_db(&path); +} + #[tokio::test] async fn test_expired_vectorset_reads_as_missing() { let test_db_path = unique_test_db_path(); @@ -612,3 +735,865 @@ async fn test_vector_storage_proposes_binlog_in_cluster_mode() { storage.shutdown().await; safe_cleanup_test_db(&test_db_path); } + +#[test] +fn test_vadd_enforces_max_dimension() { + with_redis_vector_config( + |vector| vector.max_dimension = 4, + |redis| { + let fitting = CanonicalVector::from_values(&[1.0, 0.0, 0.0, 0.0]).expect("4d vector"); + let oversized = + CanonicalVector::from_values(&[1.0, 0.0, 0.0, 0.0, 0.0]).expect("5d vector"); + + assert!(redis.vadd(b"vectors", b"a", &fitting).expect("4d insert")); + let error = redis + .vadd(b"vectors", b"b", &oversized) + .expect_err("5d vector must be rejected"); + assert!(matches!(error, Error::RedisErr { .. })); + assert!(error.to_string().contains("max_dimension")); + assert_eq!(redis.vcard(b"vectors").expect("card"), 1); + }, + ); +} + +#[test] +fn test_vadd_enforces_max_element_bytes() { + with_redis_vector_config( + |vector| vector.max_element_bytes = 4, + |redis| { + let vector = CanonicalVector::from_values(&[1.0, 0.0]).expect("vector"); + + assert!( + redis + .vadd(b"vectors", b"abcd", &vector) + .expect("fitting element") + ); + let error = redis + .vadd(b"vectors", b"abcde", &vector) + .expect_err("oversized element must be rejected"); + assert!(matches!(error, Error::RedisErr { .. })); + assert!(error.to_string().contains("max_element_bytes")); + assert_eq!(redis.vcard(b"vectors").expect("card"), 1); + }, + ); +} + +#[test] +fn test_vadd_enforces_max_vector_bytes() { + with_redis_vector_config( + |vector| vector.max_vector_bytes = 8, + |redis| { + let fitting = CanonicalVector::from_values(&[1.0, 0.0]).expect("2d vector"); + let oversized = CanonicalVector::from_values(&[1.0, 0.0, 0.0, 0.0]).expect("4d vector"); + + assert!( + redis + .vadd(b"vectors", b"a", &fitting) + .expect("8-byte vector") + ); + let error = redis + .vadd(b"vectors", b"b", &oversized) + .expect_err("16-byte vector must be rejected"); + assert!(matches!(error, Error::RedisErr { .. })); + assert!(error.to_string().contains("max_vector_bytes")); + assert_eq!(redis.vcard(b"vectors").expect("card"), 1); + }, + ); +} + +#[test] +fn test_vsim_enforces_max_k() { + with_redis_vector_config( + |vector| vector.max_k = 2, + |redis| { + let x = populate_search_vectors(redis); + + let error = redis + .vsim( + b"search", + VectorQuery::Vector(x.clone()), + search_options(3, VectorSearchMode::Approximate), + ) + .expect_err("count above max_k must be rejected"); + assert!(matches!(error, Error::RedisErr { .. })); + assert!(error.to_string().contains("max_k")); + + let hits = redis + .vsim( + b"search", + VectorQuery::Vector(x), + search_options(2, VectorSearchMode::Approximate), + ) + .expect("count within max_k"); + assert_eq!(hits.len(), 2); + }, + ); +} + +#[test] +fn test_vsim_waits_for_flat_query_permit() { + with_redis_vector_config( + |vector| vector.max_concurrent_flat_queries = 1, + |redis| { + populate_search_vectors(redis); + let held = redis + .flat_query_gate + .acquire(Instant::now() + Duration::from_secs(5)) + .expect("acquire the only permit"); + + std::thread::scope(|scope| { + let query = scope.spawn(|| { + redis.vsim( + b"search", + VectorQuery::Element(b"x".to_vec()), + search_options(2, VectorSearchMode::Approximate), + ) + }); + std::thread::sleep(Duration::from_millis(200)); + assert!(!query.is_finished(), "query must wait for the permit"); + drop(held); + let hits = query + .join() + .expect("query thread") + .expect("search after permit release"); + assert_eq!(hits.len(), 2); + }); + assert_eq!( + redis.flat_query_gate.available_permits(), + 1, + "permit must return to the gate" + ); + }, + ); +} + +#[test] +fn test_vsim_flat_deadline_covers_queue_wait() { + with_redis_vector_config( + |vector| { + vector.max_concurrent_flat_queries = 1; + vector.flat_query_timeout_ms = 150; + }, + |redis| { + populate_search_vectors(redis); + let held = redis + .flat_query_gate + .acquire(Instant::now() + Duration::from_secs(5)) + .expect("acquire the only permit"); + + let started = Instant::now(); + let result = std::thread::scope(|scope| { + let query = scope.spawn(|| { + redis.vsim( + b"search", + VectorQuery::Element(b"x".to_vec()), + search_options(2, VectorSearchMode::Approximate), + ) + }); + // Hold the permit well past the query deadline. + std::thread::sleep(Duration::from_secs(1)); + drop(held); + query.join().expect("query thread") + }); + + let error = result.expect_err("queued query must time out, no partial results"); + assert!(matches!(error, Error::VectorFlatQueryTimeout { .. })); + assert!( + started.elapsed() >= Duration::from_millis(150), + "deadline must include queue wait" + ); + assert_eq!( + redis.flat_query_gate.available_permits(), + 1, + "permit must be released after the timeout" + ); + + // The gate is fully usable again after the timed-out query. + let hits = redis + .vsim( + b"search", + VectorQuery::Element(b"x".to_vec()), + search_options(2, VectorSearchMode::Approximate), + ) + .expect("search after timeout"); + assert_eq!(hits.len(), 2); + }, + ); +} + +#[test] +fn test_vsim_scan_budget_entries_aborts_without_partial_results() { + with_redis_vector_config( + |vector| vector.flat_scan_max_entries = 50, + |redis| { + let vector = CanonicalVector::from_values(&[1.0, 0.0]).expect("vector"); + for i in 0..200 { + redis + .vadd(b"budget", format!("member-{i}").as_bytes(), &vector) + .expect("insert"); + } + + let error = redis + .vsim( + b"budget", + VectorQuery::Vector(vector.clone()), + search_options(10, VectorSearchMode::Approximate), + ) + .expect_err("scan budget must abort the query, no partial results"); + assert!(matches!(error, Error::VectorFlatScanBudgetExceeded { .. })); + assert_eq!( + redis.flat_query_gate.available_permits(), + 4, + "permit must be released after the abort" + ); + + // A set within budget still answers after the aborted scan. + redis.vadd(b"small", b"a", &vector).expect("insert small"); + let hits = redis + .vsim( + b"small", + VectorQuery::Vector(vector), + search_options(1, VectorSearchMode::Approximate), + ) + .expect("small search"); + assert_eq!(hits.len(), 1); + }, + ); +} + +#[test] +fn test_vsim_scan_budget_bytes_aborts_without_partial_results() { + with_redis_vector_config( + |vector| vector.flat_scan_max_bytes = 128, + |redis| { + let vector = CanonicalVector::from_values(&[1.0, 0.0]).expect("vector"); + for i in 0..50 { + redis + .vadd(b"bytes", format!("member-{i}").as_bytes(), &vector) + .expect("insert"); + } + + let error = redis + .vsim( + b"bytes", + VectorQuery::Vector(vector), + search_options(10, VectorSearchMode::Approximate), + ) + .expect_err("byte budget must abort the query, no partial results"); + assert!(matches!(error, Error::VectorFlatScanBudgetExceeded { .. })); + assert_eq!(redis.flat_query_gate.available_permits(), 4); + }, + ); +} + +#[test] +fn test_vsim_precancelled_token_aborts_immediately() { + with_redis(|redis| { + let x = populate_search_vectors(redis); + let cancel = FlatQueryCancel::default(); + cancel.cancel(); + + let error = redis + .vsim_with_cancel( + b"search", + VectorQuery::Vector(x), + search_options(2, VectorSearchMode::Approximate), + &cancel, + ) + .expect_err("cancelled token must abort the query"); + assert!(matches!(error, Error::VectorFlatQueryCancelled { .. })); + assert_eq!(redis.flat_query_gate.available_permits(), 4); + }); +} + +#[test] +fn test_vsim_cancel_mid_scan_aborts_and_releases_resources() { + with_redis_vector_config( + |vector| vector.flat_cancel_check_interval = 16, + |redis| { + let vector = CanonicalVector::from_values(&[1.0, 0.0]).expect("vector"); + for i in 0..50_000 { + redis + .vadd(b"big", format!("member-{i}").as_bytes(), &vector) + .expect("insert"); + } + + let cancel = FlatQueryCancel::default(); + let token = cancel.clone(); + let result = std::thread::scope(|scope| { + let query = scope.spawn(|| { + redis.vsim_with_cancel( + b"big", + VectorQuery::Vector(vector), + search_options(10, VectorSearchMode::Approximate), + &token, + ) + }); + std::thread::sleep(Duration::from_millis(20)); + cancel.cancel(); + query.join().expect("query thread") + }); + + let error = result.expect_err("mid-scan cancel must abort, no partial results"); + assert!(matches!(error, Error::VectorFlatQueryCancelled { .. })); + assert_eq!( + redis.flat_query_gate.available_permits(), + 4, + "permit must be released after cancellation" + ); + }, + ); +} + +#[test] +fn test_vinfo_missing_key_and_wrongtype() { + with_redis(|redis| { + assert_eq!(redis.vinfo(b"missing").expect("missing vinfo"), None); + + redis.set(b"string", b"value").expect("set string"); + let error = redis.vinfo(b"string").expect_err("string key must fail"); + assert!(error.to_string().contains("WRONGTYPE")); + }); +} + +#[test] +fn test_vinfo_reports_meta_and_generation_changes_on_recreate() { + with_redis(|redis| { + let vector3 = CanonicalVector::from_values(&[1.0, 0.0, 0.0]).expect("vector"); + assert!(redis.vadd(b"vectors", b"a", &vector3).expect("insert a")); + assert!(redis.vadd(b"vectors", b"b", &vector3).expect("insert b")); + + let first = redis.vinfo(b"vectors").expect("vinfo").expect("set exists"); + assert_eq!(first.dimension, 3); + assert_eq!(first.size, 2); + let first_generation = first.generation; + + // Removing the last member deletes the meta; recreating the set + // allocates a fresh generation sequence. + assert!(redis.vrem(b"vectors", b"a").expect("remove a")); + assert!(redis.vrem(b"vectors", b"b").expect("remove b")); + assert_eq!(redis.vinfo(b"vectors").expect("vinfo after drain"), None); + + let vector2 = CanonicalVector::from_values(&[1.0, 0.0]).expect("vector"); + assert!(redis.vadd(b"vectors", b"c", &vector2).expect("recreate")); + let second = redis + .vinfo(b"vectors") + .expect("vinfo") + .expect("recreated set"); + assert_eq!(second.dimension, 2); + assert_eq!(second.size, 1); + assert_ne!( + second.generation, first_generation, + "recreated set must have a new generation sequence" + ); + }); +} + +#[test] +fn test_vector_metrics_count_successful_queries() { + with_redis(|redis| { + let x = populate_search_vectors(redis); + let hits = redis + .vsim( + b"search", + VectorQuery::Vector(x), + search_options(2, VectorSearchMode::Approximate), + ) + .expect("search"); + assert_eq!(hits.len(), 2); + + let metrics = redis.vector_metrics.snapshot(); + assert_eq!(metrics.flat_queries_total, 1); + assert_eq!(metrics.flat_query_duration_count, 1); + assert_eq!(metrics.flat_query_timeouts_total, 0); + assert_eq!(metrics.flat_query_errors_total, 0); + assert_eq!(metrics.capacity_rejected_total, 0); + }); +} + +#[test] +fn test_vector_metrics_count_scan_deadline_timeout() { + with_redis_vector_config( + |vector| vector.flat_query_timeout_ms = 0, + |redis| { + let x = populate_search_vectors(redis); + let error = redis + .vsim( + b"search", + VectorQuery::Vector(x), + search_options(2, VectorSearchMode::Approximate), + ) + .expect_err("zero timeout must abort the query"); + assert!(matches!(error, Error::VectorFlatQueryTimeout { .. })); + + let metrics = redis.vector_metrics.snapshot(); + assert_eq!(metrics.flat_queries_total, 1); + assert_eq!(metrics.flat_query_timeouts_total, 1); + assert_eq!(metrics.flat_query_errors_total, 0); + assert_eq!(metrics.capacity_rejected_total, 0); + assert_eq!(metrics.flat_query_duration_count, 1); + }, + ); +} + +#[test] +fn test_vector_metrics_count_budget_and_cancel_errors() { + with_redis_vector_config( + |vector| vector.flat_scan_max_entries = 50, + |redis| { + let vector = CanonicalVector::from_values(&[1.0, 0.0]).expect("vector"); + for i in 0..200 { + redis + .vadd(b"budget", format!("member-{i}").as_bytes(), &vector) + .expect("insert"); + } + let error = redis + .vsim( + b"budget", + VectorQuery::Vector(vector.clone()), + search_options(10, VectorSearchMode::Approximate), + ) + .expect_err("budget must abort the query"); + assert!(matches!(error, Error::VectorFlatScanBudgetExceeded { .. })); + + let cancel = FlatQueryCancel::default(); + cancel.cancel(); + let error = redis + .vsim_with_cancel( + b"budget", + VectorQuery::Vector(vector), + search_options(1, VectorSearchMode::Approximate), + &cancel, + ) + .expect_err("cancelled token must abort the query"); + assert!(matches!(error, Error::VectorFlatQueryCancelled { .. })); + + let metrics = redis.vector_metrics.snapshot(); + assert_eq!(metrics.flat_queries_total, 2); + assert_eq!(metrics.flat_query_errors_total, 2); + assert_eq!(metrics.flat_query_timeouts_total, 0); + assert_eq!(metrics.flat_query_duration_count, 2); + }, + ); +} + +#[test] +fn test_vector_metrics_count_capacity_rejected_queries() { + with_redis_vector_config( + |vector| { + vector.max_concurrent_flat_queries = 1; + vector.flat_query_timeout_ms = 150; + }, + |redis| { + populate_search_vectors(redis); + let held = redis + .flat_query_gate + .acquire(Instant::now() + Duration::from_secs(5)) + .expect("acquire the only permit"); + + let result = std::thread::scope(|scope| { + let query = scope.spawn(|| { + redis.vsim( + b"search", + VectorQuery::Element(b"x".to_vec()), + search_options(2, VectorSearchMode::Approximate), + ) + }); + std::thread::sleep(Duration::from_secs(1)); + drop(held); + query.join().expect("query thread") + }); + assert!(matches!( + result.expect_err("queued query must be rejected"), + Error::VectorFlatQueryTimeout { .. } + )); + + let metrics = redis.vector_metrics.snapshot(); + assert_eq!(metrics.capacity_rejected_total, 1); + // The query never entered the scan path. + assert_eq!(metrics.flat_queries_total, 0); + assert_eq!(metrics.flat_query_timeouts_total, 0); + assert_eq!(metrics.flat_query_duration_count, 0); + }, + ); +} + +fn add_mutation(element: &[u8], values: &[f32]) -> VectorSetMutationV1 { + let vector = CanonicalVector::from_values(values).expect("valid vector"); + VectorSetMutationV1::add_from_canonical(element, &vector).expect("mutation from vector") +} + +#[test] +fn test_apply_mutation_create_add_update_and_remove() { + with_redis(|redis| { + let key = b"vectors"; + + // Create: first add creates the set. + let result = redis + .apply_vector_set_mutation(key, &add_mutation(b"a", &[1.0, 0.0]), None) + .expect("create set"); + assert_eq!(result, VectorSetApplyResult::Added); + assert_eq!(redis.vcard(key).expect("card"), 1); + assert_eq!(read_stored_vector_meta(redis, key).data_revision(), 1); + + // Update: same element keeps count, bumps data_revision. + let result = redis + .apply_vector_set_mutation(key, &add_mutation(b"a", &[0.0, 1.0]), None) + .expect("update member"); + assert_eq!(result, VectorSetApplyResult::Updated); + assert_eq!(redis.vcard(key).expect("card after update"), 1); + assert_eq!(read_stored_vector_meta(redis, key).data_revision(), 2); + assert_eq!(redis.vemb(key, b"a").expect("emb"), Some(vec![0.0, 1.0])); + + // Add: a new element increases count. + let result = redis + .apply_vector_set_mutation(key, &add_mutation(b"b", &[1.0, 0.0]), None) + .expect("add member"); + assert_eq!(result, VectorSetApplyResult::Added); + assert_eq!(redis.vcard(key).expect("card after add"), 2); + assert_eq!(read_stored_vector_meta(redis, key).data_revision(), 3); + + // Remove: existing member, then a miss. + let remove_a = VectorSetMutationV1::Remove { + element: b"a".to_vec(), + }; + let result = redis + .apply_vector_set_mutation(key, &remove_a, None) + .expect("remove member"); + assert_eq!(result, VectorSetApplyResult::Removed); + assert_eq!(redis.vcard(key).expect("card after remove"), 1); + let result = redis + .apply_vector_set_mutation(key, &remove_a, None) + .expect("remove miss"); + assert_eq!(result, VectorSetApplyResult::RemoveMissed); + + // Removing the last member deletes the meta with it. + let remove_b = VectorSetMutationV1::Remove { + element: b"b".to_vec(), + }; + let result = redis + .apply_vector_set_mutation(key, &remove_b, None) + .expect("remove last member"); + assert_eq!(result, VectorSetApplyResult::Removed); + assert_eq!(redis.vcard(key).expect("missing card"), 0); + assert!(redis.get_key_type(key).is_err()); + }); +} + +#[test] +fn test_apply_mutation_remove_miss_on_missing_key() { + with_redis(|redis| { + let mutation = VectorSetMutationV1::Remove { + element: b"member".to_vec(), + }; + let result = redis + .apply_vector_set_mutation(b"missing", &mutation, None) + .expect("remove from missing key"); + assert_eq!(result, VectorSetApplyResult::RemoveMissed); + }); +} + +#[test] +fn test_apply_mutation_uses_explicit_create_generation() { + with_redis(|redis| { + let result = redis + .apply_vector_set_mutation(b"vectors", &add_mutation(b"a", &[1.0, 0.0]), Some(42)) + .expect("create with explicit generation"); + assert_eq!(result, VectorSetApplyResult::Added); + assert_eq!(read_stored_vector_meta(redis, b"vectors").version(), 42); + + // The explicit generation only applies to creation; updates keep it. + let result = redis + .apply_vector_set_mutation(b"vectors", &add_mutation(b"b", &[0.0, 1.0]), Some(7)) + .expect("add to existing set"); + assert_eq!(result, VectorSetApplyResult::Added); + assert_eq!(read_stored_vector_meta(redis, b"vectors").version(), 42); + }); +} + +#[test] +fn test_apply_mutation_wrongtype_is_a_business_error() { + with_redis(|redis| { + redis.set(b"string-key", b"value").expect("set string"); + + let add_error = redis + .apply_vector_set_mutation(b"string-key", &add_mutation(b"a", &[1.0, 0.0]), None) + .expect_err("add on string key"); + assert!(matches!( + add_error, + VectorSetApplyError::Business(VectorSetBusinessError::WrongType) + )); + + let remove = VectorSetMutationV1::Remove { + element: b"a".to_vec(), + }; + let remove_error = redis + .apply_vector_set_mutation(b"string-key", &remove, None) + .expect_err("remove on string key"); + assert!(matches!( + remove_error, + VectorSetApplyError::Business(VectorSetBusinessError::WrongType) + )); + + // The business error maps back to the exact storage error VADD/VREM + // returned before the apply entry existed. + let mapped: Error = add_error.into(); + assert!( + mapped + .to_string() + .contains("WRONGTYPE Operation against a key holding the wrong kind of value") + ); + }); +} + +#[test] +fn test_apply_mutation_dimension_mismatch_is_a_business_error() { + with_redis(|redis| { + let result = redis + .apply_vector_set_mutation(b"vectors", &add_mutation(b"a", &[1.0, 0.0]), None) + .expect("create set"); + assert_eq!(result, VectorSetApplyResult::Added); + + let error = redis + .apply_vector_set_mutation(b"vectors", &add_mutation(b"b", &[1.0, 0.0, 0.0]), None) + .expect_err("dimension mismatch"); + assert!(matches!( + error, + VectorSetApplyError::Business(VectorSetBusinessError::DimensionMismatch { + expected: 2, + got: 3 + }) + )); + // A rejected mutation leaves the set untouched. + assert_eq!(redis.vcard(b"vectors").expect("card"), 1); + }); +} + +#[test] +fn test_apply_mutation_codec_round_trip_through_apply() { + with_redis(|redis| { + let mutation = add_mutation(b"member", &[3.0, 4.0]); + let decoded = VectorSetMutationV1::decode(&mutation.encode()).expect("decode mutation"); + assert_eq!(decoded, mutation); + + let result = redis + .apply_vector_set_mutation(b"vectors", &decoded, None) + .expect("apply decoded mutation"); + assert_eq!(result, VectorSetApplyResult::Added); + let restored = redis + .vemb(b"vectors", b"member") + .expect("embedding") + .expect("member exists"); + assert!((restored[0] - 3.0).abs() < 1e-6); + assert!((restored[1] - 4.0).abs() < 1e-6); + }); +} + +#[test] +fn test_vector_meta_read_fault_blocks_reads_and_writes() { + with_redis(|redis| { + let vector = CanonicalVector::from_values(&[1.0, 0.0]).expect("valid vector"); + redis.vadd(b"vectors", b"a", &vector).expect("insert a"); + + redis.vector_fault_hooks.set_fail_meta_read(true); + + // Reads against MetaCF all surface the injected failure as a + // complete error, never a partial result. + let error = redis.vcard(b"vectors").expect_err("vcard must fail"); + assert!(error.to_string().contains("injected fault")); + assert!(redis.vinfo(b"vectors").is_err()); + assert!(redis.vemb(b"vectors", b"a").is_err()); + assert!(redis.vismember(b"vectors", b"a").is_err()); + let error = redis + .vsim( + b"vectors", + VectorQuery::Vector(vector.clone()), + search_options(1, VectorSearchMode::Truth), + ) + .expect_err("vsim must fail"); + assert!(error.to_string().contains("injected fault")); + + // Mutations read the meta before deciding the outcome, so both an + // update on an existing set and the creation of a new set fail + // without any partial write. + let error = redis + .vadd(b"vectors", b"b", &vector) + .expect_err("vadd on existing set must fail"); + assert!(error.to_string().contains("injected fault")); + assert!(redis.vadd(b"other", b"a", &vector).is_err()); + assert!(redis.vrem(b"vectors", b"a").is_err()); + + // The failed mutations left no trace behind. + assert_eq!(count_cf_entries(redis, ColumnFamilyIndex::MetaCF), 1); + assert_eq!(count_cf_entries(redis, ColumnFamilyIndex::VectorDataCF), 1); + + redis.vector_fault_hooks.set_fail_meta_read(false); + + // After disarming, the pre-fault state is intact and writes work. + assert_eq!(redis.vcard(b"vectors").expect("card"), 1); + assert!(redis.vadd(b"vectors", b"b", &vector).expect("retry vadd")); + assert_eq!(redis.vcard(b"vectors").expect("card after retry"), 2); + assert!(redis.vrem(b"vectors", b"b").expect("retry vrem")); + assert_eq!(redis.vcard(b"vectors").expect("card after vrem"), 1); + }); +} + +#[test] +fn test_vector_member_read_fault_fails_vsim_without_partial_results() { + with_redis(|redis| { + let query = populate_search_vectors(redis); + + let baseline = redis + .vsim( + b"search", + VectorQuery::Vector(query.clone()), + search_options(2, VectorSearchMode::Truth), + ) + .expect("baseline vsim"); + assert_eq!(baseline.len(), 2); + + redis.vector_fault_hooks.set_fail_member_read(true); + + // An iterator-level member read failure aborts the whole query + // instead of skipping the damaged member and returning a partial + // top-K. + let error = redis + .vsim( + b"search", + VectorQuery::Vector(query.clone()), + search_options(2, VectorSearchMode::Truth), + ) + .expect_err("vsim must fail on member read fault"); + assert!(error.to_string().contains("injected fault")); + + // Writes and meta reads do not go through the scan path and keep + // working while the member read fault is armed. + assert_eq!(redis.vcard(b"search").expect("card during fault"), 4); + + redis.vector_fault_hooks.set_fail_member_read(false); + + // The scan recovers and returns the same results as before. + let recovered = redis + .vsim( + b"search", + VectorQuery::Vector(query), + search_options(2, VectorSearchMode::Truth), + ) + .expect("vsim after disarm"); + assert_eq!(recovered.len(), 2); + assert_eq!(redis.vcard(b"search").expect("card after fault"), 4); + }); +} + +#[test] +fn test_vector_batch_commit_fault_preserves_atomicity() { + with_redis(|redis| { + let vector = CanonicalVector::from_values(&[1.0, 0.0]).expect("valid vector"); + redis.vadd(b"vectors", b"a", &vector).expect("insert a"); + let meta_before = read_stored_vector_meta(redis, b"vectors"); + let members_before = count_cf_entries(redis, ColumnFamilyIndex::VectorDataCF); + + redis.vector_fault_hooks.set_fail_batch_commit(true); + + // Add on an existing set: commit fails, meta (count/data_revision) + // and member stay invisible together. + let error = redis + .vadd(b"vectors", b"b", &vector) + .expect_err("vadd must fail on commit fault"); + assert!(error.to_string().contains("injected fault")); + + // Create of a new set: commit fails atomically as well. + assert!(redis.vadd(b"newset", b"x", &vector).is_err()); + + // Remove: commit fails, the member and meta survive together. + assert!(redis.vrem(b"vectors", b"a").is_err()); + + // Reads are unaffected by the commit fault and observe exactly the + // pre-fault state. + assert_eq!(redis.vcard(b"vectors").expect("card"), 1); + assert!(redis.vismember(b"vectors", b"a").expect("member a")); + assert!(!redis.vismember(b"vectors", b"b").expect("member b")); + assert_eq!(redis.vcard(b"newset").expect("newset card"), 0); + assert_eq!( + count_cf_entries(redis, ColumnFamilyIndex::VectorDataCF), + members_before + ); + let meta_after = read_stored_vector_meta(redis, b"vectors"); + assert_eq!(meta_after.data_revision(), meta_before.data_revision()); + assert_eq!(meta_after.version(), meta_before.version()); + + redis.vector_fault_hooks.reset(); + + // After disarming, retried mutations succeed and stay consistent. + assert!(redis.vadd(b"vectors", b"b", &vector).expect("retry vadd")); + assert_eq!(redis.vcard(b"vectors").expect("card after retry"), 2); + assert!(redis.vrem(b"vectors", b"a").expect("retry vrem")); + assert_eq!(redis.vcard(b"vectors").expect("card after vrem"), 1); + assert!(redis.vadd(b"newset", b"x", &vector).expect("retry create")); + assert_eq!(redis.vcard(b"newset").expect("newset card"), 1); + }); +} + +#[test] +fn vector_data_sample_validation_passes_on_healthy_data() { + with_redis(|redis| { + let vector = CanonicalVector::from_values(&[1.0, 0.0]).expect("vector"); + redis.vadd(b"vs", b"e1", &vector).expect("vadd e1"); + redis.vadd(b"vs", b"e2", &vector).expect("vadd e2"); + + let sample = redis.validate_vector_data_sample(16).expect("sample"); + assert_eq!(sample.members, 2); + assert_eq!(sample.metas, 1); + + // A sample size of zero samples nothing and always passes. + let empty = redis.validate_vector_data_sample(0).expect("empty sample"); + assert_eq!(empty.members, 0); + assert_eq!(empty.metas, 0); + }); +} + +#[test] +fn vector_data_sample_validation_rejects_corrupt_member() { + with_redis(|redis| { + let vector = CanonicalVector::from_values(&[1.0, 0.0]).expect("vector"); + redis.vadd(b"vs", b"e1", &vector).expect("vadd e1"); + + let db = redis.db().expect("db"); + let vector_cf = redis + .get_cf_handle(ColumnFamilyIndex::VectorDataCF) + .expect("vector cf"); + let member_key = db + .iterator_cf(&vector_cf, IteratorMode::Start) + .next() + .expect("one member entry") + .expect("member entry") + .0; + db.put_cf(&vector_cf, &member_key, b"garbage") + .expect("corrupt member value"); + + assert!(redis.validate_vector_data_sample(16).is_err()); + }); +} + +#[test] +fn vector_data_sample_validation_rejects_corrupt_meta() { + with_redis(|redis| { + let vector = CanonicalVector::from_values(&[1.0, 0.0]).expect("vector"); + redis.vadd(b"vs", b"e1", &vector).expect("vadd e1"); + + let db = redis.db().expect("db"); + let meta_cf = redis + .get_cf_handle(ColumnFamilyIndex::MetaCF) + .expect("meta cf"); + // A VectorSet-tagged meta value that fails codec decode (too short). + let corrupt_key = BaseMetaKey::new(b"zz_corrupt").encode().expect("meta key"); + db.put_cf( + &meta_cf, + &corrupt_key, + vec![storage::DataType::VectorSet as u8, 0, 0], + ) + .expect("corrupt meta value"); + + assert!(redis.validate_vector_data_sample(16).is_err()); + }); +} diff --git a/tests/python/conftest.py b/tests/python/conftest.py index cf044ad1..d3518599 100644 --- a/tests/python/conftest.py +++ b/tests/python/conftest.py @@ -42,8 +42,8 @@ def redis_client(): 作用域为 session，所有测试共享一个连接 """ client = redis.Redis( - host='localhost', - port=6379, + host=os.getenv("KIWI_HOST", "localhost"), + port=int(os.getenv("KIWI_PORT", "7379")), decode_responses=True, socket_connect_timeout=5, socket_timeout=5, @@ -54,7 +54,10 @@ def redis_client(): client.ping() except redis.RedisError as error: client.close() - message = "Redis server is not running on localhost:6379" + message = ( + "Redis server is not running on " + f"{os.getenv('KIWI_HOST', 'localhost')}:{os.getenv('KIWI_PORT', '7379')}" + ) if _enabled("KIWI_TEST_REQUIRE_SERVER"): pytest.fail(f"{message}: {error}", pytrace=False) pytest.skip(message) @@ -117,8 +120,8 @@ def redis_binary_client(redis_client): 用于测试二进制安全功能 """ client = redis.Redis( - host='localhost', - port=6379, + host=os.getenv("KIWI_HOST", "localhost"), + port=int(os.getenv("KIWI_PORT", "7379")), decode_responses=False, # 不自动解码 socket_connect_timeout=5, socket_timeout=5, @@ -128,7 +131,10 @@ def redis_binary_client(redis_client): client.ping() except redis.RedisError as error: client.close() - message = "Redis server is not running on localhost:6379" + message = ( + "Redis server is not running on " + f"{os.getenv('KIWI_HOST', 'localhost')}:{os.getenv('KIWI_PORT', '7379')}" + ) if _enabled("KIWI_TEST_REQUIRE_SERVER"): pytest.fail(f"{message}: {error}", pytrace=False) pytest.skip(message) diff --git a/tests/python/test_mset.py b/tests/python/test_mset.py index 5a5d5a9e..1162b5ad 100644 --- a/tests/python/test_mset.py +++ b/tests/python/test_mset.py @@ -26,6 +26,7 @@ python tests/python/test_mset.py # 直接运行 """ +import os import sys import pytest @@ -246,12 +247,16 @@ def run_standalone_tests(): try: # 连接测试 - r = redis.Redis(host='localhost', port=6379, decode_responses=True) + r = redis.Redis( + host=os.getenv("KIWI_HOST", "localhost"), + port=int(os.getenv("KIWI_PORT", "7379")), + decode_responses=True, + ) r.ping() print("✓ 成功连接到 Redis 服务器\n") except redis.ConnectionError: print("✗ 无法连接到 Redis 服务器") - print("请确保服务器正在运行在 localhost:6379") + print("请确保服务器正在运行在 localhost:7379（可用 KIWI_HOST/KIWI_PORT 覆盖）") sys.exit(1) test_count = 0 @@ -302,7 +307,11 @@ def run_standalone_tests(): def test_basic(): """独立测试函数 - 基本功能""" try: - r = redis.Redis(host='localhost', port=6379, decode_responses=True) + r = redis.Redis( + host=os.getenv("KIWI_HOST", "localhost"), + port=int(os.getenv("KIWI_PORT", "7379")), + decode_responses=True, + ) r.ping() cleanup_keys(r, ['test_kiwi_mset_key1', 'test_kiwi_mset_key2']) @@ -318,7 +327,11 @@ def test_basic(): def test_with_mget(): """独立测试函数 - 与 MGET 配合""" try: - r = redis.Redis(host='localhost', port=6379, decode_responses=True) + r = redis.Redis( + host=os.getenv("KIWI_HOST", "localhost"), + port=int(os.getenv("KIWI_PORT", "7379")), + decode_responses=True, + ) r.ping() cleanup_keys(r, ['test_kiwi_mset_mget1', 'test_kiwi_mset_mget2']) @@ -334,7 +347,11 @@ def test_with_mget(): def test_overwrite(): """独立测试函数 - 覆盖测试""" try: - r = redis.Redis(host='localhost', port=6379, decode_responses=True) + r = redis.Redis( + host=os.getenv("KIWI_HOST", "localhost"), + port=int(os.getenv("KIWI_PORT", "7379")), + decode_responses=True, + ) r.ping() cleanup_keys(r, ['test_kiwi_mset_over']) @@ -350,7 +367,11 @@ def test_overwrite(): def test_single_pair(): """独立测试函数 - 单个键值对""" try: - r = redis.Redis(host='localhost', port=6379, decode_responses=True) + r = redis.Redis( + host=os.getenv("KIWI_HOST", "localhost"), + port=int(os.getenv("KIWI_PORT", "7379")), + decode_responses=True, + ) r.ping() cleanup_keys(r, ['test_kiwi_mset_single']) @@ -365,7 +386,11 @@ def test_single_pair(): def test_large_batch(): """独立测试函数 - 大批量操作""" try: - r = redis.Redis(host='localhost', port=6379, decode_responses=True) + r = redis.Redis( + host=os.getenv("KIWI_HOST", "localhost"), + port=int(os.getenv("KIWI_PORT", "7379")), + decode_responses=True, + ) r.ping() # 清理可能存在的键 @@ -385,7 +410,11 @@ def test_large_batch(): def test_batch_round_trip(): """独立测试函数 - 批量往返测试""" try: - r = redis.Redis(host='localhost', port=6379, decode_responses=True) + r = redis.Redis( + host=os.getenv("KIWI_HOST", "localhost"), + port=int(os.getenv("KIWI_PORT", "7379")), + decode_responses=True, + ) r.ping() cleanup_keys(r, ['test_kiwi_mset_a1', 'test_kiwi_mset_a2']) diff --git a/tests/python/test_vector_basic.py b/tests/python/test_vector_basic.py new file mode 100644 index 00000000..f2dba547 --- /dev/null +++ b/tests/python/test_vector_basic.py @@ -0,0 +1,128 @@ +# Copyright (c) 2024-present, arana-db Community. All rights reserved. +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +#!/usr/bin/env python3 +"""向量索引基础功能验证脚本。 + +运行前请确保 Kiwi 服务已启动（默认端口 7379）。 +""" + +import redis +import struct + +DIM = 4 +INDEX = 'vtest_idx' +PREFIX = 'vdoc:' +PORT = 7379 + + +def pack_vec(values): + """将 float 列表编码为 FLOAT32 小端字节序。""" + return struct.pack(f'<{len(values)}f', *values) + + +def clean_docs(r): + for i in range(1, 4): + r.delete(f'{PREFIX}{i}') + + +def main(): + r = redis.Redis(host='localhost', port=PORT, decode_responses=True, protocol=2) + r.ping() + + clean_docs(r) + + # FT.CREATE: 创建向量索引（已存在则忽略） + try: + r.execute_command( + 'FT.CREATE', INDEX, 'ON', 'HASH', 'PREFIX', '1', PREFIX, + 'SCHEMA', 'vec', 'VECTOR', 'FLAT', '6', + 'TYPE', 'FLOAT32', 'DIM', str(DIM), 'DISTANCE_METRIC', 'L2' + ) + print('FT.CREATE OK') + except redis.ResponseError as e: + if 'already exists' not in str(e).lower(): + raise + print('FT.CREATE: index already exists') + + # HSET: 写入带向量字段的 Hash + r.hset(f'{PREFIX}1', mapping={ + 'vec': pack_vec([1.0, 0.0, 0.0, 0.0]), + 'name': 'alpha' + }) + r.hset(f'{PREFIX}2', mapping={ + 'vec': pack_vec([0.0, 1.0, 0.0, 0.0]), + 'name': 'beta' + }) + r.hset(f'{PREFIX}3', mapping={ + 'vec': pack_vec([0.0, 0.0, 1.0, 0.0]), + 'name': 'gamma' + }) + print('HSET 3 docs OK') + + query = pack_vec([1.0, 0.0, 0.0, 0.0]) + + # 先用 KNN 3 检查 3 个文档是否都被索引 + res_all = r.execute_command( + 'FT.SEARCH', INDEX, '*=>[KNN 3 @vec $q]', + 'PARAMS', '2', 'q', query, + 'RETURN', '1', 'name', + 'DIALECT', '2' + ) + print('FT.SEARCH (all docs) result:', res_all) + keys_all = {res_all[i] for i in range(1, len(res_all), 2)} + assert keys_all == {f'{PREFIX}1', f'{PREFIX}2', f'{PREFIX}3'}, ( + f'expected 3 indexed docs, got {res_all[0]} hits with keys {keys_all}' + ) + + # KNN 2 检查 top-2 排序 + res = r.execute_command( + 'FT.SEARCH', INDEX, '*=>[KNN 2 @vec $q]', + 'PARAMS', '2', 'q', query, + 'RETURN', '1', 'name', + 'DIALECT', '2' + ) + print('FT.SEARCH (top 2) result:', res) + assert res[0] == 2, f'expected 2 hits, got {res[0]}' + assert res[1] == f'{PREFIX}1', f'expected {PREFIX}1 as first hit, got {res[1]}' + assert res[2] == ['name', 'alpha'], f"expected ['name', 'alpha'], got {res[2]}" + assert res[4] in (['name', 'beta'], ['name', 'gamma']) + + # HDEL: 删除 doc1 的向量字段 + deleted = r.hdel(f'{PREFIX}1', 'vec') + print('HDEL vec =>', deleted) + assert deleted == 1, f'expected HDEL to delete 1 field, got {deleted}' + + # 再次检索，验证 doc1 不再参与向量搜索 + res2 = r.execute_command( + 'FT.SEARCH', INDEX, '*=>[KNN 3 @vec $q]', + 'PARAMS', '2', 'q', query, + 'RETURN', '1', 'name', + 'DIALECT', '2' + ) + print('FT.SEARCH after HDEL result:', res2) + + assert res2[0] == 2, f'expected 2 hits after HDEL, got {res2[0]}' + keys_after = {res2[i] for i in range(1, len(res2), 2)} + assert keys_after == {f'{PREFIX}2', f'{PREFIX}3'}, f'unexpected keys: {keys_after}' + + clean_docs(r) + print('all assertions passed') + + +if __name__ == '__main__': + main() diff --git a/tests/python/test_vector_cluster.py b/tests/python/test_vector_cluster.py new file mode 100644 index 00000000..ab2148fa --- /dev/null +++ b/tests/python/test_vector_cluster.py @@ -0,0 +1,706 @@ +# Copyright (c) 2024-present, arana-db Community. All rights reserved. +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Vector Set 三节点 Raft 集群集成测试。 + +每个用例独立拉起一套 3 节点 kiwi 集群（动态端口、临时数据目录），通过 +grpcurl 调用 RaftAdminService 初始化/变更成员，通过 redis-py 走 RESP 协议 +验证 Vector Set 命令在集群模式下的行为。 + +运行方式（默认跳过，避免拖慢常规集成测试）： + + cargo build --bin kiwi + KIWI_RUN_CLUSTER_TESTS=1 python3 -m pytest tests/python/test_vector_cluster.py -v + +依赖：`grpcurl`（集群初始化与成员管理）、`target/debug/kiwi`（可用 +KIWI_BINARY 覆盖）、redis-py。 + +已知限制（本分支现状，非测试缺陷）： +- Raft 快照的网络安装尚未实现（src/raft/src/network.rs 的 install_snapshot + 直接返回 Unsupported，src/raft/src/grpc/core.rs 的接收端丢弃分片），因此 + "follower 落后超过 leader 日志保留窗口后通过快照追平" 的场景无法测试。 + test_lagged_follower_log_replay_and_snapshot_build 只覆盖：落后未超窗的 + follower 通过日志回放追平 + leader 本地快照构建/日志清理不影响集群。 +- Vector member key 编码内嵌节点本地随机的 storage_incarnation + （src/storage/src/format_vector_member_key.rs + storage_manifest.rs）， + 而集群复制是物理 binlog 回放（逐字节复制 leader 编码后的 CF 记录）。 + 因此 follower 重放的 member 记录带有 leader 的 incarnation，用本地 + incarnation 构造 key 的读路径（VEMB/VISMEMBER/VSIM/VREM）在新 leader + 上读不到旧数据；只有不内嵌 incarnation 的 meta（VCARD/VDIM/VINFO） + 复制后可读。这是 PR0 logical mutation replay（社区 issue #332）要解决 + 的问题。因此集群模式下 Vector 命令默认被 `vector-cluster-enabled` + 门禁确定性拒绝（ERR vector commands are not supported in cluster + mode yet）；本测试在节点配置中显式开启该开关以验证集群链路行为。 + 主用例只断言当前可兑现的行为；member 数据在副本上的可读性由 + test_replica_member_data_survives_failover（xfail）锁定，PR0 落地后 + 该用例会自然转绿（XPASS）。 +""" + +import json +import os +import shutil +import signal +import socket +import subprocess +import threading +import time + +import pytest +import redis +from redis.exceptions import AskError, MovedError + +# redis-py 的 MovedError/AskError 期望 Redis Cluster 的 "MOVED " +# 格式，而 Kiwi 返回简化形式 "MOVED "（无 slot 号），原生解析会直接 +# 抛 ValueError。这里打补丁兼容两种格式，并保留 host/port 供重定向使用。 +def _redirect_error_init(self, resp, status_code=None): + redis.ResponseError.__init__(self, resp, status_code=status_code) + self.args = (resp,) + self.message = resp + parts = resp.split(" ") + if len(parts) == 1: + self.slot_id = None + host, port = parts[0].rsplit(":", 1) + else: + self.slot_id = int(parts[0]) + host, port = parts[1].rsplit(":", 1) + self.node_addr = self.host, self.port = host, int(port) + + +AskError.__init__ = _redirect_error_init + +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) +KIWI_BINARY = os.environ.get( + "KIWI_BINARY", os.path.join(REPO_ROOT, "target", "debug", "kiwi") +) +GRPCURL = shutil.which("grpcurl") + +ELECTION_TIMEOUT_MS = 600 + +pytestmark = [ + pytest.mark.integration, + pytest.mark.slow, + pytest.mark.skipif( + os.environ.get("KIWI_RUN_CLUSTER_TESTS") != "1", + reason="cluster tests are slow; set KIWI_RUN_CLUSTER_TESTS=1 to enable", + ), + pytest.mark.skipif(GRPCURL is None, reason="grpcurl is required"), + pytest.mark.skipif( + not os.path.isfile(KIWI_BINARY), + reason=f"kiwi binary not found at {KIWI_BINARY}; run `cargo build --bin kiwi`", + ), +] + + +# --------------------------------------------------------------------------- +# 工具函数 +# --------------------------------------------------------------------------- + + +def wait_until(fn, timeout, desc, interval=0.2): + """轮询 fn 直到返回真值；超时抛 AssertionError。返回 fn 的真值结果。""" + deadline = time.monotonic() + timeout + while True: + result = fn() + if result: + return result + if time.monotonic() > deadline: + raise AssertionError(f"timed out after {timeout}s waiting for: {desc}") + time.sleep(interval) + + +def grpc_call(addr, service, method, payload, timeout=10): + """通过 grpcurl 调用 gRPC（服务带反射，无需 proto 文件）。失败返回 None。""" + try: + proc = subprocess.run( + [GRPCURL, "-plaintext", "-d", json.dumps(payload), addr, f"{service}/{method}"], + capture_output=True, + text=True, + timeout=timeout, + ) + except subprocess.TimeoutExpired: + return None + if proc.returncode != 0: + return None + out = proc.stdout.strip() + if not out: + return {} + try: + return json.loads(out) + except json.JSONDecodeError: + return None + + +def grpc_metrics(node): + return grpc_call( + f"127.0.0.1:{node.raft_port}", + "kiwi.raft.v1.RaftMetricsService", + "Metrics", + {}, + ) + + +def free_ports(count): + ports = set() + while len(ports) < count: + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.bind(("127.0.0.1", 0)) + ports.add(sock.getsockname()[1]) + sock.close() + return sorted(ports) + + +def vec(i): + """第 i 个元素的 2 维向量（非零，避免零范数边界）。""" + return (1.0 + (i % 1000) * 0.001, 0.5 - (i % 1000) * 0.0005) + + +def elem(i): + return f"e{i}".encode() + + +# force_leader 的 bump 写入使用独立 key，不影响各用例主 key 的 VCARD 计数 +BUMP_KEY = b"vc:bump" + + +def vadd_args(key, i): + x, y = vec(i) + return (b"VADD", key, b"VALUES", 2, f"{x:.6f}", f"{y:.6f}", elem(i), b"NOQUANT") + + +# --------------------------------------------------------------------------- +# 集群封装 +# --------------------------------------------------------------------------- + + +class ClusterNode: + def __init__(self, node_id, raft_port, resp_port, node_dir): + self.node_id = node_id + self.raft_port = raft_port + self.resp_port = resp_port + self.dir = node_dir + self.conf_path = os.path.join(node_dir, "node.conf") + self.log_path = os.path.join(node_dir, "kiwi.log") + self.proc = None + self.log_fd = None + self.alive = False + + @property + def raft_addr(self): + return f"127.0.0.1:{self.raft_port}" + + @property + def resp_addr(self): + return f"127.0.0.1:{self.resp_port}" + + +class VectorCluster: + def __init__(self, base_dir, binary, node_count=3): + self.base_dir = str(base_dir) + self.binary = binary + ports = free_ports(node_count * 2) + self.nodes = [] + for i in range(node_count): + node_dir = os.path.join(self.base_dir, f"node{i + 1}") + os.makedirs(node_dir, exist_ok=True) + self.nodes.append( + ClusterNode(i + 1, ports[i], ports[node_count + i], node_dir) + ) + self._clients = {} + + # -- 生命周期 --------------------------------------------------------- + + def _write_config(self, node): + with open(node.conf_path, "w", encoding="utf-8") as conf: + conf.write( + "\n".join( + [ + "binding 127.0.0.1", + f"port {node.resp_port}", + "runtime-network_threads 1", + "runtime-storage_threads 2", + "db-instance-num 1", + f"data-dir {node.dir}/db", + f"log-dir {node.dir}/logs", + f"raft-node-id {node.node_id}", + f"raft-addr {node.raft_addr}", + f"raft-resp-addr {node.resp_addr}", + f"raft-data-dir {node.dir}/raft", + "raft-heartbeat-interval-ms 100", + "raft-election-timeout-min-ms 300", + f"raft-election-timeout-max-ms {ELECTION_TIMEOUT_MS}", + # 集群模式 Vector 命令默认拒绝（failover 后成员数据 + # 不可读的防护，见 docs/vector-set-operations.md）； + # 本测试显式放开以验证集群链路行为。 + "vector-cluster-enabled yes", + "", + ] + ) + ) + + def start_node(self, node, timeout=30): + self._write_config(node) + node.log_fd = open(node.log_path, "ab") + env = dict(os.environ, RUST_LOG="info") + node.proc = subprocess.Popen( + [self.binary, "--config", node.conf_path], + cwd=node.dir, + stdout=node.log_fd, + stderr=subprocess.STDOUT, + start_new_session=True, + env=env, + ) + node.alive = True + + def ping_ok(): + if node.proc.poll() is not None: + raise AssertionError( + f"node {node.node_id} exited early (rc={node.proc.returncode}); " + f"see {node.log_path}" + ) + try: + self.client(node).ping() + return True + except redis.RedisError: + return False + + wait_until(ping_ok, timeout, f"node {node.node_id} RESP ready") + + def start_all(self): + for node in self.nodes: + self.start_node(node) + + def kill(self, node): + """SIGKILL 整个进程组，模拟节点崩溃。""" + if node.proc is not None and node.proc.poll() is None: + os.killpg(node.proc.pid, signal.SIGKILL) + node.proc.wait(timeout=10) + node.alive = False + + def restart(self, node, timeout=30): + assert not node.alive + self._clients.pop(node.node_id, None) + self.start_node(node, timeout=timeout) + # 等它重新感知到 leader（或自己成为 leader） + wait_until( + lambda: (grpc_metrics(node) or {}).get("currentLeader", 0) != 0 + or (grpc_metrics(node) or {}).get("isLeader"), + timeout, + f"node {node.node_id} rejoined the cluster", + ) + + def shutdown(self): + for node in self.nodes: + if node.alive: + self.kill(node) + if node.log_fd is not None: + node.log_fd.close() + for client in self._clients.values(): + client.close() + self._clients.clear() + + # -- 客户端与重定向 ---------------------------------------------------- + + def client(self, node): + if node.node_id not in self._clients: + self._clients[node.node_id] = redis.Redis( + host="127.0.0.1", + port=node.resp_port, + decode_responses=False, + socket_connect_timeout=2, + socket_timeout=10, + ) + return self._clients[node.node_id] + + def alive_nodes(self): + return [n for n in self.nodes if n.alive] + + def node_by_resp_port(self, port): + for node in self.nodes: + if node.resp_port == port: + return node + raise AssertionError(f"MOVED target port {port} is not a cluster node") + + def execute_from(self, node, *args, max_redirects=5): + """从指定节点发命令，跟随 MOVED 重定向（Kiwi 简化版：MOVED ）。""" + current = node + for _ in range(max_redirects): + try: + return self.client(current).execute_command(*args) + except MovedError as exc: + current = self.node_by_resp_port(exc.port) + raise AssertionError(f"too many MOVED redirects for command {args[0]!r}") + + # -- Raft 操作 ---------------------------------------------------------- + + def initialize(self): + payload = { + "nodes": [ + { + "node_id": n.node_id, + "raft_addr": n.raft_addr, + "resp_addr": n.resp_addr, + } + for n in self.nodes + ] + } + result = wait_until( + lambda: grpc_call( + self.nodes[0].raft_addr, + "kiwi.raft.v1.RaftAdminService", + "Initialize", + payload, + ), + 30, + "cluster Initialize RPC", + ) + assert result.get("response", {}).get("success") is True, ( + f"Initialize failed: {result}" + ) + + def wait_leader(self, timeout=30): + def find(): + for node in self.alive_nodes(): + metrics = grpc_metrics(node) + if metrics and metrics.get("isLeader"): + return node + return None + + return wait_until(find, timeout, "a leader to be elected") + + def wait_card(self, node, key, expected, timeout=30): + """等指定节点（应为 leader）上线性一致读 VCARD 达到期望值。""" + + def card_ok(): + try: + return self.client(node).execute_command(b"VCARD", key) == expected + except redis.RedisError: + return False + + wait_until( + card_ok, + timeout, + f"node {node.node_id} VCARD({key!r}) == {expected}", + ) + + +def force_leader(cluster, target, timeout=60): + """确定性地让 target 成为 leader。 + + 利用 Raft 选举限制——voter 不会把票投给日志不如自己完整的 candidate： + 先杀掉第三节点，通过现 leader 提交一条 bump 写入（只会到达 target， + 使 target 日志严格最长），再杀掉现 leader 并重启第三节点（其日志必然 + 短于 target），于是 target 必然赢得选举。bump 写入提交成功本身也隐含 + 了 target 已追平到该写入。返回新的 leader（应为 target）。 + """ + leader = cluster.wait_leader() + if leader.node_id == target.node_id: + return leader + third = next( + n + for n in cluster.nodes + if n.node_id not in (leader.node_id, target.node_id) + ) + cluster.kill(third) + cluster.execute_from(leader, *vadd_args(BUMP_KEY, target.node_id * 1000 + leader.node_id)) + cluster.kill(leader) + cluster.restart(third) + new_leader = cluster.wait_leader(timeout=timeout) + assert new_leader.node_id == target.node_id, ( + f"expected node {target.node_id} to win the election (longest log), " + f"got node {new_leader.node_id}" + ) + return new_leader + + +# --------------------------------------------------------------------------- +# 夹具 +# --------------------------------------------------------------------------- + + +@pytest.fixture(scope="session", autouse=True) +def redis_client(): + """覆盖 conftest 的全局单节点夹具：本文件自带 3 节点集群，不依赖外部服务器。""" + yield None + + +@pytest.fixture() +def cluster(tmp_path): + instance = VectorCluster(tmp_path, KIWI_BINARY) + try: + instance.start_all() + instance.initialize() + instance.wait_leader(timeout=30) + except Exception: + instance.shutdown() + raise + yield instance + instance.shutdown() + + +def vadd_batch(cluster, key, start, count, chunk=250): + """在当前 leader 上批量写入 e[start]..e[start+count-1]，断言全部为新元素。""" + leader = cluster.wait_leader() + client = cluster.client(leader) + end = start + count + for off in range(start, end, chunk): + pipe = client.pipeline(transaction=False) + for i in range(off, min(off + chunk, end)): + pipe.execute_command(*vadd_args(key, i)) + for reply in pipe.execute(): + assert reply == 1 + return leader + + +# --------------------------------------------------------------------------- +# 用例 +# --------------------------------------------------------------------------- + + +def test_quorum_commit_and_three_replica_consistency(cluster): + """写路径只走 leader（follower 一律 MOVED），线性一致读验证数据， + 并逐副本验证三副本一致（成员收缩为单节点后直接读）。""" + key = b"vc:t1" + leader = cluster.wait_leader() + followers = [n for n in cluster.nodes if n.node_id != leader.node_id] + + # follower 上写命令与 6 个 vector 读命令都应 MOVED 到 leader + for follower in followers: + client = cluster.client(follower) + with pytest.raises(MovedError) as write_redirect: + client.execute_command(*vadd_args(key, 0)) + assert write_redirect.value.port == leader.resp_port + for read_cmd in (b"VCARD", b"VDIM", b"VINFO"): + with pytest.raises(MovedError): + client.execute_command(read_cmd, key) + + # 从 follower 连接出发、跟随 MOVED 后写成功（客户端可见的写路径） + assert cluster.execute_from(followers[0], *vadd_args(key, 0)) == 1 + + # leader 上写入 100 个向量（quorum 提交） + vadd_batch(cluster, key, 1, 99) + + # leader 上的线性一致读 + leader_client = cluster.client(cluster.wait_leader()) + assert leader_client.execute_command(b"VCARD", key) == 100 + assert leader_client.execute_command(b"VDIM", key) == 2 + assert leader_client.execute_command(b"VISMEMBER", key, elem(7)) == 1 + emb = leader_client.execute_command(b"VEMB", key, elem(42)) + assert [float(v) for v in emb] == pytest.approx(list(vec(42))) + + # 三副本一致：vector 读只在 leader 可用，因此把每个 follower 副本依次 + # "扶上 leader"（force_leader 利用日志最完整者必胜的选举规则，确定性 + # 轮换）后直接读验证；原 leader 副本已被上面的线性一致读覆盖。 + # 注：member 级数据（VEMB/VISMEMBER）因 storage_incarnation 内嵌于 + # member key，副本上读不到（见模块 docstring），这里断言 meta 级一致； + # member 级断言见 xfail 用例 test_replica_member_data_survives_failover。 + for target in followers: + forced = force_leader(cluster, target) + assert forced.node_id == target.node_id + cluster.wait_card(target, key, 100) + assert cluster.client(target).execute_command(b"VDIM", key) == 2 + + +def test_concurrent_vadd(cluster): + """多线程并发 VADD 不同元素到同一个 set，最终 VCARD == 成功新增数。""" + key = b"vc:t2" + threads, per_thread = 8, 50 + added = [0] * threads + errors = [] + + def worker(tid): + try: + node = cluster.wait_leader() + total = 0 + for i in range(per_thread): + index = tid * per_thread + i + reply = cluster.execute_from(node, *vadd_args(key, index)) + total += int(reply) + added[tid] = total + except Exception as exc: # noqa: BLE001 - 收集到主线程统一断言 + errors.append(f"thread {tid}: {exc!r}") + + workers = [threading.Thread(target=worker, args=(t,)) for t in range(threads)] + for worker_thread in workers: + worker_thread.start() + for worker_thread in workers: + worker_thread.join(timeout=120) + + assert not errors, f"concurrent VADD errors: {errors}" + expected = sum(added) + assert expected == threads * per_thread, ( + f"VADD should report every element as newly added, got {expected}" + ) + leader_client = cluster.client(cluster.wait_leader()) + assert leader_client.execute_command(b"VCARD", key) == expected + + +def test_leader_failover(cluster): + """杀掉 leader 后新 leader 选出：旧数据可见（VCARD/VSIM），新写入成功。""" + key = b"vc:t3" + leader = vadd_batch(cluster, key, 0, 200) + leader_client = cluster.client(leader) + assert leader_client.execute_command(b"VCARD", key) == 200 + + cluster.kill(leader) + + new_leader = cluster.wait_leader(timeout=30) + assert new_leader.node_id != leader.node_id + new_client = cluster.client(new_leader) + + # 旧数据在新 leader 上可见（等任期内的线性一致屏障就绪）。 + # meta 级（VCARD）可验证；member 级（VSIM/VEMB 旧数据）受 + # storage_incarnation 限制，见 xfail 用例。 + cluster.wait_card(new_leader, key, 200) + + # 新 leader 接受写入，且新写入的数据立即可读 + assert new_client.execute_command(*vadd_args(key, 5000)) == 1 + assert new_client.execute_command(b"VCARD", key) == 201 + emb = new_client.execute_command(b"VEMB", key, elem(5000)) + assert [float(v) for v in emb] == pytest.approx(list(vec(5000))) + + +def test_lagged_follower_log_replay_and_snapshot_build(cluster): + """落后 follower 通过日志回放追平；写过量超过快照阈值后 leader 构建快照。 + + 注：快照的网络安装（install_snapshot RPC）本分支未实现，落后超过日志 + 保留窗口的 follower 无法追平，该场景无法测试；这里 B 阶段写入量刻意 + 低于保留窗口（快照阈值 5000 / 保留 1000），保证走日志回放路径。 + """ + key = b"vc:t4" + + # A：全集群健康时写入 500 + leader = vadd_batch(cluster, key, 0, 500) + follower = next(n for n in cluster.nodes if n.node_id != leader.node_id) + + # B：停掉一个 follower，再写 800（远低于日志清理窗口） + cluster.kill(follower) + vadd_batch(cluster, key, 500, 800) + leader_client = cluster.client(cluster.wait_leader()) + assert leader_client.execute_command(b"VCARD", key) == 1300 + + # C：重启 follower，等它通过日志回放追平，再把它扶上 leader 直接验证副本。 + # force_leader 中的 bump 写入只有在该 follower 完整追平后才能提交成功 + # （AppendEntries 会先补齐它缺失的全部日志），因此追平被隐式确认。 + # member 级旧数据受 incarnation 限制不可读（见模块 docstring），改为验证 + # meta 计数 + 该节点成为 leader 后的新写新读闭环。 + cluster.restart(follower) + forced = force_leader(cluster, follower) + assert forced.node_id == follower.node_id + cluster.wait_card(follower, key, 1300) + forced_client = cluster.client(forced) + assert forced_client.execute_command(*vadd_args(key, 6000)) == 1 + emb = forced_client.execute_command(b"VEMB", key, elem(6000)) + assert [float(v) for v in emb] == pytest.approx(list(vec(6000))) + + # 让 phase C 中被杀的第三节点归队：phase D 全量写入会触发快照与日志清理， + # 若它持续落后超过清理窗口，由于快照网络安装未实现将永远无法追平， + # 进而卡死后续选举出的新 leader（无法凑齐多数派提交新任期条目）。 + dead = next(n for n in cluster.nodes if not n.alive) + cluster.restart(dead) + + # D：全员在线继续写，累计日志超过快照阈值（5000），等 leader 构建快照 + vadd_batch(cluster, key, 1300, 4000) + leader = cluster.wait_leader() + cluster.wait_card(leader, key, 5301, timeout=120) + snapshot_tar = os.path.join(leader.dir, "raft", "snapshots", "current_snapshot.tar") + snapshot_meta = os.path.join( + leader.dir, "raft", "snapshots", "current_snapshot_meta.json" + ) + wait_until( + lambda: os.path.isfile(snapshot_tar) + and os.path.isfile(snapshot_meta) + and os.path.getsize(snapshot_tar) > 0, + 240, + "leader to persist a snapshot (current_snapshot.tar)", + ) + + # 快照构建后集群读写正常，且 failover 后副本数据依然完整（meta 级 + + # 新 leader 的新写新读闭环；member 级旧数据受 incarnation 限制） + assert cluster.execute_from(leader, *vadd_args(key, 9000)) == 1 + cluster.wait_card(cluster.wait_leader(), key, 5302) + other = next( + n for n in cluster.nodes if n.node_id != leader.node_id and n.alive + ) + forced = force_leader(cluster, other) + cluster.wait_card(forced, key, 5302, timeout=300) + forced_client = cluster.client(forced) + assert forced_client.execute_command(*vadd_args(key, 9001)) == 1 + emb = forced_client.execute_command(b"VEMB", key, elem(9001)) + assert [float(v) for v in emb] == pytest.approx(list(vec(9001))) + + +@pytest.mark.xfail( + reason=( + "vector member key 内嵌节点本地 storage_incarnation，物理 binlog 回放后 " + "副本上的 member 数据（VEMB/VISMEMBER/VSIM）不可读；待 PR0 logical " + "mutation replay（社区 issue #332）落地后此用例应转绿" + ), + strict=False, +) +def test_replica_member_data_survives_failover(cluster): + """规范要求：failover 后旧数据在新 leader 上完全可见（member 级）。 + + 当前分支因 incarnation 问题失败（xfail）；meta 级（VCARD）可见性由 + test_leader_failover 覆盖。 + """ + key = b"vc:t6" + leader = vadd_batch(cluster, key, 0, 50) + cluster.kill(leader) + + new_leader = cluster.wait_leader(timeout=30) + assert new_leader.node_id != leader.node_id + cluster.wait_card(new_leader, key, 50) + client = cluster.client(new_leader) + + assert client.execute_command(b"VISMEMBER", key, elem(1)) == 1 + emb = client.execute_command(b"VEMB", key, elem(1)) + assert emb is not None + assert [float(v) for v in emb] == pytest.approx(list(vec(1))) + x, y = vec(1) + similar = client.execute_command( + b"VSIM", key, b"VALUES", 2, f"{x:.6f}", f"{y:.6f}", b"COUNT", 3, b"TRUTH" + ) + assert elem(1) in similar + + +def test_minority_rejects_writes(cluster): + """杀掉两个节点后，少数派节点不可写（MOVED/ERR not leader），不断言具体错误。""" + key = b"vc:t5" + leader = cluster.wait_leader() + others = [n for n in cluster.nodes if n.node_id != leader.node_id] + survivor = others[0] + + cluster.kill(leader) + cluster.kill(others[1]) + + # 给选超时间留窗口：幸存节点无法凑齐多数派，不能成为 leader + time.sleep(ELECTION_TIMEOUT_MS / 1000 * 3) + + client = cluster.client(survivor) + deadline = time.monotonic() + 6 + while time.monotonic() < deadline: + try: + client.execute_command(*vadd_args(key, 0)) + pytest.fail("VADD unexpectedly succeeded on a minority node") + except MovedError: + pass # 重定向到已死的旧 leader，同样不可写 + except redis.ResponseError as exc: + assert "not leader" in str(exc), ( + f"unexpected error from minority node: {exc}" + ) + time.sleep(0.5) diff --git a/tests/python/test_vector_set_commands.py b/tests/python/test_vector_set_commands.py index 711c0668..76bd2678 100644 --- a/tests/python/test_vector_set_commands.py +++ b/tests/python/test_vector_set_commands.py @@ -40,6 +40,7 @@ def vector_client(request): prefix + b"errors:bin", prefix + b"string", prefix + b"malformed", + prefix + b"info", ] client = redis.Redis( host=os.getenv("KIWI_HOST", "127.0.0.1"), @@ -322,3 +323,75 @@ def test_malformed_vectors_and_options(vector_client): b"COUNT", 0, ) + + +def vinfo_fields(reply, protocol): + if protocol == 3: + assert type(reply) is dict + return reply + assert type(reply) is list + return dict(zip(reply[::2], reply[1::2])) + + +def test_vinfo_reports_flat_phase1_fields(vector_client): + client, protocol, prefix = vector_client + key = prefix + b"info" + + # Missing key: null reply in both protocols (RESP2 *-1, RESP3 _). + assert client.execute_command(b"VINFO", key) is None + + vadd_values(client, key, [1, 0, 0], b"a") + vadd_values(client, key, [0, 1, 0], b"b") + info = vinfo_fields(client.execute_command(b"VINFO", key), protocol) + assert info[b"quant-type"] == b"fp32" + assert info[b"hnsw-m"] == 0 + assert info[b"vector-dim"] == 3 + assert info[b"projection-input-dim"] == 0 + assert info[b"size"] == 2 + assert info[b"max-level"] == 0 + assert info[b"attributes-count"] == 0 + assert info[b"hnsw-max-node-uid"] == 0 + assert type(info[b"vset-uid"]) is int + first_uid = info[b"vset-uid"] + + # Recreating the set allocates a fresh generation sequence. + client.delete(key) + assert client.execute_command(b"VINFO", key) is None + vadd_values(client, key, [1, 0, 0], b"c") + info = vinfo_fields(client.execute_command(b"VINFO", key), protocol) + assert info[b"size"] == 1 + assert info[b"vset-uid"] != first_uid + + +def test_vinfo_wrongtype(vector_client): + client, _protocol, prefix = vector_client + key = prefix + b"string" + client.set(key, b"value") + assert_response_error(client, "WRONGTYPE", b"VINFO", key) + + +def test_info_vector_section(vector_client): + client, _protocol, prefix = vector_client + key = prefix + b"info" + vadd_values(client, key, [1, 0], b"a") + client.execute_command(b"VSIM", key, b"VALUES", 2, 1, 0) + + body = client.execute_command(b"INFO", b"VECTOR").decode() + assert body.startswith("# Vector") + fields = dict( + line.split(":", 1) for line in body.splitlines() if ":" in line + ) + assert fields["index-kind"] == "flat" + assert int(fields["vector_flat_queries_total"]) >= 1 + for name in ( + "vector_flat_query_timeouts_total", + "vector_flat_query_errors_total", + "vector_search_capacity_rejected_total", + "vector_flat_query_duration_micros_total", + "vector_flat_query_duration_count", + ): + assert int(fields[name]) >= 0 + + full = client.execute_command(b"INFO").decode() + assert "# Vector" in full + assert "index-kind:flat" in full diff --git a/tests/python/test_vector_set_differential.py b/tests/python/test_vector_set_differential.py new file mode 100644 index 00000000..cc0630c5 --- /dev/null +++ b/tests/python/test_vector_set_differential.py @@ -0,0 +1,341 @@ +# Copyright (c) 2024-present, arana-db Community. All rights reserved. +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Differential tests comparing Kiwi vector sets against Redis 8. + +Every test issues the same commands to a Kiwi server (KIWI_HOST/KIWI_PORT, +default 127.0.0.1:7379) and to a Redis 8 reference server +(VECTOR_REDIS_HOST/VECTOR_REDIS_PORT, default 127.0.0.1:6380) and compares +the replies. Both sides are populated with an explicit NOQUANT option. + +Out of scope on purpose: +- Redis' default Q8 quantization path. Phase 1 Kiwi only supports NOQUANT; + the error contract for omitted/unsupported quantization options is covered + by test_vector_set_commands.py and not repeated here. +- VINFO field *values* such as hnsw-m, max-level and vset-uid. Redis 8 + reports real HNSW internals while Kiwi Phase 1 reports FLAT sentinels with + a different meaning; only the field-name set and value types are compared. + +All datasets use a fixed seed so runs are reproducible. +""" + +import os +import random +import struct + +import pytest +import redis + +KIWI_HOST = os.getenv("KIWI_HOST", "127.0.0.1") +KIWI_PORT = int(os.getenv("KIWI_PORT", "7379")) +REDIS8_HOST = os.getenv("VECTOR_REDIS_HOST", "127.0.0.1") +REDIS8_PORT = int(os.getenv("VECTOR_REDIS_PORT", "6380")) + +SCORE_TOLERANCE = 1e-6 + + +def _server_reachable(host, port): + client = redis.Redis( + host=host, port=port, socket_connect_timeout=1, socket_timeout=1 + ) + try: + client.ping() + return True + except (redis.RedisError, OSError): + return False + finally: + client.close() + + +if not _server_reachable(KIWI_HOST, KIWI_PORT): + pytest.skip( + f"Kiwi server not reachable at {KIWI_HOST}:{KIWI_PORT}; " + "skipping vector set differential tests", + allow_module_level=True, + ) +if not _server_reachable(REDIS8_HOST, REDIS8_PORT): + pytest.skip( + f"Redis 8 reference server not reachable at {REDIS8_HOST}:{REDIS8_PORT}; " + "skipping vector set differential tests", + allow_module_level=True, + ) + + +TEST_KEY_NAMES = (b"main", b"dense3", b"string", b"missing") + + +def _build_main_members(): + """Fixed-seed dim-4 dataset: regular, empty and binary elements plus a + tie pair sharing one identical vector (tie order must be element bytes + ascending on both servers).""" + rng = random.Random(20240701) + members = [] + for element in (b"alpha", b"beta", b"gamma", b"delta", b"", b"\x00bin\x00"): + members.append((element, [rng.uniform(-2.0, 2.0) for _ in range(4)])) + tie_vector = [0.25, -0.5, 1.0, 0.75] + members.append((b"tie-a", list(tie_vector))) + members.append((b"tie-b", list(tie_vector))) + return members + + +def _build_dense3_members(): + rng = random.Random(7) + return [ + (element, [rng.uniform(-1.0, 1.0) for _ in range(3)]) + for element in (b"x", b"y", b"z") + ] + + +MAIN_MEMBERS = _build_main_members() +DENSE3_MEMBERS = _build_dense3_members() +MAIN_QUERY = [0.5, -0.25, 1.5, 0.125] + + +@pytest.fixture(params=[2, 3], ids=["resp2", "resp3"]) +def backends(request): + protocol = request.param + prefix = f"test_vdiff:p{protocol}:".encode() + keys = [prefix + name for name in TEST_KEY_NAMES] + kiwi = redis.Redis( + host=KIWI_HOST, port=KIWI_PORT, decode_responses=False, protocol=protocol + ) + reference = redis.Redis( + host=REDIS8_HOST, port=REDIS8_PORT, decode_responses=False, protocol=protocol + ) + for client in (kiwi, reference): + client.delete(*keys) + yield kiwi, reference, protocol, prefix + for client in (kiwi, reference): + client.delete(*keys) + client.close() + + +def vadd_noquant(client, key, values, element): + return client.execute_command( + b"VADD", key, b"VALUES", len(values), *values, element, b"NOQUANT" + ) + + +def populate(client, key, members): + for element, vector in members: + vadd_noquant(client, key, vector, element) + + +def assert_same_reply(kiwi, reference, *command): + kiwi_reply = kiwi.execute_command(*command) + redis_reply = reference.execute_command(*command) + assert kiwi_reply == redis_reply, ( + f"{command!r}: kiwi={kiwi_reply!r} != redis={redis_reply!r}" + ) + return kiwi_reply + + +def normalized_vemb(reply): + if reply is None: + return None + return [float(component) for component in reply] + + +def assert_same_vemb(kiwi, reference, key, element): + kiwi_vemb = normalized_vemb(kiwi.execute_command(b"VEMB", key, element)) + redis_vemb = normalized_vemb(reference.execute_command(b"VEMB", key, element)) + if redis_vemb is None: + assert kiwi_vemb is None + return + assert kiwi_vemb is not None + assert len(kiwi_vemb) == len(redis_vemb) + assert kiwi_vemb == pytest.approx(redis_vemb, abs=SCORE_TOLERANCE) + + +def parse_vsim(reply, protocol, withscores): + """Normalize a VSIM reply to an ordered list of (element, score) pairs.""" + if not withscores: + return [(element, None) for element in reply] + if protocol == 3: + return [(element, float(score)) for element, score in reply.items()] + return [ + (reply[index], float(reply[index + 1])) + for index in range(0, len(reply), 2) + ] + + +def assert_same_vsim(kiwi, reference, protocol, query_args, *options): + withscores = b"WITHSCORES" in options + kiwi_hits = parse_vsim( + kiwi.execute_command(b"VSIM", *query_args, *options), protocol, withscores + ) + redis_hits = parse_vsim( + reference.execute_command(b"VSIM", *query_args, *options), + protocol, + withscores, + ) + assert [element for element, _ in kiwi_hits] == [ + element for element, _ in redis_hits + ] + if withscores: + assert len(kiwi_hits) == len(redis_hits) + for (_, kiwi_score), (_, redis_score) in zip(kiwi_hits, redis_hits): + assert kiwi_score == pytest.approx(redis_score, abs=SCORE_TOLERANCE) + + +def vinfo_as_dict(reply, protocol): + if protocol == 3: + assert type(reply) is dict + return reply + assert type(reply) is list + return dict(zip(reply[::2], reply[1::2])) + + +def test_vadd_vcard_vdim_vismember_match(backends): + kiwi, reference, _protocol, prefix = backends + main_key = prefix + b"main" + dense3_key = prefix + b"dense3" + + for key, members in ((main_key, MAIN_MEMBERS), (dense3_key, DENSE3_MEMBERS)): + for element, vector in members: + assert_same_reply(kiwi, reference, b"VADD", key, b"VALUES", + len(vector), *vector, element, b"NOQUANT") + assert_same_reply(kiwi, reference, b"VCARD", key) + assert_same_reply(kiwi, reference, b"VDIM", key) + for element, _ in members: + assert_same_reply(kiwi, reference, b"VISMEMBER", key, element) + assert_same_reply(kiwi, reference, b"VISMEMBER", key, b"ghost") + + # Re-adding an existing element with a new vector is an update (0/False). + updated = [0.1, 0.2, 0.3, 0.4] + assert_same_reply(kiwi, reference, b"VADD", main_key, b"VALUES", + 4, *updated, b"alpha", b"NOQUANT") + assert_same_reply(kiwi, reference, b"VCARD", main_key) + assert_same_vemb(kiwi, reference, main_key, b"alpha") + + +def test_vemb_matches_for_every_member(backends): + kiwi, reference, _protocol, prefix = backends + main_key = prefix + b"main" + populate(kiwi, main_key, MAIN_MEMBERS) + populate(reference, main_key, MAIN_MEMBERS) + + for element, _ in MAIN_MEMBERS: + assert_same_vemb(kiwi, reference, main_key, element) + assert_same_vemb(kiwi, reference, main_key, b"ghost") + + +def test_vsim_variants_match(backends): + kiwi, reference, protocol, prefix = backends + main_key = prefix + b"main" + populate(kiwi, main_key, MAIN_MEMBERS) + populate(reference, main_key, MAIN_MEMBERS) + + fp32_blob = struct.pack(f"<{len(MAIN_QUERY)}f", *MAIN_QUERY) + query_variants = [ + (main_key, b"ELE", b"alpha"), + (main_key, b"ELE", b""), + (main_key, b"VALUES", len(MAIN_QUERY), *MAIN_QUERY), + (main_key, b"FP32", fp32_blob), + ] + option_variants = [ + (b"COUNT", 16, b"WITHSCORES", b"TRUTH"), + (b"COUNT", 3, b"WITHSCORES", b"TRUTH"), + (b"COUNT", 16, b"TRUTH"), + (b"COUNT", 3, b"TRUTH"), + ] + for query_args in query_variants: + for options in option_variants: + assert_same_vsim(kiwi, reference, protocol, query_args, *options) + + +def test_vrem_replies_and_set_disappears_when_emptied(backends): + kiwi, reference, _protocol, prefix = backends + main_key = prefix + b"main" + populate(kiwi, main_key, MAIN_MEMBERS) + populate(reference, main_key, MAIN_MEMBERS) + + # Removing an absent element first: 0/False on both. + assert_same_reply(kiwi, reference, b"VREM", main_key, b"ghost") + + for element, _ in MAIN_MEMBERS: + assert_same_reply(kiwi, reference, b"VREM", main_key, element) + assert_same_reply(kiwi, reference, b"VISMEMBER", main_key, element) + + # The last removal deletes the key on both servers. + assert_same_reply(kiwi, reference, b"VCARD", main_key) + assert_same_reply(kiwi, reference, b"TYPE", main_key) + + +def test_wrongtype_errors_match(backends): + kiwi, reference, _protocol, prefix = backends + string_key = prefix + b"string" + for client in (kiwi, reference): + client.set(string_key, b"plain-string") + + commands = [ + (b"VADD", string_key, b"VALUES", 2, 1, 0, b"member", b"NOQUANT"), + (b"VCARD", string_key), + (b"VDIM", string_key), + (b"VEMB", string_key, b"member"), + (b"VISMEMBER", string_key, b"member"), + (b"VREM", string_key, b"member"), + (b"VSIM", string_key, b"VALUES", 2, 1, 0, b"TRUTH"), + (b"VINFO", string_key), + ] + for command in commands: + for name, client in (("kiwi", kiwi), ("redis", reference)): + with pytest.raises(redis.ResponseError) as excinfo: + client.execute_command(*command) + assert str(excinfo.value).startswith("WRONGTYPE"), ( + f"{name} {command!r}: {excinfo.value}" + ) + + +def test_missing_key_semantics_match(backends): + kiwi, reference, _protocol, prefix = backends + missing_key = prefix + b"missing" + + assert_same_reply(kiwi, reference, b"VCARD", missing_key) + assert_same_reply(kiwi, reference, b"VISMEMBER", missing_key, b"member") + assert_same_reply( + kiwi, reference, b"VSIM", missing_key, b"VALUES", 4, *MAIN_QUERY, b"TRUTH" + ) + assert_same_vemb(kiwi, reference, missing_key, b"member") + # VDIM on a missing key is deliberately not compared: Kiwi Phase 1 returns + # an error while Redis 8 returns 0; this known divergence is covered by + # the Kiwi-side contract tests. + + +def test_vinfo_field_names_and_types_match(backends): + kiwi, reference, protocol, prefix = backends + main_key = prefix + b"main" + + # Missing key: null reply on both servers. + assert_same_reply(kiwi, reference, b"VINFO", main_key) + + populate(kiwi, main_key, MAIN_MEMBERS) + populate(reference, main_key, MAIN_MEMBERS) + kiwi_info = vinfo_as_dict(kiwi.execute_command(b"VINFO", main_key), protocol) + redis_info = vinfo_as_dict(reference.execute_command(b"VINFO", main_key), protocol) + + # Only the field-name set and value types are compared. Values such as + # hnsw-m, max-level and vset-uid describe Redis' real HNSW index, while + # Kiwi Phase 1 reports FLAT sentinels with different semantics. + assert set(kiwi_info) == set(redis_info), ( + f"VINFO fields differ: kiwi={sorted(kiwi_info)} redis={sorted(redis_info)}" + ) + for field in kiwi_info: + assert type(kiwi_info[field]) is type(redis_info[field]), ( + f"VINFO field {field!r} type: kiwi={type(kiwi_info[field])} " + f"redis={type(redis_info[field])}" + ) diff --git a/tests/run_python_integration.sh b/tests/run_python_integration.sh index 54f8ab32..93d220bc 100755 --- a/tests/run_python_integration.sh +++ b/tests/run_python_integration.sh @@ -95,10 +95,10 @@ import socket with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: sock.settimeout(0.2) - raise SystemExit(0 if sock.connect_ex(("127.0.0.1", 6379)) == 0 else 1) + raise SystemExit(0 if sock.connect_ex(("127.0.0.1", 7379)) == 0 else 1) PY then - echo "127.0.0.1:6379 is already in use" >&2 + echo "127.0.0.1:7379 is already in use" >&2 exit 1 fi @@ -111,7 +111,7 @@ readonly LOG_DIR="${TEMP_DIR}/logs" mkdir -p -- "${DATA_DIR}" "${LOG_DIR}" cat >"${CONFIG_PATH}" <&2 + echo "Kiwi did not become ready on 127.0.0.1:7379" >&2 exit 1 fi set +e KIWI_TEST_REQUIRE_SERVER=1 \ KIWI_TEST_ISOLATED_SERVER=1 \ + KIWI_HOST=127.0.0.1 \ + KIWI_PORT=7379 \ make -C tests test-python test_status=$? set -e From 61c30ccbc69056cdd4468997307ca79b78d9d122 Mon Sep 17 00:00:00 2001 From: happy-v587 Date: Sun, 2 Aug 2026 14:01:47 +0800 Subject: [PATCH 23/32] fix: address vector PR review and CI failures --- src/cmd/src/table.rs | 2 +- src/conf/src/vector_config.rs | 36 +++++++++++++++++------- src/server/src/main.rs | 9 +++--- src/storage/src/redis.rs | 39 +++++++++++++++++++------- src/storage/src/redis_strings.rs | 2 +- src/storage/src/storage_manifest.rs | 26 +++++++++++------ src/storage/src/storage_scan.rs | 1 + src/storage/tests/redis_string_test.rs | 29 +++++++++++++++++-- src/storage/tests/redis_vector_test.rs | 22 +++++++++------ src/storage/tests/scan_test.rs | 13 ++++++++- 10 files changed, 132 insertions(+), 47 deletions(-) diff --git a/src/cmd/src/table.rs b/src/cmd/src/table.rs index c0a30160..165817d1 100644 --- a/src/cmd/src/table.rs +++ b/src/cmd/src/table.rs @@ -38,7 +38,7 @@ pub struct CommandTableGates { pub vector_enabled: GateFlagProvider, /// Whether Vector Set commands are allowed given the cluster state: /// false in cluster mode until the Raft apply-correctness contract (PR0) - /// lands, unless `vector-cluster-enabled` is set. + /// lands. pub vector_cluster_allowed: GateFlagProvider, /// Whether FLUSHDB/FLUSHALL are allowed. Disabled in cluster mode unless /// `cluster-flush-enabled` is set. diff --git a/src/conf/src/vector_config.rs b/src/conf/src/vector_config.rs index 67f95d12..e2a6313e 100644 --- a/src/conf/src/vector_config.rs +++ b/src/conf/src/vector_config.rs @@ -118,16 +118,32 @@ mod tests { #[test] fn max_dimension_bounds_are_enforced() { - let mut config = VectorConfig::default(); - - config.max_dimension = 0; - assert!(config.validate().is_err()); - - config.max_dimension = MAX_VECTOR_DIMENSION + 1; - assert!(config.validate().is_err()); - - config.max_dimension = MAX_VECTOR_DIMENSION; - assert!(config.validate().is_ok()); + assert!( + VectorConfig { + max_dimension: 0, + ..Default::default() + } + .validate() + .is_err() + ); + + assert!( + VectorConfig { + max_dimension: MAX_VECTOR_DIMENSION + 1, + ..Default::default() + } + .validate() + .is_err() + ); + + assert!( + VectorConfig { + max_dimension: MAX_VECTOR_DIMENSION, + ..Default::default() + } + .validate() + .is_ok() + ); } #[test] diff --git a/src/server/src/main.rs b/src/server/src/main.rs index e82b4c6a..d54b25f5 100644 --- a/src/server/src/main.rs +++ b/src/server/src/main.rs @@ -39,12 +39,13 @@ struct PausePermitWrapper { } /// Build the command-table feature gates from the loaded configuration: -/// vector commands follow `vector-enabled`, are additionally rejected in -/// cluster mode unless `vector-cluster-enabled` is set, and FLUSHDB/FLUSHALL -/// are only allowed outside cluster mode unless `cluster-flush-enabled` is set. +/// vector commands follow `vector-enabled` and remain rejected in cluster +/// mode until their Raft apply/replay correctness contract is complete. +/// FLUSHDB/FLUSHALL are only allowed outside cluster mode unless +/// `cluster-flush-enabled` is set. fn command_table_gates(config: &Config) -> cmd::table::CommandTableGates { let vector_enabled = config.vector.enabled; - let vector_cluster_allowed = config.raft.is_none() || config.vector.cluster_enabled; + let vector_cluster_allowed = config.raft.is_none(); let cluster_flush_allowed = config.raft.is_none() || config.cluster_flush_enabled; cmd::table::CommandTableGates::from_flags( vector_enabled, diff --git a/src/storage/src/redis.rs b/src/storage/src/redis.rs index 0ca2863e..757d4e2f 100644 --- a/src/storage/src/redis.rs +++ b/src/storage/src/redis.rs @@ -446,6 +446,27 @@ impl Redis { let mut db_opts = self.storage.options.clone(); db_opts.add_event_listener(purger); + // A legacy database has no vector column family. Refuse to open a + // database that already has that column family but lost its manifest + // before RocksDB can create or mutate any descriptors. + let manifest_path = Path::new(db_path).join(crate::storage_manifest::STORAGE_MANIFEST_FILE); + if !manifest_path.exists() + && DB::list_cf(&db_opts, db_path).is_ok_and(|column_families| { + column_families + .iter() + .any(|name| name == ColumnFamilyIndex::VectorDataCF.name()) + }) + { + return Err(InvalidFormatSnafu { + message: format!( + "storage manifest {} is missing but the database has a vector column family; \ + refusing to reinterpret existing vector data", + manifest_path.display() + ), + } + .build()); + } + const CF_CONFIGS: &[(&str, bool, Option)] = &[ ("default", true, None), // meta & string: bloom filter ("hash_data_cf", true, None), // hash: bloom filter @@ -497,17 +518,15 @@ impl Redis { // Load (or create) the per-instance storage manifest. A manifest that // is missing while the database holds data means the data predates // the incarnation mechanism; refuse to open rather than reinterpret. - let has_entries = self.handles.iter().any(|name| { - let Some(db) = self.db.as_ref() else { - return false; - }; - db.cf_handle(name).is_some_and(|cf| { - db.iterator_cf(&cf, rocksdb::IteratorMode::Start) - .next() - .is_some() - }) + let vector_data_has_entries = self.db.as_ref().is_some_and(|db| { + db.cf_handle(ColumnFamilyIndex::VectorDataCF.name()) + .is_some_and(|cf| { + db.iterator_cf(&cf, rocksdb::IteratorMode::Start) + .next() + .is_some() + }) }); - let manifest = StorageManifest::open(Path::new(db_path), has_entries)?; + let manifest = StorageManifest::open(Path::new(db_path), vector_data_has_entries)?; let _ = incarnation_cell.set(manifest.storage_incarnation()); self.manifest = Some(manifest); diff --git a/src/storage/src/redis_strings.rs b/src/storage/src/redis_strings.rs index 37fe2eec..b016ddb3 100644 --- a/src/storage/src/redis_strings.rs +++ b/src/storage/src/redis_strings.rs @@ -2318,7 +2318,7 @@ impl Redis { DataType::String => ParsedStringsValue::new(&value_bytes[..]) .map(|parsed| !parsed.is_stale()) .unwrap_or(false), - DataType::Hash | DataType::Set | DataType::ZSet => { + DataType::Hash | DataType::Set | DataType::ZSet | DataType::VectorSet => { crate::format_base_meta_value::ParsedBaseMetaValue::new(&value_bytes[..]) .map(|meta| !meta.is_stale() && meta.count() > 0) .unwrap_or(false) diff --git a/src/storage/src/storage_manifest.rs b/src/storage/src/storage_manifest.rs index 32ba6cb6..a8f2e11c 100644 --- a/src/storage/src/storage_manifest.rs +++ b/src/storage/src/storage_manifest.rs @@ -66,21 +66,23 @@ pub(crate) struct StorageManifest { } impl StorageManifest { - /// Load the manifest for the instance stored in `db_dir`, creating one - /// when the database is empty. `db_has_entries` must report whether any - /// column family of the already-open database contains at least one key. - pub(crate) fn open(db_dir: &Path, db_has_entries: bool) -> Result { + /// Load the manifest for the instance stored in `db_dir`. + /// + /// Databases created before the manifest existed can be bootstrapped when + /// they contain only legacy data. Existing vector member data is rejected + /// without a manifest because its incarnation cannot be reconstructed. + pub(crate) fn open(db_dir: &Path, vector_data_has_entries: bool) -> Result { let path = db_dir.join(STORAGE_MANIFEST_FILE); if path.exists() { return Self::read(&path); } ensure!( - !db_has_entries, + !vector_data_has_entries, InvalidFormatSnafu { message: format!( - "storage manifest {} is missing but the database is not empty; \ - refusing to reinterpret existing data", + "storage manifest {} is missing but vector data is present; \ + refusing to reinterpret existing vector members", path.display() ) } @@ -206,7 +208,7 @@ mod tests { assert_ne!(incarnation, 0); assert!(dir.path().join(STORAGE_MANIFEST_FILE).exists()); - let reopened = StorageManifest::open(dir.path(), true).expect("reopen manifest"); + let reopened = StorageManifest::open(dir.path(), false).expect("reopen manifest"); assert_eq!(reopened.storage_incarnation(), incarnation); } @@ -229,7 +231,13 @@ mod tests { } #[test] - fn missing_manifest_on_non_empty_db_is_rejected() { + fn missing_manifest_on_legacy_non_vector_db_is_bootstrapped() { + let dir = tempfile::tempdir().expect("temp dir"); + StorageManifest::open(dir.path(), false).expect("legacy manifest bootstrap"); + } + + #[test] + fn missing_manifest_on_non_empty_vector_db_is_rejected() { let dir = tempfile::tempdir().expect("temp dir"); assert!(StorageManifest::open(dir.path(), true).is_err()); } diff --git a/src/storage/src/storage_scan.rs b/src/storage/src/storage_scan.rs index e12059d5..f14d7c8b 100644 --- a/src/storage/src/storage_scan.rs +++ b/src/storage/src/storage_scan.rs @@ -81,6 +81,7 @@ fn parse_scan_type(name: &[u8]) -> DataType { b"set" => DataType::Set, b"list" => DataType::List, b"zset" => DataType::ZSet, + b"vectorset" => DataType::VectorSet, _ => DataType::None, } } diff --git a/src/storage/tests/redis_string_test.rs b/src/storage/tests/redis_string_test.rs index 2584a50a..436acfc4 100644 --- a/src/storage/tests/redis_string_test.rs +++ b/src/storage/tests/redis_string_test.rs @@ -23,9 +23,9 @@ mod redis_string_test { use kstd::lock_mgr::LockMgr; use storage::{ - BaseMetaKey, BgTaskHandler, ColumnFamilyIndex, DataType, Redis, StorageOptions, - ZsetScoreMember, format_base_meta_value::BaseMetaValue, safe_cleanup_test_db, - slot_indexer::key_to_slot_id, storage::Storage, unique_test_db_path, + BaseMetaKey, BgTaskHandler, CanonicalVector, ColumnFamilyIndex, DataType, Redis, + StorageOptions, ZsetScoreMember, format_base_meta_value::BaseMetaValue, + safe_cleanup_test_db, slot_indexer::key_to_slot_id, storage::Storage, unique_test_db_path, }; fn cleanup_redis(redis: Redis, test_db_path: &Path) { @@ -1285,6 +1285,29 @@ mod redis_string_test { cleanup_redis(redis, &test_db_path); } + #[test] + fn test_random_key_returns_vectorset_key() { + let test_db_path = unique_test_db_path(); + safe_cleanup_test_db(&test_db_path); + + let storage_options = Arc::new(StorageOptions::default()); + let (bg_task_handler, _) = BgTaskHandler::new(); + let lock_mgr = Arc::new(LockMgr::new(1000)); + let mut redis = Redis::new(storage_options, 1, Arc::new(bg_task_handler), lock_mgr); + redis.open(test_db_path.to_str().unwrap()).unwrap(); + + let vectorset_key = b"only_vectorset_key"; + let vector = CanonicalVector::from_values(&[1.0, 0.0]).unwrap(); + redis.vadd(vectorset_key, b"member", &vector).unwrap(); + + assert_eq!( + redis.random_key().unwrap().as_deref(), + Some("only_vectorset_key") + ); + + cleanup_redis(redis, &test_db_path); + } + // Regression: `ParsedStringsValue::new` accepts any value whose leading byte // is a valid DataType tag and applies no count check, so a string-first // probe would misread a collection meta value as a live string and return a diff --git a/src/storage/tests/redis_vector_test.rs b/src/storage/tests/redis_vector_test.rs index 3ade8be3..6bfdc69f 100644 --- a/src/storage/tests/redis_vector_test.rs +++ b/src/storage/tests/redis_vector_test.rs @@ -17,7 +17,7 @@ #![allow(clippy::unwrap_used)] -use std::{path::PathBuf, sync::Arc, sync::Mutex, time::Duration, time::Instant}; +use std::{path::Path, sync::Arc, sync::Mutex, time::Duration, time::Instant}; use conf::vector_config::VectorConfig; use kstd::lock_mgr::LockMgr; @@ -30,7 +30,7 @@ use storage::{ }; use storage::{slot_indexer::key_to_slot_id, storage::Storage}; -fn open_redis_with_options(path: &PathBuf, storage_options: Arc) -> Redis { +fn open_redis_with_options(path: &Path, storage_options: Arc) -> Redis { let (bg_task_handler, _) = BgTaskHandler::new(); let lock_mgr = Arc::new(LockMgr::new(1000)); let mut redis = Redis::new(storage_options, 1, Arc::new(bg_task_handler), lock_mgr); @@ -40,7 +40,7 @@ fn open_redis_with_options(path: &PathBuf, storage_options: Arc) redis } -fn open_redis(path: &PathBuf) -> Redis { +fn open_redis(path: &Path) -> Redis { open_redis_with_options(path, Arc::new(StorageOptions::default())) } @@ -60,8 +60,10 @@ fn with_redis_vector_config(edit: impl FnOnce(&mut VectorConfig), test: impl FnO safe_cleanup_test_db(&path); let mut vector = VectorConfig::default(); edit(&mut vector); - let mut storage_options = StorageOptions::default(); - storage_options.vector = vector; + let storage_options = StorageOptions { + vector, + ..Default::default() + }; let redis = open_redis_with_options(&path, Arc::new(storage_options)); test(&redis); @@ -89,7 +91,9 @@ fn count_cf_entries(redis: &Redis, cf_index: ColumnFamilyIndex) -> usize { let db = redis.db().expect("db is initialized"); let cf = redis.get_cf_handle(cf_index).expect("column family exists"); db.iterator_cf(&cf, IteratorMode::Start) - .map(|entry| entry.expect("read column family entry")) + .inspect(|entry| { + entry.as_ref().expect("read column family entry"); + }) .count() } @@ -260,7 +264,9 @@ fn test_vector_meta_and_member_are_committed_together() { ); assert_eq!( db.iterator_cf_opt(&vector_cf, vector_options, IteratorMode::Start) - .map(|entry| entry.expect("read vector entry")) + .inspect(|entry| { + entry.as_ref().expect("read vector entry"); + }) .count(), 1 ); @@ -497,7 +503,7 @@ fn test_vadd_rebuilds_expired_vectorset_with_new_generation() { let meta_key = BaseMetaKey::new(key).encode().expect("meta key"); let mut meta = previous_meta; meta.set_etime(1); - db.put_cf(&meta_cf, &meta_key, &meta.encode()) + db.put_cf(&meta_cf, &meta_key, meta.encode()) .expect("store expired vector meta"); assert_eq!(redis.vcard(key).expect("expired card"), 0); diff --git a/src/storage/tests/scan_test.rs b/src/storage/tests/scan_test.rs index 62667fcf..133db18e 100644 --- a/src/storage/tests/scan_test.rs +++ b/src/storage/tests/scan_test.rs @@ -21,7 +21,7 @@ mod scan_test { use std::{collections::HashSet, sync::Arc}; - use storage::{StorageOptions, ZsetScoreMember, storage::Storage}; + use storage::{CanonicalVector, StorageOptions, ZsetScoreMember, storage::Storage}; fn open_storage(instances: usize) -> (Storage, tempfile::TempDir) { let dir = tempfile::tempdir().unwrap(); @@ -146,6 +146,17 @@ mod scan_test { as_set(scan_all(&storage, 2, Some(b"list"), b"*")), HashSet::from([b"a_list".to_vec()]) ); + storage + .vadd( + b"a_vectorset", + b"member", + &CanonicalVector::from_values(&[1.0, 0.0]).unwrap(), + ) + .unwrap(); + assert_eq!( + as_set(scan_all(&storage, 2, Some(b"VeCtOrSeT"), b"*")), + HashSet::from([b"a_vectorset".to_vec()]) + ); // An unknown type matches nothing but still terminates. assert!(scan_all(&storage, 2, Some(b"stream"), b"*").is_empty()); } From f3af349247d02f70ef3adbdd2abd9c03ceb3b870 Mon Sep 17 00:00:00 2001 From: happy-v587 Date: Sun, 2 Aug 2026 15:32:46 +0800 Subject: [PATCH 24/32] fix: tombstone vector sets on DEL --- src/storage/src/redis_strings.rs | 36 +------------------------- src/storage/tests/redis_vector_test.rs | 13 ++++++++-- 2 files changed, 12 insertions(+), 37 deletions(-) diff --git a/src/storage/src/redis_strings.rs b/src/storage/src/redis_strings.rs index b016ddb3..6c5c3f5a 100644 --- a/src/storage/src/redis_strings.rs +++ b/src/storage/src/redis_strings.rs @@ -33,7 +33,6 @@ use crate::{ format_base_meta_value::ParsedBaseMetaValue, format_list_meta_value::ParsedListsMetaValue, format_strings_value::{ParsedStringsValue, StringValue}, - format_vector_member_key::VectorMemberDataKey, redis_sets::glob_match_bytes, }; @@ -2059,7 +2058,7 @@ impl Redis { log::warn!("failed to update key statistics for {key_str}: {error:?}"); } } - DataType::Hash | DataType::Set | DataType::ZSet => { + DataType::Hash | DataType::Set | DataType::ZSet | DataType::VectorSet => { let mut parsed = ParsedBaseMetaValue::new(&value[..])?; if !parsed.is_valid() { return Ok(false); @@ -2081,39 +2080,6 @@ impl Redis { log::warn!("failed to update key statistics for {key_str}: {error:?}"); } } - DataType::VectorSet => { - let parsed = ParsedBaseMetaValue::new(&value[..])?; - if !parsed.is_valid() { - return Ok(false); - } - let count = parsed.count(); - let prefix = VectorMemberDataKey::encode_key_prefix(key)?; - let Some(vector_cf) = self.get_cf_handle(ColumnFamilyIndex::VectorDataCF) else { - return Ok(false); - }; - let mut keys_to_delete = Vec::new(); - for item in db.iterator_cf( - &vector_cf, - rocksdb::IteratorMode::From(&prefix, rocksdb::Direction::Forward), - ) { - let (member_key, _) = item.context(RocksSnafu)?; - if !member_key.starts_with(&prefix) { - break; - } - keys_to_delete.push(member_key.to_vec()); - } - let mut batch = self.create_batch()?; - batch.delete(ColumnFamilyIndex::MetaCF, &encoded_meta_key)?; - for member_key in keys_to_delete { - batch.delete(ColumnFamilyIndex::VectorDataCF, &member_key)?; - } - batch.commit()?; - if let Err(error) = - self.update_specific_key_statistics(DataType::VectorSet, &key_str, count) - { - log::warn!("failed to update key statistics for {key_str}: {error:?}"); - } - } DataType::List => { let mut parsed = ParsedListsMetaValue::new(&value[..])?; if !parsed.is_valid() { diff --git a/src/storage/tests/redis_vector_test.rs b/src/storage/tests/redis_vector_test.rs index 6bfdc69f..9ee81c89 100644 --- a/src/storage/tests/redis_vector_test.rs +++ b/src/storage/tests/redis_vector_test.rs @@ -653,16 +653,25 @@ async fn test_expired_vectorset_reads_as_missing() { } #[test] -fn test_del_removes_vector_meta_and_members() { +fn test_del_tombstones_vector_meta_and_defers_member_cleanup() { with_redis(|redis| { let vector = CanonicalVector::from_values(&[1.0, 0.0]).expect("vector"); redis.vadd(b"vectors", b"a", &vector).expect("insert a"); redis.vadd(b"vectors", b"b", &vector).expect("insert b"); assert_eq!(count_cf_entries(redis, ColumnFamilyIndex::VectorDataCF), 2); + let previous_version = read_stored_vector_meta(redis, b"vectors").version(); assert!(redis.del_key(b"vectors").expect("delete vector set")); assert_eq!(redis.vcard(b"vectors").expect("missing card"), 0); - assert_eq!(count_cf_entries(redis, ColumnFamilyIndex::VectorDataCF), 0); + let tombstone = read_stored_vector_meta(redis, b"vectors"); + assert_eq!(tombstone.count(), 0); + assert_eq!(tombstone.etime(), 0); + assert!(tombstone.version() > previous_version); + assert_eq!( + count_cf_entries(redis, ColumnFamilyIndex::VectorDataCF), + 2, + "DEL should leave vector members for compaction" + ); }); } From f1f8d7539812efd4c6c5e99d540744976d4f6261 Mon Sep 17 00:00:00 2001 From: happy-v587 Date: Sun, 2 Aug 2026 15:51:58 +0800 Subject: [PATCH 25/32] test: use public vector metadata assertions --- src/storage/tests/redis_vector_test.rs | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/storage/tests/redis_vector_test.rs b/src/storage/tests/redis_vector_test.rs index 9ee81c89..77d46be2 100644 --- a/src/storage/tests/redis_vector_test.rs +++ b/src/storage/tests/redis_vector_test.rs @@ -664,8 +664,6 @@ fn test_del_tombstones_vector_meta_and_defers_member_cleanup() { assert!(redis.del_key(b"vectors").expect("delete vector set")); assert_eq!(redis.vcard(b"vectors").expect("missing card"), 0); let tombstone = read_stored_vector_meta(redis, b"vectors"); - assert_eq!(tombstone.count(), 0); - assert_eq!(tombstone.etime(), 0); assert!(tombstone.version() > previous_version); assert_eq!( count_cf_entries(redis, ColumnFamilyIndex::VectorDataCF), From b0db91f701c04a05d90591dd644d2123675fb9b5 Mon Sep 17 00:00:00 2001 From: happy-v587 Date: Sun, 2 Aug 2026 20:47:24 +0800 Subject: [PATCH 26/32] fix: harden snapshot tests on Windows --- src/raft/src/state_machine.rs | 46 +++++++++++------------------ src/storage/src/storage_manifest.rs | 26 ++++++++++++---- 2 files changed, 38 insertions(+), 34 deletions(-) diff --git a/src/raft/src/state_machine.rs b/src/raft/src/state_machine.rs index 4ceef3f7..cf790818 100644 --- a/src/raft/src/state_machine.rs +++ b/src/raft/src/state_machine.rs @@ -871,6 +871,8 @@ mod snapshot_gate_tests { use super::*; + const SNAPSHOT_TEST_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10); + #[derive(Default)] struct CountingPauseController { paused: AtomicBool, @@ -958,12 +960,9 @@ mod snapshot_gate_tests { let result = builder.build_snapshot().await; (builder, result) }); - tokio::time::timeout( - std::time::Duration::from_secs(1), - hook.checkpoint_completed.notified(), - ) - .await - .expect("snapshot build should reach the post-checkpoint barrier"); + tokio::time::timeout(SNAPSHOT_TEST_TIMEOUT, hook.checkpoint_completed.notified()) + .await + .expect("snapshot build should reach the post-checkpoint barrier"); assert!( !build.is_finished(), "snapshot build should remain blocked before archive and persistence" @@ -973,16 +972,13 @@ mod snapshot_gate_tests { log_id: LogId::new(LeaderId::new(1, 1), 1), payload: EntryPayload::Blank, }; - tokio::time::timeout( - std::time::Duration::from_secs(1), - state_machine.apply([blank]), - ) - .await - .expect("apply should proceed after checkpoint while builder remains alive") - .expect("blank entry should apply"); + tokio::time::timeout(SNAPSHOT_TEST_TIMEOUT, state_machine.apply([blank])) + .await + .expect("apply should proceed after checkpoint while builder remains alive") + .expect("blank entry should apply"); hook.continue_build.notify_one(); - let (builder, snapshot) = tokio::time::timeout(std::time::Duration::from_secs(1), build) + let (builder, snapshot) = tokio::time::timeout(SNAPSHOT_TEST_TIMEOUT, build) .await .expect("snapshot build should resume after the test barrier") .expect("snapshot build task should not panic"); @@ -1040,13 +1036,10 @@ mod snapshot_gate_tests { log_id: LogId::new(LeaderId::new(1, 1), 1), payload: EntryPayload::Blank, }; - tokio::time::timeout( - std::time::Duration::from_secs(1), - state_machine.apply([blank]), - ) - .await - .expect("apply should proceed after checkpoint failure while builder remains alive") - .expect("blank entry should apply"); + tokio::time::timeout(SNAPSHOT_TEST_TIMEOUT, state_machine.apply([blank])) + .await + .expect("apply should proceed after checkpoint failure while builder remains alive") + .expect("blank entry should apply"); drop(builder); drop(state_machine); @@ -1107,7 +1100,7 @@ mod snapshot_gate_tests { ) .await }); - tokio::time::timeout(std::time::Duration::from_secs(1), async { + tokio::time::timeout(SNAPSHOT_TEST_TIMEOUT, async { while controller.pause_count.load(Ordering::SeqCst) == 0 { tokio::task::yield_now().await; } @@ -1133,12 +1126,9 @@ mod snapshot_gate_tests { .expect("live storage should remain readable"), "live-value" ); - let permit = tokio::time::timeout( - std::time::Duration::from_secs(1), - Arc::clone(&controller).enter(), - ) - .await - .expect("storage access should resume after cancelling install"); + let permit = tokio::time::timeout(SNAPSHOT_TEST_TIMEOUT, Arc::clone(&controller).enter()) + .await + .expect("storage access should resume after cancelling install"); drop(permit); drop(storage_swap); diff --git a/src/storage/src/storage_manifest.rs b/src/storage/src/storage_manifest.rs index a8f2e11c..e166ede7 100644 --- a/src/storage/src/storage_manifest.rs +++ b/src/storage/src/storage_manifest.rs @@ -123,12 +123,26 @@ impl StorageManifest { /// the snapshot carries the storage identity with the data). pub(crate) fn copy_to(&self, dir: &Path) -> Result<()> { let target = dir.join(STORAGE_MANIFEST_FILE); - fs::copy(&self.path, &target).context(IoSnafu)?; - fs::File::open(&target) - .context(IoSnafu)? - .sync_all() - .context(IoSnafu)?; - Ok(()) + let mut last_error = None; + for attempt in 0..5 { + match (|| -> std::io::Result<()> { + fs::copy(&self.path, &target)?; + fs::File::open(&target)?.sync_all()?; + Ok(()) + })() { + Ok(()) => return Ok(()), + Err(error) => { + last_error = Some(error); + if attempt < 4 { + // Windows can briefly retain a handle while RocksDB + // finishes materializing the checkpoint files. + let _ = fs::remove_file(&target); + std::thread::sleep(std::time::Duration::from_millis(100 * (attempt + 1))); + } + } + } + } + Err(last_error.expect("manifest copy must record its final I/O error")).context(IoSnafu) } fn read(path: &Path) -> Result { From 6d57a5b79e0ef1768c274b028ad1c8a737ebca0e Mon Sep 17 00:00:00 2001 From: happy-v587 Date: Sun, 2 Aug 2026 21:24:17 +0800 Subject: [PATCH 27/32] fix: avoid dynamic command gate allocations --- src/cmd/src/table.rs | 48 ++++++++++++++++++++++---------------------- 1 file changed, 24 insertions(+), 24 deletions(-) diff --git a/src/cmd/src/table.rs b/src/cmd/src/table.rs index 165817d1..24bf3d90 100644 --- a/src/cmd/src/table.rs +++ b/src/cmd/src/table.rs @@ -27,30 +27,30 @@ use crate::{Cmd, CmdMeta}; pub type CmdTable = HashMap>; -/// Provider returning whether a gated command family is currently allowed. -pub type GateFlagProvider = Arc bool + Send + Sync>; - -/// Feature gates consulted when a command table is built. Gates are evaluated -/// on every command execution, so a table always reflects the injected flags. -#[derive(Clone)] +/// Feature gates captured when a command table is built. +/// +/// The values come from startup configuration and do not change while a +/// command table is alive, so storing them by value avoids an unnecessary +/// callback/Arc allocation for every gate. +#[derive(Clone, Copy)] pub struct CommandTableGates { /// Whether the Vector Set commands (VADD/VSIM/...) are enabled. - pub vector_enabled: GateFlagProvider, + pub vector_enabled: bool, /// Whether Vector Set commands are allowed given the cluster state: /// false in cluster mode until the Raft apply-correctness contract (PR0) /// lands. - pub vector_cluster_allowed: GateFlagProvider, + pub vector_cluster_allowed: bool, /// Whether FLUSHDB/FLUSHALL are allowed. Disabled in cluster mode unless /// `cluster-flush-enabled` is set. - pub cluster_flush_allowed: GateFlagProvider, + pub cluster_flush_allowed: bool, } impl Default for CommandTableGates { fn default() -> Self { Self { - vector_enabled: Arc::new(|| true), - vector_cluster_allowed: Arc::new(|| true), - cluster_flush_allowed: Arc::new(|| true), + vector_enabled: true, + vector_cluster_allowed: true, + cluster_flush_allowed: true, } } } @@ -63,9 +63,9 @@ impl CommandTableGates { cluster_flush_allowed: bool, ) -> Self { Self { - vector_enabled: Arc::new(move || vector_enabled), - vector_cluster_allowed: Arc::new(move || vector_cluster_allowed), - cluster_flush_allowed: Arc::new(move || cluster_flush_allowed), + vector_enabled, + vector_cluster_allowed, + cluster_flush_allowed, } } } @@ -76,7 +76,7 @@ impl CommandTableGates { #[derive(Clone)] struct GatedCmd { inner: Arc, - allowed: GateFlagProvider, + allowed: bool, disabled_error: String, } @@ -86,7 +86,7 @@ impl Cmd for GatedCmd { } fn do_initial(&self, client: &Client) -> bool { - if !(self.allowed)() { + if !self.allowed { client.set_reply(RespData::Error(self.disabled_error.clone().into())); return false; } @@ -107,14 +107,14 @@ impl Cmd for GatedCmd { fn register_gated_cmds( cmd_table: &mut CmdTable, cmds: Vec>, - allowed: &GateFlagProvider, + allowed: bool, disabled_error: impl Fn(&CmdMeta) -> String, ) { for cmd in cmds { let meta = cmd.meta().clone(); let gated = GatedCmd { inner: cmd, - allowed: Arc::clone(allowed), + allowed, disabled_error: disabled_error(&meta), }; cmd_table.insert(meta.name, Arc::new(gated)); @@ -125,7 +125,7 @@ fn register_gated_cmds( /// layered on top before insertion into the table. fn wrap_gated_cmds( cmds: Vec>, - allowed: &GateFlagProvider, + allowed: bool, disabled_error: impl Fn(&CmdMeta) -> String, ) -> Vec> { cmds.into_iter() @@ -133,7 +133,7 @@ fn wrap_gated_cmds( let meta = cmd.meta().clone(); Arc::new(GatedCmd { inner: cmd, - allowed: Arc::clone(allowed), + allowed, disabled_error: disabled_error(&meta), }) as Arc }) @@ -304,7 +304,7 @@ pub fn create_command_table_with_gates( register_gated_cmds( &mut cmd_table, flush_cmds, - &gates.cluster_flush_allowed, + gates.cluster_flush_allowed, |meta| { format!( "ERR {} is not supported in cluster mode yet", @@ -329,10 +329,10 @@ pub fn create_command_table_with_gates( Arc::new(crate::vector::VInfoCmd::new()), Arc::new(crate::vector::VIsMemberCmd::new()), ]; - let vector_cmds = wrap_gated_cmds(vector_cmds, &gates.vector_cluster_allowed, |_| { + let vector_cmds = wrap_gated_cmds(vector_cmds, gates.vector_cluster_allowed, |_| { "ERR vector commands are not supported in cluster mode yet".to_string() }); - register_gated_cmds(&mut cmd_table, vector_cmds, &gates.vector_enabled, |_| { + register_gated_cmds(&mut cmd_table, vector_cmds, gates.vector_enabled, |_| { "ERR vector support is disabled (vector-enabled=false)".to_string() }); From 8c294ea0f9411c16ae062c6dbbeeefdf3ee50858 Mon Sep 17 00:00:00 2001 From: happy-v587 Date: Sun, 2 Aug 2026 21:28:42 +0800 Subject: [PATCH 28/32] test: cover vector gate precedence --- src/cmd/src/table.rs | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/src/cmd/src/table.rs b/src/cmd/src/table.rs index 24bf3d90..6a3124b5 100644 --- a/src/cmd/src/table.rs +++ b/src/cmd/src/table.rs @@ -70,9 +70,9 @@ impl CommandTableGates { } } -/// Wraps a command with a deterministic pre-execution gate: when the flag -/// provider returns false the command replies with `disabled_error` and the -/// inner command never runs. +/// Wraps a command with a deterministic pre-execution gate: when `allowed` is +/// false the command replies with `disabled_error` and the inner command never +/// runs. #[derive(Clone)] struct GatedCmd { inner: Arc, @@ -514,6 +514,19 @@ mod tests { } } + #[test] + fn vector_feature_gate_precedes_cluster_gate() { + let table = create_command_table_with_gates( + Arc::new(|| None), + CommandTableGates::from_flags(false, false, true), + ); + let reply = run_command(&table, "vcard", &[b"vcard".to_vec(), b"k".to_vec()]); + assert_eq!( + error_text(&reply), + "ERR vector support is disabled (vector-enabled=false)" + ); + } + #[test] fn vector_commands_pass_gate_when_enabled() { let table = create_command_table_with_gates( From c16c3efaea297b2796a70639eeeb5bc8a5b6c589 Mon Sep 17 00:00:00 2001 From: happy-v587 Date: Sun, 2 Aug 2026 23:22:07 +0800 Subject: [PATCH 29/32] fix(cmd): avoid zero-sized requirepass providers --- src/cmd/src/auth.rs | 17 ++++++++++++++++- src/cmd/src/hello.rs | 3 ++- src/cmd/src/substr.rs | 3 ++- src/cmd/src/table.rs | 39 +++++++++++++++++++++++---------------- 4 files changed, 43 insertions(+), 19 deletions(-) diff --git a/src/cmd/src/auth.rs b/src/cmd/src/auth.rs index 2a6eb8c5..a951cccd 100644 --- a/src/cmd/src/auth.rs +++ b/src/cmd/src/auth.rs @@ -27,6 +27,14 @@ use crate::{impl_cmd_clone_box, impl_cmd_meta}; pub type RequirepassProvider = Arc Option + Send + Sync>; +fn no_requirepass() -> Option { + None +} + +pub(crate) fn no_requirepass_provider() -> RequirepassProvider { + Arc::new(no_requirepass as fn() -> Option) +} + #[derive(Clone)] pub struct AuthCmd { meta: CmdMeta, @@ -43,7 +51,7 @@ impl Default for AuthCmd { acl_category: AclCategory::CONNECTION, ..Default::default() }, - requirepass_provider: Arc::new(|| None), + requirepass_provider: no_requirepass_provider(), } } } @@ -223,4 +231,11 @@ mod tests { let cmd = AuthCmd::default(); assert!(cmd.has_flag(CmdFlags::NO_AUTH)); } + + #[test] + fn default_requirepass_provider_has_nonzero_trait_object_data_size() { + let cmd = AuthCmd::default(); + + assert_ne!(std::mem::size_of_val(cmd.requirepass_provider.as_ref()), 0); + } } diff --git a/src/cmd/src/hello.rs b/src/cmd/src/hello.rs index c1f5d202..154f5734 100644 --- a/src/cmd/src/hello.rs +++ b/src/cmd/src/hello.rs @@ -23,6 +23,7 @@ use resp::{CommandType, HelloAuthResult, RespCommand, RespData, RespError}; use storage::storage::Storage; use subtle::ConstantTimeEq; +use crate::auth::no_requirepass_provider; use crate::{ AclCategory, Cmd, CmdFlags, CmdMeta, RequirepassProvider, impl_cmd_clone_box, impl_cmd_meta, }; @@ -43,7 +44,7 @@ impl Default for HelloCmd { acl_category: AclCategory::CONNECTION | AclCategory::FAST, ..Default::default() }, - requirepass_provider: Arc::new(|| None), + requirepass_provider: no_requirepass_provider(), } } } diff --git a/src/cmd/src/substr.rs b/src/cmd/src/substr.rs index b5b82789..2889f540 100644 --- a/src/cmd/src/substr.rs +++ b/src/cmd/src/substr.rs @@ -71,6 +71,7 @@ mod tests { use storage::{StorageOptions, safe_cleanup_test_db, unique_test_db_path}; use super::*; + use crate::auth::no_requirepass_provider; use crate::table::create_command_table; struct TestStream; @@ -143,7 +144,7 @@ mod tests { b"-100".to_vec(), ]); - let command_table = create_command_table(Arc::new(|| None)); + let command_table = create_command_table(no_requirepass_provider()); command_table .get("substr") .expect("SUBSTR should be publicly registered") diff --git a/src/cmd/src/table.rs b/src/cmd/src/table.rs index 6a3124b5..a363355e 100644 --- a/src/cmd/src/table.rs +++ b/src/cmd/src/table.rs @@ -365,13 +365,15 @@ mod tests { use resp::RespData; use storage::storage::Storage; + use crate::auth::{RequirepassProvider, no_requirepass_provider}; + use super::{ CmdTable, CommandTableGates, create_command_table, create_command_table_with_gates, }; #[test] fn registers_substr_but_not_touch_until_access_metadata_exists() { - let table = create_command_table(Arc::new(|| None)); + let table = create_command_table(no_requirepass_provider()); assert!(table.contains_key("substr")); assert!(!table.contains_key("touch")); @@ -390,6 +392,11 @@ mod tests { } } + fn test_requirepass_provider(password: Option<&str>) -> RequirepassProvider { + let password = password.map(str::to_owned); + Arc::new(move || password.clone()) + } + fn run_command(table: &CmdTable, name: &str, argv: &[Vec]) -> RespData { let command = table.get(name).expect("command should be registered"); let client = Client::new(Box::new(TestStream)); @@ -408,7 +415,7 @@ mod tests { #[test] fn hello_command_returns_resp3_handshake() { - let table = create_command_table(Arc::new(|| None)); + let table = create_command_table(no_requirepass_provider()); let command = table.get("hello").expect("HELLO should be registered"); let client = Client::new(Box::new(TestStream)); client.set_cmd_name(b"hello"); @@ -453,7 +460,7 @@ mod tests { #[test] fn vector_commands_are_registered() { - let table = create_command_table(Arc::new(|| None)); + let table = create_command_table(no_requirepass_provider()); for name in [ "vadd", "vsim", @@ -471,7 +478,7 @@ mod tests { #[test] fn vector_commands_are_rejected_when_disabled() { let table = create_command_table_with_gates( - Arc::new(|| None), + no_requirepass_provider(), CommandTableGates::from_flags(false, true, true), ); let argvs: [(&str, Vec>); 8] = [ @@ -517,7 +524,7 @@ mod tests { #[test] fn vector_feature_gate_precedes_cluster_gate() { let table = create_command_table_with_gates( - Arc::new(|| None), + no_requirepass_provider(), CommandTableGates::from_flags(false, false, true), ); let reply = run_command(&table, "vcard", &[b"vcard".to_vec(), b"k".to_vec()]); @@ -530,7 +537,7 @@ mod tests { #[test] fn vector_commands_pass_gate_when_enabled() { let table = create_command_table_with_gates( - Arc::new(|| None), + no_requirepass_provider(), CommandTableGates::from_flags(true, true, true), ); // Malformed vector spec: parsing fails before storage is touched, so @@ -554,7 +561,7 @@ mod tests { #[test] fn vector_commands_are_rejected_when_cluster_gate_disallows() { let table = create_command_table_with_gates( - Arc::new(|| None), + no_requirepass_provider(), CommandTableGates::from_flags(true, false, true), ); let reply = run_command(&table, "vcard", &[b"vcard".to_vec(), b"k".to_vec()]); @@ -583,7 +590,7 @@ mod tests { #[test] fn info_vector_section_reports_flat_index_and_metrics() { - let table = create_command_table(Arc::new(|| None)); + let table = create_command_table(no_requirepass_provider()); let reply = run_command(&table, "info", &[b"info".to_vec(), b"vector".to_vec()]); let RespData::BulkString(Some(body)) = reply else { @@ -613,7 +620,7 @@ mod tests { #[test] fn flush_commands_are_rejected_when_cluster_gate_disallows() { let table = create_command_table_with_gates( - Arc::new(|| None), + no_requirepass_provider(), CommandTableGates::from_flags(true, true, false), ); let reply = run_command(&table, "flushdb", &[b"flushdb".to_vec()]); @@ -631,7 +638,7 @@ mod tests { #[test] fn flush_commands_execute_when_gate_allows() { let table = create_command_table_with_gates( - Arc::new(|| None), + no_requirepass_provider(), CommandTableGates::from_flags(true, true, true), ); let reply = run_command(&table, "flushdb", &[b"flushdb".to_vec()]); @@ -644,7 +651,7 @@ mod tests { #[test] fn flush_commands_execute_with_default_gates() { // Default gates model standalone mode: nothing is blocked. - let table = create_command_table(Arc::new(|| None)); + let table = create_command_table(no_requirepass_provider()); let reply = run_command(&table, "flushdb", &[b"flushdb".to_vec()]); assert!( matches!(reply, RespData::SimpleString(ref s) if s.as_ref() == b"OK"), @@ -654,7 +661,7 @@ mod tests { #[test] fn hello_bare_with_requirepass_returns_noauth() { - let table = create_command_table(Arc::new(|| Some("secret".to_string()))); + let table = create_command_table(test_requirepass_provider(Some("secret"))); let command = table.get("hello").expect("HELLO should be registered"); let client = Client::new(Box::new(TestStream)); client.set_cmd_name(b"hello"); @@ -673,7 +680,7 @@ mod tests { #[test] fn hello_setname_sets_client_name() { - let table = create_command_table(Arc::new(|| None)); + let table = create_command_table(no_requirepass_provider()); let command = table.get("hello").expect("HELLO should be registered"); let client = Client::new(Box::new(TestStream)); client.set_cmd_name(b"hello"); @@ -695,7 +702,7 @@ mod tests { #[test] fn hello_auth_with_correct_password_authenticates() { - let table = create_command_table(Arc::new(|| Some("secret".to_string()))); + let table = create_command_table(test_requirepass_provider(Some("secret"))); let command = table.get("hello").expect("HELLO should be registered"); let client = Client::new(Box::new(TestStream)); client.set_cmd_name(b"hello"); @@ -720,7 +727,7 @@ mod tests { #[test] fn hello_auth_with_wrong_password_returns_wrongpass() { - let table = create_command_table(Arc::new(|| Some("secret".to_string()))); + let table = create_command_table(test_requirepass_provider(Some("secret"))); let command = table.get("hello").expect("HELLO should be registered"); let client = Client::new(Box::new(TestStream)); client.set_cmd_name(b"hello"); @@ -745,7 +752,7 @@ mod tests { #[test] fn hello_auth_without_requirepass_returns_error() { - let table = create_command_table(Arc::new(|| None)); + let table = create_command_table(no_requirepass_provider()); let command = table.get("hello").expect("HELLO should be registered"); let client = Client::new(Box::new(TestStream)); client.set_cmd_name(b"hello"); From 9a446c53e99fd095ff3dae58586727a0520ca682 Mon Sep 17 00:00:00 2001 From: happy-v587 Date: Sun, 2 Aug 2026 23:28:57 +0800 Subject: [PATCH 30/32] test(cmd): cover hello provider storage --- src/cmd/src/hello.rs | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/src/cmd/src/hello.rs b/src/cmd/src/hello.rs index 154f5734..86bd505e 100644 --- a/src/cmd/src/hello.rs +++ b/src/cmd/src/hello.rs @@ -227,4 +227,11 @@ mod tests { assert_eq!(client.name().as_slice(), b"my-client"); } + + #[test] + fn default_requirepass_provider_has_nonzero_trait_object_data_size() { + let cmd = HelloCmd::default(); + + assert_ne!(std::mem::size_of_val(cmd.requirepass_provider.as_ref()), 0); + } } From e8c8fafc2e68edd37466cc4577d0aac94f1cdc4e Mon Sep 17 00:00:00 2001 From: happy-v587 Date: Sun, 2 Aug 2026 23:39:36 +0800 Subject: [PATCH 31/32] fix(storage): avoid zero-sized cache callbacks --- src/storage/src/redis.rs | 40 +++++++++++++++++++++++++++++++++++--- src/storage/src/storage.rs | 32 ++++++++++++++++++++++++++++-- 2 files changed, 67 insertions(+), 5 deletions(-) diff --git a/src/storage/src/redis.rs b/src/storage/src/redis.rs index cc29f28f..a7f8df6a 100644 --- a/src/storage/src/redis.rs +++ b/src/storage/src/redis.rs @@ -61,6 +61,14 @@ use crate::storage_manifest::StorageManifest; /// log index that created the key (wired up by the raft layer later). pub type GenerationProvider = Arc Result + Send + Sync>; +fn default_cache_weighter(_: &K, _: &V) -> usize { + 1 +} + +fn default_cache_filter(_: &K, _: &V) -> bool { + true +} + // Import logindex types for use in Storage #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -350,7 +358,16 @@ impl Redis { }; let statistics_store: Cache = - CacheBuilder::new(storage.statistics_max_size).build(); + CacheBuilder::new(storage.statistics_max_size) + .with_weighter( + default_cache_weighter:: + as fn(&String, &KeyStatistics) -> usize, + ) + .with_filter( + default_cache_filter:: + as fn(&String, &KeyStatistics) -> bool, + ) + .build(); let flat_query_gate = crate::vector_flat::FlatQueryGate::new(storage.vector.max_concurrent_flat_queries); @@ -373,8 +390,25 @@ impl Redis { vector_fault_hooks: crate::vector_fault::VectorFaultHooks::default(), statistics_store: Arc::new(statistics_store), - scan_cursors_store: Mutex::new(CacheBuilder::new(5000).build()), - spop_counts_store: Mutex::new(CacheBuilder::new(1000).build()), + scan_cursors_store: Mutex::new( + CacheBuilder::new(5000) + .with_weighter( + default_cache_weighter::, Vec> + as fn(&Vec, &Vec) -> usize, + ) + .with_filter( + default_cache_filter::, Vec> as fn(&Vec, &Vec) -> bool, + ) + .build(), + ), + spop_counts_store: Mutex::new( + CacheBuilder::new(1000) + .with_weighter( + default_cache_weighter:: as fn(&String, &u64) -> usize, + ) + .with_filter(default_cache_filter:: as fn(&String, &u64) -> bool) + .build(), + ), small_compaction_threshold: std::sync::atomic::AtomicU64::new(5000), small_compaction_duration_threshold: std::sync::atomic::AtomicU64::new(10000), diff --git a/src/storage/src/storage.rs b/src/storage/src/storage.rs index 5dfa05fb..4f7426d4 100644 --- a/src/storage/src/storage.rs +++ b/src/storage/src/storage.rs @@ -34,6 +34,14 @@ use crate::storage_scan::{SCAN_CURSOR_STATE_CAPACITY, ScanCursorState}; use crate::{ColumnFamilyIndex, Redis, StorageOptions, data_type_to_tag}; use conf::raft_type::{Binlog, OperateType}; +fn default_cache_weighter(_: &K, _: &V) -> usize { + 1 +} + +fn default_cache_filter(_: &K, _: &V) -> bool { + true +} + pub enum TaskType { None = 0, CleanAll = 1, @@ -132,8 +140,28 @@ impl Storage { is_opened: AtomicBool::new(false), lock_mgr: Arc::new(LockMgr::new(1000)), command_access_gate: Arc::new(RwLock::new(())), - cursors_store: Arc::new(CacheBuilder::new(1000).build()), - scan_cursor_states: Arc::new(CacheBuilder::new(SCAN_CURSOR_STATE_CAPACITY).build()), + cursors_store: Arc::new( + CacheBuilder::new(1000) + .with_weighter( + default_cache_weighter:: as fn(&String, &String) -> usize, + ) + .with_filter( + default_cache_filter:: as fn(&String, &String) -> bool, + ) + .build(), + ), + scan_cursor_states: Arc::new( + CacheBuilder::new(SCAN_CURSOR_STATE_CAPACITY) + .with_weighter( + default_cache_weighter:: + as fn(&u64, &ScanCursorState) -> usize, + ) + .with_filter( + default_cache_filter:: + as fn(&u64, &ScanCursorState) -> bool, + ) + .build(), + ), db_instance_num, db_id, bg_task_handler: None, From e249191952b41f280cf7b6bb8ab8dc99b1c02d18 Mon Sep 17 00:00:00 2001 From: happy-v587 Date: Sun, 2 Aug 2026 23:54:16 +0800 Subject: [PATCH 32/32] fix(cmd): avoid zero-sized test streams --- src/cmd/src/auth.rs | 6 ++++-- src/cmd/src/hello.rs | 6 ++++-- src/cmd/src/hscan.rs | 8 +++++--- src/cmd/src/keys.rs | 8 +++++--- src/cmd/src/sscan.rs | 8 +++++--- src/cmd/src/substr.rs | 12 +++++++----- src/cmd/src/table.rs | 18 +++++++++-------- src/cmd/src/zmscore.rs | 10 ++++++---- src/storage/src/redis.rs | 40 +++----------------------------------- src/storage/src/storage.rs | 32 ++---------------------------- 10 files changed, 51 insertions(+), 97 deletions(-) diff --git a/src/cmd/src/auth.rs b/src/cmd/src/auth.rs index a951cccd..479bdd2e 100644 --- a/src/cmd/src/auth.rs +++ b/src/cmd/src/auth.rs @@ -126,7 +126,9 @@ mod tests { use client::StreamTrait; use storage::storage::Storage; - struct TestStream; + struct TestStream { + _marker: u8, + } #[async_trait::async_trait] impl StreamTrait for TestStream { @@ -141,7 +143,7 @@ mod tests { fn make_client() -> Arc { // `Client::new` is fail-closed: a freshly built client is unauthenticated. - Arc::new(Client::new(Box::new(TestStream))) + Arc::new(Client::new(Box::new(TestStream { _marker: 0 }))) } fn make_storage() -> Arc { diff --git a/src/cmd/src/hello.rs b/src/cmd/src/hello.rs index 86bd505e..be0d4c61 100644 --- a/src/cmd/src/hello.rs +++ b/src/cmd/src/hello.rs @@ -139,7 +139,9 @@ mod tests { use super::*; use client::StreamTrait; - struct TestStream; + struct TestStream { + _marker: u8, + } #[async_trait::async_trait] impl StreamTrait for TestStream { @@ -153,7 +155,7 @@ mod tests { } fn make_client() -> Arc { - Arc::new(Client::new(Box::new(TestStream))) + Arc::new(Client::new(Box::new(TestStream { _marker: 0 }))) } fn make_storage() -> Arc { diff --git a/src/cmd/src/hscan.rs b/src/cmd/src/hscan.rs index 998f0b5c..3d582142 100644 --- a/src/cmd/src/hscan.rs +++ b/src/cmd/src/hscan.rs @@ -112,7 +112,9 @@ mod tests { use super::*; - struct TestStream; + struct TestStream { + _marker: u8, + } #[async_trait::async_trait] impl StreamTrait for TestStream { @@ -175,7 +177,7 @@ mod tests { .unwrap(); storage.hset(b"binary_hash", b"\xff", b"\xfe").unwrap(); let storage = Arc::new(storage); - let client = Client::new(Box::new(TestStream)); + let client = Client::new(Box::new(TestStream { _marker: 0 })); client.set_argv(&[ b"hscan".to_vec(), b"binary_hash".to_vec(), @@ -211,7 +213,7 @@ mod tests { .unwrap(); storage.hset(b"empty_field_hash", b"", b"value").unwrap(); let storage = Arc::new(storage); - let client = Client::new(Box::new(TestStream)); + let client = Client::new(Box::new(TestStream { _marker: 0 })); client.set_argv(&[ b"hscan".to_vec(), b"empty_field_hash".to_vec(), diff --git a/src/cmd/src/keys.rs b/src/cmd/src/keys.rs index 859fa345..2d92d7f8 100644 --- a/src/cmd/src/keys.rs +++ b/src/cmd/src/keys.rs @@ -89,7 +89,9 @@ mod tests { use super::KeysCmd; use crate::Cmd; - struct TestStream; + struct TestStream { + _marker: u8, + } #[async_trait::async_trait] impl StreamTrait for TestStream { @@ -142,7 +144,7 @@ mod tests { storage.set(&invalid_key_a, b"invalid-a").unwrap(); storage.set(&invalid_key_b, b"invalid-b").unwrap(); let storage = Arc::new(storage); - let client = Client::new(Box::new(TestStream)); + let client = Client::new(Box::new(TestStream { _marker: 0 })); let command = KeysCmd::new(); client.set_argv(&[b"keys".to_vec(), b"?".to_vec()]); @@ -173,7 +175,7 @@ mod tests { storage.set(b"", b"empty-key").unwrap(); storage.set(b"literal", b"literal-key").unwrap(); let storage = Arc::new(storage); - let client = Client::new(Box::new(TestStream)); + let client = Client::new(Box::new(TestStream { _marker: 0 })); let command = KeysCmd::new(); assert_eq!( diff --git a/src/cmd/src/sscan.rs b/src/cmd/src/sscan.rs index c3a1f069..ac2c909a 100644 --- a/src/cmd/src/sscan.rs +++ b/src/cmd/src/sscan.rs @@ -110,7 +110,9 @@ mod tests { use super::*; - struct TestStream; + struct TestStream { + _marker: u8, + } #[async_trait::async_trait] impl StreamTrait for TestStream { @@ -173,7 +175,7 @@ mod tests { .unwrap(); storage.sadd(b"binary_set", &[b"\xff"]).unwrap(); let storage = Arc::new(storage); - let client = Client::new(Box::new(TestStream)); + let client = Client::new(Box::new(TestStream { _marker: 0 })); client.set_argv(&[ b"sscan".to_vec(), b"binary_set".to_vec(), @@ -206,7 +208,7 @@ mod tests { .unwrap(); storage.sadd(b"empty_member_set", &[b""]).unwrap(); let storage = Arc::new(storage); - let client = Client::new(Box::new(TestStream)); + let client = Client::new(Box::new(TestStream { _marker: 0 })); client.set_argv(&[ b"sscan".to_vec(), b"empty_member_set".to_vec(), diff --git a/src/cmd/src/substr.rs b/src/cmd/src/substr.rs index 2889f540..ca77c27a 100644 --- a/src/cmd/src/substr.rs +++ b/src/cmd/src/substr.rs @@ -74,7 +74,9 @@ mod tests { use crate::auth::no_requirepass_provider; use crate::table::create_command_table; - struct TestStream; + struct TestStream { + _marker: u8, + } #[async_trait::async_trait] impl StreamTrait for TestStream { @@ -135,7 +137,7 @@ mod tests { .unwrap(); storage.set(b"key", b"Hello World").unwrap(); let storage = Arc::new(storage); - let client = Client::new(Box::new(TestStream)); + let client = Client::new(Box::new(TestStream { _marker: 0 })); client.set_cmd_name(b"substr"); client.set_argv(&[ b"substr".to_vec(), @@ -172,7 +174,7 @@ mod tests { .unwrap(); storage.hset(b"hash", b"field", b"value").unwrap(); let storage = Arc::new(storage); - let client = Client::new(Box::new(TestStream)); + let client = Client::new(Box::new(TestStream { _marker: 0 })); client.set_cmd_name(b"substr"); client.set_argv(&[ b"substr".to_vec(), @@ -206,7 +208,7 @@ mod tests { .unwrap(); storage.set(b"binary", b"\xff\x00").unwrap(); let storage = Arc::new(storage); - let client = Client::new(Box::new(TestStream)); + let client = Client::new(Box::new(TestStream { _marker: 0 })); for (key, expected) in [ ( @@ -235,7 +237,7 @@ mod tests { #[test] fn substr_rejects_out_of_range_integer_arguments_before_storage_access() { - let client = Client::new(Box::new(TestStream)); + let client = Client::new(Box::new(TestStream { _marker: 0 })); client.set_cmd_name(b"substr"); client.set_argv(&[ b"substr".to_vec(), diff --git a/src/cmd/src/table.rs b/src/cmd/src/table.rs index a363355e..9d92530b 100644 --- a/src/cmd/src/table.rs +++ b/src/cmd/src/table.rs @@ -379,7 +379,9 @@ mod tests { assert!(!table.contains_key("touch")); } - struct TestStream; + struct TestStream { + _marker: u8, + } #[async_trait::async_trait] impl StreamTrait for TestStream { @@ -399,7 +401,7 @@ mod tests { fn run_command(table: &CmdTable, name: &str, argv: &[Vec]) -> RespData { let command = table.get(name).expect("command should be registered"); - let client = Client::new(Box::new(TestStream)); + let client = Client::new(Box::new(TestStream { _marker: 0 })); client.set_cmd_name(name.as_bytes()); client.set_argv(argv); command.execute(&client, Arc::new(Storage::new(1, 0))); @@ -417,7 +419,7 @@ mod tests { fn hello_command_returns_resp3_handshake() { let table = create_command_table(no_requirepass_provider()); let command = table.get("hello").expect("HELLO should be registered"); - let client = Client::new(Box::new(TestStream)); + let client = Client::new(Box::new(TestStream { _marker: 0 })); client.set_cmd_name(b"hello"); client.set_argv(&[b"hello".to_vec(), b"3".to_vec()]); @@ -663,7 +665,7 @@ mod tests { fn hello_bare_with_requirepass_returns_noauth() { let table = create_command_table(test_requirepass_provider(Some("secret"))); let command = table.get("hello").expect("HELLO should be registered"); - let client = Client::new(Box::new(TestStream)); + let client = Client::new(Box::new(TestStream { _marker: 0 })); client.set_cmd_name(b"hello"); client.set_argv(&[b"hello".to_vec(), b"3".to_vec()]); @@ -682,7 +684,7 @@ mod tests { fn hello_setname_sets_client_name() { let table = create_command_table(no_requirepass_provider()); let command = table.get("hello").expect("HELLO should be registered"); - let client = Client::new(Box::new(TestStream)); + let client = Client::new(Box::new(TestStream { _marker: 0 })); client.set_cmd_name(b"hello"); client.set_argv(&[ b"hello".to_vec(), @@ -704,7 +706,7 @@ mod tests { fn hello_auth_with_correct_password_authenticates() { let table = create_command_table(test_requirepass_provider(Some("secret"))); let command = table.get("hello").expect("HELLO should be registered"); - let client = Client::new(Box::new(TestStream)); + let client = Client::new(Box::new(TestStream { _marker: 0 })); client.set_cmd_name(b"hello"); client.set_argv(&[ b"hello".to_vec(), @@ -729,7 +731,7 @@ mod tests { fn hello_auth_with_wrong_password_returns_wrongpass() { let table = create_command_table(test_requirepass_provider(Some("secret"))); let command = table.get("hello").expect("HELLO should be registered"); - let client = Client::new(Box::new(TestStream)); + let client = Client::new(Box::new(TestStream { _marker: 0 })); client.set_cmd_name(b"hello"); client.set_argv(&[ b"hello".to_vec(), @@ -754,7 +756,7 @@ mod tests { fn hello_auth_without_requirepass_returns_error() { let table = create_command_table(no_requirepass_provider()); let command = table.get("hello").expect("HELLO should be registered"); - let client = Client::new(Box::new(TestStream)); + let client = Client::new(Box::new(TestStream { _marker: 0 })); client.set_cmd_name(b"hello"); client.set_argv(&[ b"hello".to_vec(), diff --git a/src/cmd/src/zmscore.rs b/src/cmd/src/zmscore.rs index f62b3334..99dee2d1 100644 --- a/src/cmd/src/zmscore.rs +++ b/src/cmd/src/zmscore.rs @@ -107,7 +107,9 @@ mod tests { use super::*; - struct TestStream; + struct TestStream { + _marker: u8, + } #[async_trait::async_trait] impl StreamTrait for TestStream { @@ -205,14 +207,14 @@ mod tests { b"first".to_vec(), ]; - let resp2_client = Client::new(Box::new(TestStream)); + let resp2_client = Client::new(Box::new(TestStream { _marker: 0 })); let resp2_reply = run_zmscore(&resp2_client, &storage, &members); assert_eq!( encode(&resp2_reply, RespVersion::RESP2), b"*4\r\n$3\r\n1.5\r\n$-1\r\n$5\r\n-2.25\r\n$3\r\n1.5\r\n" ); - let resp3_client = Client::new(Box::new(TestStream)); + let resp3_client = Client::new(Box::new(TestStream { _marker: 0 })); resp3_client.set_argv(&[b"hello".to_vec(), b"3".to_vec()]); crate::hello::HelloCmd::default().do_cmd(&resp3_client, Arc::clone(&storage)); let _hello_reply = resp3_client.take_reply(); @@ -229,7 +231,7 @@ mod tests { #[tokio::test] async fn zmscore_returns_exact_wrongtype_error_and_expired_key_nils() { let (db_path, storage) = open_storage(); - let client = Client::new(Box::new(TestStream)); + let client = Client::new(Box::new(TestStream { _marker: 0 })); storage.set(b"zmscore-key", b"not-a-zset").unwrap(); let wrongtype = run_zmscore(&client, &storage, &[b"member".to_vec()]); diff --git a/src/storage/src/redis.rs b/src/storage/src/redis.rs index a7f8df6a..cc29f28f 100644 --- a/src/storage/src/redis.rs +++ b/src/storage/src/redis.rs @@ -61,14 +61,6 @@ use crate::storage_manifest::StorageManifest; /// log index that created the key (wired up by the raft layer later). pub type GenerationProvider = Arc Result + Send + Sync>; -fn default_cache_weighter(_: &K, _: &V) -> usize { - 1 -} - -fn default_cache_filter(_: &K, _: &V) -> bool { - true -} - // Import logindex types for use in Storage #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -358,16 +350,7 @@ impl Redis { }; let statistics_store: Cache = - CacheBuilder::new(storage.statistics_max_size) - .with_weighter( - default_cache_weighter:: - as fn(&String, &KeyStatistics) -> usize, - ) - .with_filter( - default_cache_filter:: - as fn(&String, &KeyStatistics) -> bool, - ) - .build(); + CacheBuilder::new(storage.statistics_max_size).build(); let flat_query_gate = crate::vector_flat::FlatQueryGate::new(storage.vector.max_concurrent_flat_queries); @@ -390,25 +373,8 @@ impl Redis { vector_fault_hooks: crate::vector_fault::VectorFaultHooks::default(), statistics_store: Arc::new(statistics_store), - scan_cursors_store: Mutex::new( - CacheBuilder::new(5000) - .with_weighter( - default_cache_weighter::, Vec> - as fn(&Vec, &Vec) -> usize, - ) - .with_filter( - default_cache_filter::, Vec> as fn(&Vec, &Vec) -> bool, - ) - .build(), - ), - spop_counts_store: Mutex::new( - CacheBuilder::new(1000) - .with_weighter( - default_cache_weighter:: as fn(&String, &u64) -> usize, - ) - .with_filter(default_cache_filter:: as fn(&String, &u64) -> bool) - .build(), - ), + scan_cursors_store: Mutex::new(CacheBuilder::new(5000).build()), + spop_counts_store: Mutex::new(CacheBuilder::new(1000).build()), small_compaction_threshold: std::sync::atomic::AtomicU64::new(5000), small_compaction_duration_threshold: std::sync::atomic::AtomicU64::new(10000), diff --git a/src/storage/src/storage.rs b/src/storage/src/storage.rs index 4f7426d4..5dfa05fb 100644 --- a/src/storage/src/storage.rs +++ b/src/storage/src/storage.rs @@ -34,14 +34,6 @@ use crate::storage_scan::{SCAN_CURSOR_STATE_CAPACITY, ScanCursorState}; use crate::{ColumnFamilyIndex, Redis, StorageOptions, data_type_to_tag}; use conf::raft_type::{Binlog, OperateType}; -fn default_cache_weighter(_: &K, _: &V) -> usize { - 1 -} - -fn default_cache_filter(_: &K, _: &V) -> bool { - true -} - pub enum TaskType { None = 0, CleanAll = 1, @@ -140,28 +132,8 @@ impl Storage { is_opened: AtomicBool::new(false), lock_mgr: Arc::new(LockMgr::new(1000)), command_access_gate: Arc::new(RwLock::new(())), - cursors_store: Arc::new( - CacheBuilder::new(1000) - .with_weighter( - default_cache_weighter:: as fn(&String, &String) -> usize, - ) - .with_filter( - default_cache_filter:: as fn(&String, &String) -> bool, - ) - .build(), - ), - scan_cursor_states: Arc::new( - CacheBuilder::new(SCAN_CURSOR_STATE_CAPACITY) - .with_weighter( - default_cache_weighter:: - as fn(&u64, &ScanCursorState) -> usize, - ) - .with_filter( - default_cache_filter:: - as fn(&u64, &ScanCursorState) -> bool, - ) - .build(), - ), + cursors_store: Arc::new(CacheBuilder::new(1000).build()), + scan_cursor_states: Arc::new(CacheBuilder::new(SCAN_CURSOR_STATE_CAPACITY).build()), db_instance_num, db_id, bg_task_handler: None,