diff --git a/lib/collection/src/shards/forward_proxy_shard.rs b/lib/collection/src/shards/forward_proxy_shard.rs index bcacdb44ec..1dd79a6e03 100644 --- a/lib/collection/src/shards/forward_proxy_shard.rs +++ b/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, + 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 { + 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>, /// 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>, } 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, batch_size: usize, hashring_filter: Option<&HashRingRouter>, merge_points: bool, runtime_handle: &Handle, - ) -> CollectionResult { + ) -> CollectionResult { 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, }) } diff --git a/lib/collection/src/shards/replica_set/shard_transfer.rs b/lib/collection/src/shards/replica_set/shard_transfer.rs index 2d04518a12..afdb57fbff 100644 --- a/lib/collection/src/shards/replica_set/shard_transfer.rs +++ b/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, batch_size: usize, hashring_filter: Option<&HashRingRouter>, merge_points: bool, - ) -> CollectionResult { + ) -> CollectionResult { 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, diff --git a/lib/collection/src/shards/transfer/resharding_stream_records.rs b/lib/collection/src/shards/transfer/resharding_stream_records.rs index 1d148c74db..0793680c61 100644 --- a/lib/collection/src/shards/transfer/resharding_stream_records.rs +++ b/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; diff --git a/lib/collection/src/shards/transfer/stream_records.rs b/lib/collection/src/shards/transfer/stream_records.rs index 40d03e99d9..f2cb6b060e 100644 --- a/lib/collection/src/shards/transfer/stream_records.rs +++ b/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; diff --git a/tests/consensus_tests/test_shard_stream_transfer.py b/tests/consensus_tests/test_shard_stream_transfer.py index 116e454c27..f9b522f492 100644 --- a/tests/consensus_tests/test_shard_stream_transfer.py +++ b/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: +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" + )