From 9ab060b565b7101dcdc19a5e0dbed6b04b4cd55e Mon Sep 17 00:00:00 2001 From: Roman Titov Date: Wed, 11 Sep 2024 11:11:05 +0200 Subject: [PATCH] Split operations into `update_all`/`update_existing` during resharding (#4928) Co-authored-by: generall --- lib/collection/src/collection/mod.rs | 19 +- lib/collection/src/collection/point_ops.rs | 45 +++- lib/collection/src/operations/mod.rs | 39 ++- lib/collection/src/operations/payload_ops.rs | 32 +++ lib/collection/src/operations/point_ops.rs | 226 +++++++++++++++++- lib/collection/src/operations/types.rs | 12 + lib/collection/src/operations/vector_ops.rs | 19 ++ .../src/shards/replica_set/update.rs | 26 +- .../shards/resharding/stage_migrate_points.rs | 2 +- .../resharding/stage_propagate_deletes.rs | 2 +- .../src/shards/shard_holder/resharding.rs | 153 +++++++++++- src/main.rs | 8 +- 12 files changed, 548 insertions(+), 35 deletions(-) diff --git a/lib/collection/src/collection/mod.rs b/lib/collection/src/collection/mod.rs index 08f791fdcf..4fc837bf8f 100644 --- a/lib/collection/src/collection/mod.rs +++ b/lib/collection/src/collection/mod.rs @@ -419,18 +419,17 @@ impl Collection { // Abort resharding, if resharding shard is marked as `Dead`. // - // This branch should only be triggered, if resharding is currently at - // `ReshardStage::MigratingPoints` stage, because resharding shard should be marked as - // `Active` when all resharding transfers are successfully completed, and so the check - // *right above* this one should be triggered. + // This branch should only be triggered, if resharding is currently at `MigratingPoints` + // stage, because target shard should be marked as `Active`, when all resharding transfers + // are successfully completed, and so the check *right above* this one would be triggered. // - // So, if resharding reached `ReshardingStage::ReadHashRingCommitted`, this branch *won't* - // be triggered, and in this case, resharding *won't* be cancelled. Though, the update - // request should *fail* with "failed to update all replicas of a shard" error. + // So, if resharding reached `ReadHashRingCommitted`, this branch *won't* be triggered, + // and resharding *won't* be cancelled. The update request should *fail* with "failed to + // update all replicas of a shard" error. // - // If resharding reached `ReshardingStage::WriteHashRingCommitted`, and this branch is - // triggered *somehow*, then `Collection::abort_resharding` call should return an error, - // so no special handling is needed for `ReshardingStage::WriteHashRingCommitted`. + // If resharding reached `ReadHashRingCommitted`, and this branch is triggered *somehow*, + // then `Collection::abort_resharding` call should return an error, so no special handling + // is needed. if current_state == Some(ReplicaState::Resharding) && state == ReplicaState::Dead { let shard_key = shard_holder .get_shard_id_to_key_mapping() diff --git a/lib/collection/src/collection/point_ops.rs b/lib/collection/src/collection/point_ops.rs index a6c06fe3f6..8d6dd4f664 100644 --- a/lib/collection/src/collection/point_ops.rs +++ b/lib/collection/src/collection/point_ops.rs @@ -108,7 +108,7 @@ impl Collection { } shard - .update_with_consistency(operation.operation, wait, ordering) + .update_with_consistency(operation.operation, wait, ordering, false) .await .map(Some) } @@ -145,13 +145,42 @@ impl Collection { let mut results = tokio::task::spawn(async move { let _update_lock = update_lock; - let updates: FuturesUnordered<_> = shard_holder - .split_by_shard(operation, &shard_keys_selection)? - .into_iter() - .map(move |(shard, operation)| { - shard.update_with_consistency(operation, wait, ordering) - }) - .collect(); + let updates = FuturesUnordered::new(); + let operations = shard_holder.split_by_shard(operation, &shard_keys_selection)?; + + for (shard, operation) in operations { + let operation = shard_holder.split_by_mode(shard.shard_id, operation); + + updates.push(async move { + let mut result = UpdateResult { + operation_id: None, + status: UpdateStatus::Acknowledged, + clock_tag: None, + }; + + for operation in operation.update_all { + result = shard + .update_with_consistency(operation, wait, ordering, false) + .await?; + } + + for operation in operation.update_only_existing { + let res = shard + .update_with_consistency(operation, wait, ordering, true) + .await; + + if let Err(err) = &res { + if err.is_missing_point() { + continue; + } + } + + result = res?; + } + + CollectionResult::Ok(result) + }); + } let results: Vec<_> = updates.collect().await; diff --git a/lib/collection/src/operations/mod.rs b/lib/collection/src/operations/mod.rs index e9937e006b..2722ecb944 100644 --- a/lib/collection/src/operations/mod.rs +++ b/lib/collection/src/operations/mod.rs @@ -21,7 +21,7 @@ pub mod verification; use std::collections::HashMap; use segment::json_path::JsonPath; -use segment::types::{ExtendedPointId, PayloadFieldSchema}; +use segment::types::{ExtendedPointId, PayloadFieldSchema, PointIdType}; use serde::{Deserialize, Serialize}; use strum::{EnumDiscriminants, EnumIter}; use validator::Validate; @@ -146,6 +146,43 @@ pub enum CollectionUpdateOperations { FieldIndexOperation(FieldIndexOperations), } +impl CollectionUpdateOperations { + pub fn is_upsert_points(&self) -> bool { + matches!( + self, + Self::PointOperation(point_ops::PointOperations::UpsertPoints(_)) + ) + } + + pub fn is_delete_points(&self) -> bool { + matches!( + self, + Self::PointOperation(point_ops::PointOperations::DeletePoints { .. }) + ) + } + + pub fn point_ids(&self) -> Vec { + match self { + Self::PointOperation(op) => op.point_ids(), + Self::VectorOperation(op) => op.point_ids(), + Self::PayloadOperation(op) => op.point_ids(), + Self::FieldIndexOperation(_) => Vec::new(), + } + } + + pub fn retain_point_ids(&mut self, filter: F) + where + F: Fn(&PointIdType) -> bool, + { + match self { + Self::PointOperation(op) => op.retain_point_ids(filter), + Self::VectorOperation(op) => op.retain_point_ids(filter), + Self::PayloadOperation(op) => op.retain_point_ids(filter), + Self::FieldIndexOperation(_) => (), + } + } +} + /// A mapping of operation to shard. /// Is a result of splitting one operation into several shards by corresponding PointIds pub enum OperationToShard { diff --git a/lib/collection/src/operations/payload_ops.rs b/lib/collection/src/operations/payload_ops.rs index 0e61da650b..465a2de66c 100644 --- a/lib/collection/src/operations/payload_ops.rs +++ b/lib/collection/src/operations/payload_ops.rs @@ -160,6 +160,38 @@ impl PayloadOps { PayloadOps::OverwritePayload(_) => true, } } + + pub fn point_ids(&self) -> Vec { + match self { + Self::SetPayload(op) => op.points.clone().unwrap_or(Vec::new()), + Self::DeletePayload(op) => op.points.clone().unwrap_or(Vec::new()), + Self::ClearPayload { points } => points.clone(), + Self::ClearPayloadByFilter(_) => Vec::new(), + Self::OverwritePayload(op) => op.points.clone().unwrap_or(Vec::new()), + } + } + + pub fn retain_point_ids(&mut self, filter: F) + where + F: Fn(&PointIdType) -> bool, + { + match self { + Self::SetPayload(op) => retain_opt(op.points.as_mut(), filter), + Self::DeletePayload(op) => retain_opt(op.points.as_mut(), filter), + Self::ClearPayload { points } => points.retain(filter), + Self::ClearPayloadByFilter(_) => (), + Self::OverwritePayload(op) => retain_opt(op.points.as_mut(), filter), + } + } +} + +fn retain_opt(vec: Option<&mut Vec>, filter: F) +where + F: Fn(&T) -> bool, +{ + if let Some(vec) = vec { + vec.retain(filter); + } } impl Validate for PayloadOps { diff --git a/lib/collection/src/operations/point_ops.rs b/lib/collection/src/operations/point_ops.rs index ec04c1d01d..7e143098d8 100644 --- a/lib/collection/src/operations/point_ops.rs +++ b/lib/collection/src/operations/point_ops.rs @@ -1,5 +1,6 @@ use std::borrow::Cow; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; +use std::iter; use api::rest::{BatchVectorStruct, ShardKeySelector, VectorStruct}; use itertools::izip; @@ -12,9 +13,15 @@ use serde::{Deserialize, Serialize}; use strum::{EnumDiscriminants, EnumIter}; use validator::Validate; -use super::{point_to_shards, split_iter_by_shard, OperationToShard, SplitByShard}; +use super::payload_ops::SetPayloadOp; +use super::vector_ops::{PointVectors, UpdateVectorsOp}; +use super::{ + point_to_shards, split_iter_by_shard, CollectionUpdateOperations, OperationToShard, + SplitByShard, +}; use crate::hash_ring::HashRingRouter; use crate::operations::types::Record; +use crate::operations::{payload_ops, vector_ops}; use crate::shards::shard::ShardId; /// Defines write ordering guarantees for collection operations @@ -216,6 +223,200 @@ pub enum PointInsertOperationsInternal { PointsList(Vec), } +impl PointInsertOperationsInternal { + pub fn point_ids(&self) -> Vec { + match self { + Self::PointsBatch(batch) => batch.ids.clone(), + Self::PointsList(points) => points.iter().map(|point| point.id).collect(), + } + } + + pub fn retain_point_ids(&mut self, filter: F) + where + F: Fn(&PointIdType) -> bool, + { + match self { + Self::PointsBatch(batch) => { + let mut retain_indices = HashSet::new(); + + retain_with_index(&mut batch.ids, |index, id| { + if filter(id) { + retain_indices.insert(index); + true + } else { + false + } + }); + + match &mut batch.vectors { + BatchVectorStruct::Single(vectors) => { + retain_with_index(vectors, |index, _| retain_indices.contains(&index)); + } + + BatchVectorStruct::MultiDense(vectors) => { + retain_with_index(vectors, |index, _| retain_indices.contains(&index)); + } + + BatchVectorStruct::Named(vectors) => { + for (_, vectors) in vectors.iter_mut() { + retain_with_index(vectors, |index, _| retain_indices.contains(&index)); + } + } + + BatchVectorStruct::Document(documents) => { + retain_with_index(documents, |index, _| retain_indices.contains(&index)); + } + } + + if let Some(payload) = &mut batch.payloads { + retain_with_index(payload, |index, _| retain_indices.contains(&index)); + } + } + + Self::PointsList(points) => points.retain(|point| filter(&point.id)), + } + } + + pub fn into_update_only(self) -> Vec { + let mut operations = Vec::new(); + + match self { + Self::PointsBatch(batch) => { + let mut update_vectors = UpdateVectorsOp { points: Vec::new() }; + + match batch.vectors { + BatchVectorStruct::Single(vectors) => { + let ids = batch.ids.iter().copied(); + let vectors = vectors.into_iter().map(VectorStruct::Single); + + update_vectors.points = ids + .zip(vectors) + .map(|(id, vector)| PointVectors { id, vector }) + .collect(); + } + + BatchVectorStruct::MultiDense(vectors) => { + let ids = batch.ids.iter().copied(); + let vectors = vectors.into_iter().map(VectorStruct::MultiDense); + + update_vectors.points = ids + .zip(vectors) + .map(|(id, vector)| PointVectors { id, vector }) + .collect(); + } + + BatchVectorStruct::Named(batch_vectors) => { + let ids = batch.ids.iter().copied(); + + let mut batch_vectors: HashMap<_, _> = batch_vectors + .into_iter() + .map(|(name, vectors)| (name, vectors.into_iter())) + .collect(); + + let vectors = iter::repeat(()).filter_map(move |_| { + let mut point_vectors = + HashMap::with_capacity(batch_vectors.capacity()); + + for (vector_name, vectors) in batch_vectors.iter_mut() { + point_vectors.insert(vector_name.clone(), vectors.next()?); + } + + Some(VectorStruct::Named(point_vectors)) + }); + + update_vectors.points = ids + .zip(vectors) + .map(|(id, vector)| PointVectors { id, vector }) + .collect(); + } + + BatchVectorStruct::Document(documents) => { + let ids = batch.ids.iter().copied(); + let documents = documents.into_iter().map(VectorStruct::Document); + + update_vectors.points = ids + .zip(documents) + .map(|(id, vector)| PointVectors { id, vector }) + .collect(); + } + } + + let update_vectors = vector_ops::VectorOperations::UpdateVectors(update_vectors); + let update_vectors = CollectionUpdateOperations::VectorOperation(update_vectors); + + operations.push(update_vectors); + + if let Some(payloads) = batch.payloads { + let ids = batch.ids.iter().copied(); + + for (id, payload) in ids.zip(payloads) { + if let Some(payload) = payload { + let set_payload = SetPayloadOp { + points: Some(vec![id]), + payload, + filter: None, + key: None, + }; + + let set_payload = + payload_ops::PayloadOps::OverwritePayload(set_payload); + let set_payload = + CollectionUpdateOperations::PayloadOperation(set_payload); + + operations.push(set_payload); + } + } + } + } + + Self::PointsList(points) => { + let mut update_vectors = UpdateVectorsOp { points: Vec::new() }; + + for point in points { + update_vectors.points.push(PointVectors { + id: point.id, + vector: point.vector, + }); + + if let Some(payload) = point.payload { + let set_payload = SetPayloadOp { + points: Some(vec![point.id]), + payload, + filter: None, + key: None, + }; + + let set_payload = payload_ops::PayloadOps::OverwritePayload(set_payload); + let set_payload = CollectionUpdateOperations::PayloadOperation(set_payload); + + operations.push(set_payload); + } + } + + let update_vectors = vector_ops::VectorOperations::UpdateVectors(update_vectors); + let update_vectors = CollectionUpdateOperations::VectorOperation(update_vectors); + + operations.insert(0, update_vectors); + } + } + + operations + } +} + +fn retain_with_index(vec: &mut Vec, mut filter: F) +where + F: FnMut(usize, &T) -> bool, +{ + let mut index = 0; + + vec.retain(|item| { + let retain = filter(index, item); + index += 1; + retain + }); +} + impl Validate for PointInsertOperationsInternal { fn validate(&self) -> Result<(), validator::ValidationErrors> { match self { @@ -355,6 +556,27 @@ impl PointOperations { PointOperations::SyncPoints(_) => true, } } + + pub fn point_ids(&self) -> Vec { + match self { + Self::UpsertPoints(op) => op.point_ids(), + Self::DeletePoints { ids } => ids.clone(), + Self::DeletePointsByFilter(_) => Vec::new(), + Self::SyncPoints(op) => op.points.iter().map(|point| point.id).collect(), + } + } + + pub fn retain_point_ids(&mut self, filter: F) + where + F: Fn(&PointIdType) -> bool, + { + match self { + Self::UpsertPoints(op) => op.retain_point_ids(filter), + Self::DeletePoints { ids } => ids.retain(filter), + Self::DeletePointsByFilter(_) => (), + Self::SyncPoints(op) => op.points.retain(|point| filter(&point.id)), + } + } } impl Validate for PointOperations { diff --git a/lib/collection/src/operations/types.rs b/lib/collection/src/operations/types.rs index dc8df6cece..5f97658665 100644 --- a/lib/collection/src/operations/types.rs +++ b/lib/collection/src/operations/types.rs @@ -1019,6 +1019,18 @@ impl CollectionError { Self::StrictMode { .. } => false, } } + + pub fn is_pre_condition_failed(&self) -> bool { + matches!(self, Self::PreConditionFailed { .. }) + } + + pub fn is_missing_point(&self) -> bool { + match self { + CollectionError::NotFound { what } => what.contains("No point with id"), + CollectionError::PointNotFound { .. } => true, + _ => false, + } + } } impl From for CollectionError { diff --git a/lib/collection/src/operations/vector_ops.rs b/lib/collection/src/operations/vector_ops.rs index 9378861e83..8480dc5e8e 100644 --- a/lib/collection/src/operations/vector_ops.rs +++ b/lib/collection/src/operations/vector_ops.rs @@ -88,6 +88,25 @@ impl VectorOperations { VectorOperations::DeleteVectorsByFilter(..) => false, } } + + pub fn point_ids(&self) -> Vec { + match self { + Self::UpdateVectors(op) => op.points.iter().map(|point| point.id).collect(), + Self::DeleteVectors(points, _) => points.points.clone(), + Self::DeleteVectorsByFilter(_, _) => Vec::new(), + } + } + + pub fn retain_point_ids(&mut self, filter: F) + where + F: Fn(&PointIdType) -> bool, + { + match self { + Self::UpdateVectors(op) => op.points.retain(|point| filter(&point.id)), + Self::DeleteVectors(points, _) => points.points.retain(filter), + Self::DeleteVectorsByFilter(_, _) => (), + } + } } impl Validate for VectorOperations { diff --git a/lib/collection/src/shards/replica_set/update.rs b/lib/collection/src/shards/replica_set/update.rs index e2dd755f91..04b0a7b3bd 100644 --- a/lib/collection/src/shards/replica_set/update.rs +++ b/lib/collection/src/shards/replica_set/update.rs @@ -71,6 +71,7 @@ impl ShardReplicaSet { operation: CollectionUpdateOperations, wait: bool, ordering: WriteOrdering, + update_only_existing: bool, ) -> CollectionResult { // `ShardReplicaSet::update` is not cancel safe, so this method is not cancel safe. @@ -91,7 +92,7 @@ impl ShardReplicaSet { WriteOrdering::Weak => None, }; - self.update(operation, wait).await + self.update(operation, wait, update_only_existing).await } else { // Forward the update to the designated leader self.forward_update(leader_peer, operation, wait, ordering) @@ -128,7 +129,7 @@ impl ShardReplicaSet { peer_ids .into_iter() - .filter(|peer_id| self.peer_is_active(peer_id)) // re-acquire replica_state read lock + .filter(|peer_id| self.peer_is_active_or_resharding(peer_id)) // re-acquire replica_state read lock .max() } @@ -143,6 +144,7 @@ impl ShardReplicaSet { &self, operation: CollectionUpdateOperations, wait: bool, + update_only_existing: bool, ) -> CollectionResult { // `ShardRepilcaSet::update_impl` is not cancel safe, so this method is not cancel safe. @@ -158,7 +160,7 @@ impl ShardReplicaSet { let is_non_zero_tick = clock.current_tick().is_some(); let res = self - .update_impl(operation.clone(), wait, &mut clock) + .update_impl(operation.clone(), wait, &mut clock, update_only_existing) .await?; if let Some(res) = res { @@ -190,6 +192,7 @@ impl ShardReplicaSet { operation: CollectionUpdateOperations, wait: bool, clock: &mut clock_set::ClockGuard, + update_only_existing: bool, ) -> CollectionResult> { // `LocalShard::update` is not guaranteed to be cancel safe and it's impossible to cancel // multiple parallel updates in a way that is *guaranteed* not to introduce inconsistencies @@ -330,8 +333,11 @@ impl ShardReplicaSet { }; if successes.len() >= minimal_success_count { - let wait_for_deactivation = - self.handle_failed_replicas(&failures, &self.replica_state.read()); + let wait_for_deactivation = self.handle_failed_replicas( + &failures, + &self.replica_state.read(), + update_only_existing, + ); // report all failing peers to consensus if wait && wait_for_deactivation && !failures.is_empty() { @@ -373,6 +379,7 @@ impl ShardReplicaSet { .iter() .filter(|(peer_id, _)| self.peer_is_resharding(peer_id)), &self.replica_state.read(), + update_only_existing, ); // completely failed - report error to user @@ -431,6 +438,7 @@ impl ShardReplicaSet { &self, failures: impl IntoIterator, state: &ReplicaSetState, + update_only_existing: bool, ) -> bool { let mut wait_for_deactivation = false; @@ -450,14 +458,18 @@ impl ShardReplicaSet { _ => continue, } - if peer_state == ReplicaState::Partial - && matches!(err, CollectionError::PreConditionFailed { .. }) + if matches!(peer_state, ReplicaState::Partial | ReplicaState::Resharding) + && err.is_pre_condition_failed() { // Handles a special case where transfer receiver haven't created a shard yet. // In this case update should be handled by source shard and forward proxy. continue; } + if update_only_existing && err.is_missing_point() { + continue; + } + if err.is_transient() || peer_state == ReplicaState::Initializing { // If the error is transient, we should not deactivate the peer // before allowing other operations to continue. diff --git a/lib/collection/src/shards/resharding/stage_migrate_points.rs b/lib/collection/src/shards/resharding/stage_migrate_points.rs index 9b3bb2ca44..cc36c1decc 100644 --- a/lib/collection/src/shards/resharding/stage_migrate_points.rs +++ b/lib/collection/src/shards/resharding/stage_migrate_points.rs @@ -353,7 +353,7 @@ async fn drive_down( // Wait on all updates here, not just the last batch // If we don't wait on all updates it somehow results in inconsistent results target_replica_set - .update_with_consistency(operation, true, WriteOrdering::Weak) + .update_with_consistency(operation, true, WriteOrdering::Weak, false) .await?; if offset.is_none() { diff --git a/lib/collection/src/shards/resharding/stage_propagate_deletes.rs b/lib/collection/src/shards/resharding/stage_propagate_deletes.rs index 1da26679dc..ffa7a81d42 100644 --- a/lib/collection/src/shards/resharding/stage_propagate_deletes.rs +++ b/lib/collection/src/shards/resharding/stage_propagate_deletes.rs @@ -95,7 +95,7 @@ pub(super) async fn drive( // Wait on all updates here, not just the last batch // If we don't wait on all updates it somehow results in inconsistent deletes replica_set - .update_with_consistency(operation, true, WriteOrdering::Weak) + .update_with_consistency(operation, true, WriteOrdering::Weak, false) .await?; if offset.is_none() { diff --git a/lib/collection/src/shards/shard_holder/resharding.rs b/lib/collection/src/shards/shard_holder/resharding.rs index fc01d95acb..d2a16f06df 100644 --- a/lib/collection/src/shards/shard_holder/resharding.rs +++ b/lib/collection/src/shards/shard_holder/resharding.rs @@ -1,18 +1,23 @@ -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::fmt; use std::ops::Deref as _; -use segment::types::ShardKey; +use segment::types::{CustomIdCheckerCondition as _, ShardKey}; use super::ShardHolder; use crate::hash_ring::{self, HashRingRouter}; use crate::operations::cluster_ops::ReshardingDirection; use crate::operations::types::{CollectionError, CollectionResult}; +use crate::operations::{point_ops, CollectionUpdateOperations}; use crate::shards::replica_set::{ReplicaState, ShardReplicaSet}; use crate::shards::resharding::{ReshardKey, ReshardStage, ReshardState}; use crate::shards::shard::ShardId; impl ShardHolder { + pub fn resharding_state(&self) -> Option { + self.resharding_state.read().clone() + } + pub fn check_start_resharding(&mut self, resharding_key: &ReshardKey) -> CollectionResult<()> { let ReshardKey { direction, @@ -364,6 +369,117 @@ impl ShardHolder { Ok(()) } + /// Split collection update operation by "update mode": + /// - update all: + /// - "regular" operation + /// - `upsert` inserts new points and updates existing ones + /// - other update operations return error, if a point does not exist in collection + /// - update existing: + /// - `upsert` does *not* insert new points, only updates existing ones + /// - other update operations ignore points that do not exist in collection + /// + /// Depends on the current resharding state. If resharding is not active operations are not split. + pub fn split_by_mode( + &self, + shard_id: ShardId, + operation: CollectionUpdateOperations, + ) -> OperationsByMode { + let Some(state) = self.resharding_state() else { + return OperationsByMode::from(operation); + }; + + // Resharding *UP* + // ┌────────────┐ ┌──────────┐ + // │ │ │ │ + // │ Shard 1 │ │ Shard 2 │ + // │ Non-Target ├──►│ Target │ + // │ Sender │ │ Receiver │ + // │ │ │ │ + // └────────────┘ └──────────┘ + // + // Resharding *DOWN* + // ┌────────────┐ ┌──────────┐ + // │ │ │ │ + // │ Shard 1 │ │ Shard 2 │ + // │ Non-Target │◄──┤ Target │ + // │ Receiver │ │ Sender │ + // │ │ │ │ + // └────────────┘ └──────────┘ + + // Target shard of the resharding operation. This is the shard that: + // + // - *created* during resharding *up* + // - *deleted* during resharding *down* + let is_target_shard = shard_id == state.shard_id; + + // Shard that will be *receiving* migrated points during resharding: + // + // - *target* shard during resharding *up* + // - *non* target shards during resharding *down* + let is_receiver_shard = match state.direction { + ReshardingDirection::Up => is_target_shard, + ReshardingDirection::Down => !is_target_shard, + }; + + // Shard that will be *sending* migrated points during resharding: + // + // - *non* target shards during resharding *up* + // - *target* shard during resharding *down* + let is_sender_shard = !is_receiver_shard; + + // We split update operations: + // + // - on *receiver* shards during `MigratingPoints` stage (for all operations except `upsert`) + // - and on *sender* shards during `ReadHashRingCommitted` stage when resharding *up* + + let should_split_receiver = is_receiver_shard + && state.stage == ReshardStage::MigratingPoints + && !operation.is_upsert_points(); + + let should_split_sender = is_sender_shard + && state.stage >= ReshardStage::ReadHashRingCommitted + && state.direction == ReshardingDirection::Up; + + if !should_split_receiver && !should_split_sender { + return OperationsByMode::from(operation); + } + + // There's no point splitting delete operations + if operation.is_delete_points() { + return OperationsByMode::from(operation); + } + + let Some(filter) = self.resharding_filter() else { + return OperationsByMode::from(operation); + }; + + let point_ids = operation.point_ids(); + + if point_ids.is_empty() { + return OperationsByMode::from(operation); + } + + let target_point_ids: HashSet<_> = point_ids + .iter() + .copied() + .filter(|&point_id| filter.check(point_id)) + .collect(); + + if target_point_ids.is_empty() { + OperationsByMode::from(operation) + } else if target_point_ids.len() == point_ids.len() { + OperationsByMode::default().with_update_only_existing(operation) + } else { + let mut update_all = operation.clone(); + update_all.retain_point_ids(|point_id| !target_point_ids.contains(point_id)); + + let mut update_only_existing = operation; + update_only_existing.retain_point_ids(|point_id| target_point_ids.contains(point_id)); + + OperationsByMode::from(update_all).with_update_only_existing(update_only_existing) + } + } + pub fn resharding_filter(&self) -> Option { let shard_id = self.resharding_state.read().as_ref()?.shard_id; self.hash_ring_filter(shard_id) @@ -391,6 +507,39 @@ impl ShardHolder { } } +#[derive(Clone, Debug, Default)] +pub struct OperationsByMode { + pub update_all: Vec, + pub update_only_existing: Vec, +} + +impl OperationsByMode { + pub fn with_update_only_existing(mut self, operation: CollectionUpdateOperations) -> Self { + match operation { + CollectionUpdateOperations::PointOperation( + point_ops::PointOperations::UpsertPoints(operation), + ) => { + self.update_only_existing = operation.into_update_only(); + } + + operation => { + self.update_only_existing = vec![operation]; + } + } + + self + } +} + +impl From for OperationsByMode { + fn from(operation: CollectionUpdateOperations) -> Self { + Self { + update_all: vec![operation], + update_only_existing: Vec::new(), + } + } +} + fn get_ring<'a>( rings: &'a mut HashMap, HashRingRouter>, shard_key: &'_ Option, diff --git a/src/main.rs b/src/main.rs index bf22ac6dc6..8bb1ef291a 100644 --- a/src/main.rs +++ b/src/main.rs @@ -348,9 +348,11 @@ fn main() -> anyhow::Result<()> { } }); - runtime_handle.block_on(async { - toc_arc.resume_resharding_tasks().await; - }); + // TODO(resharding): Remove resharding driver? + // + // runtime_handle.block_on(async { + // toc_arc.resume_resharding_tasks().await; + // }); let collections_to_recover_in_consensus = if is_new_deployment { let existing_collections =