fix(heal): preserve cancellation and retry only failed listing pages

Co-Authored-By: heihutu <heihutu@gmail.com>
Co-Authored-By: zhi22915 <qiuzgang@gmail.com>

crates/heal/src/error.rs

@@ -54,6 +54,15 @@ pub enum Error {
     #[error("Heal task execution failed: {message}")]
     TaskExecutionFailed { message: String },
 
+    /// The current page already exhausted its local retry budget. Retrying
+    /// the enclosing bucket would replay pages whose results were counted.
+    #[error("Heal listing failed for bucket {bucket}: {source}")]
+    HealListingFailed {
+        bucket: String,
+        #[source]
+        source: Box<Error>,
+    },
+
     #[error("Invalid heal type: {heal_type}")]
     InvalidHealType { heal_type: String },
 

crates/heal/src/heal/manager/tests.rs

@@ -288,6 +288,59 @@ pub(super) async fn pause_completed_retention_before_publish(task_id: &str, stat
     }
 }
 
+#[tokio::test]
+async fn canonical_outcome_cancel_wins_before_worker_finalizes_success() {
+    use crate::heal::outcome::{HealAbortReason, HealExecutionOutcome};
+    use crate::heal::task::{OUTCOME_FINISH_TEST_HOOK, OutcomeFinishTestHook};
+    let bucket = "canonical-outcome-cancel-before-finish";
+    let manager = HealManager::new(Arc::new(MockStorage), None);
+    let request = HealRequest::object(bucket.to_string(), "object".to_string(), None);
+    let task_id = request.id.clone();
+    let duplicate = HealRequest::object(bucket.to_string(), "object".to_string(), None);
+    let alias = duplicate.id.clone();
+    let retention_hook = Arc::new(CompletedRetentionHook::default());
+    {
+        let mut hooks = COMPLETED_RETENTION_HOOKS.lock().await;
+        hooks.insert(bucket.to_string(), retention_hook.clone());
+        hooks.insert(task_id.clone(), retention_hook.clone());
+    }
+    let finish_hook = Arc::new(OutcomeFinishTestHook {
+        task_id: task_id.clone(),
+        reached: Notify::new(),
+        release: Notify::new(),
+    });
+    *OUTCOME_FINISH_TEST_HOOK.lock().await = Some(finish_hook.clone());
+    manager.submit_heal_request(request).await.expect("admit original");
+    manager.submit_heal_request(duplicate).await.expect("admit alias");
+    process_manager_queue_once(&manager).await;
+    tokio::time::timeout(Duration::from_secs(5), retention_hook.started.notified())
+        .await
+        .expect("storage started");
+    retention_hook.execute.notify_one();
+    tokio::time::timeout(Duration::from_secs(5), finish_hook.reached.notified())
+        .await
+        .expect("storage returned before outcome finalization");
+    manager.cancel_task(&alias).await.expect("cancel wins publication");
+    finish_hook.release.notify_one();
+    tokio::time::timeout(Duration::from_secs(5), retention_hook.handoff.notified())
+        .await
+        .expect("scheduler completes cancelled handoff");
+    for token in [&task_id, &alias] {
+        let report = manager.get_task_report(token).await.expect("cancelled token retained");
+        assert_eq!(report.status, HealTaskStatus::Cancelled);
+        assert_eq!(
+            report.outcome.as_ref().expect("frozen outcome").execution,
+            HealExecutionOutcome::Aborted(HealAbortReason::Cancelled)
+        );
+    }
+    retention_hook.finish.notify_one();
+    *OUTCOME_FINISH_TEST_HOOK.lock().await = None;
+    COMPLETED_RETENTION_HOOKS
+        .lock()
+        .await
+        .retain(|key, _| key != bucket && key != &task_id);
+}
+
 #[tokio::test]
 async fn completed_retention_cancel_wins_over_a_prepared_retry_snapshot() {
     let bucket = "completed-retention-retry-cancel";

crates/heal/src/heal/outcome.rs

@@ -141,7 +141,9 @@ pub struct HealTaskOutcome {
 
 impl HealTaskOutcome {
     pub(crate) fn start(&mut self) {
-        self.execution = HealExecutionOutcome::Running;
+        if self.execution != HealExecutionOutcome::Aborted(HealAbortReason::Cancelled) {
+            self.execution = HealExecutionOutcome::Running;
+        }
         self.coverage = HealTraversalCoverage::Partial;
     }
 
@@ -155,6 +157,9 @@ impl HealTaskOutcome {
     }
 
     pub(crate) fn finish(&mut self, abort: Option<HealAbortReason>) {
+        if self.execution == HealExecutionOutcome::Aborted(HealAbortReason::Cancelled) {
+            return;
+        }
         let abort = abort.or(self.untraversable.then_some(HealAbortReason::Untraversable));
         self.execution = match abort {
             Some(reason) => HealExecutionOutcome::Aborted(reason),

crates/heal/src/heal/task.rs

@@ -47,6 +47,26 @@ use uuid::Uuid;
 
 use super::{BUCKET_META_PREFIX, DATA_USAGE_CACHE_NAME, RUSTFS_META_BUCKET};
 
+#[cfg(test)]
+pub(crate) struct OutcomeFinishTestHook {
+    pub(crate) task_id: String,
+    pub(crate) reached: tokio::sync::Notify,
+    pub(crate) release: tokio::sync::Notify,
+}
+
+#[cfg(test)]
+pub(crate) static OUTCOME_FINISH_TEST_HOOK: std::sync::LazyLock<tokio::sync::Mutex<Option<Arc<OutcomeFinishTestHook>>>> =
+    std::sync::LazyLock::new(|| tokio::sync::Mutex::new(None));
+
+#[cfg(test)]
+async fn pause_outcome_finish(task_id: &str) {
+    let hook = OUTCOME_FINISH_TEST_HOOK.lock().await.clone();
+    if let Some(hook) = hook.filter(|hook| hook.task_id == task_id) {
+        hook.reached.notify_one();
+        hook.release.notified().await;
+    }
+}
+
 const LOG_COMPONENT_HEAL: &str = "heal";
 const LOG_SUBSYSTEM_TASK: &str = "task";
 const LOG_SUBSYSTEM_OBJECT: &str = "object";
@@ -932,6 +952,8 @@ impl HealTask {
             HealType::ErasureSet { buckets, set_disk_id } => self.heal_erasure_set(buckets.clone(), set_disk_id.clone()).await,
         };
 
+        #[cfg(test)]
+        pause_outcome_finish(&self.id).await;
         {
             let mut outcome = self.outcome.write().await;
             if outcome.counters.processed == 0

crates/heal/src/heal/task/heal_bucket.rs

@@ -306,23 +306,47 @@ impl HealTask {
             let mut continuation_token: Option<String> = None;
             loop {
                 self.check_control_flags().await?;
-                let (objects, next_token, is_truncated) = if let Some(set_disk_id) = set_disk_id.as_deref() {
-                    self.await_with_control(self.storage.list_versions_for_heal_page_disk_walk(
-                        set_disk_id,
-                        bucket,
-                        prefix,
-                        continuation_token.as_deref(),
-                        false,
-                    ))
-                    .await?
-                } else {
-                    self.await_with_control(self.storage.list_objects_for_heal_page(
-                        bucket,
-                        prefix,
-                        continuation_token.as_deref(),
-                        false,
-                    ))
-                    .await?
+                let mut listing_attempt = 0;
+                let (objects, next_token, is_truncated) = loop {
+                    let page = if let Some(set_disk_id) = set_disk_id.as_deref() {
+                        self.await_with_control(self.storage.list_versions_for_heal_page_disk_walk(
+                            set_disk_id,
+                            bucket,
+                            prefix,
+                            continuation_token.as_deref(),
+                            false,
+                        ))
+                        .await
+                    } else {
+                        self.await_with_control(self.storage.list_objects_for_heal_page(
+                            bucket,
+                            prefix,
+                            continuation_token.as_deref(),
+                            false,
+                        ))
+                        .await
+                    };
+                    match page {
+                        Ok(page) => break page,
+                        Err(error @ (Error::TaskCancelled | Error::TaskTimeout)) => return Err(error),
+                        Err(error) => {
+                            self.outcome.write().await.attempt_failed();
+                            if error.is_recoverable_heal() && listing_attempt < MAX_BUCKET_OBJECT_HEAL_RETRIES {
+                                listing_attempt += 1;
+                                self.await_with_control(async {
+                                    tokio::time::sleep(self.bucket_object_retry_delay(listing_attempt)).await;
+                                    Ok(())
+                                })
+                                .await?;
+                                continue;
+                            }
+                            self.outcome.write().await.mark_untraversable();
+                            return Err(Error::HealListingFailed {
+                                bucket: bucket.to_string(),
+                                source: Box::new(error),
+                            });
+                        }
+                    }
                 };
 
                 let mut pending = objects;

crates/heal/src/heal/task/tests.rs

@@ -36,6 +36,70 @@ mod canonical_outcome {
         )
     }
 
+    #[tokio::test(start_paused = true)]
+    async fn cluster_retries_only_the_failed_listing_page() {
+        let storage = Arc::new(MockStorage {
+            recoverable_second_page_failures: Mutex::new(Some(1)),
+            ..Default::default()
+        });
+        let task = HealTask::from_request(
+            HealRequest::new(
+                HealType::Cluster,
+                HealOptions {
+                    recursive: true,
+                    timeout: None,
+                    ..Default::default()
+                },
+                HealPriority::Normal,
+            ),
+            storage.clone(),
+        );
+        task.execute().await.expect("second-page retry succeeds");
+        let outcome = task.get_outcome().await;
+        assert_eq!(outcome.execution, HealExecutionOutcome::Completed);
+        assert_eq!(outcome.coverage, HealTraversalCoverage::Complete);
+        assert_eq!(outcome.counters.processed, 2);
+        assert_eq!(outcome.counters.attempt_failures, 1);
+        assert_eq!(task.get_progress().await.objects_scanned, 2);
+        assert_eq!(
+            storage.heal_object_calls.lock().expect("object calls").as_slice(),
+            ["object-a", "object-b"]
+        );
+        assert_eq!(
+            storage.listing_tokens.lock().expect("listing tokens").as_slice(),
+            [None, Some("second".to_string()), Some("second".to_string())]
+        );
+    }
+
+    #[tokio::test(start_paused = true)]
+    async fn exhausted_listing_page_cannot_restart_the_bucket() {
+        let storage = Arc::new(MockStorage {
+            recoverable_second_page_failures: Mutex::new(Some(4)),
+            ..Default::default()
+        });
+        let task = HealTask::from_request(
+            HealRequest::new(
+                HealType::Cluster,
+                HealOptions {
+                    recursive: true,
+                    timeout: None,
+                    ..Default::default()
+                },
+                HealPriority::Normal,
+            ),
+            storage.clone(),
+        );
+        task.execute().await.expect_err("listing page budget exhausted");
+        let outcome = task.get_outcome().await;
+        assert_eq!(outcome.execution, HealExecutionOutcome::Aborted(HealAbortReason::Untraversable));
+        assert_eq!(outcome.coverage, HealTraversalCoverage::Partial);
+        assert_eq!(outcome.counters.processed, 1);
+        assert_eq!(outcome.counters.attempt_failures, 4);
+        assert_eq!(task.get_progress().await.objects_scanned, 1);
+        assert_eq!(storage.heal_object_calls.lock().expect("object calls").as_slice(), ["object-a"]);
+        assert_eq!(storage.bucket_heal_calls.lock().expect("bucket calls").as_slice(), ["bucket-a"]);
+    }
+
     #[tokio::test]
     async fn listing_failure_preserves_processed_objects_and_partial_coverage() {
         let storage = Arc::new(MockStorage {
@@ -872,6 +936,8 @@ struct MockStorage {
     listed: Mutex<bool>,
     list_each_bucket: bool,
     fail_second_listing_page: bool,
+    recoverable_second_page_failures: Mutex<Option<usize>>,
+    listing_tokens: Mutex<Vec<Option<String>>>,
     healed_objects: Mutex<Vec<String>>,
     heal_object_calls: Mutex<Vec<String>>,
     heal_object_version_ids: Mutex<Vec<Option<String>>>,
@@ -1285,6 +1351,28 @@ impl HealStorageAPI for MockStorage {
         _include_lifecycle_object_info: bool,
     ) -> Result<(Vec<HealListItem>, Option<String>, bool)> {
         self.listed_prefixes.lock().unwrap().push(prefix.to_string());
+        self.listing_tokens
+            .lock()
+            .expect("listing tokens")
+            .push(continuation_token.map(ToOwned::to_owned));
+        if let Some(remaining) = self
+            .recoverable_second_page_failures
+            .lock()
+            .expect("listing failures")
+            .as_mut()
+        {
+            if continuation_token.is_none() {
+                return Ok((vec![heal_item("object-a")], Some("second".to_string()), true));
+            }
+            if *remaining > 0 {
+                *remaining -= 1;
+                return Err(Error::Storage(EcstoreError::InsufficientReadQuorum(
+                    bucket.to_string(),
+                    "page".to_string(),
+                )));
+            }
+            return Ok((vec![heal_item("object-b")], None, false));
+        }
         if self.fail_second_listing_page {
             return if continuation_token.is_none() {
                 Ok((vec![heal_item("object-a")], Some("next-page".to_string()), true))