fix(rebalance): align activation locks and preserve retryable causes (#7551)

crates/ecstore/src/core/pools.rs

@@ -3463,6 +3463,11 @@ impl PoolRebalanceActivationFence {
     }
 }
 
+#[cfg(test)]
+tokio::task_local! {
+    pub(crate) static REBALANCE_ACTIVATION_LOCK_ATTEMPT: Arc<tokio::sync::Notify>;
+}
+
 pub(crate) async fn acquire_pool_rebalance_activation_locks<S>(
     pool: Arc<S>,
     fleet_proof: Option<crate::services::notification_sys::CrossPoolFenceFleetProofToken>,
@@ -3473,17 +3478,21 @@ where
             NamespaceLock = rustfs_lock::NamespaceLockWrapper,
         >,
 {
-    // Activation lock order is always pool.bin -> rebalance.bin.
+    // Match entry admission: rebalance.bin -> pool.bin. An entry retains its
+    // run read fence while target mutations acquire the pool metadata fence;
+    // activation must not hold pool.bin while waiting for that entry to drain.
+    let rebalance_meta_lock = pool.new_ns_lock(RUSTFS_META_BUCKET, REBAL_META_NAME).await?;
+    #[cfg(test)]
+    let _ = REBALANCE_ACTIVATION_LOCK_ATTEMPT.try_with(|attempted| attempted.notify_one());
+    let rebalance_meta_guard = rebalance_meta_lock
+        .get_write_lock(get_lock_acquire_timeout())
+        .await
+        .map_err(activation_rebalance_meta_lock_error)?;
     let pool_meta_lock = pool.new_ns_lock(RUSTFS_META_BUCKET, POOL_META_NAME).await?;
     let pool_meta_guard = pool_meta_lock
         .get_write_lock(get_lock_acquire_timeout())
         .await
         .map_err(activation_pool_meta_lock_error)?;
-    let rebalance_meta_lock = pool.new_ns_lock(RUSTFS_META_BUCKET, REBAL_META_NAME).await?;
-    let rebalance_meta_guard = rebalance_meta_lock
-        .get_write_lock(get_lock_acquire_timeout())
-        .await
-        .map_err(activation_rebalance_meta_lock_error)?;
 
     Ok(PoolRebalanceActivationFence {
         pool_meta_guard,
@@ -22094,7 +22103,7 @@ mod pools_tests {
                 .resources
                 .lock()
                 .expect("activation lock recorder should not be poisoned"),
-            vec![POOL_META_NAME.to_string(), REBAL_META_NAME.to_string()]
+            vec![REBAL_META_NAME.to_string(), POOL_META_NAME.to_string()]
         );
 
         let mut second_acquire = Box::pin(acquire_pool_rebalance_activation_locks(second.clone(), None));
@@ -22110,10 +22119,50 @@ mod pools_tests {
                 .resources
                 .lock()
                 .expect("activation lock recorder should not be poisoned"),
-            vec![POOL_META_NAME.to_string(), REBAL_META_NAME.to_string()]
+            vec![REBAL_META_NAME.to_string(), POOL_META_NAME.to_string()]
         );
     }
 
+    #[tokio::test]
+    async fn test_activation_cancellation_releases_rebalance_fence_while_pool_fence_is_contended() {
+        use crate::storage_api_contracts::namespace::NamespaceLocking as _;
+        let pool = Arc::new(ActivationLockRecorder {
+            lock_manager: Arc::new(rustfs_lock::GlobalLockManager::new()),
+            owner: "activation-cancellation",
+            resources: StdMutex::new(Vec::new()),
+        });
+        let pool_lock = pool
+            .new_ns_lock(crate::disk::RUSTFS_META_BUCKET, POOL_META_NAME)
+            .await
+            .expect("pool lock should be created");
+        let pool_reader = pool_lock
+            .get_read_lock(std::time::Duration::from_secs(5))
+            .await
+            .expect("ordinary mutation should hold the pool read fence");
+        pool.resources.lock().expect("recorder should not be poisoned").clear();
+        let mut activation = Box::pin(acquire_pool_rebalance_activation_locks(Arc::clone(&pool), None));
+        assert!(matches!(futures::poll!(&mut activation), Poll::Pending));
+        assert_eq!(
+            *pool.resources.lock().expect("recorder should not be poisoned"),
+            vec![REBAL_META_NAME.to_string(), POOL_META_NAME.to_string()],
+            "activation must hold the run fence before waiting for the pool fence",
+        );
+        drop(activation);
+        let rebalance_lock = pool
+            .new_ns_lock(crate::disk::RUSTFS_META_BUCKET, REBAL_META_NAME)
+            .await
+            .expect("run lock should be created");
+        let run_writer = rebalance_lock
+            .get_write_lock(std::time::Duration::from_secs(5))
+            .await
+            .expect("cancelling activation must release its already-acquired run fence");
+        assert!(
+            !pool_reader.is_released(),
+            "cancelling activation must not release another caller's pool fence"
+        );
+        assert!(!run_writer.is_lock_lost());
+    }
+
     #[test]
     fn decommission_receipt_run_token_changes_with_persisted_start_time() {
         let first = OffsetDateTime::from_unix_timestamp(1_000).expect("first run timestamp should be valid");

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

@@ -572,7 +572,7 @@ impl ECStore {
     where
         S: EcstoreObjectIO + StorageNamespaceLocking<Error = Error, NamespaceLock = rustfs_lock::NamespaceLockWrapper>,
     {
-        // Lock order: pool_meta_save_gate -> pool.bin -> rebalance.bin.
+        // Lock order: pool_meta_save_gate -> rebalance.bin -> pool.bin.
         let mut pool_meta_guard = self.pool_meta_save_gate.lock().await;
         pool_meta_guard.ensure_write_safe("rebalance worker activation")?;
         // Classify the durable rebalance record while holding both namespace

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

@@ -50,6 +50,11 @@ fn ensure_rebalance_entry_active(cancel: &CancellationToken) -> Result<()> {
     Ok(())
 }
 
+#[cfg(test)]
+tokio::task_local! {
+    static REBALANCE_ENTRY_RUN_FENCE_BARRIER: (Arc<tokio::sync::Notify>, Arc<tokio::sync::Notify>);
+}
+
 #[derive(Debug)]
 struct RebalanceEntryTarget {
     bucket: String,
@@ -256,9 +261,15 @@ impl ECStore {
             .sort_by_key(|v| (v.mod_time.is_none(), std::cmp::Reverse(v.mod_time)));
 
         // Entry lock order is bucket incarnation -> activation_gate -> rebalance.bin -> movement gate.
+        // Target capacity admission can then acquire pool.bin under the run fence.
         // Stop waits for in-flight entries through cleanup, but not for entries admitted later.
         ensure_rebalance_entry_active(&cancel)?;
         let run_guard = self.rebalance_run_guard(rebalance_id.as_ref(), "rebalance entry").await?;
+        #[cfg(test)]
+        if let Ok((arrived, release)) = REBALANCE_ENTRY_RUN_FENCE_BARRIER.try_with(Clone::clone) {
+            arrived.notify_one();
+            release.notified().await;
+        }
         let lock_lost_signal = run_guard.lock_lost_signal();
         #[cfg(test)]
         let _run_signal_test_fence = lock_lost_signal
@@ -1237,6 +1248,130 @@ mod tests {
         assert_eq!(pool_stats.cleanup_warnings.count, 1, "deferred cleanup must not add a permanent warning");
     }
 
+    #[tokio::test]
+    #[serial_test::serial]
+    async fn real_rebalance_entry_progresses_while_peer_activation_waits_for_run_fence() {
+        const REBALANCE_ID: &str = "rebalance-peer-activation-lock-order";
+        let (_temp_dirs, store, peer) = crate::services::rebalance::test_two_pool_stores_with_isolated_node_contexts(Some(
+            active_rebalance_meta(REBALANCE_ID),
+        ))
+        .await;
+        assert!(!Arc::ptr_eq(&store.ctx, &peer.ctx), "node-local movement gates must be independent");
+        {
+            let mut meta = peer.rebalance_meta.write().await;
+            let meta = meta.as_mut().expect("peer should know the durable run");
+            meta.activation_gate = Arc::default();
+            meta.cancel = None;
+        }
+        let bucket = crate::disk::RUSTFS_META_BUCKET;
+        let object = "rebalance-peer-activation-object";
+        let version_id = uuid::Uuid::new_v4();
+        let payload = b"entry must drain before peer activation takes the pool fence".repeat(1024);
+        let source_set = store.pools[0].get_disks_by_key(object);
+        let target_set = store.pools[1].get_disks_by_key(object);
+        let opts = ObjectOptions {
+            versioned: true,
+            version_id: Some(version_id.to_string()),
+            ..Default::default()
+        };
+        let mut writer = PutObjReader::from_vec(payload.clone());
+        let source_before = source_set
+            .put_object(bucket, object, &mut writer, &opts)
+            .await
+            .expect("source version should be written");
+        let entry = metacache_entry_from_source(&source_set, bucket, object).await;
+        let arrived = Arc::new(tokio::sync::Notify::new());
+        let release = Arc::new(tokio::sync::Notify::new());
+        // JoinSet aborts both scoped tasks if an assertion or timeout fails.
+        let mut tasks = tokio::task::JoinSet::new();
+        let entry_store = Arc::clone(&store);
+        tasks.spawn(
+            REBALANCE_ENTRY_RUN_FENCE_BARRIER.scope((Arc::clone(&arrived), Arc::clone(&release)), async move {
+                entry_store
+                    .rebalance_entry(
+                        RebalanceEntryTarget {
+                            bucket: bucket.to_string(),
+                            pool_index: 0,
+                        },
+                        entry,
+                        source_set,
+                        Arc::new(RebalanceBucketConfigs::default()),
+                        Arc::from(REBALANCE_ID),
+                        CancellationToken::new(),
+                    )
+                    .await
+            }),
+        );
+        tokio::time::timeout(StdDuration::from_secs(30), arrived.notified())
+            .await
+            .expect("real entry must acquire its persisted run read fence");
+
+        let attempted = Arc::new(tokio::sync::Notify::new());
+        let peer_pool = Arc::clone(&peer.pools[0]);
+        let (activation_done, activation_result) = tokio::sync::oneshot::channel();
+        tasks.spawn(
+            crate::core::pools::REBALANCE_ACTIVATION_LOCK_ATTEMPT.scope(Arc::clone(&attempted), async move {
+                let result = peer.fence_rebalance_worker_activation(peer_pool, REBALANCE_ID).await;
+                let result = result.map(|fence| match fence {
+                    super::super::control::RebalanceWorkerActivationFence::Ready(fence) => {
+                        fence.ensure_held().expect("peer activation must retain both fences");
+                    }
+                    super::super::control::RebalanceWorkerActivationFence::NotStartedTerminal => {
+                        panic!("the paused entry's run must still require activation");
+                    }
+                });
+                activation_done.send(result).expect("activation receiver should remain alive");
+                Ok(RebalanceEntryOutcome::Completed)
+            }),
+        );
+        tokio::time::timeout(StdDuration::from_secs(30), attempted.notified())
+            .await
+            .expect("peer activation must attempt the persisted rebalance write fence");
+        release.notify_one();
+
+        tokio::time::timeout(StdDuration::from_secs(30), async {
+            while let Some(result) = tasks.join_next().await {
+                assert!(matches!(
+                    result
+                        .expect("scoped task must not panic")
+                        .expect("entry must not fail or defer"),
+                    RebalanceEntryOutcome::Completed
+                ));
+            }
+        })
+        .await
+        .expect("entry and peer activation must both make progress");
+        activation_result
+            .await
+            .expect("peer activation result should be sent")
+            .expect("peer activation must not time out behind the entry it blocks");
+
+        let mut reader = target_set
+            .get_object_reader(bucket, object, None, HeaderMap::new(), &opts)
+            .await
+            .expect("the exact target version must be readable");
+        let mut actual = Vec::new();
+        reader
+            .stream
+            .read_to_end(&mut actual)
+            .await
+            .expect("target body should drain completely");
+        assert_eq!(actual, payload);
+        assert_eq!(reader.object_info.version_id, source_before.version_id);
+        assert_eq!(reader.object_info.etag, source_before.etag);
+        assert_eq!(reader.object_info.mod_time, source_before.mod_time);
+        let source_error = store.pools[0]
+            .get_object_info(bucket, object, &opts)
+            .await
+            .expect_err("completed entry must clean up the source version");
+        assert!(crate::error::is_err_object_not_found(&source_error) || crate::error::is_err_version_not_found(&source_error));
+        let meta = store.rebalance_meta.read().await;
+        let stats = &meta.as_ref().expect("local run must remain installed").pool_stats[0];
+        assert_eq!(stats.num_objects, 1);
+        assert_eq!(stats.num_versions, 1);
+        assert_eq!(stats.cleanup_warnings.count, 0);
+    }
+
     #[tokio::test]
     #[serial_test::serial]
     async fn real_rebalance_run_fence_loss_before_target_commit_preserves_target_and_source() {

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

@@ -1907,6 +1907,124 @@ fn test_is_transient_rebalance_error_accepts_wrapped_disk_timeout() {
     assert!(is_transient_rebalance_error(&Error::Io(std::io::Error::other(DiskError::Timeout))));
 }
 
+#[test]
+fn test_rebalance_stage_wrapped_transient_errors_remain_retryable() {
+    let cases = [
+        Error::Lock(rustfs_lock::LockError::timeout(".rustfs.sys/pool.bin@latest", Duration::from_secs(5))),
+        Error::Lock(rustfs_lock::LockError::network(
+            "peer unavailable",
+            std::io::Error::from(std::io::ErrorKind::ConnectionReset),
+        )),
+        Error::SlowDown,
+        Error::ErasureReadQuorum,
+        Error::ErasureWriteQuorum,
+        Error::Io(std::io::Error::other(DiskError::Timeout)),
+        Error::Io(std::io::Error::from(std::io::ErrorKind::TimedOut)),
+    ];
+    for mut error in cases {
+        for depth in 0..=3 {
+            assert!(is_transient_rebalance_error(&error), "transient source lost at depth {depth}: {error:?}");
+            assert!(
+                should_defer_rebalance_entry_failure(&error),
+                "exhausted transient entries must be deferred"
+            );
+            assert!(should_retry_rebalance_listing(&error, 0, 3));
+            assert!(
+                !should_retry_rebalance_listing(&error, 2, 3),
+                "wrapping must not bypass the attempt limit"
+            );
+            error = data_movement::data_movement_stage_error_for_test(
+                "rebalance_object",
+                "put_object",
+                "bucket",
+                "baseline/00042.bin",
+                error,
+            );
+        }
+    }
+}
+
+#[test]
+fn test_rebalance_stage_wrapped_terminal_errors_remain_terminal() {
+    let cases = [
+        Error::FileAccessDenied,
+        Error::FileCorrupt,
+        Error::OperationCanceled,
+        Error::DataMovementOverwriteErr("bucket".to_string(), "object".to_string(), "version".to_string()),
+        Error::Lock(rustfs_lock::LockError::already_locked("bucket/object", "owner")),
+        Error::other("permission denied"),
+    ];
+    for mut error in cases {
+        for depth in 0..=3 {
+            assert!(
+                !is_transient_rebalance_error(&error),
+                "terminal source must survive depth {depth}: {error:?}"
+            );
+            assert!(!should_defer_rebalance_entry_failure(&error));
+            // Object names are untrusted context, not evidence of a transient failure.
+            error = data_movement::data_movement_stage_error_for_test(
+                "rebalance_object",
+                "put_object",
+                "bucket",
+                "remote lock rpc timed out",
+                error,
+            );
+        }
+    }
+}
+
+#[tokio::test]
+async fn test_rebalance_stage_wrapped_lock_timeout_retries_real_migration_loop() {
+    for succeeds_on_retry in [true, false] {
+        let backend = MigrationBackendSpy::new(None, None);
+        let attempts = AtomicUsize::new(0);
+        let waits = AtomicUsize::new(0);
+        let mut transfer = |_, _, _| {
+            let attempt = attempts.fetch_add(1, Ordering::SeqCst);
+            async move {
+                if succeeds_on_retry && attempt > 0 {
+                    return Ok(());
+                }
+                Err(data_movement::data_movement_stage_error_for_test(
+                    "rebalance_object",
+                    "put_object",
+                    "bucket",
+                    "baseline/00042.bin",
+                    Error::Lock(rustfs_lock::LockError::timeout(".rustfs.sys/pool.bin@latest", Duration::from_secs(5))),
+                ))
+            }
+        };
+        let version = version_normal();
+        let result = migrate_entry_version_with_retry_wait(
+            &backend,
+            "bucket".to_string(),
+            0,
+            &version,
+            None,
+            3,
+            false,
+            &mut transfer,
+            |_: String, _: String, _: ObjectOptions| async { Ok::<_, Error>(ObjectInfo::default()) },
+            |_| {
+                waits.fetch_add(1, Ordering::SeqCst);
+                std::future::ready(())
+            },
+        )
+        .await;
+        assert_eq!(result.moved, succeeds_on_retry);
+        assert_eq!(result.failed, !succeeds_on_retry);
+        assert_eq!(attempts.load(Ordering::SeqCst), if succeeds_on_retry { 2 } else { 3 });
+        assert_eq!(backend.get_calls(), attempts.load(Ordering::SeqCst));
+        assert_eq!(waits.load(Ordering::SeqCst), attempts.load(Ordering::SeqCst) - 1);
+        if !succeeds_on_retry {
+            assert_eq!(result.stage, Some("write_target"));
+            assert!(should_defer_rebalance_entry_failure(
+                result.error.as_ref().expect("exhaustion must retain its source error")
+            ));
+        }
+    }
+}
+
 #[test]
 fn test_is_transient_rebalance_error_accepts_io_timeout_message() {
     assert!(is_transient_rebalance_error(&Error::Io(std::io::Error::other("timeout"))));

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

@@ -244,6 +244,7 @@ pub(super) fn resolve_rebalance_bucket_result(
 }
 
 pub(super) fn is_transient_rebalance_error(err: &Error) -> bool {
+    let err = rebalance_error_source(err);
     match err {
         Error::SlowDown
         | Error::ErasureReadQuorum
@@ -256,6 +257,15 @@ pub(super) fn is_transient_rebalance_error(err: &Error) -> bool {
     }
 }
 
+fn rebalance_error_source(mut err: &Error) -> &Error {
+    // Stage context contains object names, so classify the preserved source,
+    // not timeout-like text supplied by an object name. Iterate nested stages.
+    while let Some(source) = crate::data_movement::data_movement_stage_source(err) {
+        err = source;
+    }
+    err
+}
+
 fn is_rebalance_transient_lock_error(err: &rustfs_lock::LockError) -> bool {
     match err {
         rustfs_lock::LockError::Timeout { .. } | rustfs_lock::LockError::Network { .. } => true,
@@ -309,6 +319,7 @@ pub(super) fn rebalance_listing_retry_delay(attempt: usize) -> Duration {
 }
 
 fn is_rebalance_lock_or_rpc_timeout(err: &Error) -> bool {
+    let err = rebalance_error_source(err);
     match err {
         Error::Lock(rustfs_lock::LockError::Timeout { .. }) | Error::Lock(rustfs_lock::LockError::Network { .. }) => true,
         Error::Io(io_err) => is_rebalance_lock_or_rpc_timeout_message(&io_err.to_string()),
@@ -585,3 +596,48 @@ impl SetDisks {
         Ok(())
     }
 }
+
+#[cfg(test)]
+mod error_source_tests {
+    use super::*;
+
+    #[test]
+    fn stage_wrapped_errors_select_the_source_backoff_policy() {
+        let cases = [
+            (
+                Error::Lock(rustfs_lock::LockError::timeout(".rustfs.sys/pool.bin@latest", Duration::from_secs(5))),
+                true,
+            ),
+            (
+                Error::Lock(rustfs_lock::LockError::network(
+                    "peer unavailable",
+                    std::io::Error::from(std::io::ErrorKind::ConnectionReset),
+                )),
+                true,
+            ),
+            (Error::other("remote lock rpc timed out"), true),
+            (Error::SlowDown, false),
+            (Error::Io(std::io::Error::other(DiskError::Timeout)), false),
+            (Error::FileAccessDenied, false),
+        ];
+        for (mut error, lock_backoff) in cases {
+            for depth in 0..=3 {
+                assert_eq!(
+                    is_rebalance_lock_or_rpc_timeout(&error),
+                    lock_backoff,
+                    "wrong backoff at depth {depth}: {error:?}"
+                );
+                if !lock_backoff {
+                    assert_eq!(rebalance_migration_retry_delay(1, &error), REBALANCE_MIGRATION_RETRY_BASE_DELAY * 2);
+                }
+                error = crate::data_movement::data_movement_stage_error_for_test(
+                    "rebalance_object",
+                    "put_object",
+                    "bucket",
+                    "remote lock rpc timed out",
+                    error,
+                );
+            }
+        }
+    }
+}