Add gRPC API to set shard cutoff point (#3661)

* Add functions to propagate updating cutoff point from collection level

* Add gRPC endpoint to set cutoff point

* Lock highest and cutoff clock maps separately

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(&[

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

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<RecoveryPoint>,
+}
 /// 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<super::UpdateShardCutoffPointRequest>,
+        ) -> std::result::Result<
+            tonic::Response<super::CollectionOperationResponse>,
+            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<super::GetShardRecoveryPointResponse>,
             tonic::Status,
         >;
+        ///
+        /// Update shard cutoff point
+        async fn update_shard_cutoff_point(
+            &self,
+            request: tonic::Request<super::UpdateShardCutoffPointRequest>,
+        ) -> std::result::Result<
+            tonic::Response<super::CollectionOperationResponse>,
+            tonic::Status,
+        >;
     }
     #[derive(Debug)]
     pub struct CollectionsInternalServer<T: CollectionsInternal> {
@@ -7252,6 +7309,56 @@ pub mod collections_internal_server {
                     };
                     Box::pin(fut)
                 }
+                "/qdrant.CollectionsInternal/UpdateShardCutoffPoint" => {
+                    #[allow(non_camel_case_types)]
+                    struct UpdateShardCutoffPointSvc<T: CollectionsInternal>(pub Arc<T>);
+                    impl<
+                        T: CollectionsInternal,
+                    > tonic::server::UnaryService<super::UpdateShardCutoffPointRequest>
+                    for UpdateShardCutoffPointSvc<T> {
+                        type Response = super::CollectionOperationResponse;
+                        type Future = BoxFuture<
+                            tonic::Response<Self::Response>,
+                            tonic::Status,
+                        >;
+                        fn call(
+                            &mut self,
+                            request: tonic::Request<super::UpdateShardCutoffPointRequest>,
+                        ) -> Self::Future {
+                            let inner = Arc::clone(&self.0);
+                            let fut = async move {
+                                <T as CollectionsInternal>::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(

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

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 {

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

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,

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.

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<UpdateShardCutoffPointRequest>,
+    ) -> Result<Response<CollectionOperationResponse>, 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))
+    }
 }