diff --git a/crates/e2e_test/src/distributed/harness.rs b/crates/e2e_test/src/distributed/harness.rs index 364034477..3ce00fa01 100644 --- a/crates/e2e_test/src/distributed/harness.rs +++ b/crates/e2e_test/src/distributed/harness.rs @@ -860,10 +860,12 @@ pub(crate) async fn start_rebalance(cluster: &RustFSTestClusterEnvironment) -> T if status.is_success() { break response; } - if status == StatusCode::INTERNAL_SERVER_ERROR - && response.contains("requires a live fleet capability proof") - && Instant::now() < deadline - { + // The init stage answers 503 + Retry-After while the cluster capability + // proof is still missing; older binaries surfaced the same condition as + // a 500 carrying the proof marker. + let waiting_for_readiness = status == StatusCode::SERVICE_UNAVAILABLE + || (status == StatusCode::INTERNAL_SERVER_ERROR && response.contains("requires a live fleet capability proof")); + if waiting_for_readiness && Instant::now() < deadline { sleep(Duration::from_secs(1)).await; continue; } diff --git a/crates/ecstore/src/core/pools.rs b/crates/ecstore/src/core/pools.rs index f807eb603..ffca15d64 100644 --- a/crates/ecstore/src/core/pools.rs +++ b/crates/ecstore/src/core/pools.rs @@ -3691,6 +3691,55 @@ fn activation_pool_meta_lock_error(err: rustfs_lock::LockError) -> Error { } } +fn rebalance_worker_read_lock_error(err: rustfs_lock::LockError, object: &str) -> Error { + match err { + rustfs_lock::LockError::QuorumNotReached { required, achieved } => Error::NamespaceLockQuorumUnavailable { + mode: "read", + bucket: RUSTFS_META_BUCKET.to_string(), + object: object.to_string(), + required, + achieved, + }, + other => Error::Lock(other), + } +} + +/// Read-only admission to an already durable rebalance run. Keeping this +/// separate from the activation writer prevents admission from authorizing a +/// pool or rebalance metadata write under shared namespace locks. +pub(crate) struct PoolRebalanceWorkerFence { + pool_meta_guard: rustfs_lock::NamespaceLockGuard, + rebalance_meta_guard: rustfs_lock::NamespaceLockGuard, + fleet_proof: Option, +} + +impl PoolRebalanceWorkerFence { + pub(crate) fn set_fleet_proof( + &mut self, + fleet_proof: Option, + ) { + self.fleet_proof = fleet_proof; + } + + pub(crate) fn ensure_held(&self) -> Result<()> { + if self.pool_meta_guard.is_lock_lost() + || self.pool_meta_guard.is_released() + || self.rebalance_meta_guard.is_lock_lost() + || self.rebalance_meta_guard.is_released() + { + return Err(Error::other("worker admission lock lost before rebalance activation")); + } + if self + .fleet_proof + .as_ref() + .is_some_and(|proof| !crate::services::notification_sys::cross_pool_fence_fleet_proof_matches(proof)) + { + return Err(Error::other(POOL_ACTIVATION_FLEET_PROOF_EXPIRED)); + } + Ok(()) + } +} + pub(crate) struct PoolRebalanceActivationFence { pool_meta_guard: Arc, rebalance_meta_guard: Arc, @@ -3762,6 +3811,35 @@ tokio::task_local! { pub(crate) static REBALANCE_ACTIVATION_LOCK_ATTEMPT: Arc; } +pub(crate) async fn acquire_pool_rebalance_worker_locks(pool: Arc) -> Result +where + S: crate::storage_api_contracts::namespace::NamespaceLocking< + Error = Error, + NamespaceLock = rustfs_lock::NamespaceLockWrapper, + >, +{ + // Match activation and entry admission: rebalance.bin -> pool.bin. Joining + // the same durable run must coexist with entries retaining their run fence. + 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_read_lock(get_lock_acquire_timeout()) + .await + .map_err(|err| rebalance_worker_read_lock_error(err, REBAL_META_NAME))?; + let pool_meta_lock = pool.new_ns_lock(RUSTFS_META_BUCKET, POOL_META_NAME).await?; + let pool_meta_guard = pool_meta_lock + .get_read_lock(get_lock_acquire_timeout()) + .await + .map_err(|err| rebalance_worker_read_lock_error(err, POOL_META_NAME))?; + + Ok(PoolRebalanceWorkerFence { + pool_meta_guard, + rebalance_meta_guard, + fleet_proof: None, + }) +} + pub(crate) async fn acquire_pool_rebalance_activation_locks( pool: Arc, fleet_proof: Option, @@ -12430,6 +12508,18 @@ impl ECStore { Ok(pool_meta) } + pub(crate) async fn load_runtime_pool_meta_under_worker_fence( + &self, + write_state: &mut PoolMetaWriteState, + worker_fence: &PoolRebalanceWorkerFence, + operation: &str, + ) -> Result { + worker_fence.ensure_held()?; + let pool_meta = self.load_runtime_pool_meta_observing(write_state, operation).await?; + worker_fence.ensure_held()?; + Ok(pool_meta) + } + pub(crate) async fn load_runtime_pool_meta_under_activation_fence( &self, write_state: &mut PoolMetaWriteState, diff --git a/crates/ecstore/src/services/rebalance/control.rs b/crates/ecstore/src/services/rebalance/control.rs index a8fef5fd8..0ee302157 100644 --- a/crates/ecstore/src/services/rebalance/control.rs +++ b/crates/ecstore/src/services/rebalance/control.rs @@ -18,8 +18,8 @@ use super::{ encode_rebalance_stop_propagation_record, }; use crate::core::pools::{ - PoolMeta, PoolRebalanceActivationFence, acquire_pool_activation_fleet_proof, acquire_pool_rebalance_activation_locks, - pool_meta_has_active_decommission, + PoolMeta, PoolRebalanceActivationFence, PoolRebalanceWorkerFence, acquire_pool_activation_fleet_proof, + acquire_pool_rebalance_activation_locks, acquire_pool_rebalance_worker_locks, pool_meta_has_active_decommission, }; use crate::error::{Error, Result}; use crate::object_api::ObjectOptions; @@ -77,10 +77,33 @@ fn ensure_rebalance_activation_pool_meta_allowed(meta: &PoolMeta) -> Result<()> } pub(super) enum RebalanceWorkerActivationFence { - Ready(Box), + Ready(Box), NotStartedTerminal, } +async fn persisted_rebalance_requires_worker_activation(pool: Arc, expected_id: &str) -> Result +where + S: EcstoreObjectIO, +{ + let mut persisted = RebalanceMeta::new(); + persisted + .load_with_opts( + pool, + ObjectOptions { + no_lock: true, + ..Default::default() + }, + ) + .await?; + if persisted.id != expected_id { + return Err(Error::other(format!( + "rebalance metadata changed before worker activation: expected {expected_id}, found {}", + persisted.id + ))); + } + Ok(crate::services::rebalance::rebalance_requires_worker_activation(&persisted)) +} + pub(super) struct RebalanceRunGuard { _local_guard: tokio::sync::OwnedRwLockReadGuard<()>, persisted_guard: rustfs_lock::NamespaceLockGuard, @@ -577,48 +600,58 @@ impl ECStore { where S: EcstoreObjectIO + StorageNamespaceLocking, { - // Lock order: pool_meta_save_gate -> rebalance.bin -> pool.bin. + // Lock order: pool_meta_save_gate -> rebalance.bin -> pool.bin. These + // shared fences admit a peer without waiting for the run's own entries. 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 - // fences. A terminal record is a no-op and must not depend on the - // notification subsystem having published a fleet proof yet. - let mut activation_fence = acquire_pool_rebalance_activation_locks(pool.clone(), None).await?; + let mut worker_fence = acquire_pool_rebalance_worker_locks(pool.clone()).await?; let pool_meta = self - .load_runtime_pool_meta_under_activation_fence(&mut pool_meta_guard, &activation_fence, "rebalance worker activation") + .load_runtime_pool_meta_under_worker_fence(&mut pool_meta_guard, &worker_fence, "rebalance worker activation") .await?; ensure_rebalance_activation_pool_meta_allowed(&pool_meta)?; - - let mut persisted = RebalanceMeta::new(); - persisted - .load_with_opts( - pool, - ObjectOptions { - no_lock: true, - ..Default::default() - }, - ) - .await?; - if persisted.id != expected_id { - return Err(Error::other(format!( - "rebalance metadata changed before worker activation: expected {expected_id}, found {}", - persisted.id - ))); - } - - activation_fence.ensure_held()?; - if !crate::services::rebalance::rebalance_requires_worker_activation(&persisted) { + let needs_activation = persisted_rebalance_requires_worker_activation(pool, expected_id).await?; + worker_fence.ensure_held()?; + if !needs_activation { return Ok(RebalanceWorkerActivationFence::NotStartedTerminal); } - // Active worker admission still requires the fail-closed fleet proof. - // Attach it immediately before the final fence validation so expiry or - // topology changes are checked again at every later commit boundary. + let fleet_proof = acquire_pool_activation_fleet_proof(&self.ctx).await?; + worker_fence.set_fleet_proof(fleet_proof); + worker_fence.ensure_held()?; + Ok(RebalanceWorkerActivationFence::Ready(Box::new(worker_fence))) + } + + pub(super) async fn fence_rebalance_worker_metadata_update( + &self, + pool: Arc, + expected_id: &str, + ) -> Result>> + where + S: EcstoreObjectIO + StorageNamespaceLocking, + { + // The admission read fences must already be dropped. Reacquire in the + // original order and revalidate every durable boundary across that gap. + let mut pool_meta_guard = self.pool_meta_save_gate.lock().await; + pool_meta_guard.ensure_write_safe("rebalance worker metadata update")?; + let mut activation_fence = acquire_pool_rebalance_activation_locks(pool.clone(), None).await?; + let pool_meta = self + .load_runtime_pool_meta_under_activation_fence( + &mut pool_meta_guard, + &activation_fence, + "rebalance worker metadata update", + ) + .await?; + ensure_rebalance_activation_pool_meta_allowed(&pool_meta)?; + let needs_activation = persisted_rebalance_requires_worker_activation(pool, expected_id).await?; + activation_fence.ensure_held()?; + if !needs_activation { + return Ok(None); + } + let fleet_proof = acquire_pool_activation_fleet_proof(&self.ctx).await?; activation_fence.set_fleet_proof(fleet_proof); activation_fence.ensure_held()?; - - Ok(RebalanceWorkerActivationFence::Ready(Box::new(activation_fence))) + Ok(Some(Box::new(activation_fence))) } #[tracing::instrument(skip_all)] @@ -1816,6 +1849,107 @@ mod tests { assert_activation_locks_released(&store).await; } + #[tokio::test] + #[serial_test::serial] + async fn rebalance_init_missing_fleet_proof_preserves_metadata_and_rechecks_competing_run() { + let previous_id = "completed-before-capability-refresh"; + let terminal = RebalanceMeta { + id: previous_id.to_string(), + stopped_at: Some(OffsetDateTime::now_utc()), + pool_stats: vec![RebalanceStats::default(); 2], + ..Default::default() + }; + let (_temp_dirs, store, competing_store) = crate::services::rebalance::test_two_pool_stores(Some(terminal.clone())).await; + let disk_stats = vec![ + DiskStat { + total_space: 100, + available_space: 0, + }, + DiskStat { + total_space: 100, + available_space: 100, + }, + ]; + let initial_epoch = store.ctx.data_movement_operation_epoch(); + let proof_guard = crate::services::notification_sys::without_cross_pool_fence_fleet_proof_for_test(); + set_rebalance_disk_stats_override_for_test(store.id, disk_stats.clone()); + let err = store + .init_rebalance_start(vec!["waiting-request".to_string()]) + .await + .expect_err("initialization must fail closed while the live proof is missing"); + assert!( + err.to_string() + .contains("pool activation requires a live fleet capability proof") + ); + assert!( + crate::core::pools::is_pool_activation_fleet_proof_error(&err), + "the admin handler retries init only on this error shape: {err}" + ); + assert_eq!(store.ctx.data_movement_operation_epoch(), initial_epoch); + { + let local = store.rebalance_meta.read().await; + let local = local.as_ref().expect("failed initialization must preserve the previous run"); + assert_eq!(local.id, previous_id); + assert_eq!(local.stopped_at, terminal.stopped_at); + assert!(local.cancel.is_none(), "failed initialization must not start a worker"); + } + let mut persisted = RebalanceMeta::new(); + persisted + .load(store.pools[0].clone()) + .await + .expect("failed initialization must leave durable metadata readable"); + assert_eq!(persisted.id, previous_id); + assert_eq!(persisted.stopped_at, terminal.stopped_at); + + drop(proof_guard); + set_rebalance_disk_stats_override_for_test(competing_store.id, disk_stats.clone()); + let winning_id = competing_store + .init_rebalance_start(vec!["winning-request".to_string()]) + .await + .expect("a ready request must be able to initialize one run after the failed attempt"); + assert_ne!(winning_id, previous_id); + + // The waiting request still holds its old local snapshot. Its retry must + // recheck durable state under the activation locks before publishing. + set_rebalance_disk_stats_override_for_test(store.id, disk_stats); + let remote_conflict = store + .init_rebalance_start(vec!["waiting-request".to_string()]) + .await + .expect_err("the waiting request must not replace a durable competing run"); + assert!( + matches!( + crate::data_movement::data_movement_stage_source(&remote_conflict), + Some(Error::RebalanceAlreadyRunning) + ), + "the save-stage error must retain the typed competing-run conflict: {remote_conflict:?}" + ); + let local_conflict = competing_store + .init_rebalance_start(vec!["another-request".to_string()]) + .await + .expect_err("a local retry must not replace its active run"); + assert!( + matches!(local_conflict, Error::RebalanceAlreadyRunning), + "the local admission check must report the competing-run conflict: {local_conflict:?}" + ); + persisted + .load(store.pools[0].clone()) + .await + .expect("the competing run must remain durable after both rejected retries"); + assert_eq!(persisted.id, winning_id); + assert_eq!(persisted.pool_stats[0].buckets, vec!["winning-request".to_string()]); + assert_eq!( + store + .rebalance_meta + .read() + .await + .as_ref() + .expect("old local run must survive") + .id, + previous_id + ); + assert_eq!(store.ctx.data_movement_operation_epoch(), initial_epoch + 1); + } + #[tokio::test] #[serial_test::serial] async fn rebalance_worker_skips_terminal_metadata_without_fleet_proof() { @@ -1874,6 +2008,128 @@ mod tests { ); } + async fn assert_worker_write_upgrade_revalidates_durable_state(change: &str) { + let rebalance_id = "rebalance-worker-write-upgrade"; + let active = RebalanceMeta { + id: rebalance_id.to_string(), + percent_free_goal: 0.5, + pool_stats: vec![ + RebalanceStats { + participating: true, + init_capacity: 100, + info: RebalanceInfo { + status: RebalStatus::Started, + ..Default::default() + }, + ..Default::default() + }, + RebalanceStats::default(), + ], + ..Default::default() + }; + let (_temp_dirs, store, peer) = + crate::services::rebalance::test_two_pool_stores_with_isolated_node_contexts(Some(active.clone())).await; + let arrived = Arc::new(tokio::sync::Notify::new()); + let release = Arc::new(tokio::sync::Notify::new()); + let worker = Arc::clone(&store); + let mut tasks = tokio::task::JoinSet::new(); + tasks.spawn( + super::super::runtime::REBALANCE_WORKER_WRITE_UPGRADE_BARRIER + .scope((Arc::clone(&arrived), Arc::clone(&release)), async move { + worker.start_rebalance_for_id(rebalance_id).await + }), + ); + tokio::time::timeout(std::time::Duration::from_secs(30), arrived.notified()) + .await + .expect("empty-queue activation must release its read fences before requesting write authority"); + + let mut changed = active.clone(); + match change { + "stopped" => { + changed.stopped_at = Some(OffsetDateTime::now_utc()); + changed.pool_stats[0].info.status = RebalStatus::Stopped; + } + "replaced" => changed.id = "replacement-rebalance-run".to_string(), + "decommission" => { + let mut pool_meta = peer.pool_meta.read().await.clone(); + pool_meta.pools[0].decommission = Some(crate::core::pools::PoolDecommissionInfo { + start_time: Some(OffsetDateTime::now_utc()), + ..Default::default() + }); + pool_meta + .save(peer.pools.clone()) + .await + .expect("a competing pool metadata writer must acquire the released pool fence"); + } + _ => panic!("unknown durable state transition"), + } + changed + .save_with_opts( + peer.pools[0].clone(), + ObjectOptions { + write_completion: crate::object_api::WriteCompletion::TailDrained, + ..Default::default() + }, + ) + .await + .expect("a competing run writer must acquire the released rebalance fence"); + release.notify_one(); + let result = tokio::time::timeout(std::time::Duration::from_secs(30), tasks.join_next()) + .await + .expect("activation must finish after the durable transition") + .expect("activation task must be tracked") + .expect("activation must not panic"); + match change { + "stopped" => result.expect("a durably stopped run must become an activation no-op"), + "replaced" => assert!( + result + .expect_err("a replaced run must reject stale activation") + .to_string() + .contains("expected"), + ), + "decommission" => assert!(matches!(result, Err(Error::DecommissionAlreadyRunning))), + _ => unreachable!(), + } + let mut persisted = RebalanceMeta::new(); + persisted + .load(peer.pools[0].clone()) + .await + .expect("the competing durable state must remain readable"); + assert_eq!(persisted.id, changed.id, "activation must not overwrite the run identity"); + assert_eq!(persisted.stopped_at, changed.stopped_at, "activation must not overwrite the stop marker"); + assert_eq!(persisted.pool_stats[0].info.status, changed.pool_stats[0].info.status); + assert!( + store + .rebalance_meta + .read() + .await + .as_ref() + .expect("local metadata must remain present") + .cancel + .is_none(), + "rejected activation must not install a worker token", + ); + assert_activation_locks_released(&store).await; + } + + #[tokio::test] + #[serial_test::serial] + async fn rebalance_worker_write_upgrade_rechecks_stop() { + assert_worker_write_upgrade_revalidates_durable_state("stopped").await; + } + + #[tokio::test] + #[serial_test::serial] + async fn rebalance_worker_write_upgrade_rechecks_run_identity() { + assert_worker_write_upgrade_revalidates_durable_state("replaced").await; + } + + #[tokio::test] + #[serial_test::serial] + async fn rebalance_worker_write_upgrade_rechecks_decommission() { + assert_worker_write_upgrade_revalidates_durable_state("decommission").await; + } + #[tokio::test] #[serial_test::serial] async fn rebalance_activation_adopts_commit_after_post_save_fence_loss() { diff --git a/crates/ecstore/src/services/rebalance/entry.rs b/crates/ecstore/src/services/rebalance/entry.rs index 05fbe5c06..09cbdc4f4 100644 --- a/crates/ecstore/src/services/rebalance/entry.rs +++ b/crates/ecstore/src/services/rebalance/entry.rs @@ -1440,7 +1440,7 @@ mod tests { #[tokio::test] #[serial_test::serial] - async fn real_rebalance_entry_progresses_while_peer_activation_waits_for_run_fence() { + async fn real_rebalance_peer_activates_while_entry_retains_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), @@ -1456,7 +1456,7 @@ mod tests { 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 payload = b"peer activation must not wait for this entry to drain".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 { @@ -1516,7 +1516,12 @@ mod tests { ); tokio::time::timeout(StdDuration::from_secs(30), attempted.notified()) .await - .expect("peer activation must attempt the persisted rebalance write fence"); + .expect("peer activation must attempt the persisted rebalance fence"); + tokio::time::timeout(StdDuration::from_secs(30), activation_result) + .await + .expect("peer activation must finish while the entry remains paused") + .expect("peer activation result should be sent") + .expect("joining the same run must not require its active entry to release the run fence"); release.notify_one(); tokio::time::timeout(StdDuration::from_secs(30), async { @@ -1531,10 +1536,6 @@ mod tests { }) .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) diff --git a/crates/ecstore/src/services/rebalance/runtime.rs b/crates/ecstore/src/services/rebalance/runtime.rs index 39302d592..abead2485 100644 --- a/crates/ecstore/src/services/rebalance/runtime.rs +++ b/crates/ecstore/src/services/rebalance/runtime.rs @@ -27,6 +27,11 @@ use tokio::time::{Duration, Instant}; use tokio_util::sync::CancellationToken; use tracing::{debug, error, info, warn}; +#[cfg(test)] +tokio::task_local! { + pub(super) static REBALANCE_WORKER_WRITE_UPGRADE_BARRIER: (Arc, Arc); +} + pub(super) fn should_fail_repeated_rebalance_bucket_defer( deferred_buckets: &mut HashSet, bucket: &str, @@ -213,6 +218,22 @@ impl ECStore { } if must_persist { + // Shared admission grants no write authority. Release it before + // taking the exclusive metadata fences, then recheck stop/run/pool + // state before publishing the locally staged terminal candidate. + drop(activation_fence); + #[cfg(test)] + if let Ok((arrived, release)) = REBALANCE_WORKER_WRITE_UPGRADE_BARRIER.try_with(Clone::clone) { + arrived.notify_one(); + release.notified().await; + } + let Some(activation_fence) = self + .fence_rebalance_worker_metadata_update(pool.clone(), expected_id.as_ref()) + .await? + else { + cancel_tx.cancel(); + return Ok(false); + }; let save_result = resolve_rebalance_meta_save_result( self.save_rebalance_meta_under_activation_fence( pool, @@ -242,13 +263,14 @@ impl ECStore { cancel_tx.cancel(); return Err(err); } + } else { + 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 !must_persist && 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 activation_outcome != RebalanceLocalActivationOutcome::Started { return Ok(must_persist); diff --git a/rustfs/src/admin/handlers/rebalance.rs b/rustfs/src/admin/handlers/rebalance.rs index 14dc61b64..e59637b2b 100644 --- a/rustfs/src/admin/handlers/rebalance.rs +++ b/rustfs/src/admin/handlers/rebalance.rs @@ -42,12 +42,9 @@ use s3s::{ s3_error, }; use serde::{Deserialize, Serialize}; -use std::{ - future::Future, - sync::Arc, - time::{Duration, Instant}, -}; +use std::{future::Future, sync::Arc, time::Duration}; use time::OffsetDateTime; +use tokio::time::Instant; use tracing::{error, info, warn}; const LOG_COMPONENT_ADMIN: &str = "admin"; @@ -128,6 +125,13 @@ fn is_rebalance_start_fleet_proof_retryable(err: &StorageError) -> bool { && err.to_string().contains(REBALANCE_START_FLEET_PROOF_MARKER) } +fn rebalance_init_failure_error(err: &StorageError) -> S3Error { + if is_rebalance_start_fleet_proof_retryable(err) { + return rebalance_fleet_proof_unavailable_error(); + } + s3_error!(InternalError, "failed to initialize rebalance: {}", err) +} + fn rebalance_start_failure_error(start_err: &StorageError, rollback_result: &Result<(), String>) -> S3Error { if crate::storage_api::capacity::is_pool_activation_fleet_proof_error(start_err) { if rollback_result.is_ok() { @@ -611,7 +615,20 @@ impl Operation for RebalanceStart { let buckets: Vec = bucket_infos.into_iter().map(|bucket| bucket.name).collect(); - let id = match store.init_rebalance_start(buckets).await { + let fleet_proof_deadline = Instant::now() + REBALANCE_START_FLEET_PROOF_RETRY_BUDGET; + // A missing proof fails before init publishes metadata. Retrying the full + // operation reacquires its gates and revalidates competing starts. + let id = match retry_rebalance_start_operation( + "init_rebalance_start", + "", + &request_id, + &actor, + &remote_addr, + fleet_proof_deadline, + || store.init_rebalance_start(buckets.clone()), + ) + .await + { Ok(id) => id, Err(StorageError::DecommissionAlreadyRunning) => { log_rebalance_request_rejected("start", "decommission_in_progress", &request_id, &actor, &remote_addr); @@ -622,7 +639,7 @@ impl Operation for RebalanceStart { return Err(s3_error!(OperationAborted, "rebalance is already in progress")); } Err(e) => { - return Err(s3_error!(InternalError, "failed to initialize rebalance: {}", e)); + return Err(rebalance_init_failure_error(&e)); } }; @@ -639,7 +656,6 @@ impl Operation for RebalanceStart { "admin rebalance state" ); let notification_sys = current_notification_system(); - let fleet_proof_deadline = Instant::now() + REBALANCE_START_FLEET_PROOF_RETRY_BUDGET; for step in rebalance_start_steps(notification_sys.is_some()) { match step { RebalanceStartStep::PropagateFence => { @@ -1170,10 +1186,10 @@ mod rebalance_handler_tests { Body, HeaderMap, Method, Operation, Params, RebalPoolProgress, RebalanceAdminStatus, RebalancePoolStatus, RebalanceStart, RebalanceStartStep, RebalanceStatus, RebalanceStop, RebalanceStopPropagationStatus, S3ErrorCode, S3Request, StatusCode, Uri, build_rebalance_admin_status, build_rebalance_pool_statuses, build_rebalance_stop_propagation_status, - is_rebalance_start_fleet_proof_retryable, rebalance_pool_used, rebalance_query_present, rebalance_remaining_buckets, - rebalance_rollback_failure_message, rebalance_rollback_stop_failure_message, rebalance_start_failure_error, - rebalance_start_rollback_error, rebalance_start_steps, rebalance_stop_target_id, rebalance_used_pct, - retry_rebalance_start_operation, rollback_result_label, stop_rebalance_admission_first, + is_rebalance_start_fleet_proof_retryable, rebalance_init_failure_error, rebalance_pool_used, rebalance_query_present, + rebalance_remaining_buckets, rebalance_rollback_failure_message, rebalance_rollback_stop_failure_message, + rebalance_start_failure_error, rebalance_start_rollback_error, rebalance_start_steps, rebalance_stop_target_id, + rebalance_used_pct, retry_rebalance_start_operation, rollback_result_label, stop_rebalance_admission_first, }; use crate::admin::storage_api::error::StorageError; use crate::admin::storage_api::rebalance::{ @@ -1533,7 +1549,7 @@ mod rebalance_handler_tests { } #[tokio::test(start_paused = true)] - async fn test_rebalance_start_retry_waits_for_fleet_proof() { + async fn test_rebalance_init_retry_waits_for_fleet_proof() { use std::sync::{ Arc, atomic::{AtomicUsize, Ordering}, @@ -1544,12 +1560,12 @@ mod rebalance_handler_tests { let attempts = Arc::clone(&attempts); async move { retry_rebalance_start_operation( - "load_rebalance_meta(start=true)", - "rebalance-id", + "init_rebalance_start", + "", "request-id", "actor", "remote", - std::time::Instant::now() + super::REBALANCE_START_FLEET_PROOF_RETRY_BUDGET, + super::Instant::now() + super::REBALANCE_START_FLEET_PROOF_RETRY_BUDGET, move || { let attempts = Arc::clone(&attempts); async move { @@ -1557,7 +1573,7 @@ mod rebalance_handler_tests { if attempt == 0 { Err(StorageError::other(super::REBALANCE_START_FLEET_PROOF_MARKER)) } else { - Ok(()) + Ok("initialized-once") } } }, @@ -1570,12 +1586,78 @@ mod rebalance_handler_tests { assert_eq!(attempts.load(Ordering::SeqCst), 1); tokio::time::advance(super::REBALANCE_START_FLEET_PROOF_RETRY_DELAY).await; - task.await + let id = task + .await .expect("retry task should not panic") .expect("fleet proof retry should eventually succeed"); + assert_eq!(id, "initialized-once"); assert_eq!(attempts.load(Ordering::SeqCst), 2); } + #[tokio::test(start_paused = true)] + async fn test_rebalance_init_retry_exhaustion_returns_503_with_retry_after() { + let mut attempts = 0; + let started = super::Instant::now(); + let deadline = started + super::REBALANCE_START_FLEET_PROOF_RETRY_BUDGET; + let err = retry_rebalance_start_operation::<(), _, _>( + "init_rebalance_start", + "", + "request-id", + "actor", + "remote", + deadline, + || { + attempts += 1; + std::future::ready(Err(StorageError::other(super::REBALANCE_START_FLEET_PROOF_MARKER))) + }, + ) + .await + .expect_err("a missing proof must exhaust the bounded retry budget"); + + assert_eq!(super::Instant::now() - started, super::REBALANCE_START_FLEET_PROOF_RETRY_BUDGET); + assert_eq!(attempts, 7, "one immediate attempt plus six bounded readiness retries"); + let response = rebalance_init_failure_error(&err); + assert_eq!(response.status_code(), Some(StatusCode::SERVICE_UNAVAILABLE)); + assert_eq!( + response + .headers() + .and_then(|headers| headers.get(http::header::RETRY_AFTER)) + .and_then(|value| value.to_str().ok()), + Some(super::REBALANCE_START_RETRY_AFTER_SECS) + ); + } + + #[tokio::test(start_paused = true)] + async fn test_rebalance_init_does_not_retry_other_failures() { + for source in [ + StorageError::other("disk read failed"), + StorageError::other("pool activation fleet capability proof expired before commit"), + StorageError::RebalanceAlreadyRunning, + ] { + let original = source.to_string(); + let mut source = Some(source); + let started = super::Instant::now(); + let err = retry_rebalance_start_operation::<(), _, _>( + "init_rebalance_start", + "", + "request-id", + "actor", + "remote", + started + super::REBALANCE_START_FLEET_PROOF_RETRY_BUDGET, + || std::future::ready(Err(source.take().expect("non-readiness errors must not be retried"))), + ) + .await + .expect_err("non-readiness errors must be returned unchanged"); + assert_eq!(err.to_string(), original); + assert_eq!(super::Instant::now(), started, "non-readiness failures must not wait"); + if !matches!(err, StorageError::RebalanceAlreadyRunning) { + let response = rebalance_init_failure_error(&err); + assert_eq!(response.code(), &S3ErrorCode::InternalError); + assert!(response.message().is_some_and(|message| message.contains(&original))); + } + } + } + #[test] fn test_rebalance_rollback_stop_failure_message_lists_failures() { let failures = vec![