fix(rebalance): drain entry tasks before listing retry (#5820)

* fix(rebalance): drain entry tasks before listing retry

* test(rebalance): satisfy clippy in retry regression

crates/ecstore/src/services/rebalance/entry.rs

@@ -519,9 +519,23 @@ impl ECStore {
             let entry_tasks = entry_tasks.clone();
 
             let job = tokio::spawn(async move {
-                let list_result =
+                let list_rx = rx.clone();
-                    run_rebalance_listing_with_retry(set, rx, bucket.clone(), rebalance_entry, set_idx, rebalance_max_attempts())
+                let list_bucket = bucket.clone();
-                        .await;
+                let list_result = run_rebalance_listing_with_retry(
+                    rx,
+                    bucket,
+                    rebalance_entry,
+                    set_idx,
+                    rebalance_max_attempts(),
+                    entry_tasks.clone(),
+                    move |cb| {
+                        let set = set.clone();
+                        let rx = list_rx.clone();
+                        let bucket = list_bucket.clone();
+                        async move { set.list_objects_to_rebalance(rx, bucket, cb).await }
+                    },
+                )
+                .await;
                 let entry_result = wait_rebalance_entry_tasks(set_idx, entry_tasks).await;
                 let result = list_result.and(entry_result);
                 if let Err(err) = &result {

crates/ecstore/src/services/rebalance/rebalance_unit_tests.rs

@@ -40,9 +40,9 @@ use super::worker::{
     resolve_rebalance_file_info_versions_result, resolve_rebalance_meta_load_result, resolve_rebalance_meta_save_result,
     resolve_rebalance_migrate_result_error, resolve_rebalance_optional_bucket_config_result, resolve_rebalance_save_task_result,
     resolve_rebalance_stats_update_result, resolve_rebalance_terminal_error, resolve_rebalance_worker_result,
-    send_rebalance_done_signal, should_cleanup_rebalance_source_entry, should_count_rebalance_version_complete,
+    run_rebalance_listing_with_retry, send_rebalance_done_signal, should_cleanup_rebalance_source_entry,
-    should_defer_rebalance_entry_failure, should_retry_rebalance_listing, should_skip_rebalance_delete_marker,
+    should_count_rebalance_version_complete, should_defer_rebalance_entry_failure, should_retry_rebalance_listing,
-    wait_rebalance_entry_tasks, wait_rebalance_listing_retry, with_rebalance_entry_context,
+    should_skip_rebalance_delete_marker, wait_rebalance_entry_tasks, wait_rebalance_listing_retry, with_rebalance_entry_context,
 };
 use super::{
     DiskStat, GetObjectReader, ObjectInfo, ObjectOptions, RebalSaveOpt, RebalStatus, RebalanceBucketConfigs,
@@ -55,8 +55,8 @@ use crate::disk::RUSTFS_META_BUCKET;
 use crate::disk::error::DiskError;
 use crate::error::{Error, Result};
 use crate::storage_api_contracts::range::HTTPRangeSpec;
-use rustfs_filemeta::FileInfo;
 use rustfs_filemeta::TRANSITION_COMPLETE;
+use rustfs_filemeta::{FileInfo, MetaCacheEntry};
 use rustfs_rio::Index;
 use s3s::dto::ReplicationConfiguration;
 use serde::Serialize;
@@ -1809,6 +1809,109 @@ fn test_should_retry_rebalance_listing_respects_attempt_limit_and_error_type() {
     assert!(!should_retry_rebalance_listing(&Error::FileAccessDenied, 0, 3));
 }
 
+#[tokio::test(start_paused = true)]
+async fn test_rebalance_listing_retry_waits_for_scheduled_entries() {
+    let entry_tasks = Arc::new(tokio::sync::Mutex::new(Vec::new()));
+    let task_registered = Arc::new(tokio::sync::Notify::new());
+    let release_task = Arc::new(tokio::sync::Notify::new());
+    let task_finished = Arc::new(std::sync::atomic::AtomicBool::new(false));
+    let callback: crate::core::pools::ListCallback = Arc::new({
+        let entry_tasks = entry_tasks.clone();
+        let task_registered = task_registered.clone();
+        let release_task = release_task.clone();
+        let task_finished = task_finished.clone();
+        move |_| {
+            let entry_tasks = entry_tasks.clone();
+            let task_registered = task_registered.clone();
+            let release_task = release_task.clone();
+            let task_finished = task_finished.clone();
+            Box::pin(async move {
+                let task = tokio::spawn(async move {
+                    release_task.notified().await;
+                    task_finished.store(true, Ordering::SeqCst);
+                    Ok(RebalanceEntryOutcome::Completed)
+                });
+                entry_tasks.lock().await.push(task);
+                task_registered.notify_one();
+            })
+        }
+    });
+    let attempts = Arc::new(AtomicUsize::new(0));
+    let runner = tokio::spawn(run_rebalance_listing_with_retry(
+        CancellationToken::new(),
+        "bucket-a".to_string(),
+        callback,
+        0,
+        3,
+        entry_tasks,
+        {
+            let attempts = attempts.clone();
+            let task_finished = task_finished.clone();
+            move |cb| {
+                let attempt = attempts.fetch_add(1, Ordering::SeqCst);
+                let task_finished = task_finished.clone();
+                async move {
+                    if attempt == 0 {
+                        cb(MetaCacheEntry::default()).await;
+                        return Err(Error::SlowDown);
+                    }
+                    assert!(task_finished.load(Ordering::SeqCst), "retry must wait for scheduled entries");
+                    Ok(())
+                }
+            }
+        },
+    ));
+
+    task_registered.notified().await;
+    tokio::time::advance(Duration::from_secs(1)).await;
+    tokio::task::yield_now().await;
+    assert_eq!(attempts.load(Ordering::SeqCst), 1, "retry must not overlap the scheduled entry task");
+    release_task.notify_one();
+
+    runner
+        .await
+        .expect("listing retry task should join")
+        .expect("listing retry should complete after the scheduled entry");
+    assert_eq!(attempts.load(Ordering::SeqCst), 2);
+}
+
+#[tokio::test(start_paused = true)]
+async fn test_rebalance_listing_retry_propagates_scheduled_entry_failure() {
+    let entry_tasks = Arc::new(tokio::sync::Mutex::new(Vec::new()));
+    let callback: crate::core::pools::ListCallback = Arc::new({
+        let entry_tasks = entry_tasks.clone();
+        move |_| {
+            let entry_tasks = entry_tasks.clone();
+            Box::pin(async move {
+                entry_tasks
+                    .lock()
+                    .await
+                    .push(tokio::spawn(async { Err(Error::other("scheduled entry failed")) }));
+            })
+        }
+    });
+    let attempts = Arc::new(AtomicUsize::new(0));
+
+    let err = run_rebalance_listing_with_retry(CancellationToken::new(), "bucket-a".to_string(), callback, 0, 3, entry_tasks, {
+        let attempts = attempts.clone();
+        move |cb| {
+            let attempt = attempts.fetch_add(1, Ordering::SeqCst);
+            async move {
+                if attempt == 0 {
+                    cb(MetaCacheEntry::default()).await;
+                    return Err(Error::SlowDown);
+                }
+                panic!("entry failure must stop listing retries")
+            }
+        }
+    })
+    .await
+    .expect_err("scheduled entry failure must be returned before retrying the listing");
+
+    assert!(err.to_string().contains("scheduled entry failed"));
+    assert_eq!(attempts.load(Ordering::SeqCst), 1);
+}
+
 #[test]
 fn test_parse_rebalance_max_attempts_uses_positive_override_or_default() {
     assert_eq!(parse_rebalance_max_attempts(Some("5")), 5);

crates/ecstore/src/services/rebalance/worker.rs

@@ -399,19 +399,24 @@ pub(super) async fn load_rebalance_bucket_configs(api: &ECStore, bucket: &str) -
     })
 }
 
-pub(super) async fn run_rebalance_listing_with_retry(
+pub(super) async fn run_rebalance_listing_with_retry<List, ListFuture>(
-    set: Arc<SetDisks>,
     rx: CancellationToken,
     bucket: String,
     cb: ListCallback,
     set_idx: usize,
     max_attempts: usize,
-) -> Result<()> {
+    entry_tasks: Arc<tokio::sync::Mutex<Vec<RebalanceEntryTask>>>,
+    mut list: List,
+) -> Result<()>
+where
+    List: FnMut(ListCallback) -> ListFuture,
+    ListFuture: std::future::Future<Output = Result<()>>,
+{
     let max_attempts = max_attempts.max(1);
     let mut last_error = None;
 
     for attempt in 0..max_attempts {
-        match set.list_objects_to_rebalance(rx.clone(), bucket.clone(), cb.clone()).await {
+        match list(cb.clone()).await {
             Ok(()) => return Ok(()),
             Err(err) if should_retry_rebalance_listing(&err, attempt, max_attempts) => {
                 let next_attempt = attempt + 2;
@@ -426,6 +431,8 @@ pub(super) async fn run_rebalance_listing_with_retry(
                     delay
                 );
                 last_error = Some(err);
+                // The full retry re-evaluates deferred entries; only task failures block the next attempt.
+                let _ = wait_rebalance_entry_tasks(set_idx, entry_tasks.clone()).await?;
                 wait_rebalance_listing_retry(&rx, delay).await?;
                 info!(
                     "rebalance listing retrying bucket {} set {} attempt {}/{}",