From de42b2c7b951c12f766613be498e0ef3a882d275 Mon Sep 17 00:00:00 2001 From: Roman Titov Date: Thu, 17 Sep 2026 19:40:05 +0900 Subject: [PATCH] Implement `CreateShardKey`/`RemoveShardKey` for consensus state machine (#10666) --- lib/collection/src/operations/types.rs | 5 + .../content_manager/collection_meta_ops.rs | 4 + .../consensus_state_machine/action.rs | 45 +- .../consensus_state_machine/mod.rs | 47 +- .../consensus_state_machine/state/apply.rs | 64 +++ .../consensus_state_machine/state/mod.rs | 18 + .../consensus_state_machine/state/plan.rs | 150 +++++- .../consensus_state_machine/tests/mod.rs | 1 + .../consensus_state_machine/tests/ops.rs | 455 +++++++++++++++++- .../consensus_state_machine/tests/prop.rs | 146 +++++- .../consensus_state_machine/tests/state.rs | 64 +++ .../toc/collection_meta_ops.rs | 4 - 12 files changed, 969 insertions(+), 34 deletions(-) create mode 100644 lib/storage/src/content_manager/consensus_state_machine/tests/state.rs diff --git a/lib/collection/src/operations/types.rs b/lib/collection/src/operations/types.rs index ed27f6762a..c56257881d 100644 --- a/lib/collection/src/operations/types.rs +++ b/lib/collection/src/operations/types.rs @@ -1961,6 +1961,11 @@ impl PeerMetadata { pub fn is_different_version(&self) -> bool { self.version != *defaults::QDRANT_VERSION } + + /// Version reported by the peer + pub fn version(&self) -> &Version { + &self.version + } } #[cfg(test)] diff --git a/lib/storage/src/content_manager/collection_meta_ops.rs b/lib/storage/src/content_manager/collection_meta_ops.rs index 978b797cfd..90889fc6a2 100644 --- a/lib/storage/src/content_manager/collection_meta_ops.rs +++ b/lib/storage/src/content_manager/collection_meta_ops.rs @@ -3,6 +3,7 @@ #![allow(deprecated)] use std::collections::BTreeMap; +use std::sync::LazyLock; use collection::config::{ CollectionConfigInternal, CollectionParams, IdTrackerParams, PayloadStorageParams, @@ -37,6 +38,9 @@ pub use super::staging::{TestSlowDown, TestTransientError}; use crate::content_manager::errors::{StorageError, StorageResult}; use crate::content_manager::shard_distribution::ShardDistributionProposal; +pub(crate) static CREATE_CUSTOM_SHARDS_IN_INITIALIZING_STATE: LazyLock = + LazyLock::new(|| semver::Version::parse("1.14.2-dev").unwrap()); + // *Operation wrapper structure is only required for better OpenAPI generation /// Create alternative name for a collection. 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 b2066a7cfb..be4467bf15 100644 --- a/lib/storage/src/content_manager/consensus_state_machine/action.rs +++ b/lib/storage/src/content_manager/consensus_state_machine/action.rs @@ -8,9 +8,10 @@ use collection::operations::config_diff::{ }; use collection::operations::types::{PeerMetadata, SparseVectorsConfig, VectorsConfigDiff}; use collection::shards::CollectionId; -use collection::shards::shard::PeerId; +use collection::shards::replica_set::replica_set_state::ReplicaState; +use collection::shards::shard::{PeerId, ShardId}; use segment::types::{ - Payload, PayloadFieldSchema, PayloadKeyType, QuantizationConfig, StrictModeConfig, + Payload, PayloadFieldSchema, PayloadKeyType, QuantizationConfig, ShardKey, StrictModeConfig, VectorNameBuf, }; use shard::operations::vector_name_ops::VectorNameConfig; @@ -60,6 +61,39 @@ pub enum Action { field_name: PayloadKeyType, }, + /// Build a shard's replica set on disk. The shard becomes visible through `RegisterShards`. + CreateShard { + collection: CollectionId, + shard_id: ShardId, + shard_key: Option, + replicas: Vec, + init_state: ReplicaState, + }, + + /// Register built shards and record them under `shard_key` in one mapping write. + RegisterShards { + collection: CollectionId, + shard_key: Option, + shards: Vec<(ShardId, Vec, ReplicaState)>, + }, + + /// Stop cleanup tasks before their shard directories disappear + InvalidateCleanLocalShards { + collection: CollectionId, + shard_ids: Vec, + }, + + /// Persist the replay gate before dropping shard directories + RemoveShardKey { + collection: CollectionId, + shard_key: ShardKey, + }, + + DropShard { + collection: CollectionId, + shard_id: ShardId, + }, + UpdateAliases { set: BTreeMap, remove: BTreeSet, @@ -93,8 +127,13 @@ impl Action { pub fn collection(&self) -> Option<&CollectionId> { match self { Action::CreateCollection { collection, .. } - | Action::DropCollection { collection } | Action::UpdateCollectionConfig { collection, .. } + | Action::DropCollection { collection } + | Action::CreateShard { collection, .. } + | Action::RegisterShards { collection, .. } + | Action::InvalidateCleanLocalShards { collection, .. } + | Action::RemoveShardKey { collection, .. } + | Action::DropShard { collection, .. } | Action::AddNamedVector { collection, .. } | Action::DropNamedVector { collection, .. } | Action::SetPayloadIndex { 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 592ee2b5c2..e973270607 100644 --- a/lib/storage/src/content_manager/consensus_state_machine/mod.rs +++ b/lib/storage/src/content_manager/consensus_state_machine/mod.rs @@ -38,7 +38,7 @@ use collection::shards::CollectionId; use collection::shards::shard::{PeerId, ShardId}; use collection::shards::transfer::ShardTransferMethod; use segment::data_types::collection_defaults::CollectionConfigDefaults; -use segment::types::HnswConfig; +use segment::types::{HnswConfig, ShardKey}; pub use self::action::{Action, CollectionConfigDiff, apply_collection_config_diffs}; pub use self::state::ClusterState; @@ -131,11 +131,7 @@ impl ConsensusStateMachine { } CollectionMetaOperations::UpdateCollection(operation) => { - // TODO: - // - // Replica changes remove a replica *and* abort its transfers and resharding, - // so they need `Transfer::Abort`/`Resharding::Abort` to be implemented first - + // TODO: Removing replica may abort transfers and resharding, which are not implemented yet if operation.has_shard_replica_changes() { ApplyOutcome::NotCovered } else { @@ -147,16 +143,27 @@ impl ConsensusStateMachine { ApplyOutcome::Accepted(self.state.plan_delete_collection(operation)) } - CollectionMetaOperations::CreateShardKey(_) - | CollectionMetaOperations::DropShardKey(_) - | CollectionMetaOperations::SetShardReplicaState(_) - | CollectionMetaOperations::TransferShard(_, _) - | CollectionMetaOperations::Resharding(_, _) => ApplyOutcome::NotCovered, - CollectionMetaOperations::ChangeAliases(operation) => { ApplyOutcome::new(self.state.plan_change_aliases(operation)) } + CollectionMetaOperations::CreateShardKey(operation) => { + ApplyOutcome::new(self.state.plan_create_shard_key(&self.context, operation)) + } + + CollectionMetaOperations::DropShardKey(operation) => { + // TODO: Dropping shard key may abort resharding, which are not implemented yet + if self.is_reshardng(&operation.collection_name, Some(&operation.shard_key)) { + ApplyOutcome::NotCovered + } else { + ApplyOutcome::new(self.state.plan_drop_shard_key(operation)) + } + } + + CollectionMetaOperations::SetShardReplicaState(_) + | CollectionMetaOperations::TransferShard(_, _) + | CollectionMetaOperations::Resharding(_, _) => ApplyOutcome::NotCovered, + CollectionMetaOperations::CreateNamedVector(operation) => { ApplyOutcome::new(self.state.plan_create_named_vector(operation)) } @@ -183,6 +190,22 @@ impl ConsensusStateMachine { } } } + + fn is_reshardng(&self, collection: &str, shard_key: Option<&ShardKey>) -> bool { + let Ok(collection_name) = self.state.resolve_collection(collection) else { + return false; + }; + + let Some(collection) = self.state.collection(&collection_name) else { + return false; + }; + + let Some(resharding) = collection.resharding.as_ref() else { + return false; + }; + + resharding.shard_key.as_ref() == shard_key + } } /// Node-local values operations read. 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 e0747801a1..d163bbe388 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 @@ -1,4 +1,5 @@ use collection::collection::vector_name_schema; +use collection::collection_state::ShardInfo; use super::*; @@ -91,6 +92,69 @@ impl ClusterState { state.payload_index_schema.schema.remove(field_name); } + // Builds a shard directory. Shards join collection state through `RegisterShards`. + Action::CreateShard { .. } => {} + + Action::RegisterShards { + collection, + shard_key, + shards, + } => { + let Some(state) = self.collection_mut(collection) else { + return; + }; + + for &(shard_id, ref peers, init_state) in shards { + let replicas = peers.iter().map(|&peer_id| (peer_id, init_state)).collect(); + state.shards.insert(shard_id, ShardInfo { replicas }); + } + + if let Some(shard_key) = shard_key { + state + .shards_key_mapping + .entry(shard_key.clone()) + .or_default() + .extend(shards.iter().map(|&(shard_id, _, _)| shard_id)); + } + } + + // Stops node-local tasks and does not change consensus state + Action::InvalidateCleanLocalShards { .. } => {} + + Action::RemoveShardKey { + collection, + shard_key, + } => { + let Some(state) = self.collection_mut(collection) else { + return; + }; + + let Some(shard_ids) = state.shards_key_mapping.remove(shard_key) else { + return; + }; + + // The mapping is removed before the shard directories are deleted. If the node + // crashes between those steps, Qdrant ignores the leftover directories during + // startup because their shards are no longer in the mapping. + // + // Remove the shards from modeled state here to match the state after restart. + // Replaying the operation then has nothing left to do. + for shard_id in shard_ids { + state.shards.remove(&shard_id); + } + } + + Action::DropShard { + collection, + shard_id, + } => { + let Some(state) = self.collection_mut(collection) else { + return; + }; + + state.shards.remove(shard_id); + } + Action::UpdateAliases { set, remove } => { for alias in remove { self.aliases.remove(alias); diff --git a/lib/storage/src/content_manager/consensus_state_machine/state/mod.rs b/lib/storage/src/content_manager/consensus_state_machine/state/mod.rs index e8217608ec..f616a229af 100644 --- a/lib/storage/src/content_manager/consensus_state_machine/state/mod.rs +++ b/lib/storage/src/content_manager/consensus_state_machine/state/mod.rs @@ -5,6 +5,7 @@ use std::collections::HashMap; use collection::collection_state; use collection::shards::CollectionId; +use semver::Version; use super::Action; use crate::content_manager::alias_mapping::AliasMapping; @@ -49,4 +50,21 @@ impl ClusterState { Ok(resolved) } + + /// Whether every known peer runs at least `version`. + /// Implementation intentionally matches `ChannelService::all_peers_at_version`. + pub fn all_peers_at_version(&self, version: &Version) -> bool { + // TODO: + // Check that peer address map and peer metadata map contain the same peers + // and each peer matches required version + + // More peer addresses than metadata means at least one version is unknown + if self.peer_address_by_id.len() > self.peer_metadata_by_id.len() { + return false; + } + + self.peer_metadata_by_id + .values() + .all(|metadata| metadata.version() >= version) + } } 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 7cd3b031fe..f290b9cb5c 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 @@ -2,9 +2,10 @@ use std::collections::{BTreeMap, BTreeSet}; use collection::collection::vector_name_schema; 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::PeerId; +use collection::shards::shard::{PeerId, ShardId}; use super::*; use crate::content_manager::collection_meta_ops::*; @@ -278,6 +279,153 @@ impl ClusterState { }]) } + pub fn plan_create_shard_key( + &self, + context: &NodeContext, + op: &CreateShardKey, + ) -> StorageResult { + let CreateShardKey { + collection_name, + shard_key, + placement, + initial_state, + } = op; + + let collection = self.resolve_collection(collection_name)?; + let collection_state = self.collection(&collection).expect("collection exists"); + + let sharding_method = collection_state + .config + .params + .sharding_method + .unwrap_or_default(); + + if sharding_method != ShardingMethod::Custom { + return Err(StorageError::bad_request(format!( + "Shard Key {shard_key} cannot be created with Auto sharding method" + ))); + } + + if collection_state.shards_key_mapping.contains_key(shard_key) { + return Err(StorageError::bad_request(format!( + "Shard key {shard_key} already exists" + ))); + } + + // TODO: Check that nested *replica* placement lists are not empty (e.g., `[[], [], []]`) + + if placement.is_empty() { + return Err(StorageError::bad_request(format!( + "Shard key {shard_key} placement cannot be empty" + ))); + } + + let unknown_peers: Vec<_> = placement + .iter() + .flatten() + .filter(|peer_id| !self.peer_address_by_id.contains_key(peer_id)) + .collect(); + + if !unknown_peers.is_empty() { + return Err(StorageError::bad_request(format!( + "Shard Key {shard_key} placement contains unknown peers: {unknown_peers:?}" + ))); + } + + let max_id = collection_state + .shards_key_mapping + .iter_shard_ids() + .max() + .unwrap_or(0); + + let base_id = max_id + 1; + + let init_state = initial_state.unwrap_or_else(|| { + if context.is_distributed + && self.all_peers_at_version(&CREATE_CUSTOM_SHARDS_IN_INITIALIZING_STATE) + { + ReplicaState::Initializing + } else { + ReplicaState::Active + } + }); + + let shards: Vec<_> = placement + .iter() + .enumerate() + .map(|(idx, replicas)| (base_id + idx as ShardId, replicas.clone(), init_state)) + .collect(); + + let mut actions: Actions = shards + .iter() + .map(|&(shard_id, ref replicas, init_state)| { + // inhibit rustfmt + Action::CreateShard { + collection: collection.clone(), + shard_id, + shard_key: Some(shard_key.clone()), + replicas: replicas.clone(), + init_state, + } + }) + .collect(); + + actions.push(Action::RegisterShards { + collection, + shard_key: Some(shard_key.clone()), + shards, + }); + + Ok(actions) + } + + pub fn plan_drop_shard_key(&self, op: &DropShardKey) -> StorageResult { + let DropShardKey { + collection_name, + shard_key, + } = op; + + let collection = self.resolve_collection(collection_name)?; + let collection_state = self.collection(&collection).expect("collection exists"); + + let sharding_method = collection_state + .config + .params + .sharding_method + .unwrap_or_default(); + + if sharding_method != ShardingMethod::Custom { + return Err(StorageError::bad_request(format!( + "shard key {shard_key} cannot be removed with Auto sharding method" + ))); + } + + let Some(shard_ids) = collection_state.shards_key_mapping.get(shard_key) else { + return Ok(Actions::new()); + }; + + let mut shard_ids: Vec<_> = shard_ids.iter().copied().collect(); + shard_ids.sort_unstable(); + + let mut actions = vec![ + Action::InvalidateCleanLocalShards { + collection: collection.clone(), + shard_ids: shard_ids.clone(), + }, + Action::RemoveShardKey { + collection: collection.clone(), + shard_key: shard_key.clone(), + }, + ]; + + actions.extend(shard_ids.into_iter().map(|shard_id| Action::DropShard { + collection: collection.clone(), + shard_id, + })); + + Ok(actions) + } + pub fn plan_update_peer_metadata(&self, peer_id: PeerId, metadata: &PeerMetadata) -> Actions { // Check if operation is already applied if self.peer_metadata_by_id.get(&peer_id) == Some(metadata) { 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 4e2655a1e1..b5254f5340 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 @@ -4,6 +4,7 @@ mod context; mod ops; mod prop; mod replay; +mod state; use std::num::NonZeroU32; 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 e0b2d62fd4..548b2586eb 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 @@ -1,12 +1,13 @@ //! Explicit tests asserting behavior of individual consensus operations //! and tests for cases that `proptest` is unlikely to generate or reach -use std::collections::{BTreeMap, BTreeSet, HashMap}; +use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet}; use std::num::NonZeroU32; use ahash::AHashMap; use collection::collection_state::ShardInfo; use collection::config::ShardingMethod; +use collection::operations::cluster_ops::ReshardingDirection; use collection::operations::config_diff::{ CollectionParamsDiff, HnswConfigDiff, OptimizersConfigDiff, QuantizationConfigDiff, }; @@ -15,12 +16,14 @@ use collection::operations::types::{ }; use collection::shards::replica_set; use collection::shards::replica_set::replica_set_state::ReplicaState; +use collection::shards::resharding::ReshardState; use collection::shards::shard::ShardId; use segment::data_types::collection_defaults::CollectionConfigDefaults; use segment::data_types::modifier::Modifier; use segment::data_types::vector_name_config::*; use segment::types::*; use serde_json::{Value, json}; +use uuid::Uuid; use super::*; use crate::content_manager::collection_meta_ops::*; @@ -491,6 +494,378 @@ fn delete_collection_alias_name() { assert_eq!(machine.state(), &state); } +#[test] +fn create_shard_key() { + let old_key = ShardKey::from("old"); + let shard_key = ShardKey::from("north"); + let placement = vec![vec![PEER_ID, OTHER_PEER_ID], vec![OTHER_PEER_ID]]; + + let mut state = custom_sharding_state(); + add_peer(&mut state, PEER_ID, None); + add_peer(&mut state, OTHER_PEER_ID, None); + + let collection = state + .collections + .get_mut(COLLECTION) + .expect("collection exists"); + collection.shards.insert( + 4, + ShardInfo { + replicas: HashMap::from([(PEER_ID, ReplicaState::Active)]), + }, + ); + collection + .shards_key_mapping + .entry(old_key) + .or_default() + .insert(4); + + let mut machine = state_machine(state); + let outcome = machine.apply(&create_shard_key_op( + shard_key.clone(), + placement.clone(), + Some(ReplicaState::Partial), + )); + + let ApplyOutcome::Accepted(actions) = outcome else { + panic!("creating a shard key should be accepted, got {outcome:?}"); + }; + + assert_eq!( + actions, + vec![ + Action::CreateShard { + collection: COLLECTION.into(), + shard_id: 5, + shard_key: Some(shard_key.clone()), + replicas: placement[0].clone(), + init_state: ReplicaState::Partial, + }, + Action::CreateShard { + collection: COLLECTION.into(), + shard_id: 6, + shard_key: Some(shard_key.clone()), + replicas: placement[1].clone(), + init_state: ReplicaState::Partial, + }, + Action::RegisterShards { + collection: COLLECTION.into(), + shard_key: Some(shard_key.clone()), + shards: vec![ + (5, placement[0].clone(), ReplicaState::Partial), + (6, placement[1].clone(), ReplicaState::Partial), + ], + }, + ], + ); + + let collection = machine + .state() + .collection(COLLECTION) + .expect("collection exists"); + let shard_ids: BTreeSet<_> = collection.shards_key_mapping[&shard_key] + .iter() + .copied() + .collect(); + + assert_eq!(shard_ids, BTreeSet::from([5, 6])); + assert_eq!( + collection.shards[&5].replicas, + HashMap::from([ + (PEER_ID, ReplicaState::Partial), + (OTHER_PEER_ID, ReplicaState::Partial), + ]), + ); + assert_eq!( + collection.shards[&6].replicas, + HashMap::from([(OTHER_PEER_ID, ReplicaState::Partial)]), + ); +} + +#[test] +fn create_shard_key_default_initial_state() { + let cases = [ + (true, "1.14.2", ReplicaState::Initializing), + (true, "1.14.0", ReplicaState::Active), + (false, "1.14.2", ReplicaState::Active), + ]; + + for (is_distributed, version, expected) in cases { + let mut context = node_context(); + context.is_distributed = is_distributed; + + let mut state = custom_sharding_state(); + add_peer(&mut state, PEER_ID, Some(version)); + + let mut machine = ConsensusStateMachine::new(state, context); + let outcome = machine.apply(&create_shard_key_op( + "north".into(), + vec![vec![PEER_ID]], + None, + )); + + let ApplyOutcome::Accepted(actions) = outcome else { + panic!("creating a shard key should be accepted, got {outcome:?}"); + }; + + assert!( + matches!( + actions.as_slice(), + [ + Action::CreateShard { + init_state, + .. + }, + Action::RegisterShards { .. }, + ] if *init_state == expected + ), + "wrong initial state for distributed={is_distributed}, version={version}: {actions:?}" + ); + + let replica_state = machine + .state() + .collection(COLLECTION) + .expect("collection exists") + .shards[&1] + .replicas[&PEER_ID]; + + assert_eq!(replica_state, expected); + } +} + +#[test] +fn create_shard_key_replay() { + let mut state = custom_sharding_state(); + add_peer(&mut state, PEER_ID, None); + + let operation = create_shard_key_op( + "north".into(), + vec![vec![PEER_ID]], + Some(ReplicaState::Active), + ); + + let mut machine = state_machine(state); + let first = machine.apply(&operation); + assert!(matches!(first, ApplyOutcome::Accepted(_))); + + let applied = machine.state().clone(); + let replay = machine.apply(&operation); + + assert!(matches!( + replay, + ApplyOutcome::Rejected(StorageError::BadRequest { .. }) + )); + assert_eq!( + machine.state(), + &applied, + "replay should not change anything" + ); +} + +#[test] +fn create_shard_key_reject_auto_sharding() { + let mut state = cluster_state(Vec::new()); + add_peer(&mut state, PEER_ID, None); + + create_shard_key_rejects_without_change( + state, + create_shard_key_op( + "north".into(), + vec![vec![PEER_ID]], + Some(ReplicaState::Active), + ), + ); +} + +#[test] +fn create_shard_key_reject_existing() { + let shard_key = ShardKey::from("north"); + let mut state = custom_sharding_state(); + add_peer(&mut state, PEER_ID, None); + state + .collections + .get_mut(COLLECTION) + .expect("collection exists") + .shards_key_mapping + .insert(shard_key.clone(), Default::default()); + + create_shard_key_rejects_without_change( + state, + create_shard_key_op(shard_key, vec![vec![PEER_ID]], Some(ReplicaState::Active)), + ); +} + +#[test] +fn create_shard_key_reject_empty_placement() { + create_shard_key_rejects_without_change( + custom_sharding_state(), + create_shard_key_op("north".into(), Vec::new(), Some(ReplicaState::Active)), + ); +} + +#[test] +fn create_shard_key_reject_unknown_peer() { + create_shard_key_rejects_without_change( + custom_sharding_state(), + create_shard_key_op( + "north".into(), + vec![vec![PEER_ID]], + Some(ReplicaState::Active), + ), + ); +} + +#[test] +fn drop_shard_key() { + let shard_key = ShardKey::from("north"); + let other_key = ShardKey::from("south"); + let mut state = custom_sharding_state(); + add_shard_key(&mut state, shard_key.clone(), &[3, 1]); + add_shard_key(&mut state, other_key.clone(), &[5]); + set_resharding(&mut state, other_key.clone(), 5); + + let mut machine = state_machine(state); + let outcome = machine.apply(&drop_shard_key_op(shard_key.clone())); + + let ApplyOutcome::Accepted(actions) = outcome else { + panic!("dropping a shard key should be accepted, got {outcome:?}"); + }; + + assert_eq!( + actions, + vec![ + Action::InvalidateCleanLocalShards { + collection: COLLECTION.into(), + shard_ids: vec![1, 3], + }, + Action::RemoveShardKey { + collection: COLLECTION.into(), + shard_key: shard_key.clone(), + }, + Action::DropShard { + collection: COLLECTION.into(), + shard_id: 1, + }, + Action::DropShard { + collection: COLLECTION.into(), + shard_id: 3, + }, + ], + ); + + let collection = machine + .state() + .collection(COLLECTION) + .expect("collection exists"); + + assert!(!collection.shards_key_mapping.contains_key(&shard_key)); + assert!(!collection.shards.contains_key(&1)); + assert!(!collection.shards.contains_key(&3)); + assert_eq!( + collection.shards_key_mapping[&other_key], + HashSet::from([5]) + ); + assert!(collection.shards.contains_key(&5)); + assert_eq!( + collection + .resharding + .as_ref() + .and_then(|resharding| resharding.shard_key.as_ref()), + Some(&other_key), + ); +} + +#[test] +fn drop_shard_key_replay() { + let state = custom_sharding_state(); + + let mut machine = state_machine(state.clone()); + let outcome = machine.apply(&drop_shard_key_op("north".into())); + + let ApplyOutcome::Accepted(actions) = outcome else { + panic!("replay of an applied shard-key drop should be accepted, got {outcome:?}"); + }; + + assert!(actions.is_empty()); + assert_eq!(machine.state(), &state, "replay should not change anything"); +} + +#[test] +fn drop_shard_key_replay_after_mapping_removal() { + let shard_key = ShardKey::from("north"); + let mut state = custom_sharding_state(); + add_shard_key(&mut state, shard_key.clone(), &[1, 2]); + + let operation = drop_shard_key_op(shard_key); + let mut uncrashed = state_machine(state.clone()); + let outcome = uncrashed.apply(&operation); + let ApplyOutcome::Accepted(actions) = outcome else { + panic!("dropping a shard key should be accepted, got {outcome:?}"); + }; + let goal = uncrashed.state().clone(); + + assert!(matches!(&actions[1], Action::RemoveShardKey { .. })); + + let mut crashed = state; + for action in &actions[..=1] { + crashed.apply_action(action); + } + + let mut replay = state_machine(crashed); + let outcome = replay.apply(&operation); + + let ApplyOutcome::Accepted(replay_actions) = outcome else { + panic!("replay after mapping removal should be accepted, got {outcome:?}"); + }; + + assert!(replay_actions.is_empty()); + assert_eq!(replay.state(), &goal); +} + +#[test] +fn drop_shard_key_reject_auto_sharding() { + let state = cluster_state(Vec::new()); + + let mut machine = state_machine(state.clone()); + let outcome = machine.apply(&drop_shard_key_op("north".into())); + + assert!(matches!( + outcome, + ApplyOutcome::Rejected(StorageError::BadRequest { .. }) + )); + assert_eq!(machine.state(), &state); +} + +#[test] +fn drop_shard_key_resharding_same_key() { + let shard_key = ShardKey::from("north"); + let mut state = custom_sharding_state(); + add_shard_key(&mut state, shard_key.clone(), &[1]); + set_resharding(&mut state, shard_key.clone(), 1); + + let mut machine = state_machine(state.clone()); + let outcome = machine.apply(&drop_shard_key_op(shard_key)); + + assert!(matches!(outcome, ApplyOutcome::NotCovered)); + assert_eq!(machine.state(), &state); +} + +fn set_resharding(state: &mut ClusterState, shard_key: ShardKey, shard_id: ShardId) { + let resharding = ReshardState::new( + Uuid::nil(), + ReshardingDirection::Up, + PEER_ID, + shard_id, + Some(shard_key), + ); + + state + .collections + .get_mut(COLLECTION) + .expect("collection exists") + .resharding = Some(resharding); +} + #[test] fn create_alias() { let state = cluster_state(Vec::new()); @@ -1391,6 +1766,84 @@ fn create_collection_op( collection_meta_op(CollectionMetaOperations::CreateCollection(operation)) } +fn create_shard_key_op( + shard_key: ShardKey, + placement: Vec>, + initial_state: Option, +) -> ConsensusOperations { + collection_meta_op(CollectionMetaOperations::CreateShardKey(CreateShardKey { + collection_name: COLLECTION.into(), + shard_key, + placement, + initial_state, + })) +} + +fn drop_shard_key_op(shard_key: ShardKey) -> ConsensusOperations { + collection_meta_op(CollectionMetaOperations::DropShardKey(DropShardKey { + collection_name: COLLECTION.into(), + shard_key, + })) +} + +fn custom_sharding_state() -> ClusterState { + let mut state = cluster_state(Vec::new()); + state + .collections + .get_mut(COLLECTION) + .expect("collection exists") + .config + .params + .sharding_method = Some(ShardingMethod::Custom); + state +} + +fn add_shard_key(state: &mut ClusterState, shard_key: ShardKey, shard_ids: &[ShardId]) { + let collection = state + .collections + .get_mut(COLLECTION) + .expect("collection exists"); + + for &shard_id in shard_ids { + collection.shards.insert( + shard_id, + ShardInfo { + replicas: HashMap::from([(PEER_ID, ReplicaState::Active)]), + }, + ); + collection + .shards_key_mapping + .entry(shard_key.clone()) + .or_default() + .insert(shard_id); + } +} + +fn add_peer(state: &mut ClusterState, peer_id: PeerId, version: Option<&str>) { + let address = format!("http://peer-{peer_id}") + .parse() + .expect("valid peer URI"); + state.peer_address_by_id.insert(peer_id, address); + + if let Some(version) = version { + state.peer_metadata_by_id.insert( + peer_id, + PeerMetadata::new(version.parse().expect("valid version")), + ); + } +} + +fn create_shard_key_rejects_without_change(state: ClusterState, operation: ConsensusOperations) { + let mut machine = state_machine(state.clone()); + let outcome = machine.apply(&operation); + + assert!(matches!( + outcome, + ApplyOutcome::Rejected(StorageError::BadRequest { .. }) + )); + assert_eq!(machine.state(), &state); +} + /// Shards a collection is created with, one replica per peer of each placement entry fn shards(placement: Vec>) -> AHashMap { placement 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 54065b607c..a1ef4f14fe 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 @@ -2,19 +2,21 @@ use std::collections::{BTreeMap, HashMap}; -use collection::collection_state; +use collection::collection_state::{self, ShardInfo}; use collection::config::ShardingMethod; use collection::operations::config_diff::{HnswConfigDiff, QuantizationConfigDiff}; use collection::operations::types::{ PeerMetadata, SparseVectorParams, SparseVectorsConfig, VectorParamsDiff, VectorsConfigDiff, }; use collection::shards::CollectionId; +use collection::shards::replica_set::replica_set_state::ReplicaState; use collection::shards::shard::ShardId; use proptest::prelude::*; use segment::data_types::modifier::Modifier; use segment::data_types::vector_name_config::*; use segment::json_path::JsonPath; use segment::types::*; +use tonic::transport::Uri; use super::*; use crate::content_manager::alias_mapping::AliasMapping; @@ -23,7 +25,10 @@ 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; +use crate::types::{PeerAddressById, PeerMetadataById}; + +const PEER_IDS: &[PeerId] = &[PEER_ID, OTHER_PEER_ID]; +const PEER_VERSIONS: &[&str] = &["1.14.0", "1.15.0"]; const COLLECTION_NAMES: &[&str] = &["alpha", "beta", "gamma"]; const MISSING_COLLECTION_NAME: &str = "missing"; @@ -31,22 +36,20 @@ const MISSING_COLLECTION_NAME: &str = "missing"; const ALIAS_NAMES: &[&str] = &["primary", "secondary"]; const DANGLING_ALIAS_NAME: &str = "dangling"; +const SHARD_KEYS: &[&str] = &["north", "south"]; + 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, OTHER_PEER_ID]; -const PEER_VERSIONS: &[&str] = &["1.14.0", "1.15.0"]; - const METADATA_KEYS: &[&str] = &["region", "tier"]; pub fn arb_state_and_operation() -> impl Strategy { arb_cluster_state().prop_flat_map(|state| { let collections = state.collections.keys().cloned(); let aliases = state.aliases.iter().map(|(alias, _)| alias.clone()); - let names = collections.chain(aliases).collect(); - let operations = arb_consensus_operation(names); + let peers = state.peer_address_by_id.keys().copied().collect(); + let operations = arb_consensus_operation(names, peers); (Just(state), operations) }) @@ -65,21 +68,29 @@ pub fn arb_cluster_state() -> impl Strategy { let state = ( Just(collections), arb_aliases(names), + arb_peer_address_by_id(), arb_peer_metadata_by_id(), arb_cluster_metadata(), arb_quota_config(), ); state.prop_map(|state| { - let (collections, aliases, peer_metadata_by_id, cluster_metadata, quota_config) = state; + let ( + collections, + aliases, + peer_address_by_id, + peer_metadata_by_id, + cluster_metadata, + quota_config, + ) = state; ClusterState { collections, aliases, + peer_address_by_id, peer_metadata_by_id, cluster_metadata, quota_config, - ..Default::default() } }) }) @@ -88,11 +99,48 @@ pub fn arb_cluster_state() -> impl Strategy { fn arb_collection_state() -> impl Strategy { let vectors = proptest::collection::btree_map(arb_vector_name(), arb_vector_name_config(), 0..3); + + let sharding_method = proptest::option::of(prop_oneof![ + Just(ShardingMethod::Auto), + Just(ShardingMethod::Custom), + ]); + + let replicas = proptest::collection::vec(arb_peer_id(), 1..3); + let shard_key = arb_shard_key(); + let shards = proptest::collection::vec((replicas, shard_key), 0..3); + let indexes = proptest::collection::hash_map(arb_field_name(), arb_field_schema(), 0..3); - (vectors, indexes).prop_map(|(vectors, indexes)| { + let state = (vectors, sharding_method, indexes, shards); + + state.prop_map(|state| { + let (vectors, sharding_method, indexes, shards) = state; + let mut state = collection_state(vectors.into_iter().collect()); + state.config.params.sharding_method = sharding_method; state.payload_index_schema.schema = indexes; + + let is_custom_sharding = sharding_method.unwrap_or_default() == ShardingMethod::Custom; + + for (shard_id, (replicas, shard_key)) in shards.into_iter().enumerate() { + let shard_id = shard_id as ShardId; + + let replicas = replicas + .into_iter() + .map(|peer_id| (peer_id, ReplicaState::Active)) + .collect(); + + state.shards.insert(shard_id, ShardInfo { replicas }); + + if is_custom_sharding { + state + .shards_key_mapping + .entry(shard_key) + .or_default() + .insert(shard_id); + } + } + state }) } @@ -127,6 +175,15 @@ fn arb_peer_metadata_by_id() -> impl Strategy { proptest::collection::hash_map(arb_peer_id(), arb_peer_metadata(), 0..3) } +fn arb_peer_address_by_id() -> impl Strategy { + proptest::collection::hash_set(arb_peer_id(), 0..3).prop_map(|peer_ids| { + peer_ids + .into_iter() + .map(|peer_id| (peer_id, peer_address(peer_id))) + .collect() + }) +} + fn arb_peer_id() -> impl Strategy { proptest::sample::select(PEER_IDS) } @@ -136,6 +193,19 @@ fn arb_peer_metadata() -> impl Strategy { .prop_map(|version| PeerMetadata::new(version.parse().expect("valid version"))) } +fn peer_address(peer_id: PeerId) -> Uri { + format!("http://peer-{peer_id}") + .parse() + .expect("valid peer URI") +} + +fn arb_shard_key() -> impl Strategy { + prop_oneof![ + proptest::sample::select(SHARD_KEYS).prop_map(ShardKey::from), + (1_u64..=2).prop_map(ShardKey::from), + ] +} + /// Cluster metadata never holds a null value: that is how a key is removed fn arb_cluster_metadata() -> impl Strategy> { proptest::collection::hash_map(arb_metadata_key(), arb_metadata_value(), 0..2) @@ -170,13 +240,14 @@ fn arb_quota_config() -> impl Strategy { pub fn arb_consensus_operation( collection_names: Vec, + peer_ids: Vec, ) -> impl Strategy { - let collection_meta = arb_collection_meta_operation(collection_names) + let collection_meta = arb_collection_meta_operation(collection_names, peer_ids) .prop_map(|operation| ConsensusOperations::CollectionMeta(Box::new(operation))); // Weighted by how many operations each arm covers, so one operation is as likely as another prop_oneof![ - 9 => collection_meta, + 11 => collection_meta, 1 => arb_update_peer_metadata(), 1 => arb_update_cluster_metadata(), 1 => arb_quota_config().prop_map(ConsensusOperations::SetQuotaConfig), @@ -185,6 +256,7 @@ pub fn arb_consensus_operation( fn arb_collection_meta_operation( mut collection_names: Vec, + peer_ids: Vec, ) -> impl Strategy { collection_names.push(MISSING_COLLECTION_NAME.into()); @@ -194,6 +266,8 @@ fn arb_collection_meta_operation( arb_update_collection(collection_names.clone()), arb_delete_collection(collection_names.clone()), arb_change_aliases(collection_names.clone()), + arb_create_shard_key(collection_names.clone(), peer_ids), + arb_drop_shard_key(collection_names.clone()), arb_create_named_vector(collection_names.clone()), arb_delete_named_vector(collection_names.clone()), arb_create_payload_index(collection_names.clone()), @@ -201,6 +275,52 @@ fn arb_collection_meta_operation( ] } +fn arb_drop_shard_key(collections: Vec) -> impl Strategy { + let collection_name = arb_collection_name(collections); + let shard_key = arb_shard_key(); + + (collection_name, shard_key).prop_map(|(collection_name, shard_key)| { + CollectionMetaOperations::DropShardKey(DropShardKey { + collection_name, + shard_key, + }) + }) +} + +fn arb_create_shard_key( + collections: Vec, + mut peer_ids: Vec, +) -> impl Strategy { + // An empty peer map cannot produce a valid placement. Keep generating placement so those + // states exercise unknown-peer rejection as well as the empty-placement check. + if peer_ids.is_empty() { + peer_ids.extend(PEER_IDS); + } + + let collection_name = arb_collection_name(collections); + let shard_key = arb_shard_key(); + let placement = proptest::collection::vec( + proptest::collection::vec(proptest::sample::select(peer_ids), 1..3), + 0..3, + ); + let initial_state = proptest::option::of(proptest::sample::select(vec![ + ReplicaState::Active, + ReplicaState::Initializing, + ReplicaState::Partial, + ])); + + (collection_name, shard_key, placement, initial_state).prop_map( + |(collection_name, shard_key, placement, initial_state)| { + CollectionMetaOperations::CreateShardKey(CreateShardKey { + collection_name, + shard_key, + placement, + initial_state, + }) + }, + ) +} + fn arb_update_peer_metadata() -> impl Strategy { (arb_peer_id(), arb_peer_metadata()).prop_map(|(peer_id, metadata)| { ConsensusOperations::UpdatePeerMetadata { peer_id, metadata } diff --git a/lib/storage/src/content_manager/consensus_state_machine/tests/state.rs b/lib/storage/src/content_manager/consensus_state_machine/tests/state.rs new file mode 100644 index 0000000000..e227ffdd91 --- /dev/null +++ b/lib/storage/src/content_manager/consensus_state_machine/tests/state.rs @@ -0,0 +1,64 @@ +use collection::operations::types::PeerMetadata; +use semver::Version; +use tonic::transport::Uri; + +use super::*; + +#[test] +fn all_peers_at_version_empty() { + let state = ClusterState::default(); + + assert!(state.all_peers_at_version(&version("1.14.2-dev"))); +} + +#[test] +fn all_peers_at_version_all_new_enough() { + let state = cluster_state_with_peers(&[(PEER_ID, "1.14.2"), (OTHER_PEER_ID, "1.15.0")]); + + assert!(state.all_peers_at_version(&version("1.14.2-dev"))); +} + +#[test] +fn all_peers_at_version_old_peer() { + let state = cluster_state_with_peers(&[(PEER_ID, "1.14.0"), (OTHER_PEER_ID, "1.15.0")]); + + assert!(!state.all_peers_at_version(&version("1.14.2-dev"))); +} + +#[test] +fn all_peers_at_version_missing_metadata() { + let mut state = cluster_state_with_peers(&[(PEER_ID, "1.15.0")]); + state + .peer_address_by_id + .insert(OTHER_PEER_ID, peer_address(OTHER_PEER_ID)); + + assert!(!state.all_peers_at_version(&version("1.14.2-dev"))); +} + +fn cluster_state_with_peers(peers: &[(PeerId, &str)]) -> ClusterState { + let peer_address_by_id = peers + .iter() + .map(|&(peer_id, _)| (peer_id, peer_address(peer_id))) + .collect(); + + let peer_metadata_by_id = peers + .iter() + .map(|&(peer_id, version)| (peer_id, PeerMetadata::new(self::version(version)))) + .collect(); + + ClusterState { + peer_address_by_id, + peer_metadata_by_id, + ..Default::default() + } +} + +fn peer_address(peer_id: PeerId) -> Uri { + format!("http://peer-{peer_id}") + .parse() + .expect("valid peer URI") +} + +fn version(version: &str) -> Version { + version.parse().expect("valid version") +} diff --git a/lib/storage/src/content_manager/toc/collection_meta_ops.rs b/lib/storage/src/content_manager/toc/collection_meta_ops.rs index 59b4578ed2..f80e71a046 100644 --- a/lib/storage/src/content_manager/toc/collection_meta_ops.rs +++ b/lib/storage/src/content_manager/toc/collection_meta_ops.rs @@ -1,5 +1,4 @@ use std::collections::HashSet; -use std::sync::LazyLock; use collection::collection::AbortReshardingScope; use collection::collection_state; @@ -22,9 +21,6 @@ use crate::content_manager::consensus_state_machine::apply_collection_config_dif use crate::content_manager::errors::StorageError; use crate::content_manager::shard_distribution::ShardDistributionProposal; -static CREATE_CUSTOM_SHARDS_IN_INITIALIZING_STATE: LazyLock = - LazyLock::new(|| semver::Version::parse("1.14.2-dev").unwrap()); - impl TableOfContent { pub(super) fn perform_collection_meta_op_sync( &self,