Implement simple consensus operations on ConsensusStateMachine (#10280)

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Roman Titov
2026-08-24 11:58:34 +02:00
committed by GitHub
co-authored by Claude Opus 5
parent 83fe47c90a
commit 2067eb3398
8 changed files with 712 additions and 11 deletions
@@ -227,7 +227,7 @@ fn sparse_schema_matches(existing: &SparseVectorParams, config: &SparseVectorCon
&& existing_datatype == datatype.map(storage_datatype_to_collection)
}
fn remove_vector_from_config(
pub fn remove_vector_from_config(
params: &mut crate::config::CollectionParams,
vector_name: &VectorNameBuf,
) {
@@ -1,5 +1,41 @@
use collection::shards::CollectionId;
use segment::types::{PayloadFieldSchema, PayloadKeyType, VectorNameBuf};
use shard::operations::vector_name_ops::VectorNameConfig;
/// A single change a consensus operation makes
#[derive(Clone, Debug, PartialEq)]
pub enum Action {
// TODO!
AddNamedVector {
collection: CollectionId,
vector_name: VectorNameBuf,
config: Box<VectorNameConfig>,
},
DropNamedVector {
collection: CollectionId,
vector_name: VectorNameBuf,
},
SetPayloadIndex {
collection: CollectionId,
field_name: PayloadKeyType,
field_schema: PayloadFieldSchema,
},
DropPayloadIndex {
collection: CollectionId,
field_name: PayloadKeyType,
},
}
impl Action {
/// Collection this action changes, if it is scoped to one
pub fn collection(&self) -> Option<&CollectionId> {
match self {
Action::AddNamedVector { collection, .. }
| Action::DropNamedVector { collection, .. }
| Action::SetPayloadIndex { collection, .. }
| Action::DropPayloadIndex { collection, .. } => Some(collection),
}
}
}
@@ -86,7 +86,6 @@ impl ConsensusStateMachine {
}
}
#[expect(clippy::unused_self)]
fn plan_collection_meta(&self, operation: &CollectionMetaOperations) -> ApplyOutcome {
match operation {
CollectionMetaOperations::Nop { .. } => ApplyOutcome::Accepted(Vec::new()),
@@ -99,11 +98,21 @@ impl ConsensusStateMachine {
| CollectionMetaOperations::DropShardKey(_)
| CollectionMetaOperations::SetShardReplicaState(_)
| CollectionMetaOperations::TransferShard(_, _)
| CollectionMetaOperations::Resharding(_, _)
| CollectionMetaOperations::CreateNamedVector(_)
| CollectionMetaOperations::DeleteNamedVector(_)
| CollectionMetaOperations::CreatePayloadIndex(_)
| CollectionMetaOperations::DropPayloadIndex(_) => ApplyOutcome::NotCovered,
| CollectionMetaOperations::Resharding(_, _) => ApplyOutcome::NotCovered,
CollectionMetaOperations::CreateNamedVector(operation) => {
ApplyOutcome::new(self.state.plan_create_named_vector(operation))
}
CollectionMetaOperations::DeleteNamedVector(operation) => {
ApplyOutcome::new(self.state.plan_delete_named_vector(operation))
}
CollectionMetaOperations::CreatePayloadIndex(operation) => {
ApplyOutcome::new(self.state.plan_create_payload_index(operation))
}
CollectionMetaOperations::DropPayloadIndex(operation) => {
ApplyOutcome::new(self.state.plan_drop_payload_index(operation))
}
// Sleeps, or fails at random. Neither is a state change to plan.
#[cfg(feature = "staging")]
@@ -1,7 +1,86 @@
use collection::collection::vector_name_schema;
use super::*;
impl ClusterState {
pub fn apply_action(&mut self, _action: &Action) {
// TODO!
/// Apply one action. Cannot fail.
///
/// Action naming a missing collection changes nothing.
/// Correct operation never emits one, so debug builds assert.
pub fn apply_action(&mut self, action: &Action) {
match action {
Action::AddNamedVector {
collection,
vector_name,
config,
} => {
let Some(state) = self.collection_mut(collection) else {
return;
};
// Planning validates using the same function, so it should never fail here
let res = vector_name_schema::add_vector_to_config(
&mut state.config.params,
vector_name,
config,
);
if let Err(err) = res {
debug_assert!(false, "rejected named vector reached the state: {err}");
log::error!("Failed to add named vector {vector_name} to {collection}: {err}");
}
}
Action::DropNamedVector {
collection,
vector_name,
} => {
let Some(state) = self.collection_mut(collection) else {
return;
};
vector_name_schema::remove_vector_from_config(
&mut state.config.params,
vector_name,
);
}
Action::SetPayloadIndex {
collection,
field_name,
field_schema,
} => {
let Some(state) = self.collection_mut(collection) else {
return;
};
state
.payload_index_schema
.schema
.insert(field_name.clone(), field_schema.clone());
}
Action::DropPayloadIndex {
collection,
field_name,
} => {
let Some(state) = self.collection_mut(collection) else {
return;
};
state.payload_index_schema.schema.remove(field_name);
}
}
}
fn collection_mut(&mut self, collection: &str) -> Option<&mut collection_state::State> {
let state = self.collections.get_mut(collection);
debug_assert!(
state.is_some(),
"action targets collection {collection}, which is not in the state",
);
state
}
}
@@ -1,4 +1,5 @@
mod apply;
mod plan;
use std::collections::HashMap;
@@ -0,0 +1,85 @@
use collection::collection::vector_name_schema;
use super::*;
use crate::content_manager::collection_meta_ops::*;
use crate::content_manager::consensus_state_machine::Action;
type Actions = Vec<Action>;
impl ClusterState {
pub fn plan_create_named_vector(&self, op: &CreateNamedVector) -> StorageResult<Actions> {
let CreateNamedVector {
collection_name,
vector_name,
config,
} = op;
let collection = self.resolve_collection(collection_name)?;
// Reject vector that already exists with different config.
//
// Validate by adding vector to the config, so that `plan` and `apply_action`
// are always in sync.
let mut params = self
.collection(&collection)
.expect("collection exists")
.config
.params
.clone();
vector_name_schema::add_vector_to_config(&mut params, vector_name, config)?;
Ok(vec![Action::AddNamedVector {
collection,
vector_name: vector_name.clone(),
config: Box::new(config.clone()),
}])
}
pub fn plan_delete_named_vector(&self, op: &DeleteNamedVector) -> StorageResult<Actions> {
let DeleteNamedVector {
collection_name,
vector_name,
} = op;
let collection = self.resolve_collection(collection_name)?;
// Deleting vector that does not exist is a no-op, not an error
Ok(vec![Action::DropNamedVector {
collection,
vector_name: vector_name.clone(),
}])
}
pub fn plan_create_payload_index(&self, op: &CreatePayloadIndex) -> StorageResult<Actions> {
let CreatePayloadIndex {
collection_name,
field_name,
field_schema,
} = op;
let collection = self.resolve_collection(collection_name)?;
Ok(vec![Action::SetPayloadIndex {
collection,
field_name: field_name.clone(),
field_schema: field_schema.clone(),
}])
}
pub fn plan_drop_payload_index(&self, op: &DropPayloadIndex) -> StorageResult<Actions> {
let DropPayloadIndex {
collection_name,
field_name,
} = op;
let collection = self.resolve_collection(collection_name)?;
Ok(vec![Action::DropPayloadIndex {
collection,
field_name: field_name.clone(),
}])
}
}
@@ -4,11 +4,13 @@
use std::collections::HashMap;
use segment::data_types::vector_name_config::*;
use segment::types::*;
use super::*;
use crate::content_manager::collection_meta_ops::*;
use crate::content_manager::consensus_ops::ConsensusOperations;
use crate::content_manager::consensus_state_machine::*;
use crate::content_manager::errors::StorageError;
const COLLECTION: &str = "alpha";
@@ -29,6 +31,354 @@ fn nop() {
assert_eq!(machine.state(), &state);
}
#[test]
fn create_named_vector_dense() {
let machine = create_named_vector_impl(dense(4, Distance::Cosine));
let params = &machine
.state()
.collection(COLLECTION)
.expect("collection exists")
.config
.params;
assert!(params.vectors.get_params("text").is_some());
}
#[test]
fn create_named_vector_sparse() {
let machine = create_named_vector_impl(sparse());
let params = &machine
.state()
.collection(COLLECTION)
.expect("collection exists")
.config
.params;
let sparse = params
.sparse_vectors
.as_ref()
.expect("sparse vector config exists");
assert!(sparse.contains_key("text"));
}
fn create_named_vector_impl(config: VectorNameConfig) -> ConsensusStateMachine {
let state = cluster_state(Vec::new());
let mut machine = state_machine(state);
let outcome = machine.apply(&create_named_vector_op("text", config));
let ApplyOutcome::Accepted(actions) = outcome else {
panic!("creating a new named vector should be accepted, got {outcome:?}");
};
assert!(matches!(
actions.as_slice(),
[Action::AddNamedVector { .. }],
));
machine
}
#[test]
fn create_named_vector_replay() {
let config = dense(4, Distance::Cosine);
let state = cluster_state(vec![("text", config.clone())]);
let mut machine = state_machine(state.clone());
let outcome = machine.apply(&create_named_vector_op("text", config));
let ApplyOutcome::Accepted(actions) = outcome else {
panic!("replay of an applied named vector should be accepted, got {outcome:?}");
};
assert!(matches!(
actions.as_slice(),
[Action::AddNamedVector { .. }],
));
assert_eq!(machine.state(), &state, "replay should not change anything");
}
#[test]
fn create_named_vector_reject_existing_diff_dim() {
create_named_vector_reject_existing(dense(8, Distance::Cosine));
}
#[test]
fn create_named_vector_reject_existing_diff_type() {
create_named_vector_reject_existing(sparse());
}
fn create_named_vector_reject_existing(config: VectorNameConfig) {
let state = cluster_state(vec![("text", dense(4, Distance::Cosine))]);
let mut machine = state_machine(state.clone());
let outcome = machine.apply(&create_named_vector_op("text", config));
assert!(matches!(
outcome,
ApplyOutcome::Rejected(StorageError::BadInput { .. })
));
assert_eq!(machine.state(), &state);
}
#[test]
fn delete_named_vector_dense() {
delete_named_vector_impl(dense(4, Distance::Cosine));
}
#[test]
fn delete_named_vector_sparse() {
delete_named_vector_impl(sparse());
}
fn delete_named_vector_impl(config: VectorNameConfig) {
let state = cluster_state(vec![("text", config)]);
let mut machine = state_machine(state);
let outcome = machine.apply(&delete_named_vector_op("text"));
let ApplyOutcome::Accepted(actions) = outcome else {
panic!("deleting an existing vector should be accepted, got {outcome:?}");
};
assert!(matches!(
actions.as_slice(),
[Action::DropNamedVector { .. }],
));
let params = &machine
.state()
.collection(COLLECTION)
.expect("collection exists")
.config
.params;
assert!(params.vectors.get_params("text").is_none());
assert!(
params
.sparse_vectors
.as_ref()
.is_none_or(|sparse| !sparse.contains_key("text")),
);
}
#[test]
fn delete_named_vector_missing() {
let state = cluster_state(Vec::new());
let mut machine = state_machine(state.clone());
let outcome = machine.apply(&delete_named_vector_op("text"));
let ApplyOutcome::Accepted(actions) = outcome else {
panic!("deleting a vector that does not exist should be accepted, got {outcome:?}");
};
// Action is emitted even if state already matches
assert!(matches!(
actions.as_slice(),
[Action::DropNamedVector { .. }],
));
assert_eq!(machine.state(), &state);
}
#[test]
fn create_payload_index() {
let state = cluster_state(Vec::new());
let mut machine = state_machine(state);
let outcome = machine.apply(&create_payload_index_op("city", PayloadSchemaType::Keyword));
let ApplyOutcome::Accepted(actions) = outcome else {
panic!("creating payload index should be accepted, got {outcome:?}");
};
assert!(matches!(
actions.as_slice(),
[Action::SetPayloadIndex { .. }],
));
let schema = &machine
.state()
.collection(COLLECTION)
.expect("collection exists")
.payload_index_schema
.schema;
assert!(schema.contains_key(&field_name("city")));
}
#[test]
fn create_payload_index_replay() {
let state = cluster_state_with_index("city", PayloadSchemaType::Keyword);
let mut machine = state_machine(state.clone());
let outcome = machine.apply(&create_payload_index_op("city", PayloadSchemaType::Keyword));
let ApplyOutcome::Accepted(actions) = outcome else {
panic!("replay of an applied payload index should be accepted, got {outcome:?}");
};
assert!(matches!(
actions.as_slice(),
[Action::SetPayloadIndex { .. }],
));
assert_eq!(machine.state(), &state, "replay should not change anything");
}
#[test]
fn create_payload_index_replace_schema() {
let state = cluster_state_with_index("city", PayloadSchemaType::Keyword);
let mut machine = state_machine(state);
let outcome = machine.apply(&create_payload_index_op("city", PayloadSchemaType::Integer));
let ApplyOutcome::Accepted(actions) = outcome else {
panic!("indexing a field again with another schema should be accepted, got {outcome:?}");
};
assert!(matches!(
actions.as_slice(),
[Action::SetPayloadIndex { .. }],
));
let schema = &machine
.state()
.collection(COLLECTION)
.expect("collection exists")
.payload_index_schema
.schema;
// A field indexed with a different schema is replaced, where a named vector is rejected
assert_eq!(
schema.get(&field_name("city")),
Some(&PayloadFieldSchema::FieldType(PayloadSchemaType::Integer)),
);
}
#[test]
fn drop_payload_index() {
let state = cluster_state_with_index("city", PayloadSchemaType::Keyword);
let mut machine = state_machine(state);
let outcome = machine.apply(&drop_payload_index_op("city"));
let ApplyOutcome::Accepted(actions) = outcome else {
panic!("dropping an indexed field should be accepted, got {outcome:?}");
};
assert!(matches!(
actions.as_slice(),
[Action::DropPayloadIndex { .. }],
));
let schema = &machine
.state()
.collection(COLLECTION)
.expect("collection exists")
.payload_index_schema
.schema;
assert!(!schema.contains_key(&field_name("city")));
}
#[test]
fn drop_payload_index_missing() {
let state = cluster_state(Vec::new());
let mut machine = state_machine(state.clone());
let outcome = machine.apply(&drop_payload_index_op("city"));
let ApplyOutcome::Accepted(actions) = outcome else {
panic!("dropping a field that is not indexed should be accepted, got {outcome:?}");
};
// Action is emitted even if state already matches
assert!(matches!(
actions.as_slice(),
[Action::DropPayloadIndex { .. }],
));
assert_eq!(machine.state(), &state);
}
#[test]
fn reject_missing_collection() {
let mut machine = state_machine(ClusterState::default());
let outcome = machine.apply(&create_named_vector_op("text", dense(4, Distance::Cosine)));
assert!(matches!(
outcome,
ApplyOutcome::Rejected(StorageError::NotFound { .. })
));
}
#[test]
fn resolve_alias() {
let mut state = cluster_state(Vec::new());
state.aliases.insert("alias".into(), COLLECTION.into());
let mut machine = state_machine(state);
let outcome = machine.apply(&collection_meta_op(
CollectionMetaOperations::CreateNamedVector(CreateNamedVector {
collection_name: "alias".into(),
vector_name: "text".into(),
config: dense(4, Distance::Cosine),
}),
));
let ApplyOutcome::Accepted(actions) = outcome else {
panic!("an alias should resolve to its collection, got {outcome:?}");
};
assert_eq!(
actions
.first()
.and_then(Action::collection)
.expect("action modifies collection"),
&COLLECTION,
);
let vectors = &machine
.state()
.collection(COLLECTION)
.expect("collection exists")
.config
.params
.vectors;
assert!(vectors.get_params("text").is_some());
}
#[test]
fn reject_dangling_alias() {
let mut state = cluster_state(Vec::new());
state.aliases.insert("dangling".into(), "missing".into());
let mut machine = state_machine(state.clone());
let outcome = machine.apply(&collection_meta_op(
CollectionMetaOperations::CreateNamedVector(CreateNamedVector {
collection_name: "dangling".into(),
vector_name: "text".into(),
config: dense(4, Distance::Cosine),
}),
));
assert!(matches!(
outcome,
ApplyOutcome::Rejected(StorageError::NotFound { .. })
));
assert_eq!(machine.state(), &state);
}
fn cluster_state(vectors: Vec<(&str, VectorNameConfig)>) -> ClusterState {
let vectors = vectors
.into_iter()
@@ -43,6 +393,78 @@ fn cluster_state(vectors: Vec<(&str, VectorNameConfig)>) -> ClusterState {
}
}
fn cluster_state_with_index(field: &str, field_type: PayloadSchemaType) -> ClusterState {
let mut state = cluster_state(Vec::new());
state
.collections
.get_mut(COLLECTION)
.expect("collection exists")
.payload_index_schema
.schema
.insert(field_name(field), PayloadFieldSchema::FieldType(field_type));
state
}
fn collection_meta_op(op: CollectionMetaOperations) -> ConsensusOperations {
ConsensusOperations::CollectionMeta(Box::new(op))
}
fn create_named_vector_op(vector_name: &str, config: VectorNameConfig) -> ConsensusOperations {
collection_meta_op(CollectionMetaOperations::CreateNamedVector(
CreateNamedVector {
collection_name: COLLECTION.into(),
vector_name: VectorNameBuf::from(vector_name),
config,
},
))
}
fn dense(size: usize, distance: Distance) -> VectorNameConfig {
VectorNameConfig::dense(DenseVectorConfig {
size,
distance,
multivector_config: None,
datatype: None,
})
}
fn sparse() -> VectorNameConfig {
VectorNameConfig::sparse(SparseVectorConfig {
modifier: None,
datatype: None,
})
}
fn delete_named_vector_op(vector_name: &str) -> ConsensusOperations {
collection_meta_op(CollectionMetaOperations::DeleteNamedVector(
DeleteNamedVector {
collection_name: COLLECTION.into(),
vector_name: VectorNameBuf::from(vector_name),
},
))
}
fn create_payload_index_op(field: &str, field_type: PayloadSchemaType) -> ConsensusOperations {
collection_meta_op(CollectionMetaOperations::CreatePayloadIndex(
CreatePayloadIndex {
collection_name: COLLECTION.to_string(),
field_name: field_name(field),
field_schema: PayloadFieldSchema::FieldType(field_type),
},
))
}
fn drop_payload_index_op(field: &str) -> ConsensusOperations {
collection_meta_op(CollectionMetaOperations::DropPayloadIndex(
DropPayloadIndex {
collection_name: COLLECTION.to_string(),
field_name: field_name(field),
},
))
}
fn field_name(field: &str) -> PayloadKeyType {
field.parse().expect("valid field name")
}
@@ -97,12 +97,81 @@ pub fn arb_consensus_operation(
collection_names.push(MISSING_COLLECTION_NAME.into());
prop_oneof![
// TODO!
Just(CollectionMetaOperations::Nop { token: 0 }),
arb_create_named_vector(collection_names.clone()),
arb_delete_named_vector(collection_names.clone()),
arb_create_payload_index(collection_names.clone()),
arb_drop_payload_index(collection_names.clone()),
]
.prop_map(|operation| ConsensusOperations::CollectionMeta(Box::new(operation)))
}
fn arb_collection_name(names: Vec<String>) -> impl Strategy<Value = String> {
proptest::sample::select(names)
}
fn arb_create_named_vector(
collections: Vec<String>,
) -> impl Strategy<Value = CollectionMetaOperations> {
let collection_name = arb_collection_name(collections);
let vector_name = arb_vector_name();
let config = arb_vector_name_config();
(collection_name, vector_name, config).prop_map(|(collection_name, vector_name, config)| {
CollectionMetaOperations::CreateNamedVector(CreateNamedVector {
collection_name,
vector_name,
config,
})
})
}
fn arb_delete_named_vector(
collections: Vec<String>,
) -> impl Strategy<Value = CollectionMetaOperations> {
let collection_name = arb_collection_name(collections);
let vector_name = arb_vector_name();
(collection_name, vector_name).prop_map(|(collection_name, vector_name)| {
CollectionMetaOperations::DeleteNamedVector(DeleteNamedVector {
collection_name,
vector_name,
})
})
}
fn arb_create_payload_index(
collections: Vec<String>,
) -> impl Strategy<Value = CollectionMetaOperations> {
let collection_name = arb_collection_name(collections);
let field_name = arb_field_name();
let field_schema = arb_field_schema();
(collection_name, field_name, field_schema).prop_map(
|(collection_name, field_name, field_schema)| {
CollectionMetaOperations::CreatePayloadIndex(CreatePayloadIndex {
collection_name,
field_name,
field_schema,
})
},
)
}
fn arb_drop_payload_index(
collections: Vec<String>,
) -> impl Strategy<Value = CollectionMetaOperations> {
let collection_name = arb_collection_name(collections);
let field_name = arb_field_name();
(collection_name, field_name).prop_map(|(collection_name, field_name)| {
CollectionMetaOperations::DropPayloadIndex(DropPayloadIndex {
collection_name,
field_name,
})
})
}
fn arb_vector_name() -> impl Strategy<Value = VectorNameBuf> {
proptest::sample::select(VECTOR_NAMES).prop_map(VectorNameBuf::from)
}