From 2e71304e6ec84e88eaa842321709e37bf449be5f Mon Sep 17 00:00:00 2001 From: Roman Titov Date: Wed, 2 Sep 2026 08:51:03 +0200 Subject: [PATCH] Implement `CreateCollection` and `DeleteCollection` for consensus state machine (#10397) Co-authored-by: Claude Opus 5 (1M context) --- .../src/content_manager/alias_mapping.rs | 21 +- .../content_manager/collection_meta_ops.rs | 45 ++- .../consensus_state_machine/action.rs | 14 +- .../consensus_state_machine/mod.rs | 172 +++++++++- .../consensus_state_machine/state/apply.rs | 18 +- .../consensus_state_machine/state/plan.rs | 100 +++++- .../consensus_state_machine/tests/mod.rs | 78 +++-- .../consensus_state_machine/tests/ops.rs | 318 ++++++++++++++++++ .../consensus_state_machine/tests/prop.rs | 63 +++- .../content_manager/toc/create_collection.rs | 38 +-- 10 files changed, 787 insertions(+), 80 deletions(-) diff --git a/lib/storage/src/content_manager/alias_mapping.rs b/lib/storage/src/content_manager/alias_mapping.rs index 09b07d73c7..9b1dbeebc9 100644 --- a/lib/storage/src/content_manager/alias_mapping.rs +++ b/lib/storage/src/content_manager/alias_mapping.rs @@ -24,6 +24,17 @@ impl AliasMapping { Ok(atomic_save_json(path, self)?) } + /// Aliases pointing at `collection_name`. + pub fn collection_aliases<'a>( + &'a self, + collection_name: &'a str, + ) -> impl Iterator + 'a { + self.0 + .iter() + .filter(move |&(_, target)| target == collection_name) + .map(|(alias, _)| alias.clone()) + } + /// Iterate over aliases and collections they point at. pub fn iter(&self) -> impl Iterator { self.0.iter() @@ -144,13 +155,9 @@ impl AliasPersistence { } pub fn collection_aliases(&self, collection_name: &str) -> Vec { - let mut result = vec![]; - for (alias, target_collection) in self.alias_mapping.0.iter() { - if collection_name == target_collection { - result.push(alias.clone()); - } - } - result + self.alias_mapping + .collection_aliases(collection_name) + .collect() } pub fn state(&self) -> &AliasMapping { diff --git a/lib/storage/src/content_manager/collection_meta_ops.rs b/lib/storage/src/content_manager/collection_meta_ops.rs index d3123fda5a..267592d369 100644 --- a/lib/storage/src/content_manager/collection_meta_ops.rs +++ b/lib/storage/src/content_manager/collection_meta_ops.rs @@ -12,7 +12,7 @@ use collection::operations::config_diff::{ WalConfigDiff, }; use collection::operations::types::{ - SparseVectorParams, SparseVectorsConfig, VectorsConfig, VectorsConfigDiff, + SparseVectorParams, SparseVectorsConfig, VectorParams, VectorsConfig, VectorsConfigDiff, }; use collection::operations::validation; use collection::shards::replica_set::replica_set_state::ReplicaState; @@ -24,7 +24,7 @@ use schemars::JsonSchema; use segment::data_types::vectors::DEFAULT_VECTOR_NAME; use segment::types::{ Payload, PayloadFieldSchema, PayloadKeyType, QuantizationConfig, ShardKey, StrictModeConfig, - VectorNameBuf, + VectorNameBuf, VectorsConfigDefaults, }; use serde::{Deserialize, Serialize}; use uuid::Uuid; @@ -193,6 +193,43 @@ pub struct CreateCollection { pub metadata: Option, } +/// Fill exactly one placement level, by precedence: request `memory`, request legacy `on_disk`, +/// default `memory`, default `on_disk`. Filling a lower level alongside a higher one would cause +/// spurious `memory`-vs-legacy mismatch warnings at resolution time. +pub fn apply_vector_placement_defaults( + params: &mut VectorParams, + defaults: &VectorsConfigDefaults, +) { + if params.memory.is_some() || params.on_disk.is_some() { + return; + } + + let &VectorsConfigDefaults { on_disk, memory } = defaults; + + if memory.is_some() { + params.memory = memory; + } else { + params.on_disk = on_disk; + } +} + +/// Service-level `payload.memory` default applies unless the request specifies `payload.memory` +/// or the legacy `on_disk_payload` flag. +pub fn apply_payload_placement_defaults( + payload: Option, + on_disk_payload: Option, + defaults: Option, +) -> Option { + if on_disk_payload.is_some() { + return payload; + } + + match (defaults, payload) { + (Some(defaults), Some(payload)) => Some(defaults.update(&payload)), + (defaults, payload) => payload.or(defaults), + } +} + /// Operation for creating new collection and (optionally) specify index params #[derive(Debug, Deserialize, Serialize, PartialEq, Eq, Hash, Clone)] #[serde(rename_all = "snake_case")] @@ -277,6 +314,10 @@ impl CreateCollectionOperation { self.distribution.is_some() } + pub fn distribution(&self) -> Option<&ShardDistributionProposal> { + self.distribution.as_ref() + } + pub fn take_distribution(&mut self) -> Option { self.distribution.take() } diff --git a/lib/storage/src/content_manager/consensus_state_machine/action.rs b/lib/storage/src/content_manager/consensus_state_machine/action.rs index bd3acac5ea..cb2cc2885b 100644 --- a/lib/storage/src/content_manager/consensus_state_machine/action.rs +++ b/lib/storage/src/content_manager/consensus_state_machine/action.rs @@ -1,5 +1,6 @@ use std::collections::{BTreeMap, BTreeSet}; +use collection::collection_state; use collection::operations::types::PeerMetadata; use collection::shards::CollectionId; use collection::shards::shard::PeerId; @@ -13,6 +14,15 @@ use crate::quota::QuotaConfig; /// A single change a consensus operation makes #[derive(Clone, Debug, PartialEq)] pub enum Action { + CreateCollection { + collection: CollectionId, + state: Box, + }, + + DropCollection { + collection: CollectionId, + }, + AddNamedVector { collection: CollectionId, vector_name: VectorNameBuf, @@ -67,7 +77,9 @@ impl Action { /// Collection this action changes, if it is scoped to one pub fn collection(&self) -> Option<&CollectionId> { match self { - Action::AddNamedVector { collection, .. } + Action::CreateCollection { collection, .. } + | Action::DropCollection { collection } + | Action::AddNamedVector { collection, .. } | Action::DropNamedVector { collection, .. } | Action::SetPayloadIndex { collection, .. } | Action::DropPayloadIndex { collection, .. } => Some(collection), diff --git a/lib/storage/src/content_manager/consensus_state_machine/mod.rs b/lib/storage/src/content_manager/consensus_state_machine/mod.rs index 4fd02b0f9a..b4b585e9f1 100644 --- a/lib/storage/src/content_manager/consensus_state_machine/mod.rs +++ b/lib/storage/src/content_manager/consensus_state_machine/mod.rs @@ -11,7 +11,9 @@ //! //! Two rules every operation follows: //! -//! 1. Never reject an operation that is partially or fully applied. +//! 1. Never reject a partially applied operation: rejecting one makes the partial state permanent. +//! An operation that was fully applied may be rejected, since the state is already complete. +//! E.g., `CreateCollection` rejects a collection that is already there. //! 2. Emit only the actions left to reach the goal state, each idempotent, and the action that //! records the operation as applied last. //! @@ -25,14 +27,24 @@ pub mod state; #[cfg(test)] mod tests; -use collection::shards::shard::PeerId; +use std::num::NonZeroU32; + +use collection::config::{ + self, CollectionConfigInternal, CollectionParams, PayloadStorageParams, ShardingMethod, + WalConfig, +}; +use collection::operations::config_diff::DiffConfig as _; +use collection::operations::types::VectorsConfig; +use collection::optimizers_builder::OptimizersConfig; +use collection::shards::shard::{PeerId, ShardId}; use collection::shards::transfer::ShardTransferMethod; use segment::data_types::collection_defaults::CollectionConfigDefaults; +use segment::types::HnswConfig; pub use self::action::Action; pub use self::state::ClusterState; use super::errors::StorageResult; -use crate::content_manager::collection_meta_ops::CollectionMetaOperations; +use crate::content_manager::collection_meta_ops::*; use crate::content_manager::consensus_ops::ConsensusOperations; use crate::content_manager::errors::StorageError; @@ -102,9 +114,15 @@ impl ConsensusStateMachine { match operation { CollectionMetaOperations::Nop { .. } => ApplyOutcome::Accepted(Vec::new()), - CollectionMetaOperations::CreateCollection(_) - | CollectionMetaOperations::UpdateCollection(_) - | CollectionMetaOperations::DeleteCollection(_) + CollectionMetaOperations::CreateCollection(operation) => { + ApplyOutcome::new(self.state.plan_create_collection(&self.context, operation)) + } + + CollectionMetaOperations::DeleteCollection(operation) => { + ApplyOutcome::Accepted(self.state.plan_delete_collection(operation)) + } + + CollectionMetaOperations::UpdateCollection(_) | CollectionMetaOperations::CreateShardKey(_) | CollectionMetaOperations::DropShardKey(_) | CollectionMetaOperations::SetShardReplicaState(_) @@ -146,7 +164,7 @@ impl ConsensusStateMachine { /// Node-local values operations read. /// /// These come from this node's config, not from consensus, so two peers can read different values -/// for the same operation. +/// for the same operation. Values scraped from `StorageConfig` keep the names they have there. #[derive(Clone, Debug)] pub struct NodeContext { pub peer_id: PeerId, @@ -155,6 +173,146 @@ pub struct NodeContext { pub collection_defaults: Option, /// Transfer method this node picks when an operation does not name one pub default_shard_transfer_method: Option, + pub max_collections: Option, + pub wal: WalConfig, + pub optimizers: OptimizersConfig, + pub hnsw_index: HnswConfig, + pub payload: Option, + /// Mirrors the deprecated storage config flag of the same name, which `payload` overrides + pub on_disk_payload: bool, +} + +impl NodeContext { + /// Shards a new collection starts with. + /// + /// The proposer picks them and the operation carries them. An operation proposed without a + /// distribution comes from a single node, which puts every shard on itself. + pub fn shard_distribution( + &self, + op: &CreateCollectionOperation, + ) -> Vec<(ShardId, Vec)> { + if let Some(distribution) = op.distribution() { + return distribution.distribution.clone(); + } + + match op.create_collection.sharding_method.unwrap_or_default() { + ShardingMethod::Auto => { + let shard_number = op.create_collection.shard_number.or_else(|| { + let defaults = self.collection_defaults.as_ref()?; + Some(defaults.get_shard_number(1)) + }); + + (0..shard_number.unwrap_or(1)) + .map(|shard_id| (shard_id, vec![self.peer_id])) + .collect() + } + + // Custom sharding creates shards with the shard key, not with the collection + ShardingMethod::Custom => Vec::new(), + } + } + + /// Resolve the config of a new collection from the operation and this node's defaults. + /// + /// `shards` is how many shards the collection starts with, which auto sharding stores as the + /// shard number when the operation names none. + #[expect(deprecated)] + pub fn collection_config( + &self, + op: &CreateCollection, + shards: usize, + ) -> StorageResult { + let CreateCollection { + mut vectors, + shard_number, + sharding_method, + on_disk_payload, + payload, + hnsw_config: hnsw_config_diff, + wal_config: wal_config_diff, + optimizers_config: optimizers_config_diff, + replication_factor, + write_consistency_factor, + quantization_config, + sparse_vectors, + strict_mode_config, + uuid, + metadata, + } = op.clone(); + + let defaults = self.collection_defaults.as_ref(); + + let shard_number = match sharding_method.unwrap_or_default() { + ShardingMethod::Auto => shard_number.unwrap_or(shards as u32), + ShardingMethod::Custom => shard_number.unwrap_or_else(|| { + defaults + .and_then(|defaults| defaults.shard_number) + .unwrap_or_else(|| config::default_shard_number().get()) + }), + }; + + let replication_factor = replication_factor + .or_else(|| defaults.and_then(|defaults| defaults.replication_factor)) + .unwrap_or_else(|| config::default_replication_factor().get()); + + let write_consistency_factor = write_consistency_factor + .or_else(|| defaults.and_then(|defaults| defaults.write_consistency_factor)) + .unwrap_or_else(|| config::default_write_consistency_factor().get()); + + if let Some(vectors_defaults) = defaults.and_then(|defaults| defaults.vectors.as_ref()) { + match &mut vectors { + VectorsConfig::Single(params) => { + apply_vector_placement_defaults(params, vectors_defaults); + } + VectorsConfig::Multi(params) => { + for params in params.values_mut() { + apply_vector_placement_defaults(params, vectors_defaults); + } + } + } + } + + let params = CollectionParams { + vectors, + sparse_vectors, + shard_number: NonZeroU32::new(shard_number) + .ok_or_else(|| StorageError::bad_input("`shard_number` cannot be 0"))?, + sharding_method, + on_disk_payload: Some(on_disk_payload.unwrap_or(self.on_disk_payload)), + payload: apply_payload_placement_defaults(payload, on_disk_payload, self.payload), + replication_factor: NonZeroU32::new(replication_factor) + .ok_or_else(|| StorageError::bad_input("`replication_factor` cannot be 0"))?, + write_consistency_factor: NonZeroU32::new(write_consistency_factor) + .ok_or_else(|| StorageError::bad_input("`write_consistency_factor` cannot be 0"))?, + read_fan_out_factor: None, + read_fan_out_delay_ms: None, + }; + + let quantization_config = quantization_config + .or_else(|| defaults.and_then(|defaults| defaults.quantization.clone())); + + let strict_mode_config = match strict_mode_config { + Some(diff) => { + let default_config = defaults + .and_then(|defaults| defaults.strict_mode.clone()) + .unwrap_or_default(); + + Some(default_config.update(&diff)) + } + None => defaults.and_then(|defaults| defaults.strict_mode.clone()), + }; + + Ok(CollectionConfigInternal { + params, + hnsw_config: self.hnsw_index.update_opt(hnsw_config_diff.as_ref()), + optimizer_config: self.optimizers.update_opt(optimizers_config_diff.as_ref()), + wal_config: self.wal.update_opt(wal_config_diff.as_ref()), + quantization_config, + strict_mode_config, + uuid, + metadata, + }) + } } /// What the machine decided about an operation. diff --git a/lib/storage/src/content_manager/consensus_state_machine/state/apply.rs b/lib/storage/src/content_manager/consensus_state_machine/state/apply.rs index 85047026fd..579bebb04c 100644 --- a/lib/storage/src/content_manager/consensus_state_machine/state/apply.rs +++ b/lib/storage/src/content_manager/consensus_state_machine/state/apply.rs @@ -4,11 +4,19 @@ use super::*; impl ClusterState { /// 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::CreateCollection { collection, state } => { + self.collections + .insert(collection.clone(), (**state).clone()); + } + + // Missing collection is legal here: the applier also deletes the directory + // a collection that failed to load leaves behind + Action::DropCollection { collection } => { + self.collections.remove(collection); + } + Action::AddNamedVector { collection, vector_name, @@ -103,6 +111,10 @@ impl ClusterState { } } + /// State of the collection an action changes. + /// + /// An action that changes a collection is never planned against a state without it, + /// so debug builds assert. fn collection_mut(&mut self, collection: &str) -> Option<&mut collection_state::State> { let state = self.collections.get_mut(collection); diff --git a/lib/storage/src/content_manager/consensus_state_machine/state/plan.rs b/lib/storage/src/content_manager/consensus_state_machine/state/plan.rs index 94d26286da..3b0ac01f5e 100644 --- a/lib/storage/src/content_manager/consensus_state_machine/state/plan.rs +++ b/lib/storage/src/content_manager/consensus_state_machine/state/plan.rs @@ -1,16 +1,114 @@ use std::collections::{BTreeMap, BTreeSet}; use collection::collection::vector_name_schema; +use collection::collection_state::ShardInfo; use collection::operations::types::PeerMetadata; +use collection::shards::replica_set::replica_set_state::ReplicaState; use collection::shards::shard::PeerId; use super::*; use crate::content_manager::collection_meta_ops::*; -use crate::content_manager::consensus_state_machine::Action; +use crate::content_manager::consensus_state_machine::{Action, NodeContext}; type Actions = Vec; impl ClusterState { + /// One action: `Collection::new` saves the config as its last step, and a collection whose + /// config is missing does not load, so creation is atomic already. + /// + /// The config is resolved here, from the operation and node-local defaults, because + /// `TableOfContent::create_collection` resolves it before it writes anything. + pub fn plan_create_collection( + &self, + context: &NodeContext, + op: &CreateCollectionOperation, + ) -> StorageResult { + let collection = &op.collection_name; + + if self.has_collection(collection) { + return Err(StorageError::already_exists(format!( + "Collection `{collection}` already exists!" + ))); + } + + if let Some(max_collections) = context.max_collections + && self.collections.len() >= max_collections + { + return Err(StorageError::bad_request(format!( + "Can't create collection with name {collection}. \ + Max collections limit reached: {max_collections}", + ))); + } + + if self.aliases.get(collection).is_some() { + return Err(StorageError::bad_input(format!( + "Can't create collection with name {collection}. \ + Alias with the same name already exists", + ))); + } + + let distribution = context.shard_distribution(op); + let config = context.collection_config(&op.create_collection, distribution.len())?; + + // Every replica of a new collection starts `Initializing`, and the peer that has a local + // one proposes `SetShardReplicaState` once the shard is built + let shards = distribution + .into_iter() + .map(|(shard_id, peers)| { + let replicas = peers + .into_iter() + .map(|peer_id| (peer_id, ReplicaState::Initializing)) + .collect(); + + (shard_id, ShardInfo { replicas }) + }) + .collect(); + + let state = collection_state::State { + config, + shards, + resharding: None, + transfers: Default::default(), + // Shard keys are set up after the collection is created + shards_key_mapping: Default::default(), + payload_index_schema: Default::default(), + }; + + Ok(vec![Action::CreateCollection { + collection: collection.clone(), + state: Box::new(state), + }]) + } + + pub fn plan_delete_collection(&self, op: &DeleteCollectionOperation) -> Actions { + let DeleteCollectionOperation(collection) = op; + + // Collection name is *not* resolved through aliases, `DeleteCollection` must name existing + // collection directly + let remove: BTreeSet<_> = self.aliases.collection_aliases(collection).collect(); + + // Remove aliases first, then collection itself. + // Either order is fine, this one simply follows the order `ToC::delete_collection` uses. + let mut actions = Actions::new(); + + // Collection without aliases does not produce an empty `UpdateAliases` action + // (similar to `plan_change_aliases`) + if !remove.is_empty() { + actions.push(Action::UpdateAliases { + set: Default::default(), + remove, + }); + } + + // Produce `DropCollection` action, even if collection does not exist: + // it removes leftover aliases and storage directory + actions.push(Action::DropCollection { + collection: collection.clone(), + }); + + actions + } + pub fn plan_create_named_vector(&self, op: &CreateNamedVector) -> StorageResult { let CreateNamedVector { collection_name, diff --git a/lib/storage/src/content_manager/consensus_state_machine/tests/mod.rs b/lib/storage/src/content_manager/consensus_state_machine/tests/mod.rs index 5ba5537bcb..4cd1fd451e 100644 --- a/lib/storage/src/content_manager/consensus_state_machine/tests/mod.rs +++ b/lib/storage/src/content_manager/consensus_state_machine/tests/mod.rs @@ -17,17 +17,48 @@ use shard::operations::VectorNameConfig; use super::*; const PEER_ID: u64 = 42; +const OTHER_PEER_ID: u64 = 43; fn state_machine(state: ClusterState) -> ConsensusStateMachine { - ConsensusStateMachine::new( - state, - NodeContext { - peer_id: PEER_ID, - is_distributed: true, - collection_defaults: None, - default_shard_transfer_method: None, - }, - ) + ConsensusStateMachine::new(state, node_context()) +} + +/// Node config is fixed, except for the parts used in tests. Extend it as needed. +fn node_context() -> NodeContext { + NodeContext { + peer_id: PEER_ID, + is_distributed: true, + collection_defaults: None, + default_shard_transfer_method: None, + max_collections: None, + wal: Default::default(), + optimizers: optimizers_config(), + hnsw_index: Default::default(), + payload: None, + on_disk_payload: false, + } +} + +/// Request leaving everything to defaults, for a test to set the parts it covers +fn create_collection_request() -> CreateCollection { + CreateCollection { + vectors: VectorsConfig::Multi(Default::default()), + shard_number: None, + sharding_method: None, + replication_factor: None, + write_consistency_factor: None, + #[expect(deprecated)] + on_disk_payload: None, + payload: None, + hnsw_config: None, + wal_config: None, + optimizers_config: None, + quantization_config: None, + sparse_vectors: None, + strict_mode_config: None, + uuid: None, + metadata: None, + } } fn collection_state(vectors: Vec<(VectorNameBuf, VectorNameConfig)>) -> collection_state::State { @@ -53,18 +84,7 @@ fn collection_config(params: CollectionParams) -> CollectionConfigInternal { CollectionConfigInternal { params, hnsw_config: Default::default(), - optimizer_config: OptimizersConfig { - deleted_threshold: 0.1, - vacuum_min_vector_number: 1000, - default_segment_number: 0, - max_segment_size: None, - #[expect(deprecated)] - memmap_threshold: None, - indexing_threshold: Some(100_000), - flush_interval_sec: 60, - max_optimization_threads: Some(0), - prevent_unoptimized: None, - }, + optimizer_config: optimizers_config(), wal_config: Default::default(), quantization_config: None, strict_mode_config: None, @@ -73,6 +93,22 @@ fn collection_config(params: CollectionParams) -> CollectionConfigInternal { } } +/// Same config a collection created by the machine gets, since `NodeContext` carries it +fn optimizers_config() -> OptimizersConfig { + OptimizersConfig { + deleted_threshold: 0.1, + vacuum_min_vector_number: 1000, + default_segment_number: 0, + max_segment_size: None, + #[expect(deprecated)] + memmap_threshold: None, + indexing_threshold: Some(100_000), + flush_interval_sec: 60, + max_optimization_threads: Some(0), + prevent_unoptimized: None, + } +} + fn collection_params(vectors: VectorsConfig) -> CollectionParams { CollectionParams { vectors, diff --git a/lib/storage/src/content_manager/consensus_state_machine/tests/ops.rs b/lib/storage/src/content_manager/consensus_state_machine/tests/ops.rs index 2c3710c459..697e9d9e28 100644 --- a/lib/storage/src/content_manager/consensus_state_machine/tests/ops.rs +++ b/lib/storage/src/content_manager/consensus_state_machine/tests/ops.rs @@ -3,7 +3,13 @@ use std::collections::{BTreeMap, BTreeSet, HashMap}; +use ahash::AHashMap; +use collection::collection_state::ShardInfo; +use collection::config::ShardingMethod; use collection::operations::types::PeerMetadata; +use collection::shards::replica_set::replica_set_state::ReplicaState; +use collection::shards::shard::ShardId; +use segment::data_types::collection_defaults::CollectionConfigDefaults; use segment::data_types::vector_name_config::*; use segment::types::*; use serde_json::{Value, json}; @@ -13,6 +19,7 @@ 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; +use crate::content_manager::shard_distribution::ShardDistributionProposal; use crate::quota::QuotaConfig; const COLLECTION: &str = "alpha"; @@ -34,6 +41,273 @@ fn nop() { assert_eq!(machine.state(), &state); } +#[test] +fn create_collection() { + let mut machine = state_machine(ClusterState::default()); + let outcome = machine.apply(&create_collection_op( + create_collection_request(), + Some(vec![vec![PEER_ID]]), + )); + + let ApplyOutcome::Accepted(actions) = outcome else { + panic!("creating a collection should be accepted, got {outcome:?}"); + }; + + assert!(matches!( + actions.as_slice(), + [Action::CreateCollection { .. }] + )); + + // Config comes from the node, and the one shard the proposer placed is local and initializing + let mut expected = collection_state(Vec::new()); + expected.shards = shards(vec![vec![PEER_ID]]); + + assert_eq!(machine.state().collection(COLLECTION), Some(&expected)); +} + +#[test] +fn create_collection_distribution() { + let mut machine = state_machine(ClusterState::default()); + let outcome = machine.apply(&create_collection_op( + create_collection_request(), + Some(vec![vec![PEER_ID, OTHER_PEER_ID], vec![OTHER_PEER_ID]]), + )); + + let ApplyOutcome::Accepted(_) = outcome else { + panic!("creating a collection should be accepted, got {outcome:?}"); + }; + + let state = machine.state().collection(COLLECTION).expect("created"); + + assert_eq!( + state.shards, + shards(vec![vec![PEER_ID, OTHER_PEER_ID], vec![OTHER_PEER_ID],]) + ); + + // Auto sharding takes the shard number from the distribution + assert_eq!(state.config.params.shard_number.get(), 2); +} + +#[test] +fn create_collection_custom_sharding() { + let mut create_collection = create_collection_request(); + create_collection.sharding_method = Some(ShardingMethod::Custom); + + let mut machine = state_machine(ClusterState::default()); + machine.apply(&create_collection_op(create_collection, None)); + + let state = machine.state().collection(COLLECTION).expect("created"); + + // Shards are created with the shard key, and the shard number counts shards per key + assert!(state.shards.is_empty()); + assert_eq!(state.config.params.shard_number.get(), 1); +} + +#[test] +fn create_collection_defaults() { + let mut context = node_context(); + context.collection_defaults = Some(CollectionConfigDefaults { + shard_number: Some(3), + replication_factor: Some(2), + shard_number_per_node: None, + write_consistency_factor: None, + vectors: None, + quantization: None, + strict_mode: None, + }); + + // Without a distribution this node places every shard on itself, as many as it defaults to + let mut machine = ConsensusStateMachine::new(ClusterState::default(), context); + machine.apply(&create_collection_op(create_collection_request(), None)); + + let state = machine.state().collection(COLLECTION).expect("created"); + + assert_eq!( + state.shards, + shards(vec![vec![PEER_ID], vec![PEER_ID], vec![PEER_ID]]), + ); + + assert_eq!(state.config.params.shard_number.get(), 3); + assert_eq!(state.config.params.replication_factor.get(), 2); +} + +#[test] +fn create_collection_reject_existing() { + let state = cluster_state(Vec::new()); + + let mut machine = state_machine(state.clone()); + let outcome = machine.apply(&create_collection_op( + create_collection_request(), + Some(vec![vec![PEER_ID]]), + )); + + assert!(matches!( + outcome, + ApplyOutcome::Rejected(StorageError::AlreadyExists { .. }) + )); + + assert_eq!(machine.state(), &state); +} + +#[test] +fn create_collection_reject_alias_name() { + let mut state = ClusterState::default(); + state.aliases.insert(COLLECTION.into(), "beta".into()); + + let mut machine = state_machine(state.clone()); + let outcome = machine.apply(&create_collection_op( + create_collection_request(), + Some(vec![vec![PEER_ID]]), + )); + + assert!(matches!( + outcome, + ApplyOutcome::Rejected(StorageError::BadInput { .. }) + )); + + assert_eq!(machine.state(), &state); +} + +#[test] +fn create_collection_reject_max_collections() { + let mut context = node_context(); + context.max_collections = Some(1); + + let mut state = ClusterState::default(); + state + .collections + .insert("beta".into(), collection_state(Vec::new())); + + let mut machine = ConsensusStateMachine::new(state.clone(), context); + let outcome = machine.apply(&create_collection_op( + create_collection_request(), + Some(vec![vec![PEER_ID]]), + )); + + assert!(matches!( + outcome, + ApplyOutcome::Rejected(StorageError::BadRequest { .. }) + )); + + assert_eq!(machine.state(), &state); +} + +#[test] +fn create_collection_reject_zero_shards() { + let state = ClusterState::default(); + + // Auto sharding takes the shard number from the distribution, and zero is not a shard number + let mut machine = state_machine(state.clone()); + let outcome = machine.apply(&create_collection_op( + create_collection_request(), + Some(Vec::new()), + )); + + assert!(matches!( + outcome, + ApplyOutcome::Rejected(StorageError::BadInput { .. }) + )); + + assert_eq!(machine.state(), &state); +} + +#[test] +fn delete_collection() { + let state = cluster_state(Vec::new()); + + let mut machine = state_machine(state); + let outcome = machine.apply(&delete_collection_op(COLLECTION)); + + let ApplyOutcome::Accepted(actions) = outcome else { + panic!("deleting a collection should be accepted, got {outcome:?}"); + }; + + assert!(matches!( + actions.as_slice(), + [Action::DropCollection { .. }] + )); + + assert!(!machine.state().has_collection(COLLECTION)); +} + +#[test] +fn delete_collection_aliases() { + let mut state = cluster_state(Vec::new()); + state + .collections + .insert("beta".into(), collection_state(Vec::new())); + state.aliases.insert("second".into(), COLLECTION.into()); + state.aliases.insert("first".into(), COLLECTION.into()); + state.aliases.insert("other".into(), "beta".into()); + + let mut machine = state_machine(state); + let outcome = machine.apply(&delete_collection_op(COLLECTION)); + + let ApplyOutcome::Accepted(actions) = outcome else { + panic!("deleting a collection should be accepted, got {outcome:?}"); + }; + + // Aliases of the collection go first, in name order, and the collection last + assert_eq!( + actions, + vec![ + Action::UpdateAliases { + set: BTreeMap::new(), + remove: BTreeSet::from(["first".into(), "second".into()]), + }, + Action::DropCollection { + collection: COLLECTION.into(), + }, + ], + ); + + let aliases = &machine.state().aliases; + + assert_eq!(aliases.get("other").map(String::as_str), Some("beta")); +} + +#[test] +fn delete_collection_replay() { + let state = ClusterState::default(); + + let mut machine = state_machine(state.clone()); + let outcome = machine.apply(&delete_collection_op(COLLECTION)); + + let ApplyOutcome::Accepted(actions) = outcome else { + panic!("replay of an applied delete should be accepted, got {outcome:?}"); + }; + + // Action is emitted even if state already matches + assert!(matches!( + actions.as_slice(), + [Action::DropCollection { .. }] + )); + + assert_eq!(machine.state(), &state, "replay should not change anything"); +} + +#[test] +fn delete_collection_alias_name() { + let mut state = cluster_state(Vec::new()); + state.aliases.insert("alias".into(), COLLECTION.into()); + + let mut machine = state_machine(state.clone()); + let outcome = machine.apply(&delete_collection_op("alias")); + + let ApplyOutcome::Accepted(actions) = outcome else { + panic!("deleting a collection by alias should be accepted, got {outcome:?}"); + }; + + // The name is not resolved: the alias names no collection to delete, and it points at + // `alpha`, not at itself, so it stays + assert!(matches!( + actions.as_slice(), + [Action::DropCollection { .. }] + )); + + assert_eq!(machine.state(), &state); +} + #[test] fn create_alias() { let state = cluster_state(Vec::new()); @@ -912,6 +1186,50 @@ fn collection_meta_op(op: CollectionMetaOperations) -> ConsensusOperations { ConsensusOperations::CollectionMeta(Box::new(op)) } +/// `placement` names the peers of every shard, the way a proposer sets them. +/// Without it the operation reaches the machine the way a single node proposes it. +fn create_collection_op( + create_collection: CreateCollection, + placement: Option>>, +) -> ConsensusOperations { + let mut operation = CreateCollectionOperation::new(COLLECTION.into(), create_collection) + .expect("valid operation"); + + if let Some(placement) = placement { + operation.set_distribution(ShardDistributionProposal { + distribution: placement + .into_iter() + .enumerate() + .map(|(idx, peers)| (idx as ShardId, peers)) + .collect(), + }); + } + + collection_meta_op(CollectionMetaOperations::CreateCollection(operation)) +} + +/// Shards a collection is created with, one replica per peer of each placement entry +fn shards(placement: Vec>) -> AHashMap { + placement + .into_iter() + .enumerate() + .map(|(idx, peers)| { + let replicas = peers + .into_iter() + .map(|peer_id| (peer_id, ReplicaState::Initializing)) + .collect(); + + (idx as ShardId, ShardInfo { replicas }) + }) + .collect() +} + +fn delete_collection_op(collection: &str) -> ConsensusOperations { + collection_meta_op(CollectionMetaOperations::DeleteCollection( + DeleteCollectionOperation(collection.into()), + )) +} + fn change_aliases_op(actions: Vec) -> ConsensusOperations { collection_meta_op(CollectionMetaOperations::ChangeAliases( ChangeAliasesOperation { actions }, diff --git a/lib/storage/src/content_manager/consensus_state_machine/tests/prop.rs b/lib/storage/src/content_manager/consensus_state_machine/tests/prop.rs index cc013f1cc7..dbb4292f84 100644 --- a/lib/storage/src/content_manager/consensus_state_machine/tests/prop.rs +++ b/lib/storage/src/content_manager/consensus_state_machine/tests/prop.rs @@ -3,8 +3,10 @@ use std::collections::HashMap; use collection::collection_state; +use collection::config::ShardingMethod; use collection::operations::types::PeerMetadata; use collection::shards::CollectionId; +use collection::shards::shard::ShardId; use proptest::prelude::*; use segment::data_types::modifier::Modifier; use segment::data_types::vector_name_config::*; @@ -16,6 +18,7 @@ use crate::content_manager::alias_mapping::AliasMapping; 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::shard_distribution::ShardDistributionProposal; use crate::quota::QuotaConfig; use crate::types::PeerMetadataById; @@ -29,7 +32,7 @@ const VECTOR_NAMES: &[&str] = &["", "text", "image"]; const FIELD_NAMES: &[&str] = &["city", "count", "nested.key"]; /// This node, and one other peer -const PEER_IDS: &[PeerId] = &[PEER_ID, 43]; +const PEER_IDS: &[PeerId] = &[PEER_ID, OTHER_PEER_ID]; const PEER_VERSIONS: &[&str] = &["1.14.0", "1.15.0"]; const METADATA_KEYS: &[&str] = &["region", "tier"]; @@ -170,7 +173,7 @@ pub fn arb_consensus_operation( // Weighted by how many operations each arm covers, so one operation is as likely as another prop_oneof![ - 6 => collection_meta, + 8 => collection_meta, 1 => arb_update_peer_metadata(), 1 => arb_update_cluster_metadata(), 1 => arb_quota_config().prop_map(ConsensusOperations::SetQuotaConfig), @@ -184,6 +187,8 @@ fn arb_collection_meta_operation( prop_oneof![ Just(CollectionMetaOperations::Nop { token: 0 }), + arb_create_collection(collection_names.clone()), + arb_delete_collection(collection_names.clone()), arb_change_aliases(collection_names.clone()), arb_create_named_vector(collection_names.clone()), arb_delete_named_vector(collection_names.clone()), @@ -209,6 +214,60 @@ fn arb_collection_name(names: Vec) -> impl Strategy { proptest::sample::select(names) } +/// Only the sharding method and the distribution vary: everything else the operation carries goes +/// into the config unchanged, or picks up a node-local default that no generated state sets. +fn arb_create_collection( + collections: Vec, +) -> impl Strategy { + let sharding_method = proptest::option::of(prop_oneof![ + Just(ShardingMethod::Auto), + Just(ShardingMethod::Custom), + ]); + + ( + arb_collection_name(collections), + sharding_method, + arb_shard_distribution(), + ) + .prop_map(|(collection_name, sharding_method, distribution)| { + let mut create_collection = create_collection_request(); + create_collection.sharding_method = sharding_method; + + let mut operation = CreateCollectionOperation::new(collection_name, create_collection) + .expect("valid operation"); + + if let Some(distribution) = distribution { + operation.set_distribution(distribution); + } + + CollectionMetaOperations::CreateCollection(operation) + }) +} + +/// The proposer picks the distribution, and only a single node proposes without one. +/// An empty one leaves auto sharding with no shards, which is rejected. +fn arb_shard_distribution() -> impl Strategy> { + let placement = proptest::collection::vec(proptest::collection::vec(arb_peer_id(), 1..3), 0..3); + + let distribution = placement.prop_map(|placement| ShardDistributionProposal { + distribution: placement + .into_iter() + .enumerate() + .map(|(idx, peers)| (idx as ShardId, peers)) + .collect(), + }); + + proptest::option::of(distribution) +} + +fn arb_delete_collection( + collections: Vec, +) -> impl Strategy { + arb_collection_name(collections).prop_map(|collection_name| { + CollectionMetaOperations::DeleteCollection(DeleteCollectionOperation(collection_name)) + }) +} + fn arb_change_aliases(collections: Vec) -> impl Strategy { let actions = proptest::collection::vec(arb_alias_operation(collections), 1..=4); diff --git a/lib/storage/src/content_manager/toc/create_collection.rs b/lib/storage/src/content_manager/toc/create_collection.rs index d4fcaf57de..b0e7621162 100644 --- a/lib/storage/src/content_manager/toc/create_collection.rs +++ b/lib/storage/src/content_manager/toc/create_collection.rs @@ -6,15 +6,12 @@ use std::num::NonZeroU32; use std::sync::Arc; use collection::collection::Collection; -use collection::config::{ - self, CollectionConfigInternal, CollectionParams, PayloadStorageParams, ShardingMethod, -}; +use collection::config::{self, CollectionConfigInternal, CollectionParams, ShardingMethod}; use collection::operations::config_diff::DiffConfig as _; -use collection::operations::types::{CollectionResult, VectorParams, VectorsConfig}; +use collection::operations::types::{CollectionResult, VectorsConfig}; use collection::shards::collection_shard_distribution::CollectionShardDistribution; use collection::shards::replica_set::replica_set_state::ReplicaState; use collection::shards::shard::{PeerId, ShardId}; -use segment::types::VectorsConfigDefaults; use super::{COLLECTION_DELETE_SPIN_INTERVAL, COLLECTION_DELETE_WAIT_TIMEOUT, TableOfContent}; use crate::common::utils::try_unwrap_with_timeout_async; @@ -23,37 +20,6 @@ use crate::content_manager::collections_ops::Checker as _; use crate::content_manager::consensus_ops::ConsensusOperations; use crate::content_manager::errors::StorageError; -/// Fill exactly one placement level, by precedence: request `memory`, request legacy `on_disk`, -/// default `memory`, default `on_disk`. Filling a lower level alongside a higher one would cause -/// spurious `memory`-vs-legacy mismatch warnings at resolution time. -fn apply_vector_placement_defaults(params: &mut VectorParams, defaults: &VectorsConfigDefaults) { - let VectorsConfigDefaults { on_disk, memory } = defaults; - if params.memory.is_some() || params.on_disk.is_some() { - return; - } - if memory.is_some() { - params.memory = *memory; - } else { - params.on_disk = *on_disk; - } -} - -/// Service-level `payload.memory` default applies unless the request specifies `payload.memory` -/// or the legacy `on_disk_payload` flag. -fn apply_payload_placement_defaults( - payload: Option, - on_disk_payload: Option, - defaults: Option, -) -> Option { - if on_disk_payload.is_some() { - return payload; - } - match (defaults, payload) { - (Some(defaults), Some(payload)) => Some(defaults.update(&payload)), - (defaults, payload) => payload.or(defaults), - } -} - impl TableOfContent { pub(super) async fn create_collection( &self,