diff --git a/lib/api/build.rs b/lib/api/build.rs index dbf11e7424..6728bb4242 100644 --- a/lib/api/build.rs +++ b/lib/api/build.rs @@ -159,6 +159,7 @@ fn configure_validation(builder: Builder) -> Builder { ("WaitForShardStateRequest.collection_name", "length(min = 1, max = 255)"), ("WaitForShardStateRequest.timeout", "range(min = 1)"), ("GetShardRecoveryPointRequest.collection_name", "length(min = 1, max = 255)"), + ("UpdateShardCutoffPointRequest.collection_name", "length(min = 1, max = 255)"), ], &[]) // Service: points.proto .validates(&[ diff --git a/lib/api/src/grpc/proto/collections_internal_service.proto b/lib/api/src/grpc/proto/collections_internal_service.proto index 556733a1a8..3d5dc824fd 100644 --- a/lib/api/src/grpc/proto/collections_internal_service.proto +++ b/lib/api/src/grpc/proto/collections_internal_service.proto @@ -23,6 +23,10 @@ service CollectionsInternal { Get shard recovery point */ rpc GetShardRecoveryPoint (GetShardRecoveryPointRequest) returns (GetShardRecoveryPointResponse) {} + /* + Update shard cutoff point + */ + rpc UpdateShardCutoffPoint (UpdateShardCutoffPointRequest) returns (CollectionOperationResponse) {} } message GetCollectionInfoRequestInternal { @@ -61,3 +65,9 @@ message RecoveryPointClockTag { uint32 clock_id = 2; uint64 clock_tick = 3; } + +message UpdateShardCutoffPointRequest { + string collection_name = 1; // Name of the collection + uint32 shard_id = 2; // Id of the shard + RecoveryPoint cutoff = 3; // Cutoff point of the shard +} diff --git a/lib/api/src/grpc/qdrant.rs b/lib/api/src/grpc/qdrant.rs index f4558e7131..dc126164c8 100644 --- a/lib/api/src/grpc/qdrant.rs +++ b/lib/api/src/grpc/qdrant.rs @@ -6732,6 +6732,22 @@ pub struct RecoveryPointClockTag { #[prost(uint64, tag = "3")] pub clock_tick: u64, } +#[derive(validator::Validate)] +#[derive(serde::Serialize)] +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct UpdateShardCutoffPointRequest { + /// Name of the collection + #[prost(string, tag = "1")] + #[validate(length(min = 1, max = 255))] + pub collection_name: ::prost::alloc::string::String, + /// Id of the shard + #[prost(uint32, tag = "2")] + pub shard_id: u32, + /// Cutoff point of the shard + #[prost(message, optional, tag = "3")] + pub cutoff: ::core::option::Option, +} /// Generated client implementations. pub mod collections_internal_client { #![allow(unused_variables, dead_code, missing_docs, clippy::let_unit_value)] @@ -6932,6 +6948,38 @@ pub mod collections_internal_client { ); self.inner.unary(req, path, codec).await } + /// + /// Update shard cutoff point + pub async fn update_shard_cutoff_point( + &mut self, + request: impl tonic::IntoRequest, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + > { + self.inner + .ready() + .await + .map_err(|e| { + tonic::Status::new( + tonic::Code::Unknown, + format!("Service was not ready: {}", e.into()), + ) + })?; + let codec = tonic::codec::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/qdrant.CollectionsInternal/UpdateShardCutoffPoint", + ); + let mut req = request.into_request(); + req.extensions_mut() + .insert( + GrpcMethod::new( + "qdrant.CollectionsInternal", + "UpdateShardCutoffPoint", + ), + ); + self.inner.unary(req, path, codec).await + } } } /// Generated server implementations. @@ -6977,6 +7025,15 @@ pub mod collections_internal_server { tonic::Response, tonic::Status, >; + /// + /// Update shard cutoff point + async fn update_shard_cutoff_point( + &self, + request: tonic::Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + >; } #[derive(Debug)] pub struct CollectionsInternalServer { @@ -7252,6 +7309,56 @@ pub mod collections_internal_server { }; Box::pin(fut) } + "/qdrant.CollectionsInternal/UpdateShardCutoffPoint" => { + #[allow(non_camel_case_types)] + struct UpdateShardCutoffPointSvc(pub Arc); + impl< + T: CollectionsInternal, + > tonic::server::UnaryService + for UpdateShardCutoffPointSvc { + type Response = super::CollectionOperationResponse; + type Future = BoxFuture< + tonic::Response, + tonic::Status, + >; + fn call( + &mut self, + request: tonic::Request, + ) -> Self::Future { + let inner = Arc::clone(&self.0); + let fut = async move { + ::update_shard_cutoff_point( + &inner, + request, + ) + .await + }; + Box::pin(fut) + } + } + let accept_compression_encodings = self.accept_compression_encodings; + let send_compression_encodings = self.send_compression_encodings; + let max_decoding_message_size = self.max_decoding_message_size; + let max_encoding_message_size = self.max_encoding_message_size; + let inner = self.inner.clone(); + let fut = async move { + let inner = inner.0; + let method = UpdateShardCutoffPointSvc(inner); + let codec = tonic::codec::ProstCodec::default(); + let mut grpc = tonic::server::Grpc::new(codec) + .apply_compression_config( + accept_compression_encodings, + send_compression_encodings, + ) + .apply_max_message_size_config( + max_decoding_message_size, + max_encoding_message_size, + ); + let res = grpc.unary(method, req).await; + Ok(res) + }; + Box::pin(fut) + } _ => { Box::pin(async move { Ok( diff --git a/lib/collection/src/collection/mod.rs b/lib/collection/src/collection/mod.rs index b053ca0d92..c675fcbb2d 100644 --- a/lib/collection/src/collection/mod.rs +++ b/lib/collection/src/collection/mod.rs @@ -416,6 +416,23 @@ impl Collection { replica_set.shard_recovery_point().await } + pub async fn update_shard_cutoff_point( + &self, + shard_id: ShardId, + cutoff: &RecoveryPoint, + ) -> CollectionResult<()> { + let shard_holder_read = self.shards_holder.read().await; + + let shard = shard_holder_read.get_shard(&shard_id); + let Some(replica_set) = shard else { + return Err(CollectionError::NotFound { + what: "Shard {shard_id}".into(), + }); + }; + + replica_set.update_shard_cutoff_point(cutoff).await + } + pub async fn state(&self) -> State { let shards_holder = self.shards_holder.read().await; let transfers = shards_holder.shard_transfers.read().clone(); diff --git a/lib/collection/src/shards/local_shard/mod.rs b/lib/collection/src/shards/local_shard/mod.rs index 0a77dd7c01..9abb13d15d 100644 --- a/lib/collection/src/shards/local_shard/mod.rs +++ b/lib/collection/src/shards/local_shard/mod.rs @@ -907,6 +907,13 @@ impl LocalShard { pub async fn recovery_point(&self) -> RecoveryPoint { self.wal.recovery_point().await } + + /// Update the cutoff point on the current shard + /// + /// This also updates the highest seen clocks. + pub async fn update_cutoff(&self, cutoff: &RecoveryPoint) { + self.wal.update_cutoff(cutoff).await + } } impl Drop for LocalShard { diff --git a/lib/collection/src/shards/replica_set/mod.rs b/lib/collection/src/shards/replica_set/mod.rs index 2f0f462af9..2544ff5c0a 100644 --- a/lib/collection/src/shards/replica_set/mod.rs +++ b/lib/collection/src/shards/replica_set/mod.rs @@ -829,6 +829,21 @@ impl ShardReplicaSet { local_shard.shard_recovery_point().await } + + /// Update the cutoff point for the local shard. + pub(crate) async fn update_shard_cutoff_point( + &self, + cutoff: &RecoveryPoint, + ) -> CollectionResult<()> { + let local_shard = self.local.read().await; + let Some(local_shard) = local_shard.as_ref() else { + return Err(CollectionError::NotFound { + what: "Peer does not have local shard".into(), + }); + }; + + local_shard.update_cutoff(cutoff).await + } } /// Represents a replica set state diff --git a/lib/collection/src/shards/shard.rs b/lib/collection/src/shards/shard.rs index aa2423c13e..07129b0e39 100644 --- a/lib/collection/src/shards/shard.rs +++ b/lib/collection/src/shards/shard.rs @@ -163,6 +163,22 @@ impl Shard { } } + pub async fn update_cutoff(&self, cutoff: &RecoveryPoint) -> CollectionResult<()> { + match self { + Self::Local(local_shard) => local_shard.update_cutoff(cutoff).await, + Self::ForwardProxy(proxy_shard) => { + proxy_shard.wrapped_shard.update_cutoff(cutoff).await + } + Self::Proxy(_) | Self::QueueProxy(_) | Self::Dummy(_) => { + return Err(CollectionError::service_error(format!( + "Setting cutoff point not supported on {}", + self.variant_name(), + ))); + } + } + Ok(()) + } + pub async fn resolve_wal_delta( &self, recovery_point: RecoveryPoint, diff --git a/lib/collection/src/wal_delta.rs b/lib/collection/src/wal_delta.rs index 58bd273f8a..ffb172d7f0 100644 --- a/lib/collection/src/wal_delta.rs +++ b/lib/collection/src/wal_delta.rs @@ -77,19 +77,27 @@ impl RecoverableWal { wal_lock.write(operation).map(|op_num| (op_num, wal_lock)) } - /// Update the cutoff clock map based on the given recovery point. + /// Update the cutoff clock map based on the given recovery point /// /// This can only increase clock ticks in the cutoff clock map. If there already are higher /// clock ticks, they're kept. /// /// It updates the highest seen clocks alongside with it. - pub async fn update_cutoff(&self, recovery_point: &RecoveryPoint) { - let mut highest_clocks = self.highest_clocks.lock().await; - let mut cutoff_clocks = self.cutoff_clocks.lock().await; - recovery_point.clock_tag_iter().for_each(|clock_tag| { - highest_clocks.advance_clock(clock_tag); - cutoff_clocks.advance_clock(clock_tag); - }); + pub async fn update_cutoff(&self, cutoff: &RecoveryPoint) { + // Lock highest and cutoff maps separately to avoid deadlocks + { + let mut highest_clocks = self.highest_clocks.lock().await; + for clock_tag in cutoff.clock_tag_iter() { + highest_clocks.advance_clock(clock_tag); + } + } + + { + let mut cutoff_clocks = self.cutoff_clocks.lock().await; + for clock_tag in cutoff.clock_tag_iter() { + cutoff_clocks.advance_clock(clock_tag); + } + } } /// Get a recovery point for this WAL. diff --git a/src/tonic/api/collections_internal_api.rs b/src/tonic/api/collections_internal_api.rs index c001306585..30b10a55e7 100644 --- a/src/tonic/api/collections_internal_api.rs +++ b/src/tonic/api/collections_internal_api.rs @@ -5,7 +5,7 @@ use api::grpc::qdrant::collections_internal_server::CollectionsInternal; use api::grpc::qdrant::{ CollectionOperationResponse, GetCollectionInfoRequestInternal, GetCollectionInfoResponse, GetShardRecoveryPointRequest, GetShardRecoveryPointResponse, InitiateShardTransferRequest, - WaitForShardStateRequest, + UpdateShardCutoffPointRequest, WaitForShardStateRequest, }; use storage::content_manager::conversions::error_to_status; use storage::content_manager::toc::TableOfContent; @@ -155,4 +155,46 @@ impl CollectionsInternal for CollectionsInternalService { }; Ok(Response::new(response)) } + + async fn update_shard_cutoff_point( + &self, + request: Request, + ) -> Result, Status> { + validate_and_log(request.get_ref()); + + let timing = Instant::now(); + let UpdateShardCutoffPointRequest { + collection_name, + shard_id, + cutoff, + } = request.into_inner(); + + let cutoff = cutoff.ok_or_else(|| Status::invalid_argument("Missing cutoff point"))?; + + let collection_read = self + .toc + .get_collection(&collection_name) + .await + .map_err(|err| { + Status::not_found(format!( + "Collection {collection_name} could not be found: {err}" + )) + })?; + + // Set the shard cutoff point + collection_read + .update_shard_cutoff_point(shard_id, &cutoff.into()) + .await + .map_err(|err| { + Status::internal(format!( + "Failed to set shard cutoff point for shard {shard_id}: {err}" + )) + })?; + + let response = CollectionOperationResponse { + result: true, + time: timing.elapsed().as_secs_f64(), + }; + Ok(Response::new(response)) + } }