Fix transfers may lock shard holder for too long (#8373)

* Add basic test that shows shard holder is locked for too long

* Rework test, move into stream records test

* Fix typos

* Don't lock shard holder for a long time in stream records transfer (#8374)

* [ai] Detach shard holder lock from sending update batch

* Fix clippy warning

* Move stream records delay up to sleep in the middle of the batch

* remove unused code

---------

Co-authored-by: Andrey Vasnetsov <andrey@vasnetsov.com>

---------

Co-authored-by: Andrey Vasnetsov <andrey@vasnetsov.com>

lib/collection/src/shards/forward_proxy_shard.rs

@@ -20,7 +20,7 @@ use shard::scroll::ScrollRequestInternal;
 use shard::search::CoreSearchRequestBatch;
 use shard::snapshots::snapshot_manifest::SnapshotManifest;
 use tokio::runtime::Handle;
-use tokio::sync::Mutex;
+use tokio::sync::{Mutex, OwnedMutexGuard};
 
 use super::shard::ShardId;
 use super::update_tracker::UpdateTracker;
@@ -53,6 +53,49 @@ pub struct TransferBatchResult {
     pub send_duration: Duration,
 }
 
+/// A transfer batch that has been read from local storage but not yet sent to the remote shard.
+///
+/// Holds the forward proxy's update lock to prevent concurrent updates between reading and sending
+/// the batch. The lock is released when this struct is dropped.
+///
+/// Use [`PreparedTransferBatch::send`] to send the batch to a remote shard.
+pub struct PreparedTransferBatch {
+    pub operation: CollectionUpdateOperations,
+    pub next_page_offset: Option<PointIdType>,
+    pub count: usize,
+    pub read_duration: Duration,
+    /// Whether to wait for the remote to process the batch.
+    pub wait: bool,
+    /// Holds the update lock to prevent concurrent updates during the transfer.
+    _update_lock: OwnedMutexGuard<()>,
+}
+
+impl PreparedTransferBatch {
+    /// Send this batch to the given remote shard.
+    ///
+    /// The update lock is released after sending completes.
+    pub async fn send(self, remote_shard: &RemoteShard) -> CollectionResult<TransferBatchResult> {
+        let send_start = Instant::now();
+        remote_shard
+            .update(
+                // Don't add clock tag, this transfers points out of regular order (SyncPoints)
+                OperationWithClockTag::from(self.operation),
+                self.wait,
+                None,
+                HwMeasurementAcc::disposable(), // Internal operation
+            )
+            .await?;
+        let send_duration = send_start.elapsed();
+
+        Ok(TransferBatchResult {
+            next_page_offset: self.next_page_offset,
+            count: self.count,
+            read_duration: self.read_duration,
+            send_duration,
+        })
+    }
+}
+
 /// ForwardProxyShard
 ///
 /// ForwardProxyShard is a wrapper type for a LocalShard.
@@ -67,7 +110,7 @@ pub struct ForwardProxyShard {
     filter: Option<Box<Filter>>,
     /// Lock required to protect transfer-in-progress updates.
     /// It should block data updating operations while the batch is being transferred.
-    update_lock: Mutex<()>,
+    update_lock: Arc<Mutex<()>>,
 }
 
 impl ForwardProxyShard {
@@ -109,7 +152,7 @@ impl ForwardProxyShard {
             remote_shard,
             resharding_hash_ring,
             filter: filter.map(Box::new),
-            update_lock: Mutex::new(()),
+            update_lock: Arc::new(Mutex::new(())),
         })
     }
 
@@ -140,24 +183,28 @@ impl ForwardProxyShard {
         Ok(())
     }
 
-    /// Move batch of points to the remote shard
+    /// Read a batch of points to transfer, without sending it yet.
     ///
-    /// Returns new point offset and actual number of transferred points. The new point offset can
-    /// be used to start the next batch from.
+    /// Returns a [`PreparedTransferBatch`] holding the batch data and the update lock. The lock
+    /// prevents concurrent updates between reading and sending. Use
+    /// [`PreparedTransferBatch::send`] to send the batch to a remote shard.
+    ///
+    /// This allows the caller to release other locks (e.g. the shard holder lock) before sending,
+    /// which avoids holding the shard holder lock during the potentially slow network transfer.
     ///
     /// # Cancel safety
     ///
     /// This method is cancel safe.
