Fix stale snapshot transfer recovery comment on failure/cancellation (#9237)

* Fix stale snapshot transfer recovery comment on failure/cancellation

The transfer status "comment" prefers the destination-side recovery
progress over the sender task status. Recovery progress was tracked in
`active_recoveries` and only removed via `finish_shard_recovery`, which
was called on the success path only. If `recover_shard_snapshot` returned
early - clearing the local shard, downloading the snapshot, checksum
mismatch, or cancellation (the future is spawned with
`spawn_cancel_on_drop`) - the entry leaked.

A leaked entry keeps reporting its last stage with an ever-growing
elapsed time (computed live from a fixed `Instant`), so subsequent
transfers for the same shard show stale timings until the next recovery
overwrites the entry or the node restarts.

Replace the manual start/finish pair with an RAII `ShardRecoveryGuard`
that removes the progress entry on drop, covering every exit path
including early returns and cancellation. The guard only removes its own
entry (checked via `Arc::ptr_eq`) so it never clobbers a newer recovery.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Extract recovery tracking into dedicated recovery_guard module

Move ShardRecoveryGuard out of the large shard_holder/mod.rs into a
dedicated recovery_guard.rs, and replace the verbose
`Arc<Mutex<HashMap<ShardId, Arc<Mutex<RecoveryProgress>>>>>` type with a
dedicated `ActiveRecoveries` struct that encapsulates start/comment/
set_stage/remove operations. No behavior change.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Bind recovery stage updates to the guard instance, not shard id

`ActiveRecoveries::remove` is pointer-guarded so an unwinding recovery
cannot delete a newer recovery's entry, but stage updates still resolved
the progress entry by shard id. If recovery A is still unwinding after
recovery B started for the same shard, A's late `set_stage` would mutate
B's progress and resurrect incorrect transfer comments.

Thread the guard's own progress handle (`RecoveryProgressHandle`) through
`recover_shard_snapshot_impl` -> `Collection::restore_shard_snapshot` ->
`ShardHolder::restore_shard_snapshot`, so the Unpacking/Restoring stages
(and Downloading, via `ShardRecoveryGuard::set_stage`) update that exact
recovery's progress. Direct API recoveries, which are not tracked as
transfer-side recoveries, pass `None`.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Allow too_many_arguments on recover_shard_snapshot_impl

Co-authored-by: Cursor <cursoragent@cursor.com>

---------

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: root <111755117+qdrant-cloud-bot@users.noreply.github.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

lib/collection/src/collection/snapshots.rs

@@ -25,6 +25,7 @@ use crate::shards::remote_shard::RemoteShard;
 use crate::shards::replica_set::ShardReplicaSet;
 use crate::shards::shard::{PeerId, ShardId};
 use crate::shards::shard_config::{self, ShardConfig};
+use crate::shards::shard_holder::recovery_guard::RecoveryProgressHandle;
 use crate::shards::shard_holder::shard_mapping::ShardKeyMapping;
 use crate::shards::shard_holder::{SHARD_KEY_MAPPING_FILE, ShardHolder, shard_not_found_error};
 use crate::shards::shard_path;
@@ -327,6 +328,7 @@ impl Collection {
         this_peer_id: PeerId,
         is_distributed: bool,
         temp_dir: &Path,
+        recovery_progress: Option<RecoveryProgressHandle>,
         cancel: cancel::CancellationToken,
     ) -> CollectionResult<impl Future<Output = CollectionResult<()>> + 'static> {
         // `ShardHolder::validate_shard_snapshot` is cancel safe, so we explicitly cancel it
@@ -351,6 +353,7 @@ impl Collection {
                     this_peer_id,
                     is_distributed,
                     &temp_dir,
+                    recovery_progress,
                     cancel,
                 )
                 .await?;

lib/collection/src/shards/shard_holder/mod.rs

@@ -1,3 +1,4 @@
+pub mod recovery_guard;
 mod resharding;
 pub(crate) mod shard_mapping;
 pub mod shared_shard_holder;
@@ -18,7 +19,6 @@ use fs_err as fs;
 use fs_err::{File, tokio as tokio_fs};
 use futures::{Future, StreamExt, TryStreamExt as _, stream};
 use itertools::Itertools;
-use parking_lot::Mutex;
 use segment::json_path::JsonPath;
 use segment::types::{PayloadFieldSchema, ShardKey, SnapshotFormat};
 use segment::utils::fs::move_all;
@@ -34,7 +34,7 @@ pub use self::shared_shard_holder::*;
 use super::replica_set::{AbortShardTransfer, ChangePeerFromState};
 use super::resharding::{ReshardState, ReshardingStage};
 use super::transfer::RecoveryStage;
-use super::transfer::transfer_tasks_pool::{RecoveryProgress, TransferTasksPool};
+use super::transfer::transfer_tasks_pool::TransferTasksPool;
 use crate::collection::payload_index_schema::PayloadIndexSchema;
 use crate::common::adaptive_handle::AdaptiveSearchHandle;
 use crate::common::collection_size_stats::CollectionSizeStats;
@@ -55,6 +55,9 @@ use crate::shards::replica_set::ShardReplicaSet;
 use crate::shards::replica_set::replica_set_state::ReplicaState;
 use crate::shards::shard::{PeerId, ShardId};
 use crate::shards::shard_config::ShardConfig;
+use crate::shards::shard_holder::recovery_guard::{
+    ActiveRecoveries, RecoveryProgressHandle, ShardRecoveryGuard,
+};
 use crate::shards::transfer::{ShardTransfer, ShardTransferKey};
 use crate::shards::{CollectionId, check_shard_path, shard_initializing_flag_path};
 
@@ -86,7 +89,8 @@ pub struct ShardHolder {
     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>>>>,
+    /// Entries are added and removed via [`ShardRecoveryGuard`].
+    active_recoveries: ActiveRecoveries,
 }
 
 impl ShardHolder {
@@ -129,7 +133,7 @@ impl ShardHolder {
             key_mapping,
             shard_id_to_key_mapping,
             sharding_method,
-            active_recoveries: Mutex::new(HashMap::new()),
+            active_recoveries: ActiveRecoveries::default(),
         })
     }
 
