mirror of
https://github.com/qdrant/qdrant.git
synced 2026-09-21 13:37:46 -05:00
Implement CreateCollection and DeleteCollection for consensus state machine (#10397)
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 5
parent
37e4b3cc3d
commit
2e71304e6e
@@ -24,6 +24,17 @@ impl AliasMapping {
|
|||||||
Ok(atomic_save_json(path, self)?)
|
Ok(atomic_save_json(path, self)?)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Aliases pointing at `collection_name`.
|
||||||
|
pub fn collection_aliases<'a>(
|
||||||
|
&'a self,
|
||||||
|
collection_name: &'a str,
|
||||||
|
) -> impl Iterator<Item = Alias> + 'a {
|
||||||
|
self.0
|
||||||
|
.iter()
|
||||||
|
.filter(move |&(_, target)| target == collection_name)
|
||||||
|
.map(|(alias, _)| alias.clone())
|
||||||
|
}
|
||||||
|
|
||||||
/// Iterate over aliases and collections they point at.
|
/// Iterate over aliases and collections they point at.
|
||||||
pub fn iter(&self) -> impl Iterator<Item = (&Alias, &CollectionId)> {
|
pub fn iter(&self) -> impl Iterator<Item = (&Alias, &CollectionId)> {
|
||||||
self.0.iter()
|
self.0.iter()
|
||||||
@@ -144,13 +155,9 @@ impl AliasPersistence {
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub fn collection_aliases(&self, collection_name: &str) -> Vec<String> {
|
pub fn collection_aliases(&self, collection_name: &str) -> Vec<String> {
|
||||||
let mut result = vec![];
|
self.alias_mapping
|
||||||
for (alias, target_collection) in self.alias_mapping.0.iter() {
|
.collection_aliases(collection_name)
|
||||||
if collection_name == target_collection {
|
.collect()
|
||||||
result.push(alias.clone());
|
|
||||||
}
|
|
||||||
}
|
|
||||||
result
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn state(&self) -> &AliasMapping {
|
pub fn state(&self) -> &AliasMapping {
|
||||||
|
|||||||
@@ -12,7 +12,7 @@ use collection::operations::config_diff::{
|
|||||||
WalConfigDiff,
|
WalConfigDiff,
|
||||||
};
|
};
|
||||||
use collection::operations::types::{
|
use collection::operations::types::{
|
||||||
SparseVectorParams, SparseVectorsConfig, VectorsConfig, VectorsConfigDiff,
|
SparseVectorParams, SparseVectorsConfig, VectorParams, VectorsConfig, VectorsConfigDiff,
|
||||||
};
|
};
|
||||||
use collection::operations::validation;
|
use collection::operations::validation;
|
||||||
use collection::shards::replica_set::replica_set_state::ReplicaState;
|
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::data_types::vectors::DEFAULT_VECTOR_NAME;
|
||||||
use segment::types::{
|
use segment::types::{
|
||||||
Payload, PayloadFieldSchema, PayloadKeyType, QuantizationConfig, ShardKey, StrictModeConfig,
|
Payload, PayloadFieldSchema, PayloadKeyType, QuantizationConfig, ShardKey, StrictModeConfig,
|
||||||
VectorNameBuf,
|
VectorNameBuf, VectorsConfigDefaults,
|
||||||
};
|
};
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
@@ -193,6 +193,43 @@ pub struct CreateCollection {
|
|||||||
pub metadata: Option<Payload>,
|
pub metadata: Option<Payload>,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// 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<PayloadStorageParams>,
|
||||||
|
on_disk_payload: Option<bool>,
|
||||||
|
defaults: Option<PayloadStorageParams>,
|
||||||
|
) -> Option<PayloadStorageParams> {
|
||||||
|
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
|
/// Operation for creating new collection and (optionally) specify index params
|
||||||
#[derive(Debug, Deserialize, Serialize, PartialEq, Eq, Hash, Clone)]
|
#[derive(Debug, Deserialize, Serialize, PartialEq, Eq, Hash, Clone)]
|
||||||
#[serde(rename_all = "snake_case")]
|
#[serde(rename_all = "snake_case")]
|
||||||
@@ -277,6 +314,10 @@ impl CreateCollectionOperation {
|
|||||||
self.distribution.is_some()
|
self.distribution.is_some()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub fn distribution(&self) -> Option<&ShardDistributionProposal> {
|
||||||
|
self.distribution.as_ref()
|
||||||
|
}
|
||||||
|
|
||||||
pub fn take_distribution(&mut self) -> Option<ShardDistributionProposal> {
|
pub fn take_distribution(&mut self) -> Option<ShardDistributionProposal> {
|
||||||
self.distribution.take()
|
self.distribution.take()
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,5 +1,6 @@
|
|||||||
use std::collections::{BTreeMap, BTreeSet};
|
use std::collections::{BTreeMap, BTreeSet};
|
||||||
|
|
||||||
|
use collection::collection_state;
|
||||||
use collection::operations::types::PeerMetadata;
|
use collection::operations::types::PeerMetadata;
|
||||||
use collection::shards::CollectionId;
|
use collection::shards::CollectionId;
|
||||||
use collection::shards::shard::PeerId;
|
use collection::shards::shard::PeerId;
|
||||||
@@ -13,6 +14,15 @@ use crate::quota::QuotaConfig;
|
|||||||
/// A single change a consensus operation makes
|
/// A single change a consensus operation makes
|
||||||
#[derive(Clone, Debug, PartialEq)]
|
#[derive(Clone, Debug, PartialEq)]
|
||||||
pub enum Action {
|
pub enum Action {
|
||||||
|
CreateCollection {
|
||||||
|
collection: CollectionId,
|
||||||
|
state: Box<collection_state::State>,
|
||||||
|
},
|
||||||
|
|
||||||
|
DropCollection {
|
||||||
|
collection: CollectionId,
|
||||||
|
},
|
||||||
|
|
||||||
AddNamedVector {
|
AddNamedVector {
|
||||||
collection: CollectionId,
|
collection: CollectionId,
|
||||||
vector_name: VectorNameBuf,
|
vector_name: VectorNameBuf,
|
||||||
@@ -67,7 +77,9 @@ impl Action {
|
|||||||
/// Collection this action changes, if it is scoped to one
|
/// Collection this action changes, if it is scoped to one
|
||||||
pub fn collection(&self) -> Option<&CollectionId> {
|
pub fn collection(&self) -> Option<&CollectionId> {
|
||||||
match self {
|
match self {
|
||||||
Action::AddNamedVector { collection, .. }
|
Action::CreateCollection { collection, .. }
|
||||||
|
| Action::DropCollection { collection }
|
||||||
|
| Action::AddNamedVector { collection, .. }
|
||||||
| Action::DropNamedVector { collection, .. }
|
| Action::DropNamedVector { collection, .. }
|
||||||
| Action::SetPayloadIndex { collection, .. }
|
| Action::SetPayloadIndex { collection, .. }
|
||||||
| Action::DropPayloadIndex { collection, .. } => Some(collection),
|
| Action::DropPayloadIndex { collection, .. } => Some(collection),
|
||||||
|
|||||||
@@ -11,7 +11,9 @@
|
|||||||
//!
|
//!
|
||||||
//! Two rules every operation follows:
|
//! 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
|
//! 2. Emit only the actions left to reach the goal state, each idempotent, and the action that
|
||||||
//! records the operation as applied last.
|
//! records the operation as applied last.
|
||||||
//!
|
//!
|
||||||
@@ -25,14 +27,24 @@ pub mod state;
|
|||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests;
|
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 collection::shards::transfer::ShardTransferMethod;
|
||||||
use segment::data_types::collection_defaults::CollectionConfigDefaults;
|
use segment::data_types::collection_defaults::CollectionConfigDefaults;
|
||||||
|
use segment::types::HnswConfig;
|
||||||
|
|
||||||
pub use self::action::Action;
|
pub use self::action::Action;
|
||||||
pub use self::state::ClusterState;
|
pub use self::state::ClusterState;
|
||||||
use super::errors::StorageResult;
|
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::consensus_ops::ConsensusOperations;
|
||||||
use crate::content_manager::errors::StorageError;
|
use crate::content_manager::errors::StorageError;
|
||||||
|
|
||||||
@@ -102,9 +114,15 @@ impl ConsensusStateMachine {
|
|||||||
match operation {
|
match operation {
|
||||||
CollectionMetaOperations::Nop { .. } => ApplyOutcome::Accepted(Vec::new()),
|
CollectionMetaOperations::Nop { .. } => ApplyOutcome::Accepted(Vec::new()),
|
||||||
|
|
||||||
CollectionMetaOperations::CreateCollection(_)
|
CollectionMetaOperations::CreateCollection(operation) => {
|
||||||
| CollectionMetaOperations::UpdateCollection(_)
|
ApplyOutcome::new(self.state.plan_create_collection(&self.context, operation))
|
||||||
| CollectionMetaOperations::DeleteCollection(_)
|
}
|
||||||
|
|
||||||
|
CollectionMetaOperations::DeleteCollection(operation) => {
|
||||||
|
ApplyOutcome::Accepted(self.state.plan_delete_collection(operation))
|
||||||
|
}
|
||||||
|
|
||||||
|
CollectionMetaOperations::UpdateCollection(_)
|
||||||
| CollectionMetaOperations::CreateShardKey(_)
|
| CollectionMetaOperations::CreateShardKey(_)
|
||||||
| CollectionMetaOperations::DropShardKey(_)
|
| CollectionMetaOperations::DropShardKey(_)
|
||||||
| CollectionMetaOperations::SetShardReplicaState(_)
|
| CollectionMetaOperations::SetShardReplicaState(_)
|
||||||
@@ -146,7 +164,7 @@ impl ConsensusStateMachine {
|
|||||||
/// Node-local values operations read.
|
/// Node-local values operations read.
|
||||||
///
|
///
|
||||||
/// These come from this node's config, not from consensus, so two peers can read different values
|
/// 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)]
|
#[derive(Clone, Debug)]
|
||||||
pub struct NodeContext {
|
pub struct NodeContext {
|
||||||
pub peer_id: PeerId,
|
pub peer_id: PeerId,
|
||||||
@@ -155,6 +173,146 @@ pub struct NodeContext {
|
|||||||
pub collection_defaults: Option<CollectionConfigDefaults>,
|
pub collection_defaults: Option<CollectionConfigDefaults>,
|
||||||
/// Transfer method this node picks when an operation does not name one
|
/// Transfer method this node picks when an operation does not name one
|
||||||
pub default_shard_transfer_method: Option<ShardTransferMethod>,
|
pub default_shard_transfer_method: Option<ShardTransferMethod>,
|
||||||
|
pub max_collections: Option<usize>,
|
||||||
|
pub wal: WalConfig,
|
||||||
|
pub optimizers: OptimizersConfig,
|
||||||
|
pub hnsw_index: HnswConfig,
|
||||||
|
pub payload: Option<PayloadStorageParams>,
|
||||||
|
/// 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<PeerId>)> {
|
||||||
|
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<CollectionConfigInternal> {
|
||||||
|
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.
|
/// What the machine decided about an operation.
|
||||||
|
|||||||
@@ -4,11 +4,19 @@ use super::*;
|
|||||||
|
|
||||||
impl ClusterState {
|
impl ClusterState {
|
||||||
/// Apply one action. Cannot fail.
|
/// 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) {
|
pub fn apply_action(&mut self, action: &Action) {
|
||||||
match 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 {
|
Action::AddNamedVector {
|
||||||
collection,
|
collection,
|
||||||
vector_name,
|
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> {
|
fn collection_mut(&mut self, collection: &str) -> Option<&mut collection_state::State> {
|
||||||
let state = self.collections.get_mut(collection);
|
let state = self.collections.get_mut(collection);
|
||||||
|
|
||||||
|
|||||||
@@ -1,16 +1,114 @@
|
|||||||
use std::collections::{BTreeMap, BTreeSet};
|
use std::collections::{BTreeMap, BTreeSet};
|
||||||
|
|
||||||
use collection::collection::vector_name_schema;
|
use collection::collection::vector_name_schema;
|
||||||
|
use collection::collection_state::ShardInfo;
|
||||||
use collection::operations::types::PeerMetadata;
|
use collection::operations::types::PeerMetadata;
|
||||||
|
use collection::shards::replica_set::replica_set_state::ReplicaState;
|
||||||
use collection::shards::shard::PeerId;
|
use collection::shards::shard::PeerId;
|
||||||
|
|
||||||
use super::*;
|
use super::*;
|
||||||
use crate::content_manager::collection_meta_ops::*;
|
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<Action>;
|
type Actions = Vec<Action>;
|
||||||
|
|
||||||
impl ClusterState {
|
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<Actions> {
|
||||||
|
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<Actions> {
|
pub fn plan_create_named_vector(&self, op: &CreateNamedVector) -> StorageResult<Actions> {
|
||||||
let CreateNamedVector {
|
let CreateNamedVector {
|
||||||
collection_name,
|
collection_name,
|
||||||
|
|||||||
@@ -17,17 +17,48 @@ use shard::operations::VectorNameConfig;
|
|||||||
use super::*;
|
use super::*;
|
||||||
|
|
||||||
const PEER_ID: u64 = 42;
|
const PEER_ID: u64 = 42;
|
||||||
|
const OTHER_PEER_ID: u64 = 43;
|
||||||
|
|
||||||
fn state_machine(state: ClusterState) -> ConsensusStateMachine {
|
fn state_machine(state: ClusterState) -> ConsensusStateMachine {
|
||||||
ConsensusStateMachine::new(
|
ConsensusStateMachine::new(state, node_context())
|
||||||
state,
|
}
|
||||||
NodeContext {
|
|
||||||
peer_id: PEER_ID,
|
/// Node config is fixed, except for the parts used in tests. Extend it as needed.
|
||||||
is_distributed: true,
|
fn node_context() -> NodeContext {
|
||||||
collection_defaults: None,
|
NodeContext {
|
||||||
default_shard_transfer_method: None,
|
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 {
|
fn collection_state(vectors: Vec<(VectorNameBuf, VectorNameConfig)>) -> collection_state::State {
|
||||||
@@ -53,18 +84,7 @@ fn collection_config(params: CollectionParams) -> CollectionConfigInternal {
|
|||||||
CollectionConfigInternal {
|
CollectionConfigInternal {
|
||||||
params,
|
params,
|
||||||
hnsw_config: Default::default(),
|
hnsw_config: Default::default(),
|
||||||
optimizer_config: OptimizersConfig {
|
optimizer_config: optimizers_config(),
|
||||||
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,
|
|
||||||
},
|
|
||||||
wal_config: Default::default(),
|
wal_config: Default::default(),
|
||||||
quantization_config: None,
|
quantization_config: None,
|
||||||
strict_mode_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 {
|
fn collection_params(vectors: VectorsConfig) -> CollectionParams {
|
||||||
CollectionParams {
|
CollectionParams {
|
||||||
vectors,
|
vectors,
|
||||||
|
|||||||
@@ -3,7 +3,13 @@
|
|||||||
|
|
||||||
use std::collections::{BTreeMap, BTreeSet, HashMap};
|
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::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::data_types::vector_name_config::*;
|
||||||
use segment::types::*;
|
use segment::types::*;
|
||||||
use serde_json::{Value, json};
|
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_ops::ConsensusOperations;
|
||||||
use crate::content_manager::consensus_state_machine::*;
|
use crate::content_manager::consensus_state_machine::*;
|
||||||
use crate::content_manager::errors::StorageError;
|
use crate::content_manager::errors::StorageError;
|
||||||
|
use crate::content_manager::shard_distribution::ShardDistributionProposal;
|
||||||
use crate::quota::QuotaConfig;
|
use crate::quota::QuotaConfig;
|
||||||
|
|
||||||
const COLLECTION: &str = "alpha";
|
const COLLECTION: &str = "alpha";
|
||||||
@@ -34,6 +41,273 @@ fn nop() {
|
|||||||
assert_eq!(machine.state(), &state);
|
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]
|
#[test]
|
||||||
fn create_alias() {
|
fn create_alias() {
|
||||||
let state = cluster_state(Vec::new());
|
let state = cluster_state(Vec::new());
|
||||||
@@ -912,6 +1186,50 @@ fn collection_meta_op(op: CollectionMetaOperations) -> ConsensusOperations {
|
|||||||
ConsensusOperations::CollectionMeta(Box::new(op))
|
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<Vec<Vec<PeerId>>>,
|
||||||
|
) -> 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<Vec<PeerId>>) -> AHashMap<ShardId, ShardInfo> {
|
||||||
|
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<AliasOperations>) -> ConsensusOperations {
|
fn change_aliases_op(actions: Vec<AliasOperations>) -> ConsensusOperations {
|
||||||
collection_meta_op(CollectionMetaOperations::ChangeAliases(
|
collection_meta_op(CollectionMetaOperations::ChangeAliases(
|
||||||
ChangeAliasesOperation { actions },
|
ChangeAliasesOperation { actions },
|
||||||
|
|||||||
@@ -3,8 +3,10 @@
|
|||||||
use std::collections::HashMap;
|
use std::collections::HashMap;
|
||||||
|
|
||||||
use collection::collection_state;
|
use collection::collection_state;
|
||||||
|
use collection::config::ShardingMethod;
|
||||||
use collection::operations::types::PeerMetadata;
|
use collection::operations::types::PeerMetadata;
|
||||||
use collection::shards::CollectionId;
|
use collection::shards::CollectionId;
|
||||||
|
use collection::shards::shard::ShardId;
|
||||||
use proptest::prelude::*;
|
use proptest::prelude::*;
|
||||||
use segment::data_types::modifier::Modifier;
|
use segment::data_types::modifier::Modifier;
|
||||||
use segment::data_types::vector_name_config::*;
|
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::collection_meta_ops::*;
|
||||||
use crate::content_manager::consensus_ops::ConsensusOperations;
|
use crate::content_manager::consensus_ops::ConsensusOperations;
|
||||||
use crate::content_manager::consensus_state_machine::*;
|
use crate::content_manager::consensus_state_machine::*;
|
||||||
|
use crate::content_manager::shard_distribution::ShardDistributionProposal;
|
||||||
use crate::quota::QuotaConfig;
|
use crate::quota::QuotaConfig;
|
||||||
use crate::types::PeerMetadataById;
|
use crate::types::PeerMetadataById;
|
||||||
|
|
||||||
@@ -29,7 +32,7 @@ const VECTOR_NAMES: &[&str] = &["", "text", "image"];
|
|||||||
const FIELD_NAMES: &[&str] = &["city", "count", "nested.key"];
|
const FIELD_NAMES: &[&str] = &["city", "count", "nested.key"];
|
||||||
|
|
||||||
/// This node, and one other peer
|
/// 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 PEER_VERSIONS: &[&str] = &["1.14.0", "1.15.0"];
|
||||||
|
|
||||||
const METADATA_KEYS: &[&str] = &["region", "tier"];
|
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
|
// Weighted by how many operations each arm covers, so one operation is as likely as another
|
||||||
prop_oneof![
|
prop_oneof![
|
||||||
6 => collection_meta,
|
8 => collection_meta,
|
||||||
1 => arb_update_peer_metadata(),
|
1 => arb_update_peer_metadata(),
|
||||||
1 => arb_update_cluster_metadata(),
|
1 => arb_update_cluster_metadata(),
|
||||||
1 => arb_quota_config().prop_map(ConsensusOperations::SetQuotaConfig),
|
1 => arb_quota_config().prop_map(ConsensusOperations::SetQuotaConfig),
|
||||||
@@ -184,6 +187,8 @@ fn arb_collection_meta_operation(
|
|||||||
|
|
||||||
prop_oneof![
|
prop_oneof![
|
||||||
Just(CollectionMetaOperations::Nop { token: 0 }),
|
Just(CollectionMetaOperations::Nop { token: 0 }),
|
||||||
|
arb_create_collection(collection_names.clone()),
|
||||||
|
arb_delete_collection(collection_names.clone()),
|
||||||
arb_change_aliases(collection_names.clone()),
|
arb_change_aliases(collection_names.clone()),
|
||||||
arb_create_named_vector(collection_names.clone()),
|
arb_create_named_vector(collection_names.clone()),
|
||||||
arb_delete_named_vector(collection_names.clone()),
|
arb_delete_named_vector(collection_names.clone()),
|
||||||
@@ -209,6 +214,60 @@ fn arb_collection_name(names: Vec<String>) -> impl Strategy<Value = String> {
|
|||||||
proptest::sample::select(names)
|
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<String>,
|
||||||
|
) -> impl Strategy<Value = CollectionMetaOperations> {
|
||||||
|
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<Value = Option<ShardDistributionProposal>> {
|
||||||
|
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<String>,
|
||||||
|
) -> impl Strategy<Value = CollectionMetaOperations> {
|
||||||
|
arb_collection_name(collections).prop_map(|collection_name| {
|
||||||
|
CollectionMetaOperations::DeleteCollection(DeleteCollectionOperation(collection_name))
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
fn arb_change_aliases(collections: Vec<String>) -> impl Strategy<Value = CollectionMetaOperations> {
|
fn arb_change_aliases(collections: Vec<String>) -> impl Strategy<Value = CollectionMetaOperations> {
|
||||||
let actions = proptest::collection::vec(arb_alias_operation(collections), 1..=4);
|
let actions = proptest::collection::vec(arb_alias_operation(collections), 1..=4);
|
||||||
|
|
||||||
|
|||||||
@@ -6,15 +6,12 @@ use std::num::NonZeroU32;
|
|||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
|
|
||||||
use collection::collection::Collection;
|
use collection::collection::Collection;
|
||||||
use collection::config::{
|
use collection::config::{self, CollectionConfigInternal, CollectionParams, ShardingMethod};
|
||||||
self, CollectionConfigInternal, CollectionParams, PayloadStorageParams, ShardingMethod,
|
|
||||||
};
|
|
||||||
use collection::operations::config_diff::DiffConfig as _;
|
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::collection_shard_distribution::CollectionShardDistribution;
|
||||||
use collection::shards::replica_set::replica_set_state::ReplicaState;
|
use collection::shards::replica_set::replica_set_state::ReplicaState;
|
||||||
use collection::shards::shard::{PeerId, ShardId};
|
use collection::shards::shard::{PeerId, ShardId};
|
||||||
use segment::types::VectorsConfigDefaults;
|
|
||||||
|
|
||||||
use super::{COLLECTION_DELETE_SPIN_INTERVAL, COLLECTION_DELETE_WAIT_TIMEOUT, TableOfContent};
|
use super::{COLLECTION_DELETE_SPIN_INTERVAL, COLLECTION_DELETE_WAIT_TIMEOUT, TableOfContent};
|
||||||
use crate::common::utils::try_unwrap_with_timeout_async;
|
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::consensus_ops::ConsensusOperations;
|
||||||
use crate::content_manager::errors::StorageError;
|
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<PayloadStorageParams>,
|
|
||||||
on_disk_payload: Option<bool>,
|
|
||||||
defaults: Option<PayloadStorageParams>,
|
|
||||||
) -> Option<PayloadStorageParams> {
|
|
||||||
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 {
|
impl TableOfContent {
|
||||||
pub(super) async fn create_collection(
|
pub(super) async fn create_collection(
|
||||||
&self,
|
&self,
|
||||||
|
|||||||
Reference in New Issue
Block a user