mirror of
https://github.com/qdrant/qdrant.git
synced 2026-09-21 13:37:46 -05:00
Add more stages for shard transfer profiling (#8027)
* Improve stages reporting * Remove transferring records info from snapshot transfer's comment * Add detailed profiling for ReshardingStreamRecords * Address AI review comment * Remove some nesting --------- Co-authored-by: timvisee <tim@visee.me>
This commit is contained in:
tellet-q
@@ -1,6 +1,6 @@
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use ahash::HashSet;
|
||||
use async_trait::async_trait;
|
||||
@@ -43,6 +43,16 @@ use crate::shards::remote_shard::RemoteShard;
|
||||
use crate::shards::shard_trait::ShardOperation;
|
||||
use crate::shards::telemetry::LocalShardTelemetry;
|
||||
|
||||
/// Result of a single batch transfer, including timing breakdown.
|
||||
pub struct TransferBatchResult {
|
||||
pub next_page_offset: Option<PointIdType>,
|
||||
pub count: usize,
|
||||
/// Time spent reading points from local storage (scroll + retrieve).
|
||||
pub read_duration: Duration,
|
||||
/// Time spent sending points to the remote shard (gRPC upsert).
|
||||
pub send_duration: Duration,
|
||||
}
|
||||
|
||||
/// ForwardProxyShard
|
||||
///
|
||||
/// ForwardProxyShard is a wrapper type for a LocalShard.
|
||||
@@ -145,10 +155,11 @@ impl ForwardProxyShard {
|
||||
hashring_filter: Option<&HashRingRouter>,
|
||||
merge_points: bool,
|
||||
runtime_handle: &Handle,
|
||||
) -> CollectionResult<(Option<PointIdType>, usize)> {
|
||||
) -> CollectionResult<TransferBatchResult> {
|
||||
debug_assert!(batch_size > 0);
|
||||
let _update_lock = self.update_lock.lock().await;
|
||||
|
||||
let read_start = Instant::now();
|
||||
let (points, next_page_offset) = match hashring_filter {
|
||||
Some(hashring_filter) => {
|
||||
self.read_batch_with_hashring(offset, batch_size, hashring_filter, runtime_handle)
|
||||
@@ -159,6 +170,7 @@ impl ForwardProxyShard {
|
||||
.await?
|
||||
}
|
||||
};
|
||||
let read_duration = read_start.elapsed();
|
||||
|
||||
// Only wait on last batch
|
||||
let wait = next_page_offset.is_none();
|
||||
@@ -179,6 +191,7 @@ impl ForwardProxyShard {
|
||||
};
|
||||
let insert_points_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)
|
||||
@@ -188,8 +201,14 @@ impl ForwardProxyShard {
|
||||
HwMeasurementAcc::disposable(), // Internal operation
|
||||
)
|
||||
.await?;
|
||||
let send_duration = send_start.elapsed();
|
||||
|
||||
Ok((next_page_offset, count))
|
||||
Ok(TransferBatchResult {
|
||||
next_page_offset,
|
||||
count,
|
||||
read_duration,
|
||||
send_duration,
|
||||
})
|
||||
}
|
||||
|
||||
/// Read a batch of points to transfer to the remote shard
|
||||
|
||||
@@ -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;
|
||||
use crate::shards::forward_proxy_shard::{ForwardProxyShard, TransferBatchResult};
|
||||
use crate::shards::local_shard::clock_map::RecoveryPoint;
|
||||
use crate::shards::queue_proxy_shard::QueueProxyShard;
|
||||
use crate::shards::remote_shard::RemoteShard;
|
||||
@@ -355,7 +355,7 @@ impl ShardReplicaSet {
|
||||
batch_size: usize,
|
||||
hashring_filter: Option<&HashRingRouter>,
|
||||
merge_points: bool,
|
||||
) -> CollectionResult<(Option<PointIdType>, usize)> {
|
||||
) -> CollectionResult<TransferBatchResult> {
|
||||
let local = self.local.read().await;
|
||||
|
||||
let Some(Shard::ForwardProxy(proxy)) = local.deref() else {
|
||||
|
||||
@@ -19,6 +19,7 @@ use fs_err::{File, tokio as tokio_fs};
|
||||
use futures::{Future, StreamExt, TryStreamExt as _, stream};
|
||||
use io::safe_delete::sync_parent_dir_async;
|
||||
use itertools::Itertools;
|
||||
use parking_lot::Mutex;
|
||||
use segment::json_path::JsonPath;
|
||||
use segment::types::{PayloadFieldSchema, ShardKey, SnapshotFormat};
|
||||
use segment::utils::fs::move_all;
|
||||
@@ -33,7 +34,8 @@ use tokio_util::io::SyncIoBridge;
|
||||
pub use self::shared_shard_holder::*;
|
||||
use super::replica_set::{AbortShardTransfer, ChangePeerFromState};
|
||||
use super::resharding::{ReshardState, ReshardingStage};
|
||||
use super::transfer::transfer_tasks_pool::TransferTasksPool;
|
||||
use super::transfer::RecoveryStage;
|
||||
use super::transfer::transfer_tasks_pool::{RecoveryProgress, TransferTasksPool};
|
||||
use crate::collection::payload_index_schema::PayloadIndexSchema;
|
||||
use crate::common::collection_size_stats::CollectionSizeStats;
|
||||
use crate::common::snapshot_stream::SnapshotStream;
|
||||
@@ -74,6 +76,9 @@ pub struct ShardHolder {
|
||||
// Duplicates the information from `key_mapping` for faster access, does not use locking
|
||||
shard_id_to_key_mapping: AHashMap<ShardId, ShardKey>,
|
||||
sharding_method: ShardingMethod,
|
||||
/// Active snapshot recoveries on this peer (destination side of transfers).
|
||||
/// Tracks progress of downloading, unpacking, and restoring snapshots.
|
||||
active_recoveries: Mutex<HashMap<ShardId, Arc<Mutex<RecoveryProgress>>>>,
|
||||
}
|
||||
|
||||
impl ShardHolder {
|
||||
@@ -119,6 +124,7 @@ impl ShardHolder {
|
||||
key_mapping,
|
||||
shard_id_to_key_mapping,
|
||||
sharding_method,
|
||||
active_recoveries: Mutex::new(HashMap::new()),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -516,6 +522,20 @@ impl ShardHolder {
|
||||
(incoming, outgoing)
|
||||
}
|
||||
|
||||
/// Start tracking recovery progress for a shard (destination side)
|
||||
pub fn start_shard_recovery(&self, shard_id: ShardId) -> Arc<Mutex<RecoveryProgress>> {
|
||||
let progress = Arc::new(Mutex::new(RecoveryProgress::new()));
|
||||
self.active_recoveries
|
||||
.lock()
|
||||
.insert(shard_id, Arc::clone(&progress));
|
||||
progress
|
||||
}
|
||||
|
||||
/// Stop tracking recovery progress for a shard
|
||||
pub fn finish_shard_recovery(&self, shard_id: ShardId) {
|
||||
self.active_recoveries.lock().remove(&shard_id);
|
||||
}
|
||||
|
||||
pub fn get_shard_transfer_info(
|
||||
&self,
|
||||
tasks_pool: &TransferTasksPool,
|
||||
@@ -528,7 +548,21 @@ impl ShardHolder {
|
||||
let from = shard_transfer.from;
|
||||
let sync = shard_transfer.sync;
|
||||
let method = shard_transfer.method;
|
||||
let status = tasks_pool.get_task_status(&shard_transfer.key());
|
||||
|
||||
// Check for active recovery on destination shard first, then sender task status
|
||||
let target_shard = to_shard_id.unwrap_or(shard_id);
|
||||
let recovery_comment = self
|
||||
.active_recoveries
|
||||
.lock()
|
||||
.get(&target_shard)
|
||||
.and_then(|p| p.lock().format_comment());
|
||||
|
||||
let comment = recovery_comment.or_else(|| {
|
||||
tasks_pool
|
||||
.get_task_status(&shard_transfer.key())
|
||||
.map(|p| p.comment)
|
||||
});
|
||||
|
||||
shard_transfers.push(ShardTransferInfo {
|
||||
shard_id,
|
||||
to_shard_id,
|
||||
@@ -536,7 +570,7 @@ impl ShardHolder {
|
||||
to,
|
||||
sync,
|
||||
method,
|
||||
comment: status.map(|p| p.comment),
|
||||
comment,
|
||||
})
|
||||
}
|
||||
shard_transfers.sort_by_key(|k| k.shard_id);
|
||||
@@ -1253,6 +1287,11 @@ impl ShardHolder {
|
||||
.prefix(&format!("{collection_name}-shard-{shard_id}"))
|
||||
.tempdir_in(temp_dir)?;
|
||||
|
||||
// Set unpacking stage
|
||||
if let Some(progress) = self.active_recoveries.lock().get(&shard_id) {
|
||||
progress.lock().set_stage(RecoveryStage::Unpacking);
|
||||
}
|
||||
|
||||
let extract = {
|
||||
let snapshot_temp_dir = snapshot_temp_dir.path().to_path_buf();
|
||||
|
||||
@@ -1290,6 +1329,11 @@ impl ShardHolder {
|
||||
|
||||
extract.await??;
|
||||
|
||||
// Set restoring stage
|
||||
if let Some(progress) = self.active_recoveries.lock().get(&shard_id) {
|
||||
progress.lock().set_stage(RecoveryStage::Restoring);
|
||||
}
|
||||
|
||||
// `ShardHolder::recover_local_shard_from` is *not* cancel safe
|
||||
// (see `ShardReplicaSet::restore_local_replica_from`)
|
||||
let recovered = self
|
||||
|
||||
@@ -61,6 +61,31 @@ impl TransferStage {
|
||||
}
|
||||
}
|
||||
|
||||
/// Current stage of snapshot recovery on the receiver (destination) node.
|
||||
///
|
||||
/// These sub-stages break down what happens during the `Recovering` stage
|
||||
/// as seen from the destination peer's perspective.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum RecoveryStage {
|
||||
/// HTTP download of snapshot from source node
|
||||
Downloading,
|
||||
/// Extracting tar archive
|
||||
Unpacking,
|
||||
/// Applying data to shard
|
||||
Restoring,
|
||||
}
|
||||
|
||||
impl RecoveryStage {
|
||||
/// Short lowercase name for display in comment
|
||||
pub fn as_str(&self) -> &'static str {
|
||||
match self {
|
||||
Self::Downloading => "downloading",
|
||||
Self::Unpacking => "unpacking",
|
||||
Self::Restoring => "restoring",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Time between consensus confirmation retries.
|
||||
const CONSENSUS_CONFIRM_RETRY_DELAY: Duration = Duration::from_secs(1);
|
||||
|
||||
|
||||
@@ -1,9 +1,11 @@
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use common::counter::hardware_accumulator::HwMeasurementAcc;
|
||||
use parking_lot::Mutex;
|
||||
use shard::count::CountRequestInternal;
|
||||
|
||||
use super::TransferStage;
|
||||
use super::transfer_tasks_pool::TransferTaskProgress;
|
||||
use crate::hash_ring::HashRingRouter;
|
||||
use crate::operations::types::{CollectionError, CollectionResult};
|
||||
@@ -41,6 +43,7 @@ pub(crate) async fn transfer_resharding_stream_records(
|
||||
);
|
||||
|
||||
// Proxify local shard and create payload indexes on remote shard
|
||||
progress.lock().set_stage(TransferStage::Proxifying);
|
||||
{
|
||||
let shard_holder = shard_holder.read().await;
|
||||
|
||||
@@ -136,9 +139,12 @@ pub(crate) async fn transfer_resharding_stream_records(
|
||||
}
|
||||
|
||||
// Transfer contents batch by batch
|
||||
progress.lock().set_stage(TransferStage::Transferring);
|
||||
log::trace!("Transferring points to shard {shard_id} by reshard streaming records");
|
||||
|
||||
let mut offset = None;
|
||||
let mut total_read = Duration::ZERO;
|
||||
let mut total_send = Duration::ZERO;
|
||||
|
||||
loop {
|
||||
let shard_holder = shard_holder.read().await;
|
||||
@@ -151,12 +157,18 @@ pub(crate) async fn transfer_resharding_stream_records(
|
||||
)));
|
||||
};
|
||||
|
||||
let (new_offset, count) = replica_set
|
||||
let result = replica_set
|
||||
.transfer_batch(offset, TRANSFER_BATCH_SIZE, Some(&hashring), true)
|
||||
.await?;
|
||||
|
||||
offset = new_offset;
|
||||
progress.lock().add(count);
|
||||
offset = result.next_page_offset;
|
||||
total_read += result.read_duration;
|
||||
total_send += result.send_duration;
|
||||
{
|
||||
let mut p = progress.lock();
|
||||
p.add(result.count);
|
||||
p.set_batch_durations(total_read, total_send);
|
||||
}
|
||||
|
||||
// If this is the last batch, finalize
|
||||
if offset.is_none() {
|
||||
|
||||
@@ -131,12 +131,12 @@ pub(super) async fn transfer_stream_records(
|
||||
)));
|
||||
};
|
||||
|
||||
let (new_offset, count) = replica_set
|
||||
let result = replica_set
|
||||
.transfer_batch(offset, TRANSFER_BATCH_SIZE, None, merge_points)
|
||||
.await?;
|
||||
|
||||
offset = new_offset;
|
||||
progress.lock().add(count);
|
||||
offset = result.next_page_offset;
|
||||
progress.lock().add(result.count);
|
||||
|
||||
#[cfg(feature = "staging")]
|
||||
if let Some(delay) = staging_delay {
|
||||
|
||||
@@ -2,14 +2,14 @@ use std::cmp::max;
|
||||
use std::collections::HashMap;
|
||||
use std::fmt::Write as _;
|
||||
use std::sync::Arc;
|
||||
use std::time::Instant;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use parking_lot::Mutex;
|
||||
|
||||
use crate::common::eta_calculator::EtaCalculator;
|
||||
use crate::common::stoppable_task_async::CancellableAsyncTaskHandle;
|
||||
use crate::shards::CollectionId;
|
||||
use crate::shards::transfer::{ShardTransfer, ShardTransferKey, TransferStage};
|
||||
use crate::shards::transfer::{RecoveryStage, ShardTransfer, ShardTransferKey, TransferStage};
|
||||
|
||||
pub struct TransferTasksPool {
|
||||
collection_id: CollectionId,
|
||||
@@ -29,6 +29,9 @@ pub struct TransferTaskProgress {
|
||||
// Stage tracking for profiling
|
||||
current_stage: Option<TransferStage>,
|
||||
stage_started: Option<Instant>,
|
||||
// Cumulative batch timing breakdown (local read vs remote send)
|
||||
batch_read_duration: Duration,
|
||||
batch_send_duration: Duration,
|
||||
}
|
||||
|
||||
#[derive(Debug, PartialEq, Eq, Hash, Clone)]
|
||||
@@ -52,6 +55,8 @@ impl TransferTaskProgress {
|
||||
eta: EtaCalculator::new(),
|
||||
current_stage: None,
|
||||
stage_started: None,
|
||||
batch_read_duration: Duration::ZERO,
|
||||
batch_send_duration: Duration::ZERO,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -78,9 +83,61 @@ impl TransferTaskProgress {
|
||||
self.current_stage
|
||||
}
|
||||
|
||||
/// Get elapsed seconds in current stage
|
||||
pub fn stage_elapsed_secs(&self) -> Option<u64> {
|
||||
self.stage_started.map(|t| t.elapsed().as_secs())
|
||||
/// Get elapsed seconds in current stage (with decimal precision)
|
||||
pub fn stage_elapsed_secs(&self) -> Option<f64> {
|
||||
self.stage_started.map(|t| t.elapsed().as_secs_f64())
|
||||
}
|
||||
|
||||
/// Update cumulative batch timing breakdown (local read vs remote send)
|
||||
pub fn set_batch_durations(&mut self, read: Duration, send: Duration) {
|
||||
self.batch_read_duration = read;
|
||||
self.batch_send_duration = send;
|
||||
}
|
||||
}
|
||||
|
||||
/// Progress tracking for snapshot recovery on the receiver (destination) node.
|
||||
///
|
||||
/// Tracks the sub-stages of recovery: downloading, unpacking, restoring.
|
||||
pub struct RecoveryProgress {
|
||||
current_stage: Option<RecoveryStage>,
|
||||
stage_started: Option<Instant>,
|
||||
}
|
||||
|
||||
impl RecoveryProgress {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
current_stage: None,
|
||||
stage_started: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Set the current recovery stage (resets stage elapsed time)
|
||||
pub fn set_stage(&mut self, stage: RecoveryStage) {
|
||||
self.current_stage = Some(stage);
|
||||
self.stage_started = Some(Instant::now());
|
||||
}
|
||||
|
||||
/// Get the current recovery stage
|
||||
pub fn current_stage(&self) -> Option<RecoveryStage> {
|
||||
self.current_stage
|
||||
}
|
||||
|
||||
/// Get elapsed seconds in current stage (with decimal precision)
|
||||
pub fn stage_elapsed_secs(&self) -> Option<f64> {
|
||||
self.stage_started.map(|t| t.elapsed().as_secs_f64())
|
||||
}
|
||||
|
||||
/// Format a comment string showing current stage and elapsed time
|
||||
pub fn format_comment(&self) -> Option<String> {
|
||||
let stage = self.current_stage?;
|
||||
let elapsed = self.stage_elapsed_secs().unwrap_or(0.0);
|
||||
Some(format!("{} ({:.2}s)", stage.as_str(), elapsed))
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for RecoveryProgress {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -105,28 +162,40 @@ impl TransferTasksPool {
|
||||
let progress = task.progress.lock();
|
||||
let total = max(progress.points_transferred, progress.points_total);
|
||||
|
||||
// Build comment with stage prefix if available
|
||||
let mut comment = String::new();
|
||||
if let Some(stage) = progress.current_stage() {
|
||||
let elapsed = progress.stage_elapsed_secs().unwrap_or(0);
|
||||
write!(comment, "{} ({}s) | ", stage.as_str(), elapsed).unwrap();
|
||||
let elapsed = progress.stage_elapsed_secs().unwrap_or(0.0);
|
||||
write!(comment, "{} ({:.2}s)", stage.as_str(), elapsed).unwrap();
|
||||
}
|
||||
|
||||
write!(
|
||||
comment,
|
||||
"Transferring records ({}/{}), started {}s ago, ETA: ",
|
||||
progress.points_transferred,
|
||||
total,
|
||||
chrono::Utc::now()
|
||||
.signed_duration_since(task.started_at)
|
||||
.num_seconds(),
|
||||
)
|
||||
.unwrap();
|
||||
// Append records progress only when points are actually tracked
|
||||
if total > 0 {
|
||||
if !comment.is_empty() {
|
||||
comment.push_str(" | ");
|
||||
}
|
||||
write!(
|
||||
comment,
|
||||
"Transferring records ({}/{})",
|
||||
progress.points_transferred, total
|
||||
)
|
||||
.unwrap();
|
||||
if let Some(eta) = progress.eta.estimate(total) {
|
||||
write!(comment, ", ETA: {:.2}s", eta.as_secs_f64()).unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(eta) = progress.eta.estimate(total) {
|
||||
write!(comment, "{:.2}s", eta.as_secs_f64()).unwrap();
|
||||
} else {
|
||||
comment.push('-');
|
||||
// Append batch timing breakdown when available
|
||||
if !progress.batch_read_duration.is_zero() || !progress.batch_send_duration.is_zero() {
|
||||
if !comment.is_empty() {
|
||||
comment.push_str(" | ");
|
||||
}
|
||||
write!(
|
||||
comment,
|
||||
"read: {:.2}s, send: {:.2}s",
|
||||
progress.batch_read_duration.as_secs_f64(),
|
||||
progress.batch_send_duration.as_secs_f64(),
|
||||
)
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
Some(TransferTaskStatus { result, comment })
|
||||
@@ -210,7 +279,7 @@ mod tests {
|
||||
progress.set_stage(TransferStage::Proxifying);
|
||||
assert_eq!(progress.current_stage(), Some(TransferStage::Proxifying));
|
||||
assert!(progress.stage_elapsed_secs().is_some());
|
||||
assert!(progress.stage_elapsed_secs().unwrap() < 2);
|
||||
assert!(progress.stage_elapsed_secs().unwrap() < 2.00);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -238,7 +307,7 @@ mod tests {
|
||||
progress.set_stage(TransferStage::Transferring);
|
||||
assert_eq!(progress.current_stage(), Some(TransferStage::Transferring));
|
||||
// New stage should have very small elapsed time
|
||||
assert!(progress.stage_elapsed_secs().unwrap() < 1);
|
||||
assert!(progress.stage_elapsed_secs().unwrap() < 1.00);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
+33
-7
@@ -9,6 +9,7 @@ use collection::operations::snapshot_ops::{
|
||||
use collection::operations::verification::VerificationPass;
|
||||
use collection::shards::replica_set::replica_set_state::ReplicaState;
|
||||
use collection::shards::shard::ShardId;
|
||||
use collection::shards::transfer::RecoveryStage;
|
||||
use shard::snapshots::snapshot_data::SnapshotData;
|
||||
use shard::snapshots::snapshot_manifest::{RecoveryType, SnapshotManifest};
|
||||
use storage::content_manager::errors::StorageError;
|
||||
@@ -161,13 +162,26 @@ pub async fn recover_shard_snapshot(
|
||||
// - but the task is *spawned* on the runtime and won't be cancelled, if request is cancelled
|
||||
|
||||
cancel::future::spawn_cancel_on_drop(async move |cancel| {
|
||||
let cancel_safe = async {
|
||||
let pre_recovery_task = async {
|
||||
let collection = toc.get_collection(&collection_pass).await?;
|
||||
collection.assert_shard_exists(shard_id).await?;
|
||||
|
||||
// Default temporary path to storage dir, to allow faster recovery within the same volume
|
||||
let download_dir = toc.optional_temp_or_storage_temp_path()?;
|
||||
Result::<_, StorageError>::Ok((collection, download_dir))
|
||||
};
|
||||
|
||||
let (collection, download_dir) =
|
||||
cancel::future::cancel_on_token(cancel.clone(), pre_recovery_task).await??;
|
||||
|
||||
// Once recovery tracking starts, `finish_shard_recovery` must run on all paths
|
||||
let recovery_progress = collection
|
||||
.shards_holder()
|
||||
.read()
|
||||
.await
|
||||
.start_shard_recovery(shard_id);
|
||||
|
||||
let download_task = async {
|
||||
let DownloadResult {
|
||||
snapshot,
|
||||
hash
|
||||
@@ -182,8 +196,11 @@ pub async fn recover_shard_snapshot(
|
||||
return Err(StorageError::bad_input(description));
|
||||
}
|
||||
|
||||
let client = client.client(api_key.as_deref())?;
|
||||
recovery_progress
|
||||
.lock()
|
||||
.set_stage(RecoveryStage::Downloading);
|
||||
|
||||
let client = client.client(api_key.as_deref())?;
|
||||
snapshots::download::download_snapshot(
|
||||
&client,
|
||||
url,
|
||||
@@ -237,14 +254,14 @@ pub async fn recover_shard_snapshot(
|
||||
}
|
||||
}
|
||||
|
||||
Ok((collection, snapshot))
|
||||
Ok(snapshot)
|
||||
};
|
||||
|
||||
let (collection, snapshot_data) =
|
||||
cancel::future::cancel_on_token(cancel.clone(), cancel_safe).await??;
|
||||
let snapshot_data =
|
||||
cancel::future::cancel_on_token(cancel.clone(), download_task).await??;
|
||||
|
||||
// `recover_shard_snapshot_impl` is *not* cancel safe
|
||||
recover_shard_snapshot_impl(
|
||||
let result = recover_shard_snapshot_impl(
|
||||
&toc,
|
||||
&collection,
|
||||
shard_id,
|
||||
@@ -253,7 +270,16 @@ pub async fn recover_shard_snapshot(
|
||||
RecoveryType::Full,
|
||||
cancel,
|
||||
)
|
||||
.await
|
||||
.await;
|
||||
|
||||
// Finish tracking recovery progress
|
||||
collection
|
||||
.shards_holder()
|
||||
.read()
|
||||
.await
|
||||
.finish_shard_recovery(shard_id);
|
||||
|
||||
result
|
||||
})
|
||||
.await??;
|
||||
|
||||
|
||||
@@ -250,7 +250,7 @@ fn aggregate_shard_transfers(
|
||||
})
|
||||
};
|
||||
|
||||
let get_transfer_from_source = |base_transfer: &ShardTransferInfo| {
|
||||
let find_transfer = |peer_id, base_transfer: &ShardTransferInfo| {
|
||||
let ShardTransferInfo {
|
||||
shard_id,
|
||||
to_shard_id,
|
||||
@@ -261,7 +261,7 @@ fn aggregate_shard_transfers(
|
||||
comment: _,
|
||||
} = base_transfer;
|
||||
|
||||
get_transfers(*from)?.iter().find(|t| {
|
||||
get_transfers(peer_id)?.iter().find(|t| {
|
||||
t.from == *from
|
||||
&& t.to == *to
|
||||
&& t.shard_id == *shard_id
|
||||
@@ -273,14 +273,25 @@ fn aggregate_shard_transfers(
|
||||
return Vec::new();
|
||||
};
|
||||
|
||||
// Try to use the information from the source peer (transfer.from),
|
||||
// otherwise, use the one from the base telemetry
|
||||
// Prefer source peer's transfer info, merge destination's comment if available
|
||||
base_transfers
|
||||
.iter()
|
||||
.map(|base_transfer| {
|
||||
get_transfer_from_source(base_transfer)
|
||||
let mut transfer = find_transfer(base_transfer.from, base_transfer)
|
||||
.cloned()
|
||||
.unwrap_or_else(|| base_transfer.clone())
|
||||
.unwrap_or_else(|| base_transfer.clone());
|
||||
|
||||
// Append destination peer's comment if available
|
||||
if let Some(dst_transfer) = find_transfer(base_transfer.to, base_transfer)
|
||||
&& let Some(ref dst_comment) = dst_transfer.comment
|
||||
{
|
||||
transfer.comment = Some(match transfer.comment {
|
||||
Some(src_comment) => format!("{src_comment} | {dst_comment}"),
|
||||
None => dst_comment.clone(),
|
||||
});
|
||||
}
|
||||
|
||||
transfer
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
@@ -110,7 +110,7 @@ def test_shard_snapshot_transfer_shows_stage_in_comment(tmp_path: pathlib.Path):
|
||||
replicate_shard(peer_api_uris[0], COLLECTION_NAME, src['local_shards'][0]['shard_id'],
|
||||
src['peer_id'], dst['peer_id'], method="snapshot")
|
||||
|
||||
stage_re = re.compile(r"^(proxifying|creating snapshot|transferring|recovering|flushing queue|waiting consensus|finalizing) \(\d+s\) \|")
|
||||
stage_re = re.compile(r"^(proxifying|creating snapshot|transferring|recovering|flushing queue|waiting consensus|finalizing) \(\d+\.\d+s\)")
|
||||
|
||||
def get_cluster_comments():
|
||||
return [t.get("comment", "") for t in get_collection_cluster_info(peer_api_uris[0], COLLECTION_NAME).get("shard_transfers", [])]
|
||||
|
||||
@@ -531,7 +531,10 @@ def check_collection_shard_transfer_progress(peer_api_uri: str, collection_name:
|
||||
comment = transfer["comment"]
|
||||
|
||||
# Compare progress or total
|
||||
current, total = re.search(r"Transferring records \((\d+)/(\d+)\), started", comment).groups()
|
||||
m = re.search(r"Transferring records \((\d+)/(\d+)\)", comment)
|
||||
if m is None:
|
||||
continue
|
||||
current, total = m.groups()
|
||||
if current is not None and expected_transfer_progress is not None and int(
|
||||
current) >= expected_transfer_progress:
|
||||
return True
|
||||
|
||||
Reference in New Issue
Block a user