@@ -476,18 +480,13 @@ 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);
+    /// Start tracking recovery progress for a shard (destination side).
+    ///
+    /// Returns a [`ShardRecoveryGuard`] that must be held for the duration of the
+    /// recovery. Dropping the guard - on success, error, or cancellation - stops
+    /// tracking and removes the progress entry.
+    pub fn start_shard_recovery(&self, shard_id: ShardId) -> ShardRecoveryGuard {
+        self.active_recoveries.start(shard_id)
     }
 
     pub fn get_shard_transfer_info(
@@ -505,13 +504,7 @@ impl ShardHolder {
 
             // 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(|| {
+            let comment = self.active_recoveries.comment(target_shard).or_else(|| {
                 tasks_pool
                     .get_task_status(&shard_transfer.key())
                     .map(|p| p.comment)
@@ -1246,6 +1239,7 @@ impl ShardHolder {
         this_peer_id: PeerId,
         is_distributed: bool,
         temp_dir: &Path,
+        recovery_progress: Option<RecoveryProgressHandle>,
         cancel: cancel::CancellationToken,
     ) -> CollectionResult<()> {
         if !self.contains_shard(shard_id) {
@@ -1261,8 +1255,8 @@ impl ShardHolder {
             .tempdir_in(temp_dir)?;
 
         // Set unpacking stage
-        if let Some(progress) = self.active_recoveries.lock().get(&shard_id) {
-            progress.lock().set_stage(RecoveryStage::Unpacking);
+        if let Some(recovery_progress) = &recovery_progress {
+            recovery_progress.lock().set_stage(RecoveryStage::Unpacking);
         }
 
         let extract = {
@@ -1303,8 +1297,8 @@ impl ShardHolder {
         extract.await??;
 
         // Set restoring stage
-        if let Some(progress) = self.active_recoveries.lock().get(&shard_id) {
-            progress.lock().set_stage(RecoveryStage::Restoring);
+        if let Some(recovery_progress) = &recovery_progress {
+            recovery_progress.lock().set_stage(RecoveryStage::Restoring);
         }
 
         // `ShardHolder::recover_local_shard_from` is *not* cancel safe

lib/collection/src/shards/shard_holder/recovery_guard.rs

@@ -0,0 +1,94 @@
+use std::collections::HashMap;
+use std::sync::Arc;
+
+use parking_lot::Mutex;
+
+use crate::shards::shard::ShardId;
+use crate::shards::transfer::RecoveryStage;
+use crate::shards::transfer::transfer_tasks_pool::RecoveryProgress;
+
+/// Shared handle to a single recovery's progress.
+///
+/// Obtained from a [`ShardRecoveryGuard`] via [`ShardRecoveryGuard::progress_handle`]
+/// and threaded into the recovery sub-steps so they update *this* recovery's progress
+/// directly, rather than re-resolving the current map entry by shard id (which a newer
+/// recovery for the same shard may have already replaced).
+pub type RecoveryProgressHandle = Arc<Mutex<RecoveryProgress>>;
+
+/// Tracks active snapshot recoveries on a peer (destination side of transfers),
+/// keyed by shard id.
+///
+/// This is a cheaply cloneable handle around a shared map. Recoveries are started
+/// via [`ActiveRecoveries::start`], which returns a [`ShardRecoveryGuard`] that
+/// removes the entry again on drop.
+#[derive(Clone, Default)]
+pub struct ActiveRecoveries {
+    recoveries: Arc<Mutex<HashMap<ShardId, RecoveryProgressHandle>>>,
+}
+
+impl ActiveRecoveries {
+    /// Start tracking a recovery for `shard_id`, returning a guard that stops
+    /// tracking on drop.
+    pub fn start(&self, shard_id: ShardId) -> ShardRecoveryGuard {
+        let progress = Arc::new(Mutex::new(RecoveryProgress::new()));
+        self.recoveries
+            .lock()
+            .insert(shard_id, Arc::clone(&progress));
+        ShardRecoveryGuard {
+            active_recoveries: self.clone(),
+            shard_id,
+            progress,
+        }
+    }
+
+    /// Format the recovery progress comment for `shard_id`, if a recovery is active.
+    pub fn comment(&self, shard_id: ShardId) -> Option<String> {
+        self.recoveries
+            .lock()
+            .get(&shard_id)
+            .and_then(|progress| progress.lock().format_comment())
+    }
+
+    /// Remove the entry for `shard_id`, but only if it still points at `progress`,
+    /// so a newer recovery for the same shard is never clobbered.
+    fn remove(&self, shard_id: ShardId, progress: &RecoveryProgressHandle) {
+        let mut recoveries = self.recoveries.lock();
+        if let Some(current) = recoveries.get(&shard_id)
+            && Arc::ptr_eq(current, progress)
+        {
+            recoveries.remove(&shard_id);
+        }
+    }
+}
+
+/// RAII guard for tracking an active snapshot recovery on the destination side.
+///
+/// While held, the recovery progress is registered in [`ActiveRecoveries`] so it can
+/// be reported in the transfer status. On drop - including early returns, errors, and
+/// cancellation - the entry is removed again, ensuring that stale recovery comments
+/// are never reported for subsequent transfers.
+pub struct ShardRecoveryGuard {
+    active_recoveries: ActiveRecoveries,
+    shard_id: ShardId,
+    progress: RecoveryProgressHandle,
+}
+
+impl ShardRecoveryGuard {
+    /// Update the recovery stage of *this* recovery.
+    pub fn set_stage(&self, stage: RecoveryStage) {
+        self.progress.lock().set_stage(stage);
+    }
+
+    /// A shared handle to this recovery's progress, for threading into recovery
+    /// sub-steps that need to update the stage. Bound to this guard's own entry, so
+    /// updates never leak into a newer recovery for the same shard.
+    pub fn progress_handle(&self) -> RecoveryProgressHandle {
+        Arc::clone(&self.progress)
+    }
+}
+
+impl Drop for ShardRecoveryGuard {
+    fn drop(&mut self) {
+        self.active_recoveries.remove(self.shard_id, &self.progress);
+    }
+}

src/actix/api/snapshot_api.rs

@@ -560,6 +560,8 @@ async fn upload_shard_snapshot(
             snapshot_data,
             priority.unwrap_or_default(),
             RecoveryType::Full,
+            // Direct API recovery is not tracked as a transfer-side recovery
+            None,
             cancel,
         )
         .await?;
@@ -741,6 +743,8 @@ async fn recover_partial_snapshot(
             snapshot_data,
             priority.unwrap_or_default(),
             RecoveryType::Partial,
+            // Direct API recovery is not tracked as a transfer-side recovery
+            None,
             cancel,
         )
         .await?;
@@ -919,6 +923,8 @@ async fn recover_partial_snapshot_from(
             snapshot_data,
             SnapshotPriority::NoSync,
             RecoveryType::Partial,
+            // Direct API recovery is not tracked as a transfer-side recovery
+            None,
             cancel,
         )
         .await?;

src/common/snapshots.rs

@@ -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::shard_holder::recovery_guard::RecoveryProgressHandle;
 use collection::shards::transfer::RecoveryStage;
 use shard::snapshots::snapshot_data::SnapshotData;
 use shard::snapshots::snapshot_manifest::{RecoveryType, SnapshotManifest};
@@ -187,8 +188,9 @@ pub async fn recover_shard_snapshot(
         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
+        // Guard tracks recovery progress and removes it on drop, on every exit path
+        // (success, error, or cancellation), so stale timings are never reported.
+        let recovery_guard = collection
             .shards_holder()
             .read()
             .await
@@ -220,9 +222,7 @@ pub async fn recover_shard_snapshot(
                         return Err(StorageError::bad_input(description));
                     }
 
-                    recovery_progress
-                        .lock()
-                        .set_stage(RecoveryStage::Downloading);
+                    recovery_guard.set_stage(RecoveryStage::Downloading);
 
                     let client = client.client(api_key.as_deref())?;
                     snapshots::download::download_snapshot(
@@ -293,16 +293,14 @@ pub async fn recover_shard_snapshot(
             snapshot_data,
             snapshot_priority,
             RecoveryType::Full,
+            Some(recovery_guard.progress_handle()),
             cancel,
         )
         .await;
 
-        // Finish tracking recovery progress
-        collection
-            .shards_holder()
-            .read()
-            .await
-            .finish_shard_recovery(shard_id);
+        // `recovery_guard` is dropped here (and on every early return above),
+        // which stops tracking recovery progress for this shard.
+        drop(recovery_guard);
 
         result
     })
@@ -314,6 +312,7 @@ pub async fn recover_shard_snapshot(
 /// # Cancel safety
 ///
 /// This function is *not* cancel safe.
+#[allow(clippy::too_many_arguments)]
 pub async fn recover_shard_snapshot_impl(
     toc: &TableOfContent,
     collection: &Collection,
@@ -321,6 +320,7 @@ pub async fn recover_shard_snapshot_impl(
     snapshot_data: SnapshotData,
     priority: SnapshotPriority,
     recovery_type: RecoveryType,
+    recovery_progress: Option<RecoveryProgressHandle>,
     cancel: cancel::CancellationToken,
 ) -> Result<(), StorageError> {
     let _recover_tracker_guard = toc
@@ -343,6 +343,7 @@ pub async fn recover_shard_snapshot_impl(
             toc.is_distributed(),
             // Default temporary path to storage dir, to allow faster recovery within the same volume
             &toc.optional_temp_or_storage_temp_path()?,
+            recovery_progress,
             cancel,
         )
         .await?