fix(ecstore): fence lost activation locks

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

@@ -16,7 +16,9 @@ use super::{
     RebalStatus, RebalanceInfo, RebalanceMeta, RebalanceStats, RebalanceStopPropagationRecord,
     encode_rebalance_stop_propagation_record,
 };
-use crate::core::pools::{PoolMeta, acquire_pool_rebalance_activation_locks, pool_meta_has_active_decommission};
+use crate::core::pools::{
+    PoolMeta, PoolRebalanceActivationFence, acquire_pool_rebalance_activation_locks, pool_meta_has_active_decommission,
+};
 use crate::error::{Error, Result};
 use crate::object_api::ObjectOptions;
 use crate::set_disk::get_lock_acquire_timeout;
@@ -38,9 +40,20 @@ fn ensure_rebalance_activation_pool_meta_allowed(meta: &PoolMeta) -> Result<()>
     Ok(())
 }
 
-async fn merge_and_save_rebalance_meta_no_lock<S>(pool: Arc<S>, local_snapshot: &RebalanceMeta, stage: &str) -> Result<()>
+pub(super) enum RebalanceWorkerActivationFence {
+    Ready(PoolRebalanceActivationFence),
+    NotStartedTerminal,
+}
+
+async fn merge_and_save_rebalance_meta_no_lock<S, F>(
+    pool: Arc<S>,
+    local_snapshot: &RebalanceMeta,
+    stage: &str,
+    before_save: F,
+) -> Result<()>
 where
     S: EcstoreObjectIO,
+    F: FnOnce() -> Result<()>,
 {
     let opts = ObjectOptions {
         no_lock: true,
@@ -59,6 +72,7 @@ where
         Err(err) => return Err(Error::other(format!("rebalance meta load before save failed during {stage}: {err}"))),
     }
 
+    before_save()?;
     merged.save_with_opts(pool, opts).await
 }
 
@@ -123,7 +137,7 @@ impl ECStore {
             .await
             .map_err(rebalance_meta_lock_error)?;
 
-        merge_and_save_rebalance_meta_no_lock(pool, local_snapshot, stage).await
+        merge_and_save_rebalance_meta_no_lock(pool, local_snapshot, stage, || Ok(())).await
     }
 
     async fn save_rebalance_activation_meta_with_merge<S>(
@@ -135,19 +149,23 @@ impl ECStore {
     where
         S: EcstoreObjectIO + StorageNamespaceLocking<Error = Error, NamespaceLock = rustfs_lock::NamespaceLockWrapper>,
     {
-        let (_pool_meta_guard, _rebalance_meta_guard) = acquire_pool_rebalance_activation_locks(pool.clone()).await?;
+        let activation_fence = acquire_pool_rebalance_activation_locks(pool.clone()).await?;
         let mut pool_meta = PoolMeta::default();
         pool_meta.load_no_lock(pool.clone()).await?;
         ensure_rebalance_activation_pool_meta_allowed(&pool_meta)?;
 
-        merge_and_save_rebalance_meta_no_lock(pool, local_snapshot, stage).await
+        merge_and_save_rebalance_meta_no_lock(pool, local_snapshot, stage, || activation_fence.ensure_held()).await
     }
 
-    pub(super) async fn fence_rebalance_worker_activation<S>(&self, pool: Arc<S>, expected_id: &str) -> Result<bool>
+    pub(super) async fn fence_rebalance_worker_activation<S>(
+        &self,
+        pool: Arc<S>,
+        expected_id: &str,
+    ) -> Result<RebalanceWorkerActivationFence>
     where
         S: EcstoreObjectIO + StorageNamespaceLocking<Error = Error, NamespaceLock = rustfs_lock::NamespaceLockWrapper>,
     {
-        let (_pool_meta_guard, _rebalance_meta_guard) = acquire_pool_rebalance_activation_locks(pool.clone()).await?;
+        let activation_fence = acquire_pool_rebalance_activation_locks(pool.clone()).await?;
         let mut pool_meta = PoolMeta::default();
         pool_meta.load_no_lock(pool.clone()).await?;
         ensure_rebalance_activation_pool_meta_allowed(&pool_meta)?;
@@ -169,7 +187,12 @@ impl ECStore {
             )));
         }
 
-        Ok(is_rebalance_conflicting_with_decommission(&persisted))
+        activation_fence.ensure_held()?;
+        if !is_rebalance_conflicting_with_decommission(&persisted) {
+            return Ok(RebalanceWorkerActivationFence::NotStartedTerminal);
+        }
+
+        Ok(RebalanceWorkerActivationFence::Ready(activation_fence))
     }
 
     #[tracing::instrument(skip_all)]

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

@@ -32,7 +32,10 @@ use super::migration::{
     MigrationBackend, MigrationVersionResult, migrate_entry_version, migrate_entry_version_with_retry_wait,
     rebalance_delete_marker_opts,
 };
