Implement CreateShardKey/RemoveShardKey for consensus state machine (#10666)

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)]

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<semver::Version> =
+    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.

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<ShardKey>,
+        replicas: Vec<PeerId>,
+        init_state: ReplicaState,
+    },
+
+    /// Register built shards and record them under `shard_key` in one mapping write.
+    RegisterShards {
+        collection: CollectionId,
+        shard_key: Option<ShardKey>,
+        shards: Vec<(ShardId, Vec<PeerId>, ReplicaState)>,
+    },
+
+    /// Stop cleanup tasks before their shard directories disappear
+    InvalidateCleanLocalShards {
+        collection: CollectionId,
+        shard_ids: Vec<ShardId>,
+    },
+
+    /// Persist the replay gate before dropping shard directories
+    RemoveShardKey {
+        collection: CollectionId,
+        shard_key: ShardKey,
+    },
+
+    DropShard {
+        collection: CollectionId,
+        shard_id: ShardId,
+    },
+
     UpdateAliases {
         set: BTreeMap<String, CollectionId>,
         remove: BTreeSet<String>,
@@ -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, .. }

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.

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);

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)
+    }
 }

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<Actions> {
+        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<Actions> {
+        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) {

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;
 

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<Vec<PeerId>>,
+    initial_state: Option<ReplicaState>,
+) -> 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<Vec<PeerId>>) -> AHashMap<ShardId, ShardInfo> {
     placement

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<Value = (ClusterState, ConsensusOperations)> {
     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<Value = ClusterState> {
         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<Value = ClusterState> {
 fn arb_collection_state() -> impl Strategy<Value = collection_state::State> {
     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<Value = PeerMetadataById> {
     proptest::collection::hash_map(arb_peer_id(), arb_peer_metadata(), 0..3)
 }
 
+fn arb_peer_address_by_id() -> impl Strategy<Value = PeerAddressById> {
+    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<Value = PeerId> {
     proptest::sample::select(PEER_IDS)
 }
@@ -136,6 +193,19 @@ fn arb_peer_metadata() -> impl Strategy<Value = PeerMetadata> {
         .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<Value = ShardKey> {
+    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<Value = HashMap<String, serde_json::Value>> {
     proptest::collection::hash_map(arb_metadata_key(), arb_metadata_value(), 0..2)
@@ -170,13 +240,14 @@ fn arb_quota_config() -> impl Strategy<Value = QuotaConfig> {
 
 pub fn arb_consensus_operation(
     collection_names: Vec<String>,
+    peer_ids: Vec<PeerId>,
 ) -> impl Strategy<Value = ConsensusOperations> {
-    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<String>,
+    peer_ids: Vec<PeerId>,
 ) -> impl Strategy<Value = CollectionMetaOperations> {
     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<String>) -> impl Strategy<Value = CollectionMetaOperations> {
+    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<String>,
+    mut peer_ids: Vec<PeerId>,
+) -> impl Strategy<Value = CollectionMetaOperations> {
+    // 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<Value = ConsensusOperations> {
     (arb_peer_id(), arb_peer_metadata()).prop_map(|(peer_id, metadata)| {
         ConsensusOperations::UpdatePeerMetadata { peer_id, metadata }

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")
+}

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<semver::Version> =
-    LazyLock::new(|| semver::Version::parse("1.14.2-dev").unwrap());
-
 impl TableOfContent {
     pub(super) fn perform_collection_meta_op_sync(
         &self,