From ce3cd2d8904f6f8ba19861d58997b7d22571e2da Mon Sep 17 00:00:00 2001 From: overtrue Date: Sat, 22 Aug 2026 02:05:53 +0800 Subject: [PATCH] fix(ecstore): commit rebalance activation after persistence --- .../ecstore/src/services/rebalance/control.rs | 27 +- .../ecstore/src/services/rebalance/entry.rs | 151 +++++++----- .../rebalance/rebalance_unit_tests.rs | 233 +++++++++++++++--- .../ecstore/src/services/rebalance/runtime.rs | 137 +++++++--- 4 files changed, 425 insertions(+), 123 deletions(-) diff --git a/crates/ecstore/src/services/rebalance/control.rs b/crates/ecstore/src/services/rebalance/control.rs index cc87ce94f..a57c8bf3e 100644 --- a/crates/ecstore/src/services/rebalance/control.rs +++ b/crates/ecstore/src/services/rebalance/control.rs @@ -86,7 +86,7 @@ where Ok(guard) } -async fn merge_and_save_rebalance_meta_no_lock( +pub(super) async fn merge_and_save_rebalance_meta_no_lock( pool: Arc, local_snapshot: &RebalanceMeta, stage: &str, @@ -276,6 +276,31 @@ impl ECStore { .await } + pub(super) async fn save_rebalance_meta_under_activation_fence( + &self, + pool: Arc, + local_snapshot: &RebalanceMeta, + stage: &str, + activation_fence: &PoolRebalanceActivationFence, + expected_id: &str, + ) -> Result<()> + where + S: EcstoreObjectIO, + { + merge_and_save_rebalance_meta_no_lock( + pool, + local_snapshot, + stage, + ObjectOptions { + no_lock: true, + ..Default::default() + }, + Some(activation_fence), + Some(expected_id), + ) + .await + } + async fn save_rebalance_meta_with_merge_for_id( &self, pool: Arc, diff --git a/crates/ecstore/src/services/rebalance/entry.rs b/crates/ecstore/src/services/rebalance/entry.rs index 08cfc680c..de5015fd2 100644 --- a/crates/ecstore/src/services/rebalance/entry.rs +++ b/crates/ecstore/src/services/rebalance/entry.rs @@ -46,6 +46,7 @@ use tracing::{debug, error, warn}; impl ECStore { async fn finish_rebalance_entry_after_cleanup( &self, + run_guard: &super::control::RebalanceRunGuard, pool_index: usize, bucket: &str, object: &str, @@ -54,39 +55,40 @@ impl ECStore { cleanup: impl std::future::Future>, ) -> Result { // Persisted stats can complete a pool on restart, so source cleanup must resolve first. - let run_guard = self.rebalance_run_guard(expected_id, "rebalance source cleanup").await?; + run_guard.ensure_held("rebalance source cleanup")?; let cleanup_result = cleanup.await; run_guard.ensure_held("rebalance source cleanup")?; - drop(run_guard); let cleanup_result = resolve_rebalance_entry_cleanup_delete_result(cleanup_result, bucket, object); let RebalanceEntryCleanupResult::Completed { warning } = cleanup_result else { return Ok(cleanup_result); }; - if let Some(message) = warning.as_ref() - && let Err(err) = self + if let Some(message) = warning.as_ref() { + run_guard.ensure_held("record rebalance cleanup warning")?; + let warning_result = self .record_rebalance_cleanup_warning(pool_index, bucket, object, message.clone(), expected_id) - .await - { - error!( - event = EVENT_REBALANCE_ENTRY, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - pool_index, - bucket, - object, - stage = "cleanup_source", - error = ?err, - "Failed to record rebalance source cleanup warning" - ); + .await; + run_guard.ensure_held("record rebalance cleanup warning")?; + if let Err(err) = warning_result { + error!( + event = EVENT_REBALANCE_ENTRY, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + pool_index, + bucket, + object, + stage = "cleanup_source", + error = ?err, + "Failed to record rebalance source cleanup warning" + ); + } } - resolve_rebalance_stats_update_result( - self.update_pool_stats_batch_for_rebalance(pool_index, bucket.to_string(), stats_updates, expected_id) - .await, - pool_index, - bucket, - object, - )?; + run_guard.ensure_held("record rebalance entry stats")?; + let stats_result = self + .update_pool_stats_batch_for_rebalance(pool_index, bucket.to_string(), stats_updates, expected_id) + .await; + run_guard.ensure_held("record rebalance entry stats")?; + resolve_rebalance_stats_update_result(stats_result, pool_index, bucket, object)?; Ok(RebalanceEntryCleanupResult::Completed { warning }) } @@ -161,14 +163,16 @@ impl ECStore { fivs.versions .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. + // Stop waits for in-flight entries through cleanup, but not for entries admitted later. + let run_guard = self.rebalance_run_guard(rebalance_id.as_ref(), "rebalance entry").await?; + let mut rebalanced: usize = 0; let mut expired: usize = 0; let mut cleanup_preflight_allowed_missing = Vec::new(); let mut stats_updates = Vec::with_capacity(fivs.versions.len()); for version in fivs.versions.iter() { - let run_guard = self - .rebalance_run_guard(rebalance_id.as_ref(), "rebalance lifecycle mutation") - .await?; + run_guard.ensure_held("rebalance lifecycle mutation")?; let lifecycle_result = crate::core::pools::should_skip_lifecycle_for_data_movement( self.clone(), &bucket, @@ -198,7 +202,6 @@ impl ECStore { reason = "expired_by_lifecycle", "Skipped rebalance version" ); - drop(run_guard); continue; } @@ -216,7 +219,6 @@ impl ECStore { reason = "last_delete_marker_without_replication", "Skipped rebalance version" ); - drop(run_guard); continue; } @@ -236,6 +238,7 @@ impl ECStore { let store = self.clone(); async move { store.delete_object(&bucket, &object, opts).await } }; + run_guard.ensure_held("rebalance version migration")?; let result = migrate_entry_version( &RebalanceMigrationBackend::new(set.as_ref(), self.as_ref()), bucket.clone(), @@ -250,7 +253,6 @@ impl ECStore { ) .await; run_guard.ensure_held("rebalance version migration")?; - drop(run_guard); if result.ignored { if should_count_rebalance_version_complete(&result) { @@ -295,6 +297,7 @@ impl ECStore { error = %err, "Deferred rebalance entry after transient migration failure" ); + run_guard.ensure_held("record rebalance last error")?; if let Err(stats_err) = self .update_rebalance_last_error(pool_index, deferred_error.clone(), rebalance_id.as_ref()) .await @@ -304,6 +307,7 @@ impl ECStore { &bucket, &entry.name, stats_err ); } + run_guard.ensure_held("record rebalance last error")?; return Ok(RebalanceEntryOutcome::Deferred { last_error: deferred_error, }); @@ -311,20 +315,23 @@ impl ECStore { let entry_err = with_rebalance_entry_context(result.stage.unwrap_or("migrate"), bucket.as_str(), version.name.as_str(), err); - if !stats_updates.is_empty() - && let Err(stats_err) = self + if !stats_updates.is_empty() { + run_guard.ensure_held("record rebalance stats before migration error")?; + let stats_result = self .update_pool_stats_batch_for_rebalance( pool_index, bucket.clone(), stats_updates.as_slice(), rebalance_id.as_ref(), ) - .await - { - error!( - "rebalance_entry {} failed to update stats before returning migration error for {}: {}", - &bucket, &entry.name, stats_err - ); + .await; + run_guard.ensure_held("record rebalance stats before migration error")?; + if let Err(stats_err) = stats_result { + error!( + "rebalance_entry {} failed to update stats before returning migration error for {}: {}", + &bucket, &entry.name, stats_err + ); + } } return Err(entry_err); @@ -342,6 +349,7 @@ impl ECStore { } let cleanup_result = self .finish_rebalance_entry_after_cleanup( + &run_guard, pool_index, bucket.as_str(), entry.name.as_str(), @@ -420,20 +428,20 @@ impl ECStore { "Rebalance source object retained" ); - resolve_rebalance_stats_update_result( - self.update_pool_stats_batch_for_rebalance( + run_guard.ensure_held("record retained rebalance entry stats")?; + let stats_result = self + .update_pool_stats_batch_for_rebalance( pool_index, bucket.clone(), stats_updates.as_slice(), rebalance_id.as_ref(), ) - .await, - pool_index, - bucket.as_str(), - entry.name.as_str(), - )?; + .await; + run_guard.ensure_held("record retained rebalance entry stats")?; + resolve_rebalance_stats_update_result(stats_result, pool_index, bucket.as_str(), entry.name.as_str())?; } + run_guard.ensure_held("rebalance entry completion")?; Ok(RebalanceEntryOutcome::Completed) } @@ -703,11 +711,23 @@ mod tests { let finish_store = Arc::clone(&store); let finish = tokio::spawn(async move { + let run_guard = finish_store + .rebalance_run_guard(rebalance_id, "rebalance source cleanup test") + .await + .expect("rebalance source cleanup test guard should be acquired"); finish_store - .finish_rebalance_entry_after_cleanup(0, "bucket", "object.bin", &[&version], rebalance_id, async move { - cleanup_released.await.expect("cleanup release sender should remain alive"); - Ok(ObjectInfo::default()) - }) + .finish_rebalance_entry_after_cleanup( + &run_guard, + 0, + "bucket", + "object.bin", + &[&version], + rebalance_id, + async move { + cleanup_released.await.expect("cleanup release sender should remain alive"); + Ok(ObjectInfo::default()) + }, + ) .await }); @@ -750,10 +770,20 @@ mod tests { let mut meta = store.rebalance_meta.write().await; meta.as_mut().expect("rebalance metadata should exist").pool_stats[0].bytes = 0; } + let warning_guard = store + .rebalance_run_guard(rebalance_id, "rebalance cleanup warning test") + .await + .expect("rebalance cleanup warning test guard should be acquired"); let warning_result = store - .finish_rebalance_entry_after_cleanup(0, "bucket", "object.bin", &[&warning_version], rebalance_id, async { - Err(Error::SlowDown.into()) - }) + .finish_rebalance_entry_after_cleanup( + &warning_guard, + 0, + "bucket", + "object.bin", + &[&warning_version], + rebalance_id, + async { Err(Error::SlowDown.into()) }, + ) .await .expect("cleanup warnings should not fail the completed migration"); assert!(matches!(warning_result, RebalanceEntryCleanupResult::Completed { warning: Some(_) })); @@ -762,15 +792,26 @@ mod tests { assert_eq!(pool_stats.cleanup_warnings.count, 1, "cleanup warning must block pool completion"); assert!(pool_stats.bytes > 0, "completed migration bytes should still be recorded"); drop(meta); + drop(warning_guard); { let mut meta = store.rebalance_meta.write().await; meta.as_mut().expect("rebalance metadata should exist").pool_stats[0].bytes = 0; } + let deferred_guard = store + .rebalance_run_guard(rebalance_id, "rebalance cleanup deferral test") + .await + .expect("rebalance cleanup deferral test guard should be acquired"); let deferred = store - .finish_rebalance_entry_after_cleanup(0, "bucket", "object.bin", &[&warning_version], rebalance_id, async { - Err(data_movement::SourceCleanupError::SourceChanged) - }) + .finish_rebalance_entry_after_cleanup( + &deferred_guard, + 0, + "bucket", + "object.bin", + &[&warning_version], + rebalance_id, + async { Err(data_movement::SourceCleanupError::SourceChanged) }, + ) .await .expect("source changes should defer cleanup without failing the worker"); assert!(matches!(deferred, RebalanceEntryCleanupResult::Deferred { .. })); diff --git a/crates/ecstore/src/services/rebalance/rebalance_unit_tests.rs b/crates/ecstore/src/services/rebalance/rebalance_unit_tests.rs index 14ad49285..b3cd1cdbb 100644 --- a/crates/ecstore/src/services/rebalance/rebalance_unit_tests.rs +++ b/crates/ecstore/src/services/rebalance/rebalance_unit_tests.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use super::control::validate_rebalance_disk_stats_coverage; +use super::control::{merge_and_save_rebalance_meta_no_lock, validate_rebalance_disk_stats_coverage}; use super::meta::{ RebalanceMetaMergeOutcome, RebalanceTerminalEvent, apply_rebalance_save_option, apply_rebalance_terminal_event, apply_stopped_at, classify_rebalance_terminal_event, clone_arc_by_index, clone_first_arc, clone_rebalance_pool_stats, @@ -33,8 +33,9 @@ use super::migration::{ rebalance_delete_marker_opts, }; use super::runtime::{ - RebalanceLocalActivationOutcome, commit_local_rebalance_worker_activation, resolve_rebalance_pre_spawn_result, - should_fail_repeated_rebalance_bucket_defer, source_cleanup_defer_attempt, + RebalanceLocalActivationOutcome, commit_local_rebalance_worker_activation, + commit_local_rebalance_worker_activation_candidate, should_fail_repeated_rebalance_bucket_defer, + source_cleanup_defer_attempt, stage_local_rebalance_worker_activation, }; use super::worker::{ RebalanceEntryCleanupResult, ensure_rebalance_listing_disks_available, is_transient_rebalance_error, @@ -2737,10 +2738,139 @@ fn test_stopped_activation_state_prevents_worker_token_commit() { } #[test] -fn test_rebalance_start_save_failure_rolls_back_local_worker_token() { - let activation_token = tokio_util::sync::CancellationToken::new(); - let observer = activation_token.clone(); - let mut meta = RebalanceMeta { +fn test_rebalance_activation_candidate_does_not_clobber_replacement_token() { + let mut local = RebalanceMeta { + id: "rebalance-a".to_string(), + pool_stats: vec![RebalanceStats { + participating: true, + buckets: vec!["bucket-a".to_string()], + info: RebalanceInfo { + status: RebalStatus::Started, + ..Default::default() + }, + ..Default::default() + }], + ..Default::default() + }; + let (candidate, outcome, must_persist) = + stage_local_rebalance_worker_activation(&local, "rebalance-a", CancellationToken::new(), OffsetDateTime::UNIX_EPOCH) + .expect("active activation candidate should be staged"); + assert_eq!(outcome, RebalanceLocalActivationOutcome::Started); + assert!(!must_persist); + + let replacement = CancellationToken::new(); + local.cancel = Some(replacement.clone()); + let err = commit_local_rebalance_worker_activation_candidate(&mut local, "rebalance-a", None, candidate) + .expect_err("a replacement token must reject the stale activation candidate"); + assert!(err.to_string().contains("worker token changed")); + assert_eq!(local.cancel.as_ref(), Some(&replacement)); +} + +#[tokio::test] +async fn test_rebalance_start_save_failure_retries_persisted_completed_state() { + let (_temp_dirs, _disk_stores, pool) = crate::set_disk::hermetic_set_disks_isolated(4).await; + let mut local = RebalanceMeta { + id: "rebalance-a".to_string(), + percent_free_goal: 0.5, + pool_stats: vec![RebalanceStats { + participating: true, + init_free_space: 400, + init_capacity: 1_000, + bytes: 100, + info: RebalanceInfo { + status: RebalStatus::Started, + ..Default::default() + }, + ..Default::default() + }], + ..Default::default() + }; + local.save(pool.clone()).await.expect("active metadata should be persisted"); + + let (failed_candidate, outcome, must_persist) = stage_local_rebalance_worker_activation( + &local, + "rebalance-a", + CancellationToken::new(), + OffsetDateTime::from_unix_timestamp(1_000).expect("test timestamp should be valid"), + ) + .expect("terminal activation candidate should be staged"); + assert_eq!(outcome, RebalanceLocalActivationOutcome::NotStartedTerminal); + assert!(must_persist); + assert_eq!(failed_candidate.pool_stats[0].info.status, RebalStatus::Completed); + assert_eq!(local.pool_stats[0].info.status, RebalStatus::Started); + + let err = merge_and_save_rebalance_meta_no_lock( + pool.clone(), + &failed_candidate, + "failed completed candidate", + ObjectOptions { + no_lock: true, + namespace_lock_fence: Some(crate::object_api::NamespaceLockFence::lost_for_test()), + ..Default::default() + }, + None, + Some(local.id.as_str()), + ) + .await + .expect_err("lost namespace quorum must reject the terminal candidate save"); + + assert!(matches!( + err, + Error::NamespaceLockQuorumUnavailable { + required: 3, + achieved: 2, + .. + } + )); + assert_eq!(local.pool_stats[0].info.status, RebalStatus::Started); + assert!(local.cancel.is_none(), "failed persistence must not publish a worker token"); + let mut after_failure = RebalanceMeta::new(); + after_failure + .load(pool.clone()) + .await + .expect("active metadata should remain readable after the failed save"); + assert_eq!(after_failure.pool_stats[0].info.status, RebalStatus::Started); + + let expected_cancel = local.cancel.clone(); + let (retry_candidate, retry_outcome, retry_must_persist) = stage_local_rebalance_worker_activation( + &local, + "rebalance-a", + CancellationToken::new(), + OffsetDateTime::from_unix_timestamp(2_000).expect("test timestamp should be valid"), + ) + .expect("retry terminal activation candidate should be staged"); + assert_eq!(retry_outcome, RebalanceLocalActivationOutcome::NotStartedTerminal); + assert!(retry_must_persist); + merge_and_save_rebalance_meta_no_lock( + pool.clone(), + &retry_candidate, + "retry completed candidate", + ObjectOptions { + no_lock: true, + ..Default::default() + }, + None, + Some(local.id.as_str()), + ) + .await + .expect("retry must persist the completed candidate"); + commit_local_rebalance_worker_activation_candidate(&mut local, "rebalance-a", expected_cancel.as_ref(), retry_candidate) + .expect("successful persistence should publish the completed candidate locally"); + + let mut persisted = RebalanceMeta::new(); + persisted + .load(pool) + .await + .expect("retry-persisted completed metadata should be readable"); + assert_eq!(persisted.pool_stats[0].info.status, RebalStatus::Completed); + assert_eq!(local.pool_stats[0].info.status, RebalStatus::Completed); + assert!(local.cancel.is_none()); +} + +#[tokio::test] +async fn test_rebalance_start_save_failure_retries_persisted_stopped_state() { + let (_temp_dirs, _disk_stores, pool) = crate::set_disk::hermetic_set_disks_isolated(4).await; + let active = RebalanceMeta { id: "rebalance-a".to_string(), pool_stats: vec![RebalanceStats { participating: true, @@ -2752,36 +2882,71 @@ fn test_rebalance_start_save_failure_rolls_back_local_worker_token() { }], ..Default::default() }; - assert_eq!( - commit_local_rebalance_worker_activation(&mut meta, "rebalance-a", activation_token.clone()) - .expect("active metadata should accept the worker token"), - RebalanceLocalActivationOutcome::Started - ); + active.save(pool.clone()).await.expect("active metadata should be persisted"); + let stopped_at = OffsetDateTime::from_unix_timestamp(1_000).expect("test timestamp should be valid"); + let mut local = active.clone(); + local.stopped_at = Some(stopped_at); + local.pool_stats[0].info.status = RebalStatus::Stopped; + local.pool_stats[0].info.end_time = Some(stopped_at); - let err = resolve_rebalance_pre_spawn_result( - Some(&mut meta), - "rebalance-a", - &activation_token, - Err(Error::NamespaceLockQuorumUnavailable { - mode: "write", - bucket: crate::disk::RUSTFS_META_BUCKET.to_string(), - object: super::REBAL_META_NAME.to_string(), - required: 3, - achieved: 2, - }), + let (failed_candidate, outcome, must_persist) = + stage_local_rebalance_worker_activation(&local, "rebalance-a", CancellationToken::new(), stopped_at) + .expect("stopped activation candidate should be staged"); + assert_eq!(outcome, RebalanceLocalActivationOutcome::NotStartedTerminal); + assert!(must_persist, "an already-terminal local state must still be persisted on retry"); + let err = merge_and_save_rebalance_meta_no_lock( + pool.clone(), + &failed_candidate, + "failed stopped candidate", + ObjectOptions { + no_lock: true, + namespace_lock_fence: Some(crate::object_api::NamespaceLockFence::lost_for_test()), + ..Default::default() + }, + None, + Some(local.id.as_str()), ) - .expect_err("metadata save failure must abort local worker activation"); + .await + .expect_err("lost namespace quorum must reject the stopped candidate save"); + assert!(matches!(err, Error::NamespaceLockQuorumUnavailable { .. })); + let mut after_failure = RebalanceMeta::new(); + after_failure + .load(pool.clone()) + .await + .expect("active metadata should remain readable after the failed save"); + assert_eq!(after_failure.pool_stats[0].info.status, RebalStatus::Started); + assert_eq!(local.pool_stats[0].info.status, RebalStatus::Stopped); - assert!(matches!( - err, - Error::NamespaceLockQuorumUnavailable { - required: 3, - achieved: 2, - .. - } - )); - assert!(observer.is_cancelled(), "the failed activation token must be canceled"); - assert!(meta.cancel.is_none(), "a retry must not see a phantom active worker token"); + let expected_cancel = local.cancel.clone(); + let (retry_candidate, retry_outcome, retry_must_persist) = + stage_local_rebalance_worker_activation(&local, "rebalance-a", CancellationToken::new(), stopped_at) + .expect("retry stopped activation candidate should be staged"); + assert_eq!(retry_outcome, RebalanceLocalActivationOutcome::NotStartedTerminal); + assert!(retry_must_persist); + merge_and_save_rebalance_meta_no_lock( + pool.clone(), + &retry_candidate, + "retry stopped candidate", + ObjectOptions { + no_lock: true, + ..Default::default() + }, + None, + Some(local.id.as_str()), + ) + .await + .expect("retry must persist the stopped candidate"); + commit_local_rebalance_worker_activation_candidate(&mut local, "rebalance-a", expected_cancel.as_ref(), retry_candidate) + .expect("successful persistence should preserve the stopped candidate locally"); + + let mut persisted = RebalanceMeta::new(); + persisted + .load(pool) + .await + .expect("retry-persisted stopped metadata should be readable"); + assert_eq!(persisted.stopped_at, Some(stopped_at)); + assert_eq!(persisted.pool_stats[0].info.status, RebalStatus::Stopped); + assert_eq!(local.pool_stats[0].info.status, RebalStatus::Stopped); } #[tokio::test] diff --git a/crates/ecstore/src/services/rebalance/runtime.rs b/crates/ecstore/src/services/rebalance/runtime.rs index 6a5b12a0f..5f03b46c2 100644 --- a/crates/ecstore/src/services/rebalance/runtime.rs +++ b/crates/ecstore/src/services/rebalance/runtime.rs @@ -64,6 +64,47 @@ pub(super) fn commit_local_rebalance_worker_activation( Ok(RebalanceLocalActivationOutcome::Started) } +pub(super) fn stage_local_rebalance_worker_activation( + meta: &super::RebalanceMeta, + expected_id: &str, + cancel: CancellationToken, + now: OffsetDateTime, +) -> Result<(super::RebalanceMeta, RebalanceLocalActivationOutcome, bool)> { + let mut candidate = meta.clone(); + let completed_at_goal = complete_rebalance_pools_at_goal(&mut candidate, now); + let completed_empty_queue = complete_rebalance_pools_with_empty_queue(&mut candidate, now); + let outcome = commit_local_rebalance_worker_activation(&mut candidate, expected_id, cancel)?; + let must_persist = + completed_at_goal || completed_empty_queue || outcome == RebalanceLocalActivationOutcome::NotStartedTerminal; + Ok((candidate, outcome, must_persist)) +} + +pub(super) fn commit_local_rebalance_worker_activation_candidate( + current: &mut super::RebalanceMeta, + expected_id: &str, + expected_cancel: Option<&CancellationToken>, + candidate: super::RebalanceMeta, +) -> Result<()> { + if current.id != expected_id || candidate.id != expected_id { + return Err(Error::other(format!( + "rebalance metadata changed before local worker activation commit: expected {expected_id}, found {}", + current.id + ))); + } + if !Arc::ptr_eq(¤t.activation_gate, &candidate.activation_gate) { + return Err(Error::other(format!( + "rebalance activation gate changed before local worker activation commit: {expected_id}" + ))); + } + if current.cancel.as_ref() != expected_cancel { + return Err(Error::other(format!( + "rebalance worker token changed before local worker activation commit: {expected_id}" + ))); + } + *current = candidate; + Ok(()) +} + pub(super) fn rollback_local_rebalance_worker_activation( meta: Option<&mut super::RebalanceMeta>, expected_id: &str, @@ -82,18 +123,6 @@ pub(super) fn rollback_local_rebalance_worker_activation( false } -pub(super) fn resolve_rebalance_pre_spawn_result( - meta: Option<&mut super::RebalanceMeta>, - expected_id: &str, - activation_token: &CancellationToken, - result: Result<()>, -) -> Result<()> { - if result.is_err() { - rollback_local_rebalance_worker_activation(meta, expected_id, activation_token); - } - result -} - impl ECStore { #[tracing::instrument(skip_all)] pub async fn start_rebalance(self: &Arc) -> Result<()> { @@ -126,8 +155,10 @@ impl ECStore { let cancel_tx = CancellationToken::new(); let rx = cancel_tx.clone(); - let mut meta_to_save = None; let activation_outcome; + let candidate; + let expected_cancel; + let must_persist; { let mut rebalance_meta = self.rebalance_meta.write().await; @@ -147,32 +178,72 @@ impl ECStore { ); return Ok(()); } - let now = OffsetDateTime::now_utc(); - if complete_rebalance_pools_at_goal(meta, now) { - meta_to_save = Some(meta.clone()); + expected_cancel = meta.cancel.clone(); + (candidate, activation_outcome, must_persist) = stage_local_rebalance_worker_activation( + meta, + expected_id.as_ref(), + cancel_tx.clone(), + OffsetDateTime::now_utc(), + )?; + if let Err(err) = activation_fence.ensure_held() { + cancel_tx.cancel(); + return Err(err); } - if complete_rebalance_pools_with_empty_queue(meta, now) { - meta_to_save = Some(meta.clone()); + if !must_persist { + if let Err(err) = commit_local_rebalance_worker_activation_candidate( + meta, + expected_id.as_ref(), + expected_cancel.as_ref(), + candidate.clone(), + ) { + cancel_tx.cancel(); + return Err(err); + } } - activation_fence.ensure_held()?; - activation_outcome = commit_local_rebalance_worker_activation(meta, expected_id.as_ref(), cancel_tx)?; + } - drop(rebalance_meta); + if must_persist { + let save_result = resolve_rebalance_meta_save_result( + self.save_rebalance_meta_under_activation_fence( + pool, + &candidate, + "start_rebalance persist activation candidate", + activation_fence.as_ref(), + expected_id.as_ref(), + ) + .await, + "start_rebalance persist activation candidate", + ); + if let Err(err) = save_result { + cancel_tx.cancel(); + return Err(err); + } + if let Err(err) = activation_fence.ensure_held() { + cancel_tx.cancel(); + return Err(err); + } + let mut rebalance_meta = self.rebalance_meta.write().await; + let Some(meta) = rebalance_meta.as_mut() else { + cancel_tx.cancel(); + return Err(Error::ConfigNotFound); + }; + if let Err(err) = commit_local_rebalance_worker_activation_candidate( + meta, + expected_id.as_ref(), + expected_cancel.as_ref(), + candidate, + ) { + cancel_tx.cancel(); + return Err(err); + } + } + if let Err(err) = activation_fence.ensure_held() { + let mut rebalance_meta = self.rebalance_meta.write().await; + rollback_local_rebalance_worker_activation(rebalance_meta.as_mut(), expected_id.as_ref(), &rx); + return Err(err); } drop(activation_fence); - if let Some(meta) = meta_to_save { - let save_result = resolve_rebalance_meta_save_result( - self.save_rebalance_meta_with_merge(pool, &meta, "start_rebalance complete pools at goal") - .await, - "start_rebalance complete pools at goal", - ); - if save_result.is_err() { - let mut rebalance_meta = self.rebalance_meta.write().await; - return resolve_rebalance_pre_spawn_result(rebalance_meta.as_mut(), expected_id.as_ref(), &rx, save_result); - } - } - if activation_outcome != RebalanceLocalActivationOutcome::Started { return Ok(()); }