-    pub async fn transfer_batch(
+    pub async fn read_transfer_batch(
         &self,
         offset: Option<PointIdType>,
         batch_size: usize,
         hashring_filter: Option<&HashRingRouter>,
         merge_points: bool,
         runtime_handle: &Handle,
-    ) -> CollectionResult<TransferBatchResult> {
+    ) -> CollectionResult<PreparedTransferBatch> {
         debug_assert!(batch_size > 0);
-        let _update_lock = self.update_lock.lock().await;
+        let update_lock = self.update_lock.clone().lock_owned().await;
 
         let read_start = Instant::now();
         let (points, next_page_offset) = match hashring_filter {
@@ -176,10 +223,6 @@ impl ForwardProxyShard {
         let wait = next_page_offset.is_none();
         let count = points.len();
 
-        // Use sync API to leverage potentially existing points
-        // Normally use SyncPoints, to completely replace everything in the target shard
-        // For resharding we need to merge points from multiple transfers, requiring a different operation
-        // Same when there is a filter, as we are only transferring a subset of points
         let point_operation = if !merge_points {
             PointOperations::SyncPoints(PointSyncOperation {
                 from_id: offset,
@@ -189,25 +232,15 @@ impl ForwardProxyShard {
         } else {
             PointOperations::UpsertPoints(PointInsertOperationsInternal::PointsList(points))
         };
-        let insert_points_operation = CollectionUpdateOperations::PointOperation(point_operation);
+        let operation = CollectionUpdateOperations::PointOperation(point_operation);
 
-        let send_start = Instant::now();
-        self.remote_shard
-            .update(
-                // Don't add clock tag, this transfers points out of regular order (SyncPoints)
-                OperationWithClockTag::from(insert_points_operation),
-                wait,
-                None,
-                HwMeasurementAcc::disposable(), // Internal operation
-            )
-            .await?;
-        let send_duration = send_start.elapsed();
-
-        Ok(TransferBatchResult {
+        Ok(PreparedTransferBatch {
+            operation,
             next_page_offset,
             count,
             read_duration,
-            send_duration,
+            wait,
+            _update_lock: update_lock,
         })
     }
 

lib/collection/src/shards/replica_set/shard_transfer.rs

@@ -7,7 +7,7 @@ use segment::types::{Filter, PointIdType};
 use super::ShardReplicaSet;
 use crate::hash_ring::HashRingRouter;
 use crate::operations::types::{CollectionError, CollectionResult};
-use crate::shards::forward_proxy_shard::{ForwardProxyShard, TransferBatchResult};
+use crate::shards::forward_proxy_shard::{ForwardProxyShard, PreparedTransferBatch};
 use crate::shards::local_shard::clock_map::RecoveryPoint;
 use crate::shards::queue_proxy_shard::QueueProxyShard;
 use crate::shards::remote_shard::RemoteShard;
@@ -342,20 +342,21 @@ impl ShardReplicaSet {
         let _ = local.insert(Shard::Local(local_shard));
     }
 
-    /// Custom operation for transferring data from one shard to another during transfer
+    /// Read a transfer batch without sending it yet.
     ///
-    /// Returns new point offset and transferred count
+    /// Returns a [`PreparedTransferBatch`] that holds the update lock. The caller can then drop
+    /// other locks (e.g. the shard holder lock) before calling [`PreparedTransferBatch::send`].
     ///
     /// # Cancel safety
     ///
     /// This method is cancel safe.
-    pub async fn transfer_batch(
+    pub async fn read_transfer_batch(
         &self,
         offset: Option<PointIdType>,
         batch_size: usize,
         hashring_filter: Option<&HashRingRouter>,
         merge_points: bool,
-    ) -> CollectionResult<TransferBatchResult> {
+    ) -> CollectionResult<PreparedTransferBatch> {
         let local = self.local.read().await;
 
         let Some(Shard::ForwardProxy(proxy)) = local.deref() else {
@@ -366,7 +367,7 @@ impl ShardReplicaSet {
         };
 
         proxy
-            .transfer_batch(
+            .read_transfer_batch(
                 offset,
                 batch_size,
                 hashring_filter,

lib/collection/src/shards/transfer/resharding_stream_records.rs

@@ -160,19 +160,28 @@ pub(crate) async fn transfer_resharding_stream_records(
     let mut total_send = Duration::ZERO;
 
     loop {
-        let shard_holder = shard_holder.read().await;
+        // Read batch under shard holder lock
+        let prepared_batch = {
+            let shard_holder = shard_holder.read().await;
 
-        let Some(replica_set) = shard_holder.get_shard(shard_id) else {
-            // Forward proxy gone?!
-            // That would be a programming error.
-            return Err(CollectionError::service_error(format!(
-                "Shard {shard_id} is not found"
-            )));
+            let Some(replica_set) = shard_holder.get_shard(shard_id) else {
+                // Forward proxy gone?!
+                // That would be a programming error.
+                return Err(CollectionError::service_error(format!(
+                    "Shard {shard_id} is not found"
+                )));
+            };
+
+            replica_set
+                .read_transfer_batch(offset, TRANSFER_BATCH_SIZE, Some(&hashring), true)
+                .await?
+
+            // Shard holder lock is dropped here, but the forward proxy update lock is still
+            // held inside the prepared batch.
         };
 
-        let result = replica_set
-            .transfer_batch(offset, TRANSFER_BATCH_SIZE, Some(&hashring), true)
-            .await?;
+        // Send batch to remote shard without holding the shard holder lock.
+        let result = prepared_batch.send(&remote_shard).await?;
 
         offset = result.next_page_offset;
         total_read += result.read_duration;

lib/collection/src/shards/transfer/stream_records.rs

@@ -134,28 +134,41 @@ pub(super) async fn transfer_stream_records(
     let mut offset = None;
 
     loop {
-        let shard_holder = shard_holder.read().await;
+        // Read batch under shard holder lock
+        let prepared_batch = {
+            let shard_holder = shard_holder.read().await;
 
-        let Some(replica_set) = shard_holder.get_shard(shard_id) else {
-            // Forward proxy gone?!
-            // That would be a programming error.
-            return Err(CollectionError::service_error(format!(
-                "Shard {shard_id} is not found"
-            )));
+            let Some(replica_set) = shard_holder.get_shard(shard_id) else {
+                // Forward proxy gone?!
+                // That would be a programming error.
+                return Err(CollectionError::service_error(format!(
+                    "Shard {shard_id} is not found"
+                )));
+            };
+
+            replica_set
+                .read_transfer_batch(offset, TRANSFER_BATCH_SIZE, None, merge_points)
+                .await?
+
+            // Shard holder lock is dropped here, but the forward proxy update lock is still
+            // held inside the prepared batch, preventing concurrent updates between reading
+            // and sending.
         };
 
-        let result = replica_set
-            .transfer_batch(offset, TRANSFER_BATCH_SIZE, None, merge_points)
-            .await?;
-
-        offset = result.next_page_offset;
-        progress.lock().add(result.count);
-
         #[cfg(feature = "staging")]
         if let Some(delay) = staging_delay {
             tokio::time::sleep(delay).await;
         }
 
+        // Send batch to remote shard without holding the shard holder lock.
+        // This is important because sending can take a very long time (especially the last
+        // batch which waits for the remote to fully process). Holding the shard holder lock
+        // during this time would block all other operations on the collection.
+        let result = prepared_batch.send(&remote_shard).await?;
+
+        offset = result.next_page_offset;
+        progress.lock().add(result.count);
+
         // If this is the last batch, finalize
         if offset.is_none() {
             break;

tests/consensus_tests/test_shard_stream_transfer.py

@@ -478,3 +478,103 @@ def test_shard_stream_transfer_pending_queue_data_race(tmp_path: pathlib.Path):
         assert_http_ok(r)
         counts.append(r.json()["result"]['count'])
     assert counts[0] == counts[1] == counts[2]
+
+
+# Transfer a shard, and assert that it doesn't block the shard holder for a long
+# time.
+#
+# A stream records transfer sends batches with wait=false. The last batch uses
+# wait=true to ensure all the transferred data is visible once the call returns.
+# Because of this the last batch may hang for a very long time. In older
+# versions such batch would hold a read lock on the shards holder for the entire
+# duration. It could cascade into freezing the entire cluster for a long time.
+#
+# This test ensures that doesn't happen anymore. It triggers a shard transfer
+# with artificial delay in batches to emulate a long running batch. Then it sends
+# a resharding start operation and confirms consensus is still responsive.
+# Previously this would be blocked and it would time out causing the test to
+# fail.
+#
+# See: <https://github.com/qdrant/qdrant/pull/8373>
+def test_shard_transfer_blocking_shard_holder(tmp_path: pathlib.Path):
+    assert_project_root()
+
+    # Prevent optimizers from interfering during the test
+    env = {
+        "QDRANT__STORAGE__OPTIMIZERS__INDEXING_THRESHOLD_KB": "0",
+        # Artificially make stream records transfer very slow
+        # Would hold shard holder read lock for a long time in previous version
+        "QDRANT_STAGING_SHARD_TRANSFER_DELAY_SEC": "10",
+    }
+
+    # Start a 3-node cluster
+    peer_api_uris, peer_dirs, bootstrap_uri = start_cluster(tmp_path, N_PEERS, extra_env=env)
+
+    # Staging feature must be enabled
+    skip_if_no_feature(peer_api_uris[0], "staging")
+
+    # Collect peer IDs
+    peer_ids = [get_cluster_info(uri)["peer_id"] for uri in peer_api_uris]
+
+    # Create collection with 3 shards and replication factor 2
+    create_collection(
+        peer_api_uris[0],
+        COLLECTION_NAME,
+        shard_number=N_PEERS,
+        replication_factor=1,
+    )
+    wait_collection_exists_and_active_on_all_peers(COLLECTION_NAME, peer_api_uris)
+    upsert_random_points(
+        peer_api_uris[0],
+        1000,
+        collection_name=COLLECTION_NAME,
+        wait="true",
+        with_sparse_vector=False,
+    )
+
+    # Find what shard is on the first peer
+    collection_cluster_info = get_collection_cluster_info(peer_api_uris[0], COLLECTION_NAME)
+    local_shard = collection_cluster_info["local_shards"][0]
+    shard_id = local_shard["shard_id"]
+
+    # Trigger shard move transfer (non-blocking)
+    # This would previously hold the shard holder read lock for a long time
+    r = requests.post(
+        f"{peer_api_uris[0]}/collections/{COLLECTION_NAME}/cluster",
+        params={
+            "timeout": "3", # Operation is supposed to be accepted quickly
+        },
+        json={
+            "move_shard": {
+                "shard_id": shard_id,
+                "from_peer_id": peer_ids[0],
+                "to_peer_id": peer_ids[1],
+                "method": "stream_records",
+            }
+        },
+    )
+    assert_http_ok(r)
+
+    # Immediately trigger resharding (start_resharding action)
+    r = requests.post(
+        f"{peer_api_uris[0]}/collections/{COLLECTION_NAME}/cluster",
+        params={
+            "timeout": "3", # Operation is supposed to be accepted quickly
+        },
+        json={
+            "start_resharding": {
+                "direction": "up",
+            }
+        },
+    )
+    assert_http_ok(r)
+
+    sleep(3)
+
+    # Check that all peers have 0 pending consensus operations
+    for i, uri in enumerate(peer_api_uris):
+        cluster_info = get_cluster_info(uri)
+        pending = cluster_info["raft_info"]["pending_operations"]
+        assert pending == 0, (
+            f"Peer {peer_ids[i]} at {uri} has {pending} pending operations, expected 0"
+        )