-use super::runtime::{should_fail_repeated_rebalance_bucket_defer, source_cleanup_defer_attempt};
+use super::runtime::{
+    RebalanceLocalActivationOutcome, commit_local_rebalance_worker_activation, should_fail_repeated_rebalance_bucket_defer,
+    source_cleanup_defer_attempt,
+};
 use super::worker::{
     RebalanceEntryCleanupResult, ensure_rebalance_listing_disks_available, is_transient_rebalance_error,
     parse_rebalance_max_attempts, rebalance_listing_retry_delay, rebalance_migration_retry_delay,
@@ -2708,6 +2711,43 @@ async fn test_start_rebalance_for_id_rejects_stopped_metadata() {
     assert!(err.to_string().contains("was stopped before start"));
 }
 
+#[tokio::test]
+async fn test_stop_at_activation_barrier_prevents_worker_token_commit() {
+    let meta = Arc::new(tokio::sync::RwLock::new(RebalanceMeta {
+        id: "rebalance-a".to_string(),
+        pool_stats: vec![RebalanceStats {
+            participating: true,
+            info: RebalanceInfo {
+                status: RebalStatus::Started,
+                ..Default::default()
+            },
+            ..Default::default()
+        }],
+        ..Default::default()
+    }));
+    let fence_reached = Arc::new(tokio::sync::Barrier::new(2));
+    let stop_committed = Arc::new(tokio::sync::Barrier::new(2));
+
+    let stop_meta = Arc::clone(&meta);
+    let stop_fence_reached = Arc::clone(&fence_reached);
+    let stop_committed_signal = Arc::clone(&stop_committed);
+    let stop = tokio::spawn(async move {
+        stop_fence_reached.wait().await;
+        stop_meta.write().await.stopped_at = Some(OffsetDateTime::now_utc());
+        stop_committed_signal.wait().await;
+    });
+
+    fence_reached.wait().await;
+    stop_committed.wait().await;
+    stop.await.expect("stop barrier task should finish");
+
+    let mut meta = meta.write().await;
+    let outcome = commit_local_rebalance_worker_activation(&mut meta, "rebalance-a", tokio_util::sync::CancellationToken::new())
+        .expect("stopped metadata should produce a non-start outcome");
+    assert_eq!(outcome, RebalanceLocalActivationOutcome::NotStartedTerminal);
+    assert!(meta.cancel.is_none(), "stopped rebalance must not receive a worker token");
+}
+
 fn test_store_with_rebalance_meta(meta: RebalanceMeta) -> Arc<crate::store::ECStore> {
     let endpoint_pools: crate::layout::endpoints::EndpointServerPools = Vec::new().into();
     Arc::new(crate::store::ECStore {

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

@@ -1,3 +1,4 @@
+use super::control::RebalanceWorkerActivationFence;
 use super::meta::{
     apply_rebalance_save_option, apply_rebalance_terminal_event, classify_rebalance_terminal_event, clone_first_arc,
     complete_rebalance_pools_at_goal, complete_rebalance_pools_with_empty_queue, ensure_valid_rebalance_pool_index,
@@ -39,6 +40,30 @@ pub(super) fn source_cleanup_defer_attempt(deferred_attempts: &mut HashMap<Strin
     *attempts
 }
 
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub(super) enum RebalanceLocalActivationOutcome {
+    Started,
+    NotStartedTerminal,
+}
+
+pub(super) fn commit_local_rebalance_worker_activation(
+    meta: &mut super::RebalanceMeta,
+    expected_id: &str,
+    cancel: CancellationToken,
+) -> Result<RebalanceLocalActivationOutcome> {
+    if meta.id != expected_id {
+        return Err(Error::other(format!(
+            "rebalance metadata changed before local worker activation: expected {expected_id}, found {}",
+            meta.id
+        )));
+    }
+    if meta.stopped_at.is_some() || !is_rebalance_in_progress(meta) {
+        return Ok(RebalanceLocalActivationOutcome::NotStartedTerminal);
+    }
+    meta.cancel = Some(cancel);
+    Ok(RebalanceLocalActivationOutcome::Started)
+}
+
 impl ECStore {
     #[tracing::instrument(skip_all)]
     pub async fn start_rebalance(self: &Arc<Self>) -> Result<()> {
@@ -59,15 +84,17 @@ impl ECStore {
             rebalance_meta.as_ref().ok_or(Error::ConfigNotFound)?.id.clone()
         };
         let pool = clone_first_arc(self.pools.as_slice(), "start_rebalance: no pools available")?;
-        if !self.fence_rebalance_worker_activation(pool, &expected_id).await? {
-            return Ok(());
-        }
+        let activation_fence = match self.fence_rebalance_worker_activation(pool, &expected_id).await? {
+            RebalanceWorkerActivationFence::Ready(fence) => fence,
+            RebalanceWorkerActivationFence::NotStartedTerminal => return Ok(()),
+        };
 
         let decommission_running = self.is_decommission_running().await;
 
         let cancel_tx = CancellationToken::new();
         let rx = cancel_tx.clone();
         let mut meta_to_save = None;
+        let activation_outcome;
 
         {
             let mut rebalance_meta = self.rebalance_meta.write().await;
@@ -94,10 +121,12 @@ impl ECStore {
             if complete_rebalance_pools_with_empty_queue(meta, now) {
                 meta_to_save = Some(meta.clone());
             }
-            meta.cancel = Some(cancel_tx);
+            activation_fence.ensure_held()?;
+            activation_outcome = commit_local_rebalance_worker_activation(meta, &expected_id, cancel_tx)?;
 
             drop(rebalance_meta);
         }
+        drop(activation_fence);
 
         if let Some(meta) = meta_to_save {
             let pool = clone_first_arc(self.pools.as_slice(), "start_rebalance: no pools available")?;
@@ -108,6 +137,10 @@ impl ECStore {
             )?;
         }
 
+        if activation_outcome != RebalanceLocalActivationOutcome::Started {
+            return Ok(());
+        }
+
         let participants = if let Some(ref meta) = *self.rebalance_meta.read().await {
             resolve_rebalance_participants(meta.pool_stats.as_slice(), self.pools.len())
         } else {