From d933d5e938659bfd2f585e8be9a9fb69acd332b1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=AE=89=E6=AD=A3=E8=B6=85?= Date: Sat, 20 Jun 2026 17:56:11 +0800 Subject: [PATCH] refactor: move ecstore rebalance flow modules (#3658) --- crates/ecstore/src/rebalance.rs | 4975 +---------------- crates/ecstore/src/rebalance/control.rs | 410 ++ crates/ecstore/src/rebalance/entry.rs | 513 ++ crates/ecstore/src/rebalance/meta.rs | 193 +- .../src/rebalance/rebalance_unit_tests.rs | 3230 +++++++++++ crates/ecstore/src/rebalance/runtime.rs | 574 ++ crates/ecstore/src/rebalance/types.rs | 119 + crates/utils/src/os/windows.rs | 2 +- docs/architecture/migration-progress.md | 179 +- 9 files changed, 5212 insertions(+), 4983 deletions(-) create mode 100644 crates/ecstore/src/rebalance/control.rs create mode 100644 crates/ecstore/src/rebalance/entry.rs create mode 100644 crates/ecstore/src/rebalance/rebalance_unit_tests.rs create mode 100644 crates/ecstore/src/rebalance/runtime.rs create mode 100644 crates/ecstore/src/rebalance/types.rs diff --git a/crates/ecstore/src/rebalance.rs b/crates/ecstore/src/rebalance.rs index 46d063530..19ff68e5a 100644 --- a/crates/ecstore/src/rebalance.rs +++ b/crates/ecstore/src/rebalance.rs @@ -12,30 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. -use crate::config::com::{read_config_with_metadata, save_config_with_opts}; -use crate::data_movement; use crate::error::{Error, Result}; -use crate::global::get_global_endpoints; use crate::object_api::{GetObjectReader, ObjectInfo, ObjectOptions, PutObjReader}; -use crate::pools::ListCallback; -use crate::set_disk::{SetDisks, get_lock_acquire_timeout}; -use crate::storage_api_contracts::EcstoreObjectIO; -use crate::store::ECStore; -use http::HeaderMap; -use rustfs_filemeta::{FileInfo, MetaCacheEntry}; -use rustfs_storage_api::{ - HTTPRangeSpec, NamespaceLocking as StorageNamespaceLocking, ObjectIO, ObjectOperations as _, StorageAdminApi, -}; -use rustfs_utils::path::encode_dir_object; -use serde::{Deserialize, Serialize}; -use std::collections::HashSet; -use std::fmt; -use std::io::Cursor; -use std::sync::Arc; -use time::OffsetDateTime; -use tokio::time::{Duration, Instant}; -use tokio_util::sync::CancellationToken; -use tracing::{debug, error, info, warn}; +use tokio::time::Duration; const LOG_COMPONENT_ECSTORE: &str = "ecstore"; const LOG_SUBSYSTEM_REBALANCE: &str = "rebalance"; @@ -43,7 +22,6 @@ const EVENT_REBALANCE_STATE: &str = "rebalance_state"; const EVENT_REBALANCE_BUCKET: &str = "rebalance_bucket"; const EVENT_REBALANCE_ENTRY: &str = "rebalance_entry"; const EVENT_REBALANCE_LISTING: &str = "rebalance_listing"; -use uuid::Uuid; const REBAL_META_FMT: u16 = 1; // Replace with actual format value const REBAL_META_VER: u16 = 1; // Replace with actual version value @@ -55,4951 +33,18 @@ const REBALANCE_MIGRATION_RETRY_BASE_DELAY: Duration = Duration::from_millis(250 const REBALANCE_MIGRATION_LOCK_RETRY_CAP: Duration = Duration::from_secs(10); const REBALANCE_DEFERRED_ENTRY_ERROR_PREFIX: &str = "deferred transient rebalance entry failure:"; +mod control; +mod entry; mod meta; mod migration; +mod runtime; +mod types; mod worker; -use meta::{ - apply_rebalance_save_option, apply_rebalance_terminal_event, classify_rebalance_terminal_event, clone_arc_by_index, - clone_first_arc, clone_rebalance_pool_stats, complete_rebalance_pools_at_goal, complete_rebalance_pools_with_empty_queue, - defer_bucket_in_rebalance_queue, ensure_valid_rebalance_pool_index, has_deferred_rebalance_error, - invalid_rebalance_pool_index_error, is_rebalance_conflicting_with_decommission, is_rebalance_in_progress, - mark_rebalance_bucket_done, merge_rebalance_meta, percent_free_ratio, rebalance_goal_reached, - rebalance_meta_load_no_data_error, rebalance_meta_load_unknown_format_error, rebalance_meta_load_unknown_version_error, - rebalance_metadata_not_initialized_error, record_rebalance_cleanup_warning_in_meta, resolve_next_rebalance_bucket, - resolve_rebalance_participants, should_accept_rebalance_stats_update, should_ignore_rebalance_data_usage_cache, - should_pool_participate, should_preserve_rebalance_stopped_state, should_skip_start_rebalance, stop_rebalance_meta_snapshot, - validate_start_rebalance_state, + +pub use types::{ + DiskStat, RStats, RebalSaveOpt, RebalStatus, RebalanceCleanupWarnings, RebalanceInfo, RebalanceMeta, RebalanceStats, }; -use migration::migrate_entry_version; -use worker::{ - RebalanceEntryTask, load_rebalance_bucket_configs, rebalance_max_attempts, rebalance_meta_lock_error, - resolve_load_rebalance_stats_update_result, resolve_rebalance_bucket_error, resolve_rebalance_bucket_result, - resolve_rebalance_entry_cleanup_delete_result, resolve_rebalance_file_info_versions_result, - resolve_rebalance_meta_load_result, resolve_rebalance_meta_save_result, resolve_rebalance_migrate_result_error, - resolve_rebalance_save_task_result, resolve_rebalance_stats_update_result, resolve_rebalance_terminal_error, - resolve_rebalance_worker_result, run_rebalance_listing_with_retry, send_rebalance_done_signal, - should_cleanup_rebalance_source_entry, should_count_rebalance_version_complete, should_defer_rebalance_entry_failure, - should_skip_rebalance_delete_marker, wait_rebalance_entry_tasks, with_rebalance_entry_context, -}; - -#[derive(Debug, Default, Clone, Serialize, Deserialize)] -pub struct RebalanceStats { - #[serde(rename = "ifs")] - pub init_free_space: u64, // Pool free space at the start of rebalance - #[serde(rename = "ic")] - pub init_capacity: u64, // Pool capacity at the start of rebalance - #[serde(rename = "bus")] - pub buckets: Vec, // Buckets being rebalanced or to be rebalanced - #[serde(rename = "rbs")] - pub rebalanced_buckets: Vec, // Buckets rebalanced - #[serde(rename = "bu")] - pub bucket: String, // Last rebalanced bucket - #[serde(rename = "ob")] - pub object: String, // Last rebalanced object - #[serde(rename = "no")] - pub num_objects: u64, // Number of objects rebalanced - #[serde(rename = "nv")] - pub num_versions: u64, // Number of versions rebalanced - #[serde(rename = "bs")] - pub bytes: u64, // Number of bytes rebalanced - #[serde(rename = "par")] - pub participating: bool, // Whether the pool is participating in rebalance - #[serde(rename = "inf")] - pub info: RebalanceInfo, // Rebalance operation info - #[serde(rename = "cw", default)] - pub cleanup_warnings: RebalanceCleanupWarnings, -} - -impl RebalanceStats { - pub fn update(&mut self, bucket: String, fi: &FileInfo) { - if fi.is_latest { - self.num_objects += 1; - } - - self.num_versions += 1; - let on_disk_size = if fi.deleted || fi.erasure.data_blocks == 0 || fi.size <= 0 { - 0 - } else { - let data_blocks = fi.erasure.data_blocks as i64; - let total_blocks = fi.erasure.data_blocks.saturating_add(fi.erasure.parity_blocks) as i64; - fi.size - .saturating_mul(total_blocks) - .checked_div(data_blocks) - .unwrap_or(0) - .max(0) as u64 - }; - self.bytes = self.bytes.saturating_add(on_disk_size); - self.bucket = bucket; - self.object = fi.name.clone(); - } - - pub fn update_batch(&mut self, bucket: String, versions: &[&FileInfo]) { - for version in versions { - self.update(bucket.clone(), version); - } - } -} - -pub type RStats = Vec>; - -#[derive(Debug, Default)] -struct RebalanceBucketConfigs { - lifecycle_config: Option, - lock_retention: Option, - replication_config: Option<(s3s::dto::ReplicationConfiguration, OffsetDateTime)>, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -enum RebalanceBucketOutcome { - Completed, - Deferred { last_error: String }, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -enum RebalanceEntryOutcome { - Completed, - Deferred { last_error: String }, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)] -pub enum RebalStatus { - #[default] - None, - Started, - Completed, - Stopped, - Failed, -} - -impl fmt::Display for RebalStatus { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - let status = match self { - RebalStatus::None => "None", - RebalStatus::Started => "Started", - RebalStatus::Completed => "Completed", - RebalStatus::Stopped => "Stopped", - RebalStatus::Failed => "Failed", - }; - write!(f, "{status}") - } -} - -impl From for RebalStatus { - fn from(value: u8) -> Self { - match value { - 1 => RebalStatus::Started, - 2 => RebalStatus::Completed, - 3 => RebalStatus::Stopped, - 4 => RebalStatus::Failed, - _ => RebalStatus::None, - } - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)] -pub enum RebalSaveOpt { - #[default] - Stats, - StoppedAt, -} - -#[derive(Debug, Default, Clone, Serialize, Deserialize)] -pub struct RebalanceInfo { - #[serde(rename = "startTs")] - pub start_time: Option, // Time at which rebalance-start was issued - #[serde(rename = "stopTs")] - pub end_time: Option, // Time at which rebalance operation completed or rebalance-stop was called - #[serde(rename = "err")] - pub last_error: Option, // Last rebalance error message - #[serde(rename = "status")] - pub status: RebalStatus, // Current state of rebalance operation -} - -#[derive(Debug, Default, Clone, Serialize, Deserialize, PartialEq, Eq)] -pub struct RebalanceCleanupWarnings { - #[serde(rename = "count", default)] - pub count: u64, - #[serde(rename = "lastMsg", default)] - pub last_message: Option, - #[serde(rename = "lastBucket", default)] - pub last_bucket: Option, - #[serde(rename = "lastObject", default)] - pub last_object: Option, - #[serde(rename = "lastAt", default)] - pub last_at: Option, -} - -#[allow(dead_code)] -#[derive(Debug, Clone, Default)] -pub struct DiskStat { - pub total_space: u64, - pub available_space: u64, -} - -#[derive(Debug, Default, Serialize, Deserialize, Clone)] -pub struct RebalanceMeta { - #[serde(skip)] - pub cancel: Option, // To be invoked on rebalance-stop - #[serde(skip)] - pub last_refreshed_at: Option, - #[serde(rename = "stopTs")] - pub stopped_at: Option, // Time when rebalance-stop was issued - #[serde(rename = "id")] - pub id: String, // ID of the ongoing rebalance operation - #[serde(rename = "pf")] - pub percent_free_goal: f64, // Computed from total free space and capacity at the start of rebalance - #[serde(rename = "rss")] - pub pool_stats: Vec, // Per-pool rebalance stats keyed by pool index -} - -impl RebalanceMeta { - pub fn new() -> Self { - Self::default() - } - pub async fn load(&mut self, store: Arc) -> Result<()> - where - S: ObjectIO< - Error = Error, - RangeSpec = HTTPRangeSpec, - HeaderMap = HeaderMap, - ObjectOptions = ObjectOptions, - ObjectInfo = ObjectInfo, - GetObjectReader = GetObjectReader, - PutObjectReader = PutObjReader, - >, - { - self.load_with_opts(store, ObjectOptions::default()).await - } - - pub async fn load_with_opts(&mut self, store: Arc, opts: ObjectOptions) -> Result<()> - where - S: ObjectIO< - Error = Error, - RangeSpec = HTTPRangeSpec, - HeaderMap = HeaderMap, - ObjectOptions = ObjectOptions, - ObjectInfo = ObjectInfo, - GetObjectReader = GetObjectReader, - PutObjectReader = PutObjReader, - >, - { - let (data, _) = read_config_with_metadata(store, REBAL_META_NAME, &opts).await?; - if data.is_empty() { - debug!( - event = EVENT_REBALANCE_STATE, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - state = "metadata_empty", - "Rebalance metadata is empty" - ); - return Ok(()); - } - if data.len() <= 4 { - return Err(rebalance_meta_load_no_data_error()); - } - - // Read header - match u16::from_le_bytes([data[0], data[1]]) { - REBAL_META_FMT => {} - fmt => return Err(rebalance_meta_load_unknown_format_error(fmt)), - } - match u16::from_le_bytes([data[2], data[3]]) { - REBAL_META_VER => {} - ver => return Err(rebalance_meta_load_unknown_version_error(ver)), - } - - let meta: Self = rmp_serde::from_read(Cursor::new(&data[4..]))?; - *self = meta; - - self.last_refreshed_at = Some(OffsetDateTime::now_utc()); - - debug!( - event = EVENT_REBALANCE_STATE, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - state = "metadata_loaded", - "Loaded rebalance metadata" - ); - Ok(()) - } - - pub async fn save(&self, store: Arc) -> Result<()> - where - S: ObjectIO< - Error = Error, - RangeSpec = HTTPRangeSpec, - HeaderMap = HeaderMap, - ObjectOptions = ObjectOptions, - ObjectInfo = ObjectInfo, - GetObjectReader = GetObjectReader, - PutObjectReader = PutObjReader, - >, - { - self.save_with_opts(store, ObjectOptions::default()).await - } - - pub async fn save_with_opts(&self, store: Arc, opts: ObjectOptions) -> Result<()> - where - S: ObjectIO< - Error = Error, - RangeSpec = HTTPRangeSpec, - HeaderMap = HeaderMap, - ObjectOptions = ObjectOptions, - ObjectInfo = ObjectInfo, - GetObjectReader = GetObjectReader, - PutObjectReader = PutObjReader, - >, - { - if self.pool_stats.is_empty() { - debug!( - event = EVENT_REBALANCE_STATE, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - state = "metadata_save_skipped", - reason = "no_pool_stats", - "Skipped rebalance metadata save" - ); - return Ok(()); - } - - let mut data = Vec::new(); - - // Initialize the header - data.extend(&REBAL_META_FMT.to_le_bytes()); - data.extend(&REBAL_META_VER.to_le_bytes()); - - let msg = rmp_serde::to_vec(self)?; - data.extend(msg); - - save_config_with_opts(store, REBAL_META_NAME, data, &opts).await?; - - Ok(()) - } -} - -impl ECStore { - async fn save_rebalance_meta_with_merge(&self, pool: Arc, local_snapshot: &RebalanceMeta, stage: &str) -> Result<()> - where - S: EcstoreObjectIO + StorageNamespaceLocking, - { - let ns_lock = pool.new_ns_lock(crate::disk::RUSTFS_META_BUCKET, REBAL_META_NAME).await?; - let _guard = ns_lock - .get_write_lock(get_lock_acquire_timeout()) - .await - .map_err(rebalance_meta_lock_error)?; - - let opts = ObjectOptions { - no_lock: true, - ..Default::default() - }; - let mut merged = RebalanceMeta::new(); - match merged.load_with_opts(pool.clone(), opts.clone()).await { - Ok(()) => { - merge_rebalance_meta(&mut merged, local_snapshot); - } - Err(Error::ConfigNotFound) => { - merged = local_snapshot.clone(); - } - Err(err) => return Err(Error::other(format!("rebalance meta load before save failed during {stage}: {err}"))), - } - - merged.save_with_opts(pool, opts).await - } - - #[tracing::instrument(skip_all)] - pub async fn load_rebalance_meta(&self) -> Result<()> { - let mut meta = RebalanceMeta::new(); - debug!( - event = EVENT_REBALANCE_STATE, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - state = "metadata_loading", - "Loading rebalance metadata" - ); - let pool = clone_first_arc(&self.pools, "rebalanceMeta: no pools available")?; - if resolve_rebalance_meta_load_result(meta.load(pool).await)? { - { - let mut rebalance_meta = self.rebalance_meta.write().await; - - *rebalance_meta = Some(meta); - - drop(rebalance_meta); - } - - resolve_load_rebalance_stats_update_result(self.update_rebalance_stats().await)?; - debug!( - event = EVENT_REBALANCE_STATE, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - state = "metadata_loaded", - "Loaded rebalance metadata" - ); - } else { - debug!( - event = EVENT_REBALANCE_STATE, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - state = "metadata_missing", - reason = "rebalance_not_started", - "Rebalance metadata not found" - ); - } - - Ok(()) - } - - #[tracing::instrument(skip_all)] - pub async fn update_rebalance_stats(&self) -> Result<()> { - let mut ok = false; - - let pool_stats = { - let rebalance_meta = self.rebalance_meta.read().await; - clone_rebalance_pool_stats(rebalance_meta.as_ref())? - }; - - debug!( - event = EVENT_REBALANCE_STATE, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - pool_count = pool_stats.len(), - "Refreshing rebalance stats snapshot" - ); - - for i in 0..self.pools.len() { - if pool_stats.get(i).is_none() { - let mut rebalance_meta = self.rebalance_meta.write().await; - debug!( - event = EVENT_REBALANCE_STATE, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - pool_index = i, - state = "pool_stat_missing", - "Adding missing rebalance pool stats entry" - ); - if let Some(meta) = rebalance_meta.as_mut() { - meta.pool_stats.push(RebalanceStats::default()); - } - ok = true; - drop(rebalance_meta); - } - } - - if ok { - debug!( - event = EVENT_REBALANCE_STATE, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - state = "metadata_saving", - "Saving rebalance metadata after stats refresh" - ); - - let rebalance_meta = self.rebalance_meta.read().await; - if let Some(meta) = rebalance_meta.as_ref() { - let pool = clone_first_arc(&self.pools, "update_rebalance_stats: no pools available")?; - resolve_rebalance_meta_save_result( - self.save_rebalance_meta_with_merge(pool, meta, "update_rebalance_stats") - .await, - "update_rebalance_stats", - )?; - } - } - - Ok(()) - } - - // async fn find_index(&self, index: usize) -> Option { - // if let Some(meta) = self.rebalance_meta.read().await.as_ref() { - // return meta.pool_stats.get(index).map(|_v| index); - // } - - // None - // } - - #[tracing::instrument(skip(self))] - pub async fn init_rebalance_meta(&self, bucktes: Vec) -> Result { - info!( - event = EVENT_REBALANCE_STATE, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - state = "initializing", - bucket_count = bucktes.len(), - "Initializing rebalance metadata" - ); - let si = StorageAdminApi::storage_info(self).await; - - let mut disk_stats = vec![DiskStat::default(); self.pools.len()]; - - let mut total_cap = 0; - let mut total_free = 0; - for disk in si.disks.iter() { - if disk.pool_index < 0 || disk_stats.len() <= disk.pool_index as usize { - continue; - } - - total_cap += disk.total_space; - total_free += disk.available_space; - - disk_stats[disk.pool_index as usize].total_space += disk.total_space; - disk_stats[disk.pool_index as usize].available_space += disk.available_space; - } - - let percent_free_goal = percent_free_ratio(total_free, total_cap); - - let mut pool_stats = Vec::with_capacity(self.pools.len()); - - let now = OffsetDateTime::now_utc(); - - for disk_stat in disk_stats.iter() { - let mut pool_stat = RebalanceStats { - init_free_space: disk_stat.available_space, - init_capacity: disk_stat.total_space, - buckets: bucktes.clone(), - rebalanced_buckets: Vec::with_capacity(bucktes.len()), - ..Default::default() - }; - - if should_pool_participate(disk_stat.available_space, disk_stat.total_space, percent_free_goal) { - pool_stat.participating = true; - pool_stat.info = RebalanceInfo { - start_time: Some(now), - status: RebalStatus::Started, - ..Default::default() - }; - } - - pool_stats.push(pool_stat); - } - - let meta = RebalanceMeta { - id: Uuid::new_v4().to_string(), - percent_free_goal, - pool_stats, - ..Default::default() - }; - - let pool = clone_first_arc(&self.pools, "init_rebalance_meta: no pools available")?; - resolve_rebalance_meta_save_result( - self.save_rebalance_meta_with_merge(pool, &meta, "init_rebalance_meta").await, - "init_rebalance_meta", - )?; - - info!( - event = EVENT_REBALANCE_STATE, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - state = "metadata_initialized", - bucket_count = bucktes.len(), - "Rebalance metadata initialized" - ); - - let id = meta.id.clone(); - - { - let mut rebalance_meta = self.rebalance_meta.write().await; - *rebalance_meta = Some(meta); - drop(rebalance_meta); - } - - Ok(id) - } - - #[tracing::instrument(skip(self, fi))] - pub async fn update_pool_stats(&self, pool_index: usize, bucket: String, fi: &FileInfo) -> Result<()> { - self.update_pool_stats_batch(pool_index, bucket, &[fi]).await - } - - #[tracing::instrument(skip(self, versions))] - pub async fn update_pool_stats_batch(&self, pool_index: usize, bucket: String, versions: &[&FileInfo]) -> Result<()> { - if versions.is_empty() { - return Ok(()); - } - - let mut rebalance_meta = self.rebalance_meta.write().await; - if let Some(meta) = rebalance_meta.as_mut() { - if !should_accept_rebalance_stats_update(meta, pool_index) { - return Ok(()); - } - - if let Some(pool_stat) = meta.pool_stats.get_mut(pool_index) { - pool_stat.update_batch(bucket, versions); - } - } - - Ok(()) - } - - #[tracing::instrument(skip(self))] - pub async fn next_rebal_bucket(&self, pool_index: usize) -> Result> { - let rebalance_meta = self.rebalance_meta.read().await; - debug!( - event = EVENT_REBALANCE_BUCKET, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - pool_index, - has_meta = rebalance_meta.is_some(), - state = "next_bucket_lookup", - "Rebalance next bucket lookup" - ); - resolve_next_rebalance_bucket(rebalance_meta.as_ref(), pool_index) - } - - #[tracing::instrument(skip(self))] - pub async fn bucket_rebalance_done(&self, pool_index: usize, bucket: String) -> Result<()> { - let mut rebalance_meta = self.rebalance_meta.write().await; - mark_rebalance_bucket_done(rebalance_meta.as_mut(), pool_index, &bucket) - } - - async fn record_rebalance_cleanup_warning( - &self, - pool_index: usize, - bucket: &str, - object: &str, - message: String, - ) -> Result<()> { - let mut rebalance_meta = self.rebalance_meta.write().await; - record_rebalance_cleanup_warning_in_meta( - rebalance_meta.as_mut(), - pool_index, - bucket, - object, - message, - OffsetDateTime::now_utc(), - ) - } - - async fn defer_rebalance_bucket(&self, pool_index: usize, bucket: String, last_error: String) -> Result<()> { - let mut rebalance_meta = self.rebalance_meta.write().await; - let Some(meta) = rebalance_meta.as_mut() else { - return Err(rebalance_metadata_not_initialized_error("defer rebalance bucket")); - }; - let pool_count = meta.pool_stats.len(); - ensure_valid_rebalance_pool_index(pool_count, pool_index)?; - let Some(pool_stat) = meta.pool_stats.get_mut(pool_index) else { - return Err(invalid_rebalance_pool_index_error(pool_index, pool_count)); - }; - - defer_bucket_in_rebalance_queue(pool_stat, &bucket)?; - pool_stat.info.last_error = Some(last_error); - meta.last_refreshed_at = Some(OffsetDateTime::now_utc()); - Ok(()) - } - - pub async fn is_rebalance_started(&self) -> bool { - let rebalance_meta = self.rebalance_meta.read().await; - if let Some(meta) = rebalance_meta.as_ref() { - meta.pool_stats.iter().enumerate().for_each(|(i, v)| { - debug!( - event = EVENT_REBALANCE_STATE, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - pool_index = i, - participating = v.participating, - status = ?v.info.status, - state = "status_inspected", - "Rebalance status inspected" - ); - }); - - let started = is_rebalance_conflicting_with_decommission(meta); - if started { - debug!( - event = EVENT_REBALANCE_STATE, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - state = "running", - "Rebalance is running" - ); - return true; - } - } - - debug!( - event = EVENT_REBALANCE_STATE, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - state = "not_running", - "Rebalance is not running" - ); - false - } - - pub async fn is_rebalance_conflicting_with_decommission(&self) -> bool { - let rebalance_meta = self.rebalance_meta.read().await; - rebalance_meta - .as_ref() - .is_some_and(is_rebalance_conflicting_with_decommission) - } - - pub async fn is_pool_rebalancing(&self, pool_index: usize) -> bool { - let rebalance_meta = self.rebalance_meta.read().await; - if let Some(ref meta) = *rebalance_meta { - if meta.stopped_at.is_some() { - return false; - } - - if let Some(pool_stat) = meta.pool_stats.get(pool_index) { - return pool_stat.participating && pool_stat.info.status == RebalStatus::Started; - } - } - - false - } - - #[tracing::instrument(skip(self))] - pub async fn stop_rebalance(self: &Arc) -> Result<()> { - let meta_to_save = { - let mut rebalance_meta = self.rebalance_meta.write().await; - stop_rebalance_meta_snapshot(rebalance_meta.as_mut(), OffsetDateTime::now_utc()) - }; - - if let Some(meta_to_save) = meta_to_save { - let pool = clone_first_arc(self.pools.as_slice(), "stop_rebalance: no pools available")?; - resolve_rebalance_meta_save_result( - self.save_rebalance_meta_with_merge(pool, &meta_to_save, "stop_rebalance") - .await, - "stop_rebalance", - )?; - } - - Ok(()) - } - - #[tracing::instrument(skip_all)] - pub async fn start_rebalance(self: &Arc) -> Result<()> { - info!( - event = EVENT_REBALANCE_STATE, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - state = "starting", - "Starting rebalance" - ); - let decommission_running = self.is_decommission_running().await; - // let rebalance_meta = self.rebalance_meta.read().await; - - let cancel_tx = CancellationToken::new(); - let rx = cancel_tx.clone(); - let mut meta_to_save = None; - - { - let mut rebalance_meta = self.rebalance_meta.write().await; - validate_start_rebalance_state(decommission_running, rebalance_meta.is_some())?; - - let Some(meta) = rebalance_meta.as_mut() else { - return Err(Error::ConfigNotFound); - }; - if should_skip_start_rebalance(meta.cancel.is_some(), is_rebalance_in_progress(meta)) { - debug!( - event = EVENT_REBALANCE_STATE, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - state = "start_skipped", - reason = "already_in_progress", - "Skipped duplicate rebalance start" - ); - return Ok(()); - } - let now = OffsetDateTime::now_utc(); - if complete_rebalance_pools_at_goal(meta, now) { - meta_to_save = Some(meta.clone()); - } - if complete_rebalance_pools_with_empty_queue(meta, now) { - meta_to_save = Some(meta.clone()); - } - meta.cancel = Some(cancel_tx); - - drop(rebalance_meta); - } - - if let Some(meta) = meta_to_save { - let pool = clone_first_arc(self.pools.as_slice(), "start_rebalance: no pools available")?; - 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", - )?; - } - - let participants = if let Some(ref meta) = *self.rebalance_meta.read().await { - resolve_rebalance_participants(meta.pool_stats.as_slice(), self.pools.len()) - } else { - debug!( - event = EVENT_REBALANCE_STATE, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - state = "start_skipped", - reason = "metadata_missing", - "Skipped rebalance start because metadata is unavailable" - ); - Vec::new() - }; - - if !participants.iter().any(|participating| *participating) { - debug!( - event = EVENT_REBALANCE_STATE, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - state = "start_skipped", - reason = "no_participants", - "Skipped rebalance start because no pools are participating" - ); - return Ok(()); - } - - let mut workers_started = 0usize; - for (idx, participating) in participants.iter().enumerate() { - if !*participating { - debug!( - event = EVENT_REBALANCE_STATE, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - pool_index = idx, - state = "pool_skipped", - reason = "not_participating", - "Skipped rebalance pool" - ); - continue; - } - - if !get_global_endpoints() - .as_ref() - .get(idx) - .is_some_and(|v| v.endpoints.as_ref().first().is_some_and(|e| e.is_local)) - { - debug!( - event = EVENT_REBALANCE_STATE, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - pool_index = idx, - state = "pool_skipped", - reason = "not_local", - "Skipped non-local rebalance pool" - ); - continue; - } - - let pool_idx = idx; - let store = self.clone(); - let rx_clone = rx.clone(); - workers_started += 1; - tokio::spawn(async move { - if let Err(err) = store.rebalance_buckets(rx_clone, pool_idx).await { - error!( - event = EVENT_REBALANCE_STATE, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - pool_index = pool_idx, - state = "pool_failed", - error = %err, - "Rebalance pool failed" - ); - } else { - debug!( - event = EVENT_REBALANCE_STATE, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - pool_index = pool_idx, - state = "completed", - "Rebalance pool completed" - ); - } - }); - } - - if workers_started == 0 { - debug!( - event = EVENT_REBALANCE_STATE, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - state = "start_skipped", - reason = "no_local_participants", - "Skipped rebalance start because no local pools are participating" - ); - return Ok(()); - } - - info!( - event = EVENT_REBALANCE_STATE, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - state = "started", - worker_count = workers_started, - "Rebalance started" - ); - Ok(()) - } - - #[tracing::instrument(skip(self, rx))] - async fn rebalance_buckets(self: &Arc, rx: CancellationToken, pool_index: usize) -> Result<()> { - ensure_valid_rebalance_pool_index(self.pools.len(), pool_index)?; - - let (done_tx, mut done_rx) = tokio::sync::mpsc::channel::>(1); - - // Save rebalance metadata periodically - let store = self.clone(); - let save_task = tokio::spawn(async move { - let mut timer = tokio::time::interval_at(Instant::now() + Duration::from_secs(30), Duration::from_secs(10)); - let mut msg: String; - let mut quit = false; - - loop { - tokio::select! { - result = done_rx.recv() => { - quit = true; - let now = OffsetDateTime::now_utc(); - let terminal_event = classify_rebalance_terminal_event(result, now); - msg = terminal_event.message().to_string(); - let mut rebalance_meta = store.rebalance_meta.write().await; - if let Some(meta) = rebalance_meta.as_mut() { - let meta_stopped = meta.stopped_at.is_some(); - if let Some(pool_stat) = meta.pool_stats.get_mut(pool_index) { - if should_preserve_rebalance_stopped_state( - meta_stopped, - pool_stat.info.status, - &terminal_event, - ) { - debug!( - event = EVENT_REBALANCE_STATE, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - pool_index, - state = "stopped_preserved", - "Preserved stopped rebalance status" - ); - } else { - apply_rebalance_terminal_event( - &mut pool_stat.info.status, - &mut pool_stat.info.end_time, - &mut pool_stat.info.last_error, - terminal_event, - now, - ); - } - } - } - } - _ = timer.tick() => { - let now = OffsetDateTime::now_utc(); - msg = format!("Saving rebalance metadata at {now:?}"); - } - } - - if let Err(err) = store.save_rebalance_stats(pool_index, RebalSaveOpt::Stats).await { - let wrapped = Error::other(format!("rebalance save_task stats save failed for pool {pool_index}: {err}")); - error!("{} err: {:?}", msg, wrapped); - if quit { - return Err(wrapped); - } - } else { - debug!( - event = EVENT_REBALANCE_STATE, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - pool_index, - state = "metadata_saved", - message = %msg, - "Saved rebalance metadata" - ); - } - - if quit { - debug!( - event = EVENT_REBALANCE_STATE, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - pool_index, - state = "save_task_exiting", - message = %msg, - "Exiting rebalance save task" - ); - return Ok(()); - } - - timer.reset(); - } - }); - - info!( - event = EVENT_REBALANCE_STATE, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - pool_index, - state = "pool_started", - "Rebalance worker started" - ); - let mut final_result: Result<()> = Ok(()); - let mut deferred_buckets = HashSet::new(); - - loop { - if rx.is_cancelled() { - info!( - event = EVENT_REBALANCE_STATE, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - pool_index, - state = "pool_stopped", - reason = "cancelled", - "Stopped rebalance worker" - ); - let err = Error::OperationCanceled; - final_result = Err(resolve_rebalance_terminal_error( - err.clone(), - send_rebalance_done_signal(&done_tx, Err(err.clone()), pool_index).await, - )); - break; - } - - let next_bucket = match self.next_rebal_bucket(pool_index).await { - Ok(bucket) => bucket, - Err(err) => { - error!( - event = EVENT_REBALANCE_BUCKET, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - pool_index, - state = "next_bucket_failed", - error = ?err, - "Rebalance next bucket lookup failed" - ); - final_result = Err(resolve_rebalance_terminal_error( - err.clone(), - send_rebalance_done_signal(&done_tx, Err(err.clone()), pool_index).await, - )); - break; - } - }; - - if let Some(bucket) = next_bucket { - debug!( - event = EVENT_REBALANCE_BUCKET, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - pool_index, - bucket = %bucket, - state = "started", - "Starting rebalance bucket" - ); - - let outcome = match resolve_rebalance_bucket_result( - self.rebalance_bucket(rx.clone(), bucket.clone(), pool_index).await, - pool_index, - &bucket, - ) { - Ok(outcome) => outcome, - Err(err) => { - error!( - event = EVENT_REBALANCE_BUCKET, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - pool_index, - bucket = %bucket, - state = "bucket_failed", - error = ?err, - "Rebalance bucket failed" - ); - final_result = Err(resolve_rebalance_terminal_error( - err.clone(), - send_rebalance_done_signal(&done_tx, Err(err.clone()), pool_index).await, - )); - break; - } - }; - - if let RebalanceBucketOutcome::Deferred { last_error } = outcome { - if !deferred_buckets.insert(bucket.clone()) { - let err = Error::other(format!( - "rebalance bucket {bucket} deferred repeatedly due to transient object failures: {last_error}" - )); - error!( - event = EVENT_REBALANCE_BUCKET, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - pool_index, - bucket = %bucket, - state = "bucket_deferred_repeatedly", - error = ?err, - "Rebalance bucket failed after repeated deferral" - ); - final_result = Err(resolve_rebalance_terminal_error( - err.clone(), - send_rebalance_done_signal(&done_tx, Err(err.clone()), pool_index).await, - )); - break; - } - - warn!( - event = EVENT_REBALANCE_BUCKET, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - pool_index, - bucket = %bucket, - state = "deferred", - error = %last_error, - "Deferred rebalance bucket after transient object failures" - ); - if let Err(err) = self.defer_rebalance_bucket(pool_index, bucket.clone(), last_error).await { - error!( - event = EVENT_REBALANCE_BUCKET, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - pool_index, - bucket = %bucket, - state = "defer_failed", - error = ?err, - "Rebalance bucket defer failed" - ); - final_result = Err(resolve_rebalance_terminal_error( - err.clone(), - send_rebalance_done_signal(&done_tx, Err(err.clone()), pool_index).await, - )); - break; - } - continue; - } - - debug!( - event = EVENT_REBALANCE_BUCKET, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - pool_index, - bucket = %bucket, - state = "completed", - "Completed rebalance bucket" - ); - if let Err(err) = self.bucket_rebalance_done(pool_index, bucket).await { - error!( - event = EVENT_REBALANCE_BUCKET, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - pool_index, - state = "bucket_done_mark_failed", - error = ?err, - "Rebalance bucket completion mark failed" - ); - final_result = Err(resolve_rebalance_terminal_error( - err.clone(), - send_rebalance_done_signal(&done_tx, Err(err.clone()), pool_index).await, - )); - break; - } - } else { - debug!( - event = EVENT_REBALANCE_BUCKET, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - pool_index, - state = "idle", - reason = "no_bucket_to_rebalance", - "No rebalance bucket available" - ); - break; - } - } - - info!( - event = EVENT_REBALANCE_STATE, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - pool_index, - state = "pool_done", - "Rebalance worker finished" - ); - - if final_result.is_ok() - && let Err(err) = send_rebalance_done_signal(&done_tx, Ok(()), pool_index).await - { - final_result = Err(err); - } - drop(done_tx); - if let Err(err) = resolve_rebalance_save_task_result(pool_index, save_task.await) - && final_result.is_ok() - { - final_result = Err(err); - } - debug!( - event = EVENT_REBALANCE_STATE, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - pool_index, - state = "pool_result_returned", - "Rebalance worker result returned" - ); - final_result - } - - async fn check_if_rebalance_done(&self, pool_index: usize) -> bool { - let mut rebalance_meta = self.rebalance_meta.write().await; - - if let Some(meta) = rebalance_meta.as_mut() - && let Some(pool_stat) = meta.pool_stats.get_mut(pool_index) - { - // Check if the pool's rebalance status is already completed - if pool_stat.info.status == RebalStatus::Completed { - debug!( - event = EVENT_REBALANCE_STATE, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - pool_index, - state = "already_completed", - "Rebalance pool is already completed" - ); - return true; - } - - // Mark pool rebalance as done only after it reaches the PercentFreeGoal. - let pfi = if pool_stat.init_capacity == 0 { - 0.0 - } else { - (pool_stat.init_free_space + pool_stat.bytes) as f64 / pool_stat.init_capacity as f64 - }; - - if !has_deferred_rebalance_error(pool_stat) - && rebalance_goal_reached( - pool_stat.init_free_space, - pool_stat.init_capacity, - pool_stat.bytes, - meta.percent_free_goal, - ) - { - pool_stat.info.status = RebalStatus::Completed; - pool_stat.info.end_time = Some(OffsetDateTime::now_utc()); - info!( - event = EVENT_REBALANCE_STATE, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - pool_index, - state = "completed", - percent_free = pfi, - "Marked rebalance pool completed" - ); - return true; - } - } - - false - } -} - -impl ECStore { - #[allow(unused_assignments)] - #[tracing::instrument(skip(self, set))] - async fn rebalance_entry( - self: Arc, - bucket: String, - pool_index: usize, - entry: MetaCacheEntry, - set: Arc, - bucket_configs: Arc, - // wk: Arc, - ) -> Result { - debug!( - event = EVENT_REBALANCE_ENTRY, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - bucket = %bucket, - object = %entry.name, - pool_index, - state = "started", - "Starting rebalance entry" - ); - - // defer!(|| async { - // warn!("rebalance_entry: defer give worker start"); - // wk.give().await; - // warn!("rebalance_entry: defer give worker done"); - // }); - - if entry.is_dir() { - debug!( - event = EVENT_REBALANCE_ENTRY, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - bucket = %bucket, - object = %entry.name, - pool_index, - state = "skipped", - reason = "directory_entry", - "Skipped rebalance entry" - ); - return Ok(RebalanceEntryOutcome::Completed); - } - - if self.check_if_rebalance_done(pool_index).await { - debug!( - event = EVENT_REBALANCE_ENTRY, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - pool_index, - bucket = %bucket, - object = %entry.name, - state = "skipped", - reason = "pool_completed", - "Skipped rebalance entry" - ); - return Ok(RebalanceEntryOutcome::Completed); - } - - let mut fivs = - resolve_rebalance_file_info_versions_result(entry.file_info_versions(&bucket), bucket.as_str(), entry.name.as_str())?; - - fivs.versions - .sort_by_key(|v| (v.mod_time.is_none(), std::cmp::Reverse(v.mod_time))); - - let mut rebalanced: usize = 0; - let mut expired: usize = 0; - let mut stats_updates = Vec::with_capacity(fivs.versions.len()); - for version in fivs.versions.iter() { - if crate::pools::should_skip_lifecycle_for_data_movement( - self.clone(), - &bucket, - version, - bucket_configs.lifecycle_config.as_ref(), - bucket_configs.lock_retention.clone(), - bucket_configs.replication_config.clone(), - true, - &crate::bucket::lifecycle::bucket_lifecycle_audit::LcEventSrc::Rebal, - ) - .await - { - expired += 1; - debug!( - event = EVENT_REBALANCE_ENTRY, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - pool_index, - bucket = %bucket, - object = %version.name, - state = "skipped", - reason = "expired_by_lifecycle", - "Skipped rebalance version" - ); - continue; - } - - let remaining_versions = fivs.versions.len() - expired; - if should_skip_rebalance_delete_marker(version, remaining_versions, bucket_configs.replication_config.is_some()) { - rebalanced += 1; - debug!( - event = EVENT_REBALANCE_ENTRY, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - pool_index, - bucket = %bucket, - object = %version.name, - state = "skipped", - reason = "last_delete_marker_without_replication", - "Skipped rebalance version" - ); - continue; - } - - let version_id = version.version_id.map(|v| v.to_string()); - let mut transfer = |src_pool_idx: usize, bucket: String, rd: GetObjectReader| { - let store = self.clone(); - async move { store.rebalance_object(src_pool_idx, bucket, rd).await } - }; - let result = migrate_entry_version( - set.as_ref(), - bucket.clone(), - pool_index, - version, - version_id.clone(), - rebalance_max_attempts(), - should_ignore_rebalance_data_usage_cache(bucket.as_str()), - &mut transfer, - ) - .await; - - if result.ignored { - if should_count_rebalance_version_complete(&result) { - rebalanced += 1; - } - debug!( - event = EVENT_REBALANCE_ENTRY, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - pool_index, - bucket = %bucket, - object = %version.name, - state = "skipped", - reason = "already_deleted", - "Skipped rebalance version" - ); - continue; - } - - if result.failed { - let err = resolve_rebalance_migrate_result_error( - result.error, - pool_index, - bucket.as_str(), - version.name.as_str(), - version_id.as_deref(), - ); - error!( - "rebalance_entry {} Error rebalancing entry {}/{:?}: {:?}", - &bucket, &version.name, &version.version_id, err - ); - if should_defer_rebalance_entry_failure(&err) { - let deferred_error = format!("{REBALANCE_DEFERRED_ENTRY_ERROR_PREFIX} {err}"); - warn!( - event = EVENT_REBALANCE_ENTRY, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - pool_index, - bucket = %bucket, - object = %version.name, - state = "deferred", - error = %err, - "Deferred rebalance entry after transient migration failure" - ); - if let Err(stats_err) = self.update_rebalance_last_error(pool_index, deferred_error.clone()).await { - error!( - "rebalance_entry {} failed to record deferred transient failure for {}: {}", - &bucket, &entry.name, stats_err - ); - } - return Ok(RebalanceEntryOutcome::Deferred { - last_error: deferred_error, - }); - } - 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 - .update_pool_stats_batch(pool_index, bucket.clone(), stats_updates.as_slice()) - .await - { - error!( - "rebalance_entry {} failed to update stats before returning migration error for {}: {}", - &bucket, &entry.name, stats_err - ); - } - - return Err(entry_err); - } - - stats_updates.push(version); - if should_count_rebalance_version_complete(&result) { - rebalanced += 1; - } - } - - resolve_rebalance_stats_update_result( - self.update_pool_stats_batch(pool_index, bucket.clone(), stats_updates.as_slice()) - .await, - pool_index, - bucket.as_str(), - entry.name.as_str(), - )?; - - if should_cleanup_rebalance_source_entry(rebalanced, fivs.versions.len()) { - let cleanup_warning = resolve_rebalance_entry_cleanup_delete_result( - set.delete_object( - bucket.as_str(), - &encode_dir_object(&entry.name), - ObjectOptions { - delete_prefix: true, - delete_prefix_object: true, - - ..Default::default() - }, - ) - .await, - bucket.as_str(), - entry.name.as_str(), - )?; - if let Some(message) = cleanup_warning { - warn!( - event = EVENT_REBALANCE_ENTRY, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - pool_index, - bucket = %bucket, - object = %entry.name, - stage = "cleanup_source", - cleanup_status = "failed_ignored", - error = %message, - "Ignored rebalance source cleanup failure" - ); - if let Err(err) = self - .record_rebalance_cleanup_warning(pool_index, bucket.as_str(), entry.name.as_str(), message) - .await - { - error!( - event = EVENT_REBALANCE_ENTRY, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - pool_index, - bucket = %bucket, - object = %entry.name, - stage = "cleanup_source", - error = ?err, - "Failed to record rebalance source cleanup warning" - ); - } - } else { - debug!( - event = EVENT_REBALANCE_ENTRY, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - pool_index, - bucket = %bucket, - object = %entry.name, - state = "source_deleted", - "Deleted rebalance source entry" - ); - } - } - - Ok(RebalanceEntryOutcome::Completed) - } - - #[tracing::instrument(skip(self, rd))] - async fn rebalance_object(self: Arc, pool_idx: usize, bucket: String, rd: GetObjectReader) -> Result<()> { - data_movement::migrate_object(self, pool_idx, bucket, rd, "rebalance_object").await - } - - async fn update_rebalance_last_error(&self, pool_idx: usize, message: String) -> Result<()> { - let mut rebalance_meta = self.rebalance_meta.write().await; - let Some(meta) = rebalance_meta.as_mut() else { - return Err(rebalance_metadata_not_initialized_error("record rebalance last error")); - }; - let pool_count = meta.pool_stats.len(); - ensure_valid_rebalance_pool_index(pool_count, pool_idx)?; - let Some(pool_stat) = meta.pool_stats.get_mut(pool_idx) else { - return Err(invalid_rebalance_pool_index_error(pool_idx, pool_count)); - }; - - pool_stat.info.last_error = Some(message); - meta.last_refreshed_at = Some(OffsetDateTime::now_utc()); - Ok(()) - } - - #[tracing::instrument(skip(self, rx))] - async fn rebalance_bucket( - self: &Arc, - rx: CancellationToken, - bucket: String, - pool_index: usize, - ) -> Result { - ensure_valid_rebalance_pool_index(self.pools.len(), pool_index)?; - - // Placeholder for actual bucket rebalance logic - debug!( - event = EVENT_REBALANCE_BUCKET, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - pool_index, - bucket = %bucket, - state = "entry_scan_started", - "Rebalance bucket entry scan started" - ); - - // TODO: other config - // if bucket != RUSTFS_META_BUCKET{ - - // } - - let pool = clone_arc_by_index(self.pools.as_slice(), pool_index, "invalid rebalance pool index")?; - let bucket_configs = Arc::new(load_rebalance_bucket_configs(&bucket).await?); - - let mut jobs = Vec::new(); - let entry_error = Arc::new(tokio::sync::Mutex::new(None::)); - let entry_workers = Arc::new(tokio::sync::Semaphore::new(pool.disk_set.len().max(1))); - - for (set_idx, set) in pool.disk_set.iter().enumerate() { - let entry_tasks = Arc::new(tokio::sync::Mutex::new(Vec::::new())); - let rebalance_entry: ListCallback = Arc::new({ - let this = Arc::clone(self); - let bucket = bucket.clone(); - let entry_error = entry_error.clone(); - let callback_rx = rx.clone(); - let set = set.clone(); - let bucket_configs = bucket_configs.clone(); - let entry_tasks = entry_tasks.clone(); - let entry_workers = entry_workers.clone(); - move |entry: MetaCacheEntry| { - let this = this.clone(); - let bucket = bucket.clone(); - let entry_error = entry_error.clone(); - let callback_rx = callback_rx.clone(); - let set = set.clone(); - let bucket_configs = bucket_configs.clone(); - let entry_tasks = entry_tasks.clone(); - let entry_workers = entry_workers.clone(); - Box::pin(async move { - if callback_rx.is_cancelled() { - return; - } - if entry_error.lock().await.is_some() { - return; - } - - let permit = tokio::select! { - _ = callback_rx.cancelled() => return, - permit = entry_workers.clone().acquire_owned() => match permit { - Ok(permit) => permit, - Err(err) => { - error!("rebalance_entry: worker semaphore closed: {err}"); - return; - } - }, - }; - - if entry_error.lock().await.is_some() { - return; - } - - let task = tokio::spawn(async move { - let _permit = permit; - debug!( - event = EVENT_REBALANCE_ENTRY, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - set_index = set_idx, - state = "task_started", - "Started rebalance entry task" - ); - let result = this.rebalance_entry(bucket, pool_index, entry, set, bucket_configs).await; - if let Err(err) = &result { - error!("rebalance_entry: rebalance entry failed: {err}"); - let mut first_err = entry_error.lock().await; - if first_err.is_none() { - *first_err = Some(err.clone()); - callback_rx.cancel(); - } - } - debug!( - event = EVENT_REBALANCE_ENTRY, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - set_index = set_idx, - state = "task_completed", - "Completed rebalance entry task" - ); - result - }); - - entry_tasks.lock().await.push(task); - }) - } - }); - - let set = set.clone(); - let rx = rx.clone(); - let bucket = bucket.clone(); - let entry_tasks = entry_tasks.clone(); - - let job = tokio::spawn(async move { - let list_result = - run_rebalance_listing_with_retry(set, rx, bucket.clone(), rebalance_entry, set_idx, rebalance_max_attempts()) - .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 { - error!("Rebalance worker {} error: {}", set_idx, err); - } else { - debug!( - event = EVENT_REBALANCE_STATE, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - set_index = set_idx, - state = "worker_completed", - "Completed rebalance worker" - ); - } - result - }); - - jobs.push((set_idx, job)); - } - - let mut worker_error: Option = None; - let mut deferred_error: Option = None; - for (set_idx, job) in jobs { - match resolve_rebalance_worker_result(set_idx, job.await) { - Ok(Some(last_error)) if deferred_error.is_none() => { - deferred_error = Some(last_error); - } - Ok(_) => {} - Err(err) if worker_error.is_none() => { - worker_error = Some(err); - } - Err(_) => {} - } - } - let entry_error = entry_error.lock().await.clone(); - resolve_rebalance_bucket_error(entry_error, worker_error)?; - if let Some(last_error) = deferred_error { - return Ok(RebalanceBucketOutcome::Deferred { last_error }); - } - - debug!( - event = EVENT_REBALANCE_BUCKET, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - pool_index, - bucket = %bucket, - state = "completed", - "Finished rebalance bucket" - ); - Ok(RebalanceBucketOutcome::Completed) - } - - #[tracing::instrument(skip(self))] - pub async fn save_rebalance_stats(&self, pool_idx: usize, opt: RebalSaveOpt) -> Result<()> { - let meta_to_save = { - let mut rebalance_meta = self.rebalance_meta.write().await; - let Some(meta) = rebalance_meta.as_mut() else { - return Ok(()); - }; - - let now = OffsetDateTime::now_utc(); - apply_rebalance_save_option(meta, pool_idx, opt, now); - meta.clone() - }; - - let pool = clone_first_arc(&self.pools, "save_rebalance_stats: no pools available")?; - - debug!( - event = EVENT_REBALANCE_STATE, - component = LOG_COMPONENT_ECSTORE, - subsystem = LOG_SUBSYSTEM_REBALANCE, - pool_index = pool_idx, - save_opt = ?opt, - state = "metadata_save_requested", - "Rebalance metadata save requested" - ); - let stage = format!("save_rebalance_stats for pool {pool_idx} opt {opt:?}"); - resolve_rebalance_meta_save_result( - self.save_rebalance_meta_with_merge(pool, &meta_to_save, stage.as_str()).await, - stage.as_str(), - )?; - - Ok(()) - } -} +use types::{RebalanceBucketConfigs, RebalanceBucketOutcome, RebalanceEntryOutcome}; #[cfg(test)] -mod rebalance_unit_tests { - use super::REBALANCE_DEFERRED_ENTRY_ERROR_PREFIX; - use super::meta::{ - 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, - complete_rebalance_pools_at_goal, complete_rebalance_pools_with_empty_queue, defer_bucket_in_rebalance_queue, - ensure_rebalance_not_decommissioning, ensure_valid_rebalance_pool_index, first_rebalance_bucket, - has_deferred_rebalance_error, is_rebalance_actively_running, is_rebalance_conflicting_with_decommission, - is_rebalance_in_progress, is_rebalance_stopped_terminal_event, mark_rebalance_bucket_done, merge_rebalance_bucket_lists, - merge_rebalance_meta, next_rebal_bucket_from_stat, percent_free_ratio, rebalance_goal_reached, - rebalance_meta_load_no_data_error, rebalance_meta_load_unknown_format_error, rebalance_meta_load_unknown_version_error, - record_rebalance_cleanup_warning_in_meta, remove_rebalanced_buckets_from_queue, resolve_next_rebalance_bucket, - resolve_rebalance_participants, should_accept_rebalance_stats_update, should_ignore_rebalance_data_usage_cache, - should_pool_participate, should_preserve_rebalance_stopped_state, should_skip_start_rebalance, - stop_rebalance_meta_snapshot, stop_rebalance_state, take_bucket_from_rebalance_queue, validate_start_rebalance_state, - }; - use super::migration::{ - MigrationBackend, MigrationVersionResult, migrate_entry_version, migrate_entry_version_with_retry_wait, - rebalance_delete_marker_opts, - }; - use super::worker::{ - ensure_rebalance_listing_disks_available, is_transient_rebalance_error, load_rebalance_bucket_configs, - parse_rebalance_max_attempts, rebalance_listing_retry_delay, rebalance_migration_retry_delay, - resolve_load_rebalance_stats_update_result, resolve_rebalance_bucket_error, resolve_rebalance_bucket_result, - resolve_rebalance_entry_cleanup_delete_result, 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, - should_defer_rebalance_entry_failure, should_retry_rebalance_listing, should_skip_rebalance_delete_marker, - wait_rebalance_listing_retry, with_rebalance_entry_context, - }; - use super::{ - GetObjectReader, ObjectInfo, ObjectOptions, RebalSaveOpt, RebalStatus, RebalanceBucketOutcome, RebalanceCleanupWarnings, - RebalanceEntryOutcome, RebalanceInfo, RebalanceMeta, RebalanceStats, - }; - use crate::data_movement; - use crate::data_usage::DATA_USAGE_CACHE_NAME; - use crate::disk::RUSTFS_META_BUCKET; - use crate::disk::error::DiskError; - use crate::error::{Error, Result}; - use rustfs_filemeta::FileInfo; - use rustfs_filemeta::TRANSITION_COMPLETE; - use rustfs_rio::Index; - use rustfs_storage_api::HTTPRangeSpec; - use s3s::dto::ReplicationConfiguration; - use serde::Serialize; - use std::io::Cursor; - use std::sync::Arc; - use std::sync::Mutex; - use std::sync::atomic::{AtomicUsize, Ordering}; - use time::OffsetDateTime; - use tokio::sync::mpsc; - use tokio::time::Duration; - use tokio_util::sync::CancellationToken; - - #[derive(Debug, Default, Serialize)] - struct LegacyRebalanceStats { - #[serde(rename = "ifs")] - init_free_space: u64, - #[serde(rename = "ic")] - init_capacity: u64, - #[serde(rename = "bus")] - buckets: Vec, - #[serde(rename = "rbs")] - rebalanced_buckets: Vec, - #[serde(rename = "bu")] - bucket: String, - #[serde(rename = "ob")] - object: String, - #[serde(rename = "no")] - num_objects: u64, - #[serde(rename = "nv")] - num_versions: u64, - #[serde(rename = "bs")] - bytes: u64, - #[serde(rename = "par")] - participating: bool, - #[serde(rename = "inf")] - info: RebalanceInfo, - } - - #[derive(Debug, Default, Serialize)] - struct LegacyRebalanceMeta { - #[serde(rename = "stopTs")] - stopped_at: Option, - #[serde(rename = "id")] - id: String, - #[serde(rename = "pf")] - percent_free_goal: f64, - #[serde(rename = "rss")] - pool_stats: Vec, - } - - struct MigrationBackendSpy { - get_object_reader: Mutex>>, - delete_object: Mutex>>, - move_remote: Mutex>>, - get_calls: AtomicUsize, - delete_calls: AtomicUsize, - move_remote_calls: AtomicUsize, - } - - impl MigrationBackendSpy { - fn new( - get_object_reader: Option>, - delete_object: Option>, - move_remote: Option>, - ) -> Self { - Self { - get_object_reader: Mutex::new(get_object_reader), - delete_object: Mutex::new(delete_object), - move_remote: Mutex::new(move_remote), - get_calls: AtomicUsize::new(0), - delete_calls: AtomicUsize::new(0), - move_remote_calls: AtomicUsize::new(0), - } - } - - fn get_calls(&self) -> usize { - self.get_calls.load(Ordering::SeqCst) - } - - fn delete_calls(&self) -> usize { - self.delete_calls.load(Ordering::SeqCst) - } - - fn move_remote_calls(&self) -> usize { - self.move_remote_calls.load(Ordering::SeqCst) - } - - fn make_reader() -> GetObjectReader { - GetObjectReader { - stream: Box::new(Cursor::new(vec![0_u8; 3])), - object_info: ObjectInfo::default(), - } - } - } - - #[async_trait::async_trait] - impl MigrationBackend for MigrationBackendSpy { - async fn get_object_reader_for_migration( - &self, - _bucket: &str, - _object: &str, - _range: Option, - _h: http::HeaderMap, - _opts: &ObjectOptions, - ) -> Result { - self.get_calls.fetch_add(1, Ordering::SeqCst); - if let Some(result) = self.get_object_reader.lock().unwrap().take() { - return result; - } - - Ok(Self::make_reader()) - } - - async fn delete_object_for_migration(&self, _bucket: &str, _object: &str, _opts: ObjectOptions) -> Result { - self.delete_calls.fetch_add(1, Ordering::SeqCst); - if let Some(result) = self.delete_object.lock().unwrap().take() { - return result; - } - - Ok(ObjectInfo::default()) - } - - async fn move_remote_version_for_migration( - &self, - _bucket: &str, - _object: &str, - _fi: &FileInfo, - _opts: &ObjectOptions, - ) -> Result<()> { - self.move_remote_calls.fetch_add(1, Ordering::SeqCst); - if let Some(result) = self.move_remote.lock().unwrap().take() { - return result; - } - - Ok(()) - } - } - - fn version_deleted() -> FileInfo { - let mut version = FileInfo::new("object.bin", 4, 2); - version.name = "object.bin".to_string(); - version.deleted = true; - version - } - - fn version_normal() -> FileInfo { - let mut version = FileInfo::new("object.bin", 4, 2); - version.name = "object.bin".to_string(); - version.size = 64; - version - } - - fn version_remote() -> FileInfo { - let mut version = FileInfo::new("object.bin", 4, 2); - version.name = "object.bin".to_string(); - version.transition_status = TRANSITION_COMPLETE.to_string(); - version - } - - #[test] - fn test_rebalance_delete_marker_opts_preserves_replication_state() { - let mod_time = OffsetDateTime::now_utc(); - let version = FileInfo { - mod_time: Some(mod_time), - replication_state_internal: Some(rustfs_filemeta::ReplicationState { - replica_status: rustfs_filemeta::ReplicationStatusType::Replica, - delete_marker: true, - replicate_decision_str: "existing".to_string(), - ..Default::default() - }), - ..version_deleted() - }; - - let opts = rebalance_delete_marker_opts(&version, Some("version-id".to_string()), 7); - let replication = opts.delete_replication.expect("replication state should be preserved"); - - assert!(opts.versioned); - assert!(opts.data_movement); - assert!(opts.delete_marker); - assert!(opts.skip_decommissioned); - assert_eq!(opts.src_pool_idx, 7); - assert_eq!(opts.version_id.as_deref(), Some("version-id")); - assert_eq!(opts.mod_time, Some(mod_time)); - assert_eq!(replication.replica_status, rustfs_filemeta::ReplicationStatusType::Replica); - assert!(replication.delete_marker); - assert_eq!(replication.replicate_decision_str, "existing"); - } - - #[tokio::test] - async fn test_migrate_entry_version_remote_version_is_moved_without_transfer() { - let backend = MigrationBackendSpy::new(None, Some(Ok(ObjectInfo::default())), Some(Ok(()))); - let version = version_remote(); - let transfer_count = Arc::new(AtomicUsize::new(0)); - let mut transfer = { - let transfer_count = transfer_count.clone(); - move |_, _, _| { - let transfer_count = transfer_count.clone(); - async move { - transfer_count.fetch_add(1, Ordering::SeqCst); - Ok(()) - } - } - }; - - let result = migrate_entry_version( - &backend, - "bucket".to_string(), - 0, - &version, - version.version_id.map(|v| v.to_string()), - 3, - false, - &mut transfer, - ) - .await; - - assert!(!result.ignored); - assert!(!result.cleanup_ignored); - assert!(result.moved); - assert!(!result.failed); - assert!(result.error.is_none()); - assert_eq!(transfer_count.load(Ordering::SeqCst), 0); - assert_eq!(backend.move_remote_calls(), 1); - assert_eq!(backend.get_calls(), 0); - assert_eq!(backend.delete_calls(), 0); - } - - #[tokio::test] - async fn test_migrate_entry_version_remote_not_found_is_cleanup_ignored() { - let backend = MigrationBackendSpy::new( - None, - Some(Ok(ObjectInfo::default())), - Some(Err(Error::ObjectNotFound("bucket".to_string(), "object.bin".to_string()))), - ); - let version = version_remote(); - let transfer_count = Arc::new(AtomicUsize::new(0)); - let mut transfer = { - let transfer_count = transfer_count.clone(); - move |_, _, _| { - let transfer_count = transfer_count.clone(); - async move { - transfer_count.fetch_add(1, Ordering::SeqCst); - Ok(()) - } - } - }; - - let result = migrate_entry_version( - &backend, - "bucket".to_string(), - 0, - &version, - version.version_id.map(|v| v.to_string()), - 3, - false, - &mut transfer, - ) - .await; - - assert!(result.ignored); - assert!(result.cleanup_ignored); - assert!(!result.moved); - assert!(!result.failed); - assert!(result.error.is_none()); - assert_eq!(transfer_count.load(Ordering::SeqCst), 0); - assert_eq!(backend.move_remote_calls(), 1); - assert_eq!(backend.get_calls(), 0); - assert_eq!(backend.delete_calls(), 0); - } - - #[tokio::test] - async fn test_migrate_entry_version_remote_overwrite_is_not_ignored() { - let backend = MigrationBackendSpy::new( - None, - None, - Some(Err(Error::DataMovementOverwriteErr( - "bucket".to_string(), - "object.bin".to_string(), - "vid-1".to_string(), - ))), - ); - let version = version_remote(); - let mut transfer = |_, _, _| async move { Ok(()) }; - - let result = migrate_entry_version( - &backend, - "bucket".to_string(), - 0, - &version, - Some("vid-1".to_string()), - 3, - false, - &mut transfer, - ) - .await; - - assert!(result.failed); - assert!(!result.ignored); - assert!(!result.cleanup_ignored); - assert!(!result.moved); - assert_eq!(result.stage, Some("move_remote_version")); - assert!(matches!(result.error, Some(Error::DataMovementOverwriteErr(_, _, _)))); - assert_eq!(backend.move_remote_calls(), 1); - } - - #[tokio::test] - async fn test_migrate_entry_version_remote_failure_is_reported() { - let backend = MigrationBackendSpy::new(None, Some(Ok(ObjectInfo::default())), Some(Err(Error::SlowDown))); - let version = version_remote(); - let transfer_count = Arc::new(AtomicUsize::new(0)); - let mut transfer = { - let transfer_count = transfer_count.clone(); - move |_, _, _| { - let transfer_count = transfer_count.clone(); - async move { - transfer_count.fetch_add(1, Ordering::SeqCst); - Ok(()) - } - } - }; - - let result = migrate_entry_version( - &backend, - "bucket".to_string(), - 0, - &version, - version.version_id.map(|v| v.to_string()), - 3, - false, - &mut transfer, - ) - .await; - - assert!(!result.ignored); - assert!(!result.cleanup_ignored); - assert!(!result.moved); - assert!(result.failed); - assert!(matches!(result.error, Some(Error::SlowDown))); - assert_eq!(transfer_count.load(Ordering::SeqCst), 0); - assert_eq!(backend.move_remote_calls(), 1); - assert_eq!(backend.get_calls(), 0); - assert_eq!(backend.delete_calls(), 0); - } - - #[tokio::test] - async fn test_migrate_entry_version_deleted_version_calls_delete_and_moved() { - let backend = MigrationBackendSpy::new(None, Some(Ok(ObjectInfo::default())), None); - let version = version_deleted(); - let mut transfer = |_, _, _| async move { Ok(()) }; - - let result = migrate_entry_version( - &backend, - "bucket".to_string(), - 1, - &version, - version.version_id.map(|v| v.to_string()), - 3, - false, - &mut transfer, - ) - .await; - - assert!(!result.ignored); - assert!(!result.cleanup_ignored); - assert!(result.moved); - assert!(!result.failed); - assert!(result.error.is_none()); - assert_eq!(backend.get_calls(), 0); - assert_eq!(backend.delete_calls(), 1); - } - - #[tokio::test] - async fn test_migrate_entry_version_deleted_version_not_found_is_ignored() { - let backend = MigrationBackendSpy::new( - None, - Some(Err(Error::ObjectNotFound("bucket".to_string(), "object.bin".to_string()))), - None, - ); - let version = version_deleted(); - let mut transfer = |_, _, _| async move { Ok(()) }; - - let result = migrate_entry_version( - &backend, - "bucket".to_string(), - 1, - &version, - version.version_id.map(|v| v.to_string()), - 3, - false, - &mut transfer, - ) - .await; - - assert!(result.ignored); - assert!(result.cleanup_ignored); - assert!(!result.moved); - assert!(!result.failed); - assert!(result.error.is_none()); - assert_eq!(backend.delete_calls(), 1); - } - - #[tokio::test] - async fn test_migrate_entry_version_deleted_version_overwrite_is_not_ignored() { - let backend = MigrationBackendSpy::new( - None, - Some(Err(Error::DataMovementOverwriteErr( - "bucket".to_string(), - "object.bin".to_string(), - "vid-1".to_string(), - ))), - None, - ); - let version = version_deleted(); - let mut transfer = |_, _, _| async move { Ok(()) }; - - let result = migrate_entry_version( - &backend, - "bucket".to_string(), - 1, - &version, - Some("vid-1".to_string()), - 3, - false, - &mut transfer, - ) - .await; - - assert!(result.failed); - assert!(!result.ignored); - assert!(!result.cleanup_ignored); - assert!(!result.moved); - assert_eq!(result.stage, Some("delete_marker")); - assert!(matches!(result.error, Some(Error::DataMovementOverwriteErr(_, _, _)))); - assert_eq!(backend.delete_calls(), 1); - } - - #[tokio::test] - async fn test_migrate_entry_version_reader_not_found_is_ignored() { - let backend = MigrationBackendSpy::new( - Some(Err(Error::ObjectNotFound("bucket".to_string(), "object.bin".to_string()))), - None, - None, - ); - let version = version_normal(); - let mut transfer = |_, _, _| async move { Ok(()) }; - - let result = migrate_entry_version( - &backend, - "bucket".to_string(), - 1, - &version, - version.version_id.map(|v| v.to_string()), - 3, - false, - &mut transfer, - ) - .await; - - assert!(result.ignored); - assert!(result.cleanup_ignored); - assert!(!result.moved); - assert!(!result.failed); - assert!(result.error.is_none()); - assert_eq!(backend.get_calls(), 1); - assert_eq!(backend.delete_calls(), 0); - } - - #[tokio::test] - async fn test_migrate_entry_version_reader_retries_before_success() { - let backend = MigrationBackendSpy::new(Some(Err(Error::SlowDown)), None, None); - let transfer_count = Arc::new(AtomicUsize::new(0)); - let wait_count = Arc::new(AtomicUsize::new(0)); - let mut transfer = { - let transfer_count = transfer_count.clone(); - move |_, _, _| { - let transfer_count = transfer_count.clone(); - async move { - transfer_count.fetch_add(1, Ordering::SeqCst); - Ok(()) - } - } - }; - - let version = version_normal(); - let result = migrate_entry_version_with_retry_wait( - &backend, - "bucket".to_string(), - 1, - &version, - version.version_id.map(|v| v.to_string()), - 3, - false, - &mut transfer, - { - let wait_count = wait_count.clone(); - move |_| { - let wait_count = wait_count.clone(); - async move { - wait_count.fetch_add(1, Ordering::SeqCst); - } - } - }, - ) - .await; - - assert!(!result.ignored); - assert!(!result.cleanup_ignored); - assert!(result.moved); - assert!(!result.failed); - assert!(result.error.is_none()); - assert_eq!(backend.get_calls(), 2); - assert_eq!(backend.delete_calls(), 0); - assert_eq!(transfer_count.load(Ordering::SeqCst), 1); - assert_eq!(wait_count.load(Ordering::SeqCst), 1); - } - - struct AlwaysFailGetBackend { - get_calls: AtomicUsize, - } - - impl AlwaysFailGetBackend { - fn new() -> Self { - Self { - get_calls: AtomicUsize::new(0), - } - } - - fn get_calls(&self) -> usize { - self.get_calls.load(Ordering::SeqCst) - } - } - - #[async_trait::async_trait] - impl MigrationBackend for AlwaysFailGetBackend { - async fn get_object_reader_for_migration( - &self, - _bucket: &str, - _object: &str, - _range: Option, - _h: http::HeaderMap, - _opts: &ObjectOptions, - ) -> Result { - self.get_calls.fetch_add(1, Ordering::SeqCst); - Err(Error::SlowDown) - } - - async fn delete_object_for_migration(&self, _bucket: &str, _object: &str, _opts: ObjectOptions) -> Result { - Ok(ObjectInfo::default()) - } - - async fn move_remote_version_for_migration( - &self, - _bucket: &str, - _object: &str, - _fi: &FileInfo, - _opts: &ObjectOptions, - ) -> Result<()> { - Ok(()) - } - } - - #[tokio::test] - async fn test_migrate_entry_version_reader_fails_after_retries() { - let backend = AlwaysFailGetBackend::new(); - let transfer_count = Arc::new(AtomicUsize::new(0)); - let mut transfer = { - let transfer_count = transfer_count.clone(); - move |_, _, _| { - let transfer_count = transfer_count.clone(); - async move { - transfer_count.fetch_add(1, Ordering::SeqCst); - Ok(()) - } - } - }; - - let version = version_normal(); - let result = migrate_entry_version( - &backend, - "bucket".to_string(), - 1, - &version, - version.version_id.map(|v| v.to_string()), - 3, - false, - &mut transfer, - ) - .await; - - assert!(!result.ignored); - assert!(!result.cleanup_ignored); - assert!(!result.moved); - assert!(result.failed); - assert!(matches!(result.error, Some(Error::SlowDown))); - assert_eq!(backend.get_calls(), 3); - assert_eq!(transfer_count.load(Ordering::SeqCst), 0); - } - - #[tokio::test] - async fn test_migrate_entry_version_zero_max_attempts_still_attempts_once() { - let backend = AlwaysFailGetBackend::new(); - let transfer_count = Arc::new(AtomicUsize::new(0)); - let mut transfer = { - let transfer_count = transfer_count.clone(); - move |_, _, _| { - let transfer_count = transfer_count.clone(); - async move { - transfer_count.fetch_add(1, Ordering::SeqCst); - Ok(()) - } - } - }; - - let version = version_normal(); - let result = migrate_entry_version( - &backend, - "bucket".to_string(), - 1, - &version, - version.version_id.map(|v| v.to_string()), - 0, - false, - &mut transfer, - ) - .await; - - assert!(!result.ignored); - assert!(!result.cleanup_ignored); - assert!(!result.moved); - assert!(result.failed); - assert!(matches!(result.error, Some(Error::SlowDown))); - assert_eq!(backend.get_calls(), 1); - assert_eq!(transfer_count.load(Ordering::SeqCst), 0); - } - - #[tokio::test] - async fn test_migrate_entry_version_transfer_retries_before_success() { - let backend = MigrationBackendSpy::new(Some(Ok(MigrationBackendSpy::make_reader())), None, None); - let transfer_count = Arc::new(AtomicUsize::new(0)); - let wait_count = Arc::new(AtomicUsize::new(0)); - let mut transfer = { - let transfer_count = transfer_count.clone(); - move |_, _, _| { - let transfer_count = transfer_count.clone(); - async move { - let attempt = transfer_count.fetch_add(1, Ordering::SeqCst); - if attempt == 0 { - return Err(Error::SlowDown); - } - Ok(()) - } - } - }; - - let version = version_normal(); - let result = migrate_entry_version_with_retry_wait( - &backend, - "bucket".to_string(), - 1, - &version, - version.version_id.map(|v| v.to_string()), - 3, - false, - &mut transfer, - { - let wait_count = wait_count.clone(); - move |_| { - let wait_count = wait_count.clone(); - async move { - wait_count.fetch_add(1, Ordering::SeqCst); - } - } - }, - ) - .await; - - assert!(!result.ignored); - assert!(!result.cleanup_ignored); - assert!(result.moved); - assert!(!result.failed); - assert_eq!(backend.get_calls(), 2); - assert_eq!(transfer_count.load(Ordering::SeqCst), 2); - assert_eq!(wait_count.load(Ordering::SeqCst), 1); - } - - #[tokio::test] - async fn test_migrate_entry_version_transfer_non_transient_fails_without_retry() { - let backend = MigrationBackendSpy::new(Some(Ok(MigrationBackendSpy::make_reader())), None, None); - let transfer_count = Arc::new(AtomicUsize::new(0)); - let wait_count = Arc::new(AtomicUsize::new(0)); - let mut transfer = { - let transfer_count = transfer_count.clone(); - move |_, _, _| { - let transfer_count = transfer_count.clone(); - async move { - transfer_count.fetch_add(1, Ordering::SeqCst); - Err(Error::FileAccessDenied) - } - } - }; - - let version = version_normal(); - let result = migrate_entry_version_with_retry_wait( - &backend, - "bucket".to_string(), - 1, - &version, - version.version_id.map(|v| v.to_string()), - 3, - false, - &mut transfer, - { - let wait_count = wait_count.clone(); - move |_| { - let wait_count = wait_count.clone(); - async move { - wait_count.fetch_add(1, Ordering::SeqCst); - } - } - }, - ) - .await; - - assert!(result.failed); - assert!(!result.ignored); - assert_eq!(result.stage, Some("write_target")); - assert!(matches!(result.error, Some(Error::FileAccessDenied))); - assert_eq!(backend.get_calls(), 1); - assert_eq!(transfer_count.load(Ordering::SeqCst), 1); - assert_eq!(wait_count.load(Ordering::SeqCst), 0); - } - - #[tokio::test] - async fn test_migrate_entry_version_transfer_fails_after_retries() { - let backend = MigrationBackendSpy::new(Some(Ok(MigrationBackendSpy::make_reader())), None, None); - let transfer_count = Arc::new(AtomicUsize::new(0)); - let mut transfer = { - let transfer_count = transfer_count.clone(); - move |_, _, _| { - let transfer_count = transfer_count.clone(); - async move { - transfer_count.fetch_add(1, Ordering::SeqCst); - Err(Error::SlowDown) - } - } - }; - - let version = version_normal(); - let result = migrate_entry_version( - &backend, - "bucket".to_string(), - 1, - &version, - version.version_id.map(|v| v.to_string()), - 2, - false, - &mut transfer, - ) - .await; - - assert!(result.failed); - assert!(!result.ignored); - assert!(!result.cleanup_ignored); - assert!(!result.moved); - assert!(result.error.is_some()); - assert_eq!(backend.get_calls(), 2); - assert_eq!(transfer_count.load(Ordering::SeqCst), 2); - } - - #[tokio::test] - async fn test_migrate_entry_version_transfer_not_found_is_ignored() { - let backend = MigrationBackendSpy::new(Some(Ok(MigrationBackendSpy::make_reader())), None, None); - let transfer_count = Arc::new(AtomicUsize::new(0)); - let mut transfer = { - let transfer_count = transfer_count.clone(); - move |_, _, _| { - let transfer_count = transfer_count.clone(); - async move { - transfer_count.fetch_add(1, Ordering::SeqCst); - Err(Error::ObjectNotFound("bucket".to_string(), "object.bin".to_string())) - } - } - }; - - let version = version_normal(); - let result = migrate_entry_version( - &backend, - "bucket".to_string(), - 1, - &version, - version.version_id.map(|v| v.to_string()), - 3, - false, - &mut transfer, - ) - .await; - - assert!(result.ignored); - assert!(result.cleanup_ignored); - assert!(!result.moved); - assert!(!result.failed); - assert!(result.error.is_none()); - assert_eq!(transfer_count.load(Ordering::SeqCst), 1); - } - - #[tokio::test] - async fn test_migrate_entry_version_transfer_overwrite_is_not_ignored() { - let backend = MigrationBackendSpy::new(Some(Ok(MigrationBackendSpy::make_reader())), None, None); - let transfer_count = Arc::new(AtomicUsize::new(0)); - let mut transfer = { - let transfer_count = transfer_count.clone(); - move |_, _, _| { - let transfer_count = transfer_count.clone(); - async move { - transfer_count.fetch_add(1, Ordering::SeqCst); - Err(Error::DataMovementOverwriteErr( - "bucket".to_string(), - "object.bin".to_string(), - "vid-1".to_string(), - )) - } - } - }; - - let version = version_normal(); - let result = migrate_entry_version( - &backend, - "bucket".to_string(), - 1, - &version, - Some("vid-1".to_string()), - 3, - false, - &mut transfer, - ) - .await; - - assert!(result.failed); - assert!(!result.ignored); - assert!(!result.cleanup_ignored); - assert!(!result.moved); - assert_eq!(result.stage, Some("write_target")); - assert!(matches!(result.error, Some(Error::DataMovementOverwriteErr(_, _, _)))); - assert_eq!(transfer_count.load(Ordering::SeqCst), 1); - } - - #[tokio::test] - async fn test_migrate_entry_version_ignores_data_usage_cache_when_enabled() { - let backend = MigrationBackendSpy::new(Some(Ok(MigrationBackendSpy::make_reader())), None, None); - let version = { - let mut version = version_normal(); - version.name = format!("{}.{}", DATA_USAGE_CACHE_NAME, version.name); - version - }; - let transfer_count = Arc::new(AtomicUsize::new(0)); - let mut transfer = { - let transfer_count = transfer_count.clone(); - move |_, _, _| { - let transfer_count = transfer_count.clone(); - async move { - transfer_count.fetch_add(1, Ordering::SeqCst); - Ok(()) - } - } - }; - - let result = migrate_entry_version( - &backend, - RUSTFS_META_BUCKET.to_string(), - 1, - &version, - version.version_id.map(|v| v.to_string()), - 2, - true, - &mut transfer, - ) - .await; - - assert!(result.ignored); - assert!(!result.cleanup_ignored); - assert!(!result.moved); - assert!(!result.failed); - assert!(result.error.is_none()); - assert_eq!(transfer_count.load(Ordering::SeqCst), 0); - assert_eq!(backend.get_calls(), 0); - assert_eq!(backend.delete_calls(), 0); - } - - #[tokio::test] - async fn test_migrate_entry_version_data_usage_cache_moves_when_ignore_disabled() { - let backend = MigrationBackendSpy::new(Some(Ok(MigrationBackendSpy::make_reader())), None, None); - let version = { - let mut version = version_normal(); - version.name = format!("{}.{}", DATA_USAGE_CACHE_NAME, version.name); - version - }; - let transfer_count = Arc::new(AtomicUsize::new(0)); - let mut transfer = { - let transfer_count = transfer_count.clone(); - move |_, _, _| { - let transfer_count = transfer_count.clone(); - async move { - transfer_count.fetch_add(1, Ordering::SeqCst); - Ok(()) - } - } - }; - - let result = migrate_entry_version( - &backend, - RUSTFS_META_BUCKET.to_string(), - 1, - &version, - version.version_id.map(|v| v.to_string()), - 2, - false, - &mut transfer, - ) - .await; - - assert!(!result.ignored); - assert!(!result.cleanup_ignored); - assert!(result.moved); - assert!(!result.failed); - assert!(result.error.is_none()); - assert_eq!(transfer_count.load(Ordering::SeqCst), 1); - assert_eq!(backend.get_calls(), 1); - assert_eq!(backend.delete_calls(), 0); - } - - #[test] - fn test_should_ignore_rebalance_data_usage_cache_true_for_meta_bucket() { - assert!(should_ignore_rebalance_data_usage_cache(RUSTFS_META_BUCKET)); - } - - #[test] - fn test_should_ignore_rebalance_data_usage_cache_false_for_regular_bucket() { - assert!(!should_ignore_rebalance_data_usage_cache("bucket-a")); - } - - #[test] - fn test_rebalance_goal_reached_at_target() { - let init_free_space = 200_u64; - let init_capacity = 1_000_u64; - let goal = 0.45_f64; - - assert!(rebalance_goal_reached(init_free_space, init_capacity, 250, goal)); - assert!(!rebalance_goal_reached(init_free_space, init_capacity, 249, goal)); - } - - #[test] - fn test_rebalance_goal_reached_above_target() { - let init_free_space = 200_u64; - let init_capacity = 1_000_u64; - let bytes = 251_u64; - let goal = 0.45_f64; - - assert!(rebalance_goal_reached(init_free_space, init_capacity, bytes, goal)); - } - - #[test] - fn test_rebalance_goal_not_reached_outside_tolerance() { - let init_free_space = 100_u64; - let init_capacity = 1_000_u64; - let bytes = 80_u64; - let goal = 0.5_f64; - - assert!(!rebalance_goal_reached(init_free_space, init_capacity, bytes, goal)); - } - - #[test] - fn test_rebalance_goal_zero_capacity_is_false() { - assert!(!rebalance_goal_reached(100, 0, 50, 0.5)); - } - - #[test] - fn test_rebalance_goal_above_one_is_true_when_reached() { - assert!(rebalance_goal_reached(950, 1_000, 100, 1.0)); - } - - #[test] - fn test_rebalance_goal_below_zero_is_true_when_reached() { - assert!(rebalance_goal_reached(10, 1_000, 0, -0.01)); - } - - #[test] - fn test_resolve_rebalance_worker_result_passthrough() { - assert!(resolve_rebalance_worker_result(0, Ok(Ok(()))).is_ok()); - - let err = resolve_rebalance_worker_result::<()>(0, Ok(Err(Error::OperationCanceled))).unwrap_err(); - assert!(matches!(err, Error::OperationCanceled)); - } - - #[tokio::test] - async fn test_resolve_rebalance_worker_result_join_error_keeps_context() { - let join_error = tokio::spawn(async { - panic!("rebalance worker panic"); - }) - .await - .expect_err("panic task should return JoinError"); - - let err = resolve_rebalance_worker_result::<()>(7, Err(join_error)).unwrap_err(); - assert!(err.to_string().contains("rebalance worker 7 task join error")); - } - - #[tokio::test] - async fn test_wait_rebalance_entry_tasks_returns_ok_for_successful_tasks() { - let tasks = Arc::new(tokio::sync::Mutex::new(vec![tokio::spawn(async { - Ok(RebalanceEntryOutcome::Completed) - })])); - - let result = super::wait_rebalance_entry_tasks(1, tasks) - .await - .expect("successful entry tasks should pass"); - - assert!(result.is_none()); - } - - #[tokio::test] - async fn test_wait_rebalance_entry_tasks_returns_first_task_error() { - let tasks = Arc::new(tokio::sync::Mutex::new(vec![ - tokio::spawn(async { Ok(RebalanceEntryOutcome::Completed) }), - tokio::spawn(async { Err(Error::other("entry failed")) }), - ])); - - let err = super::wait_rebalance_entry_tasks(1, tasks) - .await - .expect_err("entry task failure should be returned"); - - assert!(err.to_string().contains("entry failed")); - } - - #[tokio::test] - async fn test_wait_rebalance_entry_tasks_returns_deferred_error_without_failing() { - let tasks = Arc::new(tokio::sync::Mutex::new(vec![ - tokio::spawn(async { Ok(RebalanceEntryOutcome::Completed) }), - tokio::spawn(async { - Ok(RebalanceEntryOutcome::Deferred { - last_error: "deferred transient rebalance entry failure: timeout".to_string(), - }) - }), - ])); - - let result = super::wait_rebalance_entry_tasks(1, tasks) - .await - .expect("deferred transient entry should not fail worker"); - - assert_eq!(result.as_deref(), Some("deferred transient rebalance entry failure: timeout")); - } - - #[test] - fn test_resolve_rebalance_save_task_result_passthrough() { - assert!(resolve_rebalance_save_task_result(0, Ok(Ok(()))).is_ok()); - } - - #[test] - fn test_resolve_rebalance_save_task_result_wraps_inner_error_context() { - let err = resolve_rebalance_save_task_result(1, Ok(Err(Error::SlowDown))) - .expect_err("inner save-task error should include pool context"); - assert!(err.to_string().contains("rebalance save_task failed for pool 1")); - } - - #[test] - fn test_resolve_rebalance_meta_save_result_passthrough() { - assert!(resolve_rebalance_meta_save_result(Ok(()), "stop_rebalance").is_ok()); - } - - #[test] - fn test_resolve_rebalance_meta_save_result_wraps_error_context() { - let err = resolve_rebalance_meta_save_result(Err(Error::SlowDown), "init_rebalance_meta") - .expect_err("meta save failure should include stage context"); - let message = err.to_string(); - assert!(message.contains("rebalance meta save failed during init_rebalance_meta")); - assert!(message.contains(Error::SlowDown.to_string().as_str())); - } - - #[test] - fn test_merge_rebalance_meta_preserves_updates_from_multiple_pools() { - let start_time = OffsetDateTime::from_unix_timestamp(1_000).unwrap(); - let warning_at = OffsetDateTime::from_unix_timestamp(1_500).unwrap(); - let mut remote = RebalanceMeta { - id: "rebal-1".to_string(), - percent_free_goal: 0.5, - pool_stats: vec![ - RebalanceStats { - buckets: vec!["bucket-a".to_string()], - participating: true, - info: RebalanceInfo { - start_time: Some(start_time), - status: RebalStatus::Started, - ..Default::default() - }, - num_versions: 4, - bytes: 400, - bucket: "bucket-a".to_string(), - object: "remote-object".to_string(), - ..Default::default() - }, - RebalanceStats { - buckets: vec!["bucket-a".to_string()], - participating: true, - info: RebalanceInfo { - start_time: Some(start_time), - status: RebalStatus::Started, - ..Default::default() - }, - ..Default::default() - }, - ], - ..Default::default() - }; - let local = RebalanceMeta { - id: "rebal-1".to_string(), - percent_free_goal: 0.5, - pool_stats: vec![ - RebalanceStats { - buckets: vec!["bucket-a".to_string()], - participating: true, - info: RebalanceInfo { - start_time: Some(start_time), - status: RebalStatus::Started, - ..Default::default() - }, - ..Default::default() - }, - RebalanceStats { - buckets: Vec::new(), - rebalanced_buckets: vec!["bucket-a".to_string()], - participating: true, - info: RebalanceInfo { - start_time: Some(start_time), - status: RebalStatus::Started, - ..Default::default() - }, - num_versions: 7, - bytes: 700, - bucket: "bucket-a".to_string(), - object: "local-object".to_string(), - cleanup_warnings: RebalanceCleanupWarnings { - count: 1, - last_message: Some("cleanup failed".to_string()), - last_bucket: Some("bucket-a".to_string()), - last_object: Some("local-object".to_string()), - last_at: Some(warning_at), - }, - ..Default::default() - }, - ], - ..Default::default() - }; - - merge_rebalance_meta(&mut remote, &local); - - assert_eq!(remote.pool_stats[0].num_versions, 4); - assert_eq!(remote.pool_stats[0].object, "remote-object"); - assert_eq!(remote.pool_stats[1].num_versions, 7); - assert_eq!(remote.pool_stats[1].object, "local-object"); - assert!(remote.pool_stats[1].buckets.is_empty()); - assert_eq!(remote.pool_stats[1].rebalanced_buckets, vec!["bucket-a"]); - assert_eq!(remote.pool_stats[1].cleanup_warnings.count, 1); - assert_eq!(remote.pool_stats[1].cleanup_warnings.last_message.as_deref(), Some("cleanup failed")); - assert_eq!(remote.pool_stats[1].cleanup_warnings.last_at, Some(warning_at)); - } - - #[test] - fn test_merge_rebalance_meta_does_not_overwrite_failed_with_started_stats() { - let now = OffsetDateTime::from_unix_timestamp(2_000).unwrap(); - let mut remote = RebalanceMeta { - id: "rebal-1".to_string(), - pool_stats: vec![RebalanceStats { - info: RebalanceInfo { - status: RebalStatus::Failed, - end_time: Some(now), - last_error: Some("timeout".to_string()), - ..Default::default() - }, - num_versions: 10, - bytes: 1_000, - ..Default::default() - }], - ..Default::default() - }; - let local = RebalanceMeta { - id: "rebal-1".to_string(), - pool_stats: vec![RebalanceStats { - participating: true, - info: RebalanceInfo { - status: RebalStatus::Started, - ..Default::default() - }, - num_versions: 8, - bytes: 800, - ..Default::default() - }], - ..Default::default() - }; - - merge_rebalance_meta(&mut remote, &local); - - assert_eq!(remote.pool_stats[0].info.status, RebalStatus::Failed); - assert_eq!(remote.pool_stats[0].info.last_error.as_deref(), Some("timeout")); - assert_eq!(remote.pool_stats[0].num_versions, 10); - assert_eq!(remote.pool_stats[0].bytes, 1_000); - } - - #[test] - fn test_merge_rebalance_meta_does_not_overwrite_stopped_with_started_stats() { - let stopped_at = OffsetDateTime::from_unix_timestamp(3_000).unwrap(); - let mut remote = RebalanceMeta { - id: "rebal-1".to_string(), - stopped_at: Some(stopped_at), - pool_stats: vec![RebalanceStats { - info: RebalanceInfo { - status: RebalStatus::Stopped, - end_time: Some(stopped_at), - ..Default::default() - }, - num_versions: 5, - ..Default::default() - }], - ..Default::default() - }; - let local = RebalanceMeta { - id: "rebal-1".to_string(), - pool_stats: vec![RebalanceStats { - info: RebalanceInfo { - status: RebalStatus::Started, - ..Default::default() - }, - num_versions: 9, - ..Default::default() - }], - ..Default::default() - }; - - merge_rebalance_meta(&mut remote, &local); - - assert_eq!(remote.stopped_at, Some(stopped_at)); - assert_eq!(remote.pool_stats[0].info.status, RebalStatus::Stopped); - assert_eq!(remote.pool_stats[0].info.end_time, Some(stopped_at)); - assert_eq!(remote.pool_stats[0].num_versions, 9); - } - - #[test] - fn test_merge_rebalance_meta_preserves_failed_status_over_stopped() { - let stopped_at = OffsetDateTime::from_unix_timestamp(3_000).unwrap(); - let failed_at = OffsetDateTime::from_unix_timestamp(4_000).unwrap(); - let mut remote = RebalanceMeta { - id: "rebal-1".to_string(), - stopped_at: Some(stopped_at), - pool_stats: vec![RebalanceStats { - info: RebalanceInfo { - status: RebalStatus::Stopped, - end_time: Some(stopped_at), - ..Default::default() - }, - ..Default::default() - }], - ..Default::default() - }; - let local = RebalanceMeta { - id: "rebal-1".to_string(), - pool_stats: vec![RebalanceStats { - info: RebalanceInfo { - status: RebalStatus::Failed, - end_time: Some(failed_at), - last_error: Some("late failure".to_string()), - ..Default::default() - }, - ..Default::default() - }], - ..Default::default() - }; - - merge_rebalance_meta(&mut remote, &local); - - assert_eq!(remote.pool_stats[0].info.status, RebalStatus::Failed); - assert_eq!(remote.pool_stats[0].info.end_time, Some(failed_at)); - assert_eq!(remote.pool_stats[0].info.last_error.as_deref(), Some("late failure")); - } - - #[test] - fn test_merge_rebalance_meta_does_not_overwrite_stopped_with_completed_status() { - let stopped_at = OffsetDateTime::from_unix_timestamp(3_000).unwrap(); - let completed_at = OffsetDateTime::from_unix_timestamp(5_000).unwrap(); - let mut remote = RebalanceMeta { - id: "rebal-1".to_string(), - stopped_at: Some(stopped_at), - pool_stats: vec![RebalanceStats { - info: RebalanceInfo { - status: RebalStatus::Stopped, - end_time: Some(stopped_at), - ..Default::default() - }, - num_versions: 5, - ..Default::default() - }], - ..Default::default() - }; - let local = RebalanceMeta { - id: "rebal-1".to_string(), - pool_stats: vec![RebalanceStats { - info: RebalanceInfo { - status: RebalStatus::Completed, - end_time: Some(completed_at), - ..Default::default() - }, - num_versions: 12, - ..Default::default() - }], - ..Default::default() - }; - - merge_rebalance_meta(&mut remote, &local); - - assert_eq!(remote.stopped_at, Some(stopped_at)); - assert_eq!(remote.pool_stats[0].info.status, RebalStatus::Stopped); - assert_eq!(remote.pool_stats[0].info.end_time, Some(stopped_at)); - assert_eq!(remote.pool_stats[0].num_versions, 12); - } - - #[test] - fn test_rebalance_meta_load_no_data_error_formats_context() { - let err = rebalance_meta_load_no_data_error(); - let rendered = err.to_string(); - - assert!(rendered.contains("rebalance metadata load failed"), "{rendered}"); - assert!(rendered.contains("payload is too short"), "{rendered}"); - } - - #[test] - fn test_rebalance_meta_load_unknown_format_error_formats_context() { - let err = rebalance_meta_load_unknown_format_error(9); - let rendered = err.to_string(); - - assert!(rendered.contains("rebalance metadata load failed"), "{rendered}"); - assert!(rendered.contains("unknown format 9"), "{rendered}"); - } - - #[test] - fn test_rebalance_meta_load_unknown_version_error_formats_context() { - let err = rebalance_meta_load_unknown_version_error(3); - let rendered = err.to_string(); - - assert!(rendered.contains("rebalance metadata load failed"), "{rendered}"); - assert!(rendered.contains("unknown version 3"), "{rendered}"); - } - - #[test] - fn test_rebalance_meta_deserializes_legacy_stats_without_cleanup_warnings() { - let legacy = LegacyRebalanceMeta { - id: "rebal-legacy".to_string(), - percent_free_goal: 0.35, - pool_stats: vec![LegacyRebalanceStats { - buckets: vec!["bucket-a".to_string()], - participating: true, - info: RebalanceInfo { - status: RebalStatus::Started, - ..Default::default() - }, - ..Default::default() - }], - ..Default::default() - }; - let data = rmp_serde::to_vec(&legacy).expect("legacy rebalance metadata should serialize"); - - let decoded: RebalanceMeta = - rmp_serde::from_slice(data.as_slice()).expect("legacy rebalance metadata should deserialize"); - - assert_eq!(decoded.id, "rebal-legacy"); - assert_eq!(decoded.pool_stats.len(), 1); - assert_eq!(decoded.pool_stats[0].cleanup_warnings, RebalanceCleanupWarnings::default()); - assert_eq!(decoded.pool_stats[0].info.status, RebalStatus::Started); - } - - #[test] - fn test_resolve_rebalance_stats_update_result_passthrough() { - assert!(resolve_rebalance_stats_update_result(Ok(()), 0, "bucket", "object").is_ok()); - } - - #[test] - fn test_resolve_rebalance_stats_update_result_wraps_error_context() { - let err = resolve_rebalance_stats_update_result(Err(Error::SlowDown), 2, "bucket-a", "obj.txt") - .expect_err("stats update error should include context"); - assert!( - err.to_string() - .contains("rebalance stats update failed for pool 2 bucket bucket-a object obj.txt") - ); - } - - #[test] - fn test_resolve_load_rebalance_stats_update_result_passthrough() { - assert!(resolve_load_rebalance_stats_update_result(Ok(())).is_ok()); - } - - #[test] - fn test_resolve_load_rebalance_stats_update_result_wraps_error_context() { - let err = resolve_load_rebalance_stats_update_result(Err(Error::SlowDown)) - .expect_err("load-time stats refresh failure should include context"); - assert!(err.to_string().contains("rebalance metadata stats refresh failed after load")); - } - - #[test] - fn test_resolve_rebalance_file_info_versions_result_passthrough() { - let value = resolve_rebalance_file_info_versions_result::(Ok(7), "bucket-a", "obj.txt") - .expect("ok results should pass through"); - assert_eq!(value, 7); - } - - #[test] - fn test_resolve_rebalance_file_info_versions_result_wraps_error_context() { - let err = resolve_rebalance_file_info_versions_result::(Err(Error::SlowDown), "bucket-a", "obj.txt") - .expect_err("errors should be wrapped"); - let message = err.to_string(); - assert!(message.contains("rebalance file_info_versions failed for bucket-a/obj.txt")); - } - - #[test] - fn test_resolve_rebalance_meta_load_result_returns_true_for_loaded_meta() { - assert!(resolve_rebalance_meta_load_result(Ok(())).expect("loaded rebalance metadata should pass through")); - } - - #[test] - fn test_resolve_rebalance_meta_load_result_returns_false_for_missing_meta() { - assert!( - !resolve_rebalance_meta_load_result(Err(Error::ConfigNotFound)) - .expect("missing rebalance metadata should be treated as not started") - ); - } - - #[test] - fn test_resolve_rebalance_meta_load_result_wraps_error_context() { - let err = - resolve_rebalance_meta_load_result(Err(Error::SlowDown)).expect_err("unexpected load failures should be wrapped"); - let message = err.to_string(); - assert!(message.contains("rebalance metadata load failed during load_rebalance_meta")); - } - - #[test] - fn test_resolve_rebalance_entry_cleanup_delete_result_passthrough() { - let result = resolve_rebalance_entry_cleanup_delete_result(Ok(ObjectInfo::default()), "bucket-a", "obj.txt"); - assert_eq!(result.expect("successful cleanup should pass through"), None); - } - - #[test] - fn test_resolve_rebalance_entry_cleanup_delete_result_ignores_not_found() { - let result = resolve_rebalance_entry_cleanup_delete_result( - Err(Error::ObjectNotFound("bucket-a".to_string(), "obj.txt".to_string())), - "bucket-a", - "obj.txt", - ); - assert_eq!(result.expect("missing cleanup source should be ignored"), None); - } - - #[test] - fn test_resolve_rebalance_entry_cleanup_delete_result_returns_warning_for_failures() { - let warning = resolve_rebalance_entry_cleanup_delete_result(Err(Error::SlowDown), "bucket-a", "obj.txt") - .expect("cleanup delete failures should be downgraded to warnings") - .expect("cleanup delete failure should return warning"); - let message = warning.as_str(); - assert!(message.contains("rebalance cleanup delete failed for bucket-a/obj.txt")); - } - - #[test] - fn test_resolve_rebalance_migrate_result_error_preserves_inner_error() { - let err = resolve_rebalance_migrate_result_error(Some(Error::SlowDown), 2, "bucket-a", "obj.txt", Some("vid-1")); - assert!(matches!(err, Error::SlowDown)); - } - - #[test] - fn test_resolve_rebalance_migrate_result_error_wraps_missing_error_context() { - let err = resolve_rebalance_migrate_result_error(None, 2, "bucket-a", "obj.txt", Some("vid-1")); - let message = err.to_string(); - assert!( - message - .contains("rebalance migration reported failure without error for pool 2 entry bucket-a/obj.txt version vid-1"), - "{message}" - ); - } - - #[test] - fn test_resolve_rebalance_bucket_result_passthrough() { - let outcome = resolve_rebalance_bucket_result(Ok(RebalanceBucketOutcome::Completed), 2, "bucket-a") - .expect("completed bucket should pass through"); - assert_eq!(outcome, RebalanceBucketOutcome::Completed); - } - - #[test] - fn test_resolve_rebalance_bucket_result_preserves_operation_canceled() { - let err = resolve_rebalance_bucket_result(Err(Error::OperationCanceled), 2, "bucket-a") - .expect_err("operation canceled should be preserved"); - assert!(matches!(err, Error::OperationCanceled)); - } - - #[test] - fn test_resolve_rebalance_bucket_result_wraps_not_initialized_with_context() { - let err = resolve_rebalance_bucket_result(Err(Error::other("errServerNotInitialized")), 2, "bucket-a") - .expect_err("not initialized should be surfaced with context"); - let message = err.to_string(); - assert!(message.contains("rebalance bucket bucket-a failed for pool 2")); - assert!(message.contains("errServerNotInitialized")); - } - - #[test] - fn test_rebalance_listing_disks_available_rejects_empty_set() { - let err = ensure_rebalance_listing_disks_available(false, "bucket-a") - .expect_err("missing online disks should be reported with bucket context"); - assert!( - err.to_string() - .contains("failed to list objects to rebalance for bucket bucket-a: no disks available") - ); - } - - #[test] - fn test_rebalance_listing_disks_available_allows_online_disks() { - assert!(ensure_rebalance_listing_disks_available(true, "bucket-a").is_ok()); - } - - #[test] - fn test_is_transient_rebalance_error_classifies_retryable_errors() { - assert!(is_transient_rebalance_error(&Error::SlowDown)); - assert!(is_transient_rebalance_error(&Error::ErasureReadQuorum)); - assert!(is_transient_rebalance_error(&Error::ErasureWriteQuorum)); - assert!(is_transient_rebalance_error(&Error::Io(std::io::Error::new( - std::io::ErrorKind::TimedOut, - "timed out", - )))); - assert!(is_transient_rebalance_error(&Error::other("i/o timeout"))); - } - - #[test] - fn test_is_transient_rebalance_error_accepts_lock_and_rpc_timeouts() { - assert!(is_transient_rebalance_error(&Error::Io(std::io::Error::other( - "Failed to acquire read lock: ns_loc: read lock acquisition timed out on bucket/object" - )))); - assert!(is_transient_rebalance_error(&Error::other( - "Remote lock RPC timed out: RPC timed out after 50ms" - ))); - assert!(is_transient_rebalance_error(&Error::other( - "Unknown hyper error: hyper::Error(Http2, KeepAliveTimedOut)" - ))); - assert!(is_transient_rebalance_error(&Error::Lock(rustfs_lock::LockError::timeout( - "bucket/object", - Duration::from_secs(5), - )))); - } - - #[test] - 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_is_transient_rebalance_error_accepts_io_timeout_message() { - assert!(is_transient_rebalance_error(&Error::Io(std::io::Error::other("timeout")))); - } - - #[test] - fn test_is_transient_rebalance_error_rejects_terminal_errors() { - assert!(!is_transient_rebalance_error(&Error::DataMovementOverwriteErr( - "bucket".to_string(), - "object".to_string(), - "version".to_string(), - ))); - assert!(!is_transient_rebalance_error(&Error::ObjectNotFound( - "bucket".to_string(), - "object".to_string(), - ))); - assert!(!is_transient_rebalance_error(&Error::FileAccessDenied)); - assert!(!is_transient_rebalance_error(&Error::Lock(rustfs_lock::LockError::already_locked( - "bucket/object", - "owner-a", - )))); - assert!(!is_transient_rebalance_error(&Error::other( - "Failed to acquire read lock: ns_loc: read lock acquisition failed on bucket/object: permission denied" - ))); - } - - #[test] - fn test_should_retry_rebalance_listing_respects_attempt_limit_and_error_type() { - assert!(should_retry_rebalance_listing(&Error::SlowDown, 0, 3)); - assert!(should_retry_rebalance_listing(&Error::SlowDown, 1, 3)); - assert!(!should_retry_rebalance_listing(&Error::SlowDown, 2, 3)); - assert!(!should_retry_rebalance_listing(&Error::FileAccessDenied, 0, 3)); - } - - #[test] - fn test_parse_rebalance_max_attempts_uses_positive_override_or_default() { - assert_eq!(parse_rebalance_max_attempts(Some("5")), 5); - assert_eq!(parse_rebalance_max_attempts(Some(" 7 ")), 7); - assert_eq!(parse_rebalance_max_attempts(Some("0")), 3); - assert_eq!(parse_rebalance_max_attempts(Some("invalid")), 3); - assert_eq!(parse_rebalance_max_attempts(None), 3); - } - - #[test] - fn test_rebalance_listing_retry_delay_scales_by_attempt() { - assert_eq!(rebalance_listing_retry_delay(0), Duration::from_millis(250)); - assert_eq!(rebalance_listing_retry_delay(1), Duration::from_millis(500)); - } - - #[test] - fn test_rebalance_migration_retry_delay_scales_for_non_lock_timeout() { - assert_eq!(rebalance_migration_retry_delay(0, &Error::SlowDown), Duration::from_millis(250)); - assert_eq!(rebalance_migration_retry_delay(1, &Error::SlowDown), Duration::from_millis(500)); - } - - #[test] - fn test_should_defer_rebalance_entry_failure_only_for_transient_errors() { - assert!(should_defer_rebalance_entry_failure(&Error::Io(std::io::Error::other( - "Failed to acquire write lock: ns_loc: write lock acquisition timed out on bucket/object" - )))); - assert!(!should_defer_rebalance_entry_failure(&Error::DataMovementOverwriteErr( - "bucket".to_string(), - "object".to_string(), - "version".to_string(), - ))); - assert!(!should_defer_rebalance_entry_failure(&Error::FileAccessDenied)); - } - - #[tokio::test] - async fn test_wait_rebalance_listing_retry_returns_canceled_without_sleeping() { - let token = CancellationToken::new(); - token.cancel(); - - let err = wait_rebalance_listing_retry(&token, Duration::from_secs(30)) - .await - .expect_err("canceled rebalance should not wait before retrying"); - - assert!(matches!(err, Error::OperationCanceled)); - } - - #[test] - fn test_with_rebalance_entry_context_formats_stage_bucket_and_object() { - let err = with_rebalance_entry_context("migrate", "bucket-a", "obj.txt", Error::SlowDown); - let message = err.to_string(); - assert!(message.contains("rebalance entry migrate failed for bucket-a/obj.txt")); - assert!(message.contains("Please reduce your request rate")); - } - - #[test] - fn test_with_rebalance_entry_context_formats_precise_stage() { - let err = with_rebalance_entry_context("write_target", "bucket-a", "obj.txt", Error::SlowDown); - let message = err.to_string(); - assert!(message.contains("rebalance entry write_target failed for bucket-a/obj.txt")); - assert!(message.contains("Please reduce your request rate")); - } - - #[tokio::test] - async fn test_migrate_entry_version_transfer_failure_reports_write_target_stage() { - let backend = MigrationBackendSpy::new(Some(Ok(MigrationBackendSpy::make_reader())), None, None); - let mut transfer = |_, _, _| async { Err(Error::SlowDown) }; - let version = version_normal(); - - let result = migrate_entry_version( - &backend, - "bucket".to_string(), - 1, - &version, - version.version_id.map(|v| v.to_string()), - 1, - false, - &mut transfer, - ) - .await; - - assert!(result.failed); - assert_eq!(result.stage, Some("write_target")); - assert!(matches!(result.error, Some(Error::SlowDown))); - } - - #[tokio::test] - async fn test_migrate_entry_version_reader_failure_reports_read_source_stage() { - let backend = AlwaysFailGetBackend::new(); - let mut transfer = |_, _, _| async { Ok(()) }; - let version = version_normal(); - - let result = migrate_entry_version( - &backend, - "bucket".to_string(), - 1, - &version, - version.version_id.map(|v| v.to_string()), - 1, - false, - &mut transfer, - ) - .await; - - assert!(result.failed); - assert_eq!(result.stage, Some("read_source")); - assert!(matches!(result.error, Some(Error::SlowDown))); - } - - #[test] - fn test_should_count_rebalance_version_complete_for_cleanup_safe_ignored_result() { - let result = MigrationVersionResult { - ignored: true, - cleanup_ignored: true, - ..Default::default() - }; - assert!(should_count_rebalance_version_complete(&result)); - } - - #[test] - fn test_should_count_rebalance_version_complete_rejects_skip_only_ignored_result() { - let result = MigrationVersionResult { - ignored: true, - cleanup_ignored: false, - ..Default::default() - }; - assert!(!should_count_rebalance_version_complete(&result)); - } - - #[test] - fn test_should_count_rebalance_version_complete_for_moved_result() { - let result = MigrationVersionResult { - moved: true, - ..Default::default() - }; - assert!(should_count_rebalance_version_complete(&result)); - } - - #[test] - fn test_should_count_rebalance_version_complete_rejects_failed_result() { - let result = MigrationVersionResult { - moved: true, - failed: true, - ..Default::default() - }; - assert!(!should_count_rebalance_version_complete(&result)); - } - - #[test] - fn test_should_count_rebalance_version_complete_rejects_incomplete_result() { - assert!(!should_count_rebalance_version_complete(&MigrationVersionResult::default())); - } - - #[test] - fn test_should_accept_rebalance_stats_update_only_for_started_pool() { - let meta = RebalanceMeta { - pool_stats: vec![RebalanceStats { - info: RebalanceInfo { - status: RebalStatus::Started, - ..Default::default() - }, - ..Default::default() - }], - ..Default::default() - }; - - assert!(should_accept_rebalance_stats_update(&meta, 0)); - } - - #[test] - fn test_should_accept_rebalance_stats_update_rejects_stopped_meta() { - let meta = RebalanceMeta { - stopped_at: Some(OffsetDateTime::UNIX_EPOCH), - pool_stats: vec![RebalanceStats { - info: RebalanceInfo { - status: RebalStatus::Started, - ..Default::default() - }, - ..Default::default() - }], - ..Default::default() - }; - - assert!(!should_accept_rebalance_stats_update(&meta, 0)); - } - - #[test] - fn test_should_accept_rebalance_stats_update_rejects_terminal_or_invalid_pool() { - for status in [RebalStatus::Completed, RebalStatus::Stopped, RebalStatus::Failed] { - let meta = RebalanceMeta { - pool_stats: vec![RebalanceStats { - info: RebalanceInfo { - status, - ..Default::default() - }, - ..Default::default() - }], - ..Default::default() - }; - - assert!(!should_accept_rebalance_stats_update(&meta, 0)); - assert!(!should_accept_rebalance_stats_update(&meta, 1)); - } - } - - #[test] - fn test_should_skip_rebalance_delete_marker_when_last_remaining_without_replication() { - assert!(should_skip_rebalance_delete_marker(&version_deleted(), 1, false)); - } - - #[test] - fn test_should_skip_rebalance_delete_marker_rejects_configured_replication() { - assert!(!should_skip_rebalance_delete_marker(&version_deleted(), 1, true)); - } - - #[test] - fn test_should_skip_rebalance_delete_marker_rejects_non_deleted_versions() { - assert!(!should_skip_rebalance_delete_marker(&version_normal(), 1, false)); - } - - #[test] - fn test_should_skip_rebalance_delete_marker_rejects_multiple_remaining_versions() { - assert!(!should_skip_rebalance_delete_marker(&version_deleted(), 2, false)); - } - - #[test] - fn test_should_cleanup_rebalance_source_entry_accepts_all_versions_completed() { - assert!(should_cleanup_rebalance_source_entry(3, 3)); - } - - #[test] - fn test_should_cleanup_rebalance_source_entry_rejects_versions_only_expired_by_lifecycle() { - assert!(!should_cleanup_rebalance_source_entry(2, 3)); - } - - #[test] - fn test_resolve_rebalance_optional_bucket_config_result_passthrough() { - let result = resolve_rebalance_optional_bucket_config_result( - "bucket-a", - "replication", - Ok((ReplicationConfiguration::default(), OffsetDateTime::UNIX_EPOCH)), - ) - .expect("bucket config should pass through"); - assert!(result.is_some()); - } - - #[test] - fn test_resolve_rebalance_optional_bucket_config_result_returns_none_for_missing_config() { - let result = resolve_rebalance_optional_bucket_config_result::<()>("bucket-a", "versioning", Err(Error::ConfigNotFound)) - .expect("missing bucket config should map to None"); - assert!(result.is_none()); - } - - #[test] - fn test_resolve_rebalance_optional_bucket_config_result_wraps_other_errors() { - let err = resolve_rebalance_optional_bucket_config_result::<()>("bucket-a", "replication", Err(Error::SlowDown)) - .expect_err("unexpected bucket config errors should be wrapped with context"); - assert!( - err.to_string() - .contains("rebalance replication config load failed for bucket bucket-a") - ); - } - - #[tokio::test] - async fn test_load_rebalance_bucket_configs_skips_meta_bucket_lookup() { - let configs = load_rebalance_bucket_configs(RUSTFS_META_BUCKET) - .await - .expect("meta bucket config loading should short-circuit"); - assert!(configs.lifecycle_config.is_none()); - assert!(configs.lock_retention.is_none()); - assert!(configs.replication_config.is_none()); - } - - #[tokio::test] - async fn test_resolve_rebalance_save_task_result_join_error_keeps_context() { - let join_error = tokio::spawn(async { - panic!("rebalance save task panic"); - }) - .await - .expect_err("panic task should return JoinError"); - - let err = resolve_rebalance_save_task_result(3, Err(join_error)).unwrap_err(); - assert!(err.to_string().contains("rebalance save_task for pool 3 join error")); - } - - #[tokio::test] - async fn test_send_rebalance_done_signal_sends_message() { - let (tx, mut rx) = mpsc::channel(1); - - send_rebalance_done_signal(&tx, Ok(()), 2) - .await - .expect("send should succeed when receiver is active"); - - let received = rx.recv().await.expect("receiver should get signal"); - assert!(received.is_ok()); - } - - #[tokio::test] - async fn test_send_rebalance_done_signal_reports_closed_channel() { - let (tx, rx) = mpsc::channel(1); - drop(rx); - - let err = send_rebalance_done_signal(&tx, Ok(()), 5) - .await - .expect_err("send should fail when receiver is closed"); - assert!(err.to_string().contains("rebalance done signal send failed for pool 5")); - } - - #[test] - fn test_resolve_rebalance_terminal_error_keeps_primary_when_signal_ok() { - let err = resolve_rebalance_terminal_error(Error::SlowDown, Ok(())); - assert!(matches!(err, Error::SlowDown)); - } - - #[test] - fn test_resolve_rebalance_terminal_error_wraps_signal_failure_context() { - let err = resolve_rebalance_terminal_error(Error::SlowDown, Err(Error::OperationCanceled)); - assert!(err.to_string().contains("rebalance terminal signal failed after error")); - } - - #[test] - fn test_resolve_rebalance_bucket_error_prefers_entry_error() { - let err = resolve_rebalance_bucket_error(Some(Error::OperationCanceled), Some(Error::SlowDown)).unwrap_err(); - assert!(matches!(err, Error::OperationCanceled)); - } - - #[test] - fn test_resolve_rebalance_bucket_error_uses_worker_error_when_entry_ok() { - let err = resolve_rebalance_bucket_error(None, Some(Error::SlowDown)).unwrap_err(); - assert!(matches!(err, Error::SlowDown)); - } - - #[test] - fn test_resolve_rebalance_bucket_error_is_ok_when_no_errors() { - assert!(resolve_rebalance_bucket_error(None, None).is_ok()); - } - - #[test] - fn test_ensure_valid_rebalance_pool_index_allows_in_range() { - assert!(ensure_valid_rebalance_pool_index(3, 2).is_ok()); - } - - #[test] - fn test_ensure_valid_rebalance_pool_index_rejects_out_of_range() { - let err = ensure_valid_rebalance_pool_index(2, 2).expect_err("out of range index should fail"); - assert!(err.to_string().contains("invalid rebalance pool index")); - } - - #[test] - fn test_clone_first_arc_returns_first_value() { - let values = vec![Arc::new(7_u8), Arc::new(9_u8)]; - let first = clone_first_arc(values.as_slice(), "empty values").expect("first value should be returned"); - assert_eq!(*first, 7_u8); - } - - #[test] - fn test_clone_first_arc_rejects_empty_values() { - let values: Vec> = Vec::new(); - let err = clone_first_arc(values.as_slice(), "empty values").expect_err("empty values should fail"); - assert!(err.to_string().contains("empty values")); - } - - #[test] - fn test_clone_arc_by_index_returns_value() { - let values = vec![Arc::new(7_u8), Arc::new(9_u8)]; - let value = - clone_arc_by_index(values.as_slice(), 1, "invalid rebalance pool index").expect("index within bounds should work"); - assert_eq!(*value, 9_u8); - } - - #[test] - fn test_clone_arc_by_index_rejects_out_of_range() { - let values = vec![Arc::new(7_u8)]; - let err = - clone_arc_by_index(values.as_slice(), 2, "invalid rebalance pool index").expect_err("out of range index should fail"); - assert!(err.to_string().contains("invalid rebalance pool index: 2")); - } - - #[test] - fn test_classify_rebalance_terminal_event_completed() { - let now = OffsetDateTime::now_utc(); - match classify_rebalance_terminal_event(Some(Ok(())), now) { - RebalanceTerminalEvent::Completed { msg } => assert!(msg.contains("Rebalance completed")), - _ => panic!("expected completed terminal event"), - } - } - - #[test] - fn test_classify_rebalance_terminal_event_stopped() { - let now = OffsetDateTime::now_utc(); - match classify_rebalance_terminal_event(Some(Err(Error::OperationCanceled)), now) { - RebalanceTerminalEvent::Stopped { msg } => assert!(msg.contains("Rebalance stopped")), - _ => panic!("expected stopped terminal event"), - } - } - - #[test] - fn test_classify_rebalance_terminal_event_failed() { - let now = OffsetDateTime::now_utc(); - match classify_rebalance_terminal_event(Some(Err(Error::SlowDown)), now) { - RebalanceTerminalEvent::Failed { msg, last_error } => { - assert!(msg.contains("Rebalance failed")); - assert!(msg.contains("with err")); - assert_eq!(last_error, Error::SlowDown.to_string()); - } - _ => panic!("expected failed terminal event"), - } - } - - #[test] - fn test_classify_rebalance_terminal_event_channel_closed() { - let now = OffsetDateTime::now_utc(); - match classify_rebalance_terminal_event(None, now) { - RebalanceTerminalEvent::ChannelClosed { msg, last_error } => { - assert!(msg.contains("channel closed")); - assert!(last_error.contains("before terminal event")); - assert!(msg.contains("at")); - assert!(last_error.contains("at")); - } - _ => panic!("expected channel closed terminal event"), - } - } - - #[test] - fn test_apply_rebalance_terminal_event_channel_closed_marks_failed() { - let now = OffsetDateTime::now_utc(); - let mut status = RebalStatus::Started; - let mut end_time = None; - let mut last_error = None; - - apply_rebalance_terminal_event( - &mut status, - &mut end_time, - &mut last_error, - RebalanceTerminalEvent::ChannelClosed { - msg: "channel closed".to_string(), - last_error: "rebalance save channel closed before terminal event".to_string(), - }, - now, - ); - - assert_eq!(status, RebalStatus::Failed); - assert_eq!(end_time, Some(now)); - assert_eq!(last_error.as_deref(), Some("rebalance save channel closed before terminal event")); - } - - #[test] - fn test_apply_rebalance_terminal_event_stopped_clears_error() { - let now = OffsetDateTime::now_utc(); - let mut status = RebalStatus::Started; - let mut end_time = None; - let mut last_error = Some("old-error".to_string()); - - apply_rebalance_terminal_event( - &mut status, - &mut end_time, - &mut last_error, - RebalanceTerminalEvent::Stopped { - msg: "rebalance stopped".to_string(), - }, - now, - ); - - assert_eq!(status, RebalStatus::Stopped); - assert_eq!(end_time, Some(now)); - assert_eq!(last_error, None); - } - - #[test] - fn test_is_rebalance_stopped_terminal_event_only_matches_stopped_variant() { - let stopped = RebalanceTerminalEvent::Stopped { - msg: "stopped".to_string(), - }; - let completed = RebalanceTerminalEvent::Completed { - msg: "completed".to_string(), - }; - - assert!(is_rebalance_stopped_terminal_event(&stopped)); - assert!(!is_rebalance_stopped_terminal_event(&completed)); - } - - #[test] - fn test_should_preserve_rebalance_stopped_state_when_meta_marked_stopped() { - let event = RebalanceTerminalEvent::Completed { - msg: "completed".to_string(), - }; - - assert!(should_preserve_rebalance_stopped_state(true, RebalStatus::Started, &event)); - } - - #[test] - fn test_should_preserve_rebalance_stopped_state_when_pool_already_stopped() { - let event = RebalanceTerminalEvent::Failed { - msg: "failed".to_string(), - last_error: "boom".to_string(), - }; - - assert!(should_preserve_rebalance_stopped_state(false, RebalStatus::Stopped, &event)); - } - - #[test] - fn test_should_preserve_rebalance_stopped_state_allows_stopped_terminal_update() { - let event = RebalanceTerminalEvent::Stopped { - msg: "stopped".to_string(), - }; - - assert!(!should_preserve_rebalance_stopped_state(true, RebalStatus::Started, &event)); - } - - #[test] - fn test_ensure_rebalance_not_decommissioning_rejects_running_decommission() { - assert!(!ensure_rebalance_not_decommissioning(true)); - } - - #[test] - fn test_ensure_rebalance_not_decommissioning_allows_idle_decommission() { - assert!(ensure_rebalance_not_decommissioning(false)); - } - - #[test] - fn test_validate_start_rebalance_state_rejects_running_decommission() { - let err = validate_start_rebalance_state(true, true).expect_err("running decommission should block rebalance start"); - assert!(matches!(err, Error::DecommissionAlreadyRunning)); - } - - #[test] - fn test_validate_start_rebalance_state_rejects_missing_meta() { - let err = validate_start_rebalance_state(false, false).expect_err("missing rebalance meta should fail"); - assert!(matches!(err, Error::ConfigNotFound)); - } - - #[test] - fn test_validate_start_rebalance_state_allows_loaded_meta() { - validate_start_rebalance_state(false, true).expect("loaded rebalance meta should allow start"); - } - - #[test] - fn test_percent_free_ratio_zero_capacity_is_zero() { - assert_eq!(percent_free_ratio(100, 0), 0.0); - } - - #[test] - fn test_percent_free_ratio_normal_case() { - assert_eq!(percent_free_ratio(250, 1_000), 0.25); - } - - #[test] - fn test_should_pool_participate_false_when_capacity_zero() { - assert!(!should_pool_participate(0, 0, 0.2)); - } - - #[test] - fn test_should_pool_participate_true_when_ratio_below_goal() { - assert!(should_pool_participate(200, 1_000, 0.3)); - } - - #[test] - fn test_should_pool_participate_false_when_ratio_meets_goal() { - assert!(!should_pool_participate(300, 1_000, 0.3)); - } - - #[test] - fn test_rebalance_goal_not_reached_for_issue_3137_initial_imbalance() { - let pool0_capacity = 10_000_u64; - let pool0_free = 9_135_u64; - let pool1_capacity = 10_000_u64; - let pool1_free = 9_800_u64; - let goal = percent_free_ratio(pool0_free + pool1_free, pool0_capacity + pool1_capacity); - - assert!(should_pool_participate(pool0_free, pool0_capacity, goal)); - assert!(!rebalance_goal_reached(pool0_free, pool0_capacity, 0, goal)); - } - - #[test] - fn test_complete_rebalance_pools_at_goal_marks_started_participants_completed() { - let now = OffsetDateTime::from_unix_timestamp(1_000).unwrap(); - let mut meta = RebalanceMeta { - 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() - }, - RebalanceStats { - participating: true, - init_free_space: 100, - init_capacity: 1_000, - bytes: 0, - info: RebalanceInfo { - status: RebalStatus::Started, - ..Default::default() - }, - ..Default::default() - }, - ], - ..Default::default() - }; - - assert!(complete_rebalance_pools_at_goal(&mut meta, now)); - assert_eq!(meta.pool_stats[0].info.status, RebalStatus::Completed); - assert_eq!(meta.pool_stats[0].info.end_time, Some(now)); - assert_eq!(meta.pool_stats[1].info.status, RebalStatus::Started); - } - - #[test] - fn test_complete_rebalance_pools_at_goal_skips_deferred_transient_error() { - let now = OffsetDateTime::from_unix_timestamp(1_000).unwrap(); - let mut meta = RebalanceMeta { - 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, - last_error: Some(format!("{REBALANCE_DEFERRED_ENTRY_ERROR_PREFIX} timeout")), - ..Default::default() - }, - ..Default::default() - }], - ..Default::default() - }; - - assert!(!complete_rebalance_pools_at_goal(&mut meta, now)); - assert_eq!(meta.pool_stats[0].info.status, RebalStatus::Started); - assert!(has_deferred_rebalance_error(&meta.pool_stats[0])); - } - - #[test] - fn test_complete_rebalance_pools_with_empty_queue_marks_started_participants_completed() { - let now = OffsetDateTime::from_unix_timestamp(1_000).unwrap(); - let mut meta = RebalanceMeta { - pool_stats: vec![ - RebalanceStats { - participating: true, - buckets: Vec::new(), - info: RebalanceInfo { - status: RebalStatus::Started, - last_error: Some("stale error".to_string()), - ..Default::default() - }, - ..Default::default() - }, - RebalanceStats { - participating: true, - buckets: vec!["bucket-a".to_string()], - info: RebalanceInfo { - status: RebalStatus::Started, - ..Default::default() - }, - ..Default::default() - }, - ], - ..Default::default() - }; - - assert!(complete_rebalance_pools_with_empty_queue(&mut meta, now)); - assert_eq!(meta.pool_stats[0].info.status, RebalStatus::Completed); - assert_eq!(meta.pool_stats[0].info.end_time, Some(now)); - assert!(meta.pool_stats[0].info.last_error.is_none()); - assert_eq!(meta.pool_stats[1].info.status, RebalStatus::Started); - } - - #[test] - fn test_should_skip_start_rebalance_only_when_running_and_cancel_attached() { - assert!(should_skip_start_rebalance(true, true)); - assert!(!should_skip_start_rebalance(true, false)); - assert!(!should_skip_start_rebalance(false, true)); - assert!(!should_skip_start_rebalance(false, false)); - } - - #[test] - fn test_new_multipart_abort_flag_defaults_to_abort_enabled() { - let flag = data_movement::new_multipart_abort_flag(); - assert!(data_movement::should_abort_multipart_upload(&flag)); - } - - #[test] - fn test_mark_multipart_upload_completed_disables_abort_cleanup() { - let flag = data_movement::new_multipart_abort_flag(); - data_movement::mark_multipart_upload_completed(&flag); - assert!(!data_movement::should_abort_multipart_upload(&flag)); - } - - #[test] - fn test_decode_part_index_returns_none_when_absent() { - assert!(data_movement::decode_part_index(None).is_none()); - } - - #[test] - fn test_decode_part_index_returns_none_for_invalid_payload() { - let invalid = bytes::Bytes::from_static(b"not-a-valid-index"); - assert!(data_movement::decode_part_index(Some(&invalid)).is_none()); - } - - #[test] - fn test_decode_part_index_returns_some_for_valid_payload() { - let mut index = Index::new(); - index.add(0, 0).expect("first index entry should be accepted"); - index - .add(2_097_152, 2_097_152) - .expect("second index entry should advance totals"); - - let encoded = index.into_vec(); - let decoded = data_movement::decode_part_index(Some(&encoded)).expect("valid index payload should decode"); - - assert_eq!(decoded.total_uncompressed, 2_097_152); - assert_eq!(decoded.total_compressed, 2_097_152); - } - - #[test] - fn test_merge_rebalance_bucket_lists_keeps_existing_order_and_skips_duplicates() { - let mut remote = vec!["bucket-a".to_string(), "bucket-b".to_string()]; - let local = vec![ - "bucket-b".to_string(), - "bucket-c".to_string(), - "bucket-c".to_string(), - "bucket-d".to_string(), - ]; - - merge_rebalance_bucket_lists(&mut remote, &local); - - assert_eq!( - remote, - vec![ - "bucket-a".to_string(), - "bucket-b".to_string(), - "bucket-c".to_string(), - "bucket-d".to_string() - ] - ); - } - - #[test] - fn test_remove_rebalanced_buckets_from_queue_filters_membership_linearly() { - let mut pool_stat = RebalanceStats { - buckets: vec![ - "bucket-a".to_string(), - "bucket-b".to_string(), - "bucket-c".to_string(), - "bucket-b".to_string(), - ], - rebalanced_buckets: vec!["bucket-b".to_string(), "bucket-d".to_string()], - ..Default::default() - }; - - remove_rebalanced_buckets_from_queue(&mut pool_stat); - - assert_eq!(pool_stat.buckets, vec!["bucket-a".to_string(), "bucket-c".to_string()]); - } - - #[test] - fn test_resolve_rebalance_participants_respects_runtime_pool_count() { - let now = OffsetDateTime::now_utc(); - let stats = vec![ - RebalanceStats { - participating: true, - info: RebalanceInfo { - status: RebalStatus::Started, - start_time: Some(now), - ..Default::default() - }, - ..Default::default() - }, - RebalanceStats { - participating: false, - info: RebalanceInfo { - status: RebalStatus::Started, - start_time: Some(now), - ..Default::default() - }, - ..Default::default() - }, - RebalanceStats { - participating: true, - info: RebalanceInfo { - status: RebalStatus::Started, - start_time: Some(now), - ..Default::default() - }, - ..Default::default() - }, - ]; - - let participants = resolve_rebalance_participants(stats.as_slice(), 2); - assert_eq!(participants, vec![true, false]); - } - - #[test] - fn test_resolve_rebalance_participants_requires_started_status() { - let now = OffsetDateTime::now_utc(); - let stats = vec![ - RebalanceStats { - participating: true, - info: RebalanceInfo { - status: RebalStatus::Completed, - start_time: Some(now), - ..Default::default() - }, - ..Default::default() - }, - RebalanceStats { - participating: true, - info: RebalanceInfo { - status: RebalStatus::Started, - start_time: Some(now), - ..Default::default() - }, - ..Default::default() - }, - ]; - - let participants = resolve_rebalance_participants(stats.as_slice(), 2); - assert_eq!(participants, vec![false, true]); - } - - #[test] - fn test_is_rebalance_actively_running_requires_cancel_and_started_state() { - let now = OffsetDateTime::now_utc(); - let mut meta = RebalanceMeta { - pool_stats: vec![RebalanceStats { - participating: true, - info: RebalanceInfo { - start_time: Some(now), - status: RebalStatus::Started, - ..Default::default() - }, - ..Default::default() - }], - ..Default::default() - }; - - assert!(!is_rebalance_actively_running(&meta)); - meta.cancel = Some(CancellationToken::new()); - assert!(is_rebalance_actively_running(&meta)); - } - - #[test] - fn test_is_rebalance_in_progress_only_started_participants() { - let now = OffsetDateTime::now_utc(); - let meta = RebalanceMeta { - stopped_at: None, - pool_stats: vec![ - RebalanceStats { - participating: true, - info: RebalanceInfo { - start_time: Some(now), - status: RebalStatus::Completed, - ..Default::default() - }, - ..Default::default() - }, - RebalanceStats { - participating: true, - info: RebalanceInfo { - start_time: Some(now), - status: RebalStatus::Started, - ..Default::default() - }, - ..Default::default() - }, - ], - ..Default::default() - }; - - assert!(is_rebalance_in_progress(&meta)); - } - - #[test] - fn test_is_rebalance_conflicting_with_decommission_true_when_in_progress() { - let now = OffsetDateTime::now_utc(); - let meta = RebalanceMeta { - stopped_at: None, - pool_stats: vec![RebalanceStats { - participating: true, - info: RebalanceInfo { - start_time: Some(now), - status: RebalStatus::Started, - ..Default::default() - }, - ..Default::default() - }], - ..Default::default() - }; - - assert!(is_rebalance_conflicting_with_decommission(&meta)); - } - - #[test] - fn test_is_rebalance_conflicting_with_decommission_false_when_stopped() { - let now = OffsetDateTime::now_utc(); - let meta = RebalanceMeta { - stopped_at: Some(now), - pool_stats: vec![RebalanceStats { - participating: true, - info: RebalanceInfo { - start_time: Some(now), - status: RebalStatus::Started, - ..Default::default() - }, - ..Default::default() - }], - ..Default::default() - }; - - assert!(!is_rebalance_conflicting_with_decommission(&meta)); - } - - #[test] - fn test_is_rebalance_in_progress_stopped_takes_precedence() { - let now = OffsetDateTime::now_utc(); - let meta = RebalanceMeta { - stopped_at: Some(now), - pool_stats: vec![RebalanceStats { - participating: true, - info: RebalanceInfo { - start_time: Some(now), - status: RebalStatus::Started, - ..Default::default() - }, - ..Default::default() - }], - ..Default::default() - }; - - assert!(!is_rebalance_in_progress(&meta)); - } - - #[test] - fn test_first_rebalance_bucket_returns_first_name() { - let pool_stat = RebalanceStats { - buckets: vec!["bucket-a".to_string(), "bucket-b".to_string()], - ..Default::default() - }; - - assert_eq!(first_rebalance_bucket(&pool_stat), Some("bucket-a".to_string())); - } - - #[test] - fn test_first_rebalance_bucket_returns_none_when_empty() { - let pool_stat = RebalanceStats::default(); - - assert_eq!(first_rebalance_bucket(&pool_stat), None); - } - - #[test] - fn test_next_rebal_bucket_from_stat_respects_empty_queue() { - let pool_stat = RebalanceStats { - buckets: vec![], - ..Default::default() - }; - - assert_eq!(next_rebal_bucket_from_stat(&pool_stat), None); - } - - #[test] - fn test_next_rebal_bucket_from_stat_returns_first_bucket() { - let now = OffsetDateTime::now_utc(); - let pool_stat = RebalanceStats { - participating: true, - info: RebalanceInfo { - status: RebalStatus::Started, - start_time: Some(now), - ..Default::default() - }, - buckets: vec!["bucket-a".to_string(), "bucket-b".to_string()], - ..Default::default() - }; - - assert_eq!(next_rebal_bucket_from_stat(&pool_stat), Some("bucket-a".to_string())); - } - - #[test] - fn test_clone_rebalance_pool_stats_rejects_missing_meta() { - let err = clone_rebalance_pool_stats(None).expect_err("missing rebalance meta should fail"); - assert!( - err.to_string() - .contains("failed to clone rebalance pool stats: rebalance metadata not initialized") - ); - } - - #[test] - fn test_clone_rebalance_pool_stats_clones_entries() { - let meta = RebalanceMeta { - pool_stats: vec![RebalanceStats::default()], - ..Default::default() - }; - - let stats = clone_rebalance_pool_stats(Some(&meta)).expect("metadata should clone pool stats"); - assert_eq!(stats.len(), 1); - } - - #[test] - fn test_resolve_next_rebalance_bucket_rejects_missing_meta() { - let err = resolve_next_rebalance_bucket(None, 0).expect_err("missing meta should fail"); - assert!( - err.to_string() - .contains("failed to resolve next rebalance bucket: rebalance metadata not initialized") - ); - } - - #[test] - fn test_resolve_next_rebalance_bucket_rejects_invalid_pool_index() { - let meta = RebalanceMeta { - pool_stats: vec![RebalanceStats::default()], - ..Default::default() - }; - - let err = resolve_next_rebalance_bucket(Some(&meta), 3).expect_err("invalid pool index should fail"); - assert!(err.to_string().contains("invalid rebalance pool index 3 for 1 pools")); - } - - #[test] - fn test_resolve_next_rebalance_bucket_returns_none_for_completed_pool() { - let meta = RebalanceMeta { - pool_stats: vec![RebalanceStats { - participating: true, - info: RebalanceInfo { - status: RebalStatus::Completed, - ..Default::default() - }, - buckets: vec!["bucket-a".to_string()], - ..Default::default() - }], - ..Default::default() - }; - - let next = resolve_next_rebalance_bucket(Some(&meta), 0).expect("completed pool should return none"); - assert!(next.is_none()); - } - - #[test] - fn test_resolve_next_rebalance_bucket_returns_first_bucket_for_active_pool() { - let now = OffsetDateTime::now_utc(); - let meta = RebalanceMeta { - pool_stats: vec![RebalanceStats { - participating: true, - info: RebalanceInfo { - status: RebalStatus::Started, - start_time: Some(now), - ..Default::default() - }, - buckets: vec!["bucket-a".to_string(), "bucket-b".to_string()], - ..Default::default() - }], - ..Default::default() - }; - - let next = resolve_next_rebalance_bucket(Some(&meta), 0).expect("active pool should return first bucket"); - assert_eq!(next.as_deref(), Some("bucket-a")); - } - - #[test] - fn test_take_bucket_from_rebalance_queue_moves_bucket_and_keeps_remaining() { - let mut pool_stat = RebalanceStats { - buckets: vec!["bucket-a".to_string(), "bucket-b".to_string(), "bucket-a".to_string()], - rebalanced_buckets: Vec::new(), - ..Default::default() - }; - - assert!(take_bucket_from_rebalance_queue(&mut pool_stat, "bucket-a")); - assert_eq!(pool_stat.buckets, vec!["bucket-b".to_string()]); - assert_eq!(pool_stat.rebalanced_buckets, vec!["bucket-a".to_string(), "bucket-a".to_string()]); - } - - #[test] - fn test_take_bucket_from_rebalance_queue_no_match_keeps_queue() { - let mut pool_stat = RebalanceStats { - buckets: vec!["bucket-a".to_string(), "bucket-b".to_string()], - rebalanced_buckets: Vec::new(), - ..Default::default() - }; - - assert!(!take_bucket_from_rebalance_queue(&mut pool_stat, "bucket-c")); - assert_eq!(pool_stat.buckets, vec!["bucket-a".to_string(), "bucket-b".to_string()]); - assert!(pool_stat.rebalanced_buckets.is_empty()); - } - - #[test] - fn test_defer_bucket_in_rebalance_queue_moves_bucket_to_back() { - let mut pool_stat = RebalanceStats { - buckets: vec!["bucket-a".to_string(), "bucket-b".to_string(), "bucket-c".to_string()], - rebalanced_buckets: Vec::new(), - ..Default::default() - }; - - defer_bucket_in_rebalance_queue(&mut pool_stat, "bucket-a").expect("queued bucket should be deferred"); - - assert_eq!( - pool_stat.buckets, - vec!["bucket-b".to_string(), "bucket-c".to_string(), "bucket-a".to_string()] - ); - assert!(pool_stat.rebalanced_buckets.is_empty()); - } - - #[test] - fn test_defer_bucket_in_rebalance_queue_rejects_missing_bucket() { - let mut pool_stat = RebalanceStats { - buckets: vec!["bucket-a".to_string(), "bucket-b".to_string()], - rebalanced_buckets: Vec::new(), - ..Default::default() - }; - - let err = defer_bucket_in_rebalance_queue(&mut pool_stat, "bucket-c").expect_err("missing bucket should not be deferred"); - - assert!(err.to_string().contains("failed to defer rebalance bucket bucket-c")); - assert_eq!(pool_stat.buckets, vec!["bucket-a".to_string(), "bucket-b".to_string()]); - assert!(pool_stat.rebalanced_buckets.is_empty()); - } - - #[test] - fn test_mark_rebalance_bucket_done_rejects_missing_meta() { - let err = mark_rebalance_bucket_done(None, 0, "bucket-a").expect_err("missing meta should fail"); - assert!( - err.to_string() - .contains("failed to mark rebalance bucket done: rebalance metadata not initialized") - ); - } - - #[test] - fn test_mark_rebalance_bucket_done_rejects_invalid_pool_index() { - let mut meta = RebalanceMeta { - pool_stats: vec![RebalanceStats::default()], - ..Default::default() - }; - - let err = mark_rebalance_bucket_done(Some(&mut meta), 3, "bucket-a").expect_err("invalid pool index should fail"); - assert!(err.to_string().contains("invalid rebalance pool index 3 for 1 pools")); - } - - #[test] - fn test_mark_rebalance_bucket_done_rejects_missing_bucket() { - let mut meta = RebalanceMeta { - pool_stats: vec![RebalanceStats { - buckets: vec!["bucket-a".to_string()], - ..Default::default() - }], - ..Default::default() - }; - - let err = mark_rebalance_bucket_done(Some(&mut meta), 0, "bucket-x").expect_err("missing bucket should fail"); - assert!( - err.to_string() - .contains("failed to mark rebalance bucket done: bucket bucket-x was not queued for pool 0") - ); - } - - #[test] - fn test_mark_rebalance_bucket_done_marks_bucket_as_rebalanced() { - let mut meta = RebalanceMeta { - pool_stats: vec![RebalanceStats { - buckets: vec!["bucket-a".to_string(), "bucket-b".to_string()], - rebalanced_buckets: Vec::new(), - ..Default::default() - }], - ..Default::default() - }; - - mark_rebalance_bucket_done(Some(&mut meta), 0, "bucket-a").expect("bucket in queue should be marked done"); - assert_eq!(meta.pool_stats[0].buckets, vec!["bucket-b".to_string()]); - assert_eq!(meta.pool_stats[0].rebalanced_buckets, vec!["bucket-a".to_string()]); - } - - #[test] - fn test_mark_rebalance_bucket_done_clears_deferred_transient_error() { - let mut meta = RebalanceMeta { - pool_stats: vec![RebalanceStats { - buckets: vec!["bucket-a".to_string()], - info: RebalanceInfo { - last_error: Some(format!("{REBALANCE_DEFERRED_ENTRY_ERROR_PREFIX} timeout")), - ..Default::default() - }, - ..Default::default() - }], - ..Default::default() - }; - - mark_rebalance_bucket_done(Some(&mut meta), 0, "bucket-a") - .expect("bucket in queue should be marked done after deferred retry succeeds"); - - assert!(meta.pool_stats[0].info.last_error.is_none()); - } - - #[test] - fn test_record_rebalance_cleanup_warning_in_meta_preserves_last_error() { - let now = OffsetDateTime::from_unix_timestamp(10_000).unwrap(); - let mut meta = RebalanceMeta { - pool_stats: vec![RebalanceStats { - info: RebalanceInfo { - last_error: Some("old-error".to_string()), - ..Default::default() - }, - ..Default::default() - }], - ..Default::default() - }; - - record_rebalance_cleanup_warning_in_meta(Some(&mut meta), 0, "bucket-a", "obj.txt", "cleanup failed".to_string(), now) - .expect("cleanup warning should be recorded"); - - assert_eq!(meta.pool_stats[0].info.last_error.as_deref(), Some("old-error")); - assert_eq!(meta.pool_stats[0].cleanup_warnings.count, 1); - assert_eq!(meta.pool_stats[0].cleanup_warnings.last_message.as_deref(), Some("cleanup failed")); - assert_eq!(meta.pool_stats[0].cleanup_warnings.last_bucket.as_deref(), Some("bucket-a")); - assert_eq!(meta.pool_stats[0].cleanup_warnings.last_object.as_deref(), Some("obj.txt")); - assert_eq!(meta.pool_stats[0].cleanup_warnings.last_at, Some(now)); - assert_eq!(meta.last_refreshed_at, Some(now)); - } - - #[test] - fn test_complete_rebalance_pools_with_empty_queue_preserves_cleanup_warnings() { - let warning_at = OffsetDateTime::from_unix_timestamp(9_000).unwrap(); - let completed_at = OffsetDateTime::from_unix_timestamp(10_000).unwrap(); - let mut meta = RebalanceMeta { - pool_stats: vec![RebalanceStats { - participating: true, - info: RebalanceInfo { - status: RebalStatus::Started, - ..Default::default() - }, - cleanup_warnings: RebalanceCleanupWarnings { - count: 1, - last_message: Some("cleanup failed".to_string()), - last_bucket: Some("bucket-a".to_string()), - last_object: Some("obj.txt".to_string()), - last_at: Some(warning_at), - }, - ..Default::default() - }], - ..Default::default() - }; - - assert!(complete_rebalance_pools_with_empty_queue(&mut meta, completed_at)); - - assert_eq!(meta.pool_stats[0].info.status, RebalStatus::Completed); - assert!(meta.pool_stats[0].info.last_error.is_none()); - assert_eq!(meta.pool_stats[0].cleanup_warnings.count, 1); - assert_eq!(meta.pool_stats[0].cleanup_warnings.last_message.as_deref(), Some("cleanup failed")); - assert_eq!(meta.pool_stats[0].cleanup_warnings.last_at, Some(warning_at)); - } - - #[test] - fn test_apply_stopped_at_transitions_started_pools_only() { - let now = OffsetDateTime::now_utc(); - let mut meta = RebalanceMeta { - pool_stats: vec![ - RebalanceStats { - info: RebalanceInfo { - status: RebalStatus::Started, - end_time: None, - last_error: Some("old".to_string()), - ..Default::default() - }, - ..Default::default() - }, - RebalanceStats { - info: RebalanceInfo { - status: RebalStatus::Failed, - end_time: Some(now), - last_error: Some("failed".to_string()), - ..Default::default() - }, - ..Default::default() - }, - ], - ..Default::default() - }; - - apply_stopped_at(&mut meta, now); - - assert_eq!(meta.stopped_at, Some(now)); - assert_eq!(meta.pool_stats[0].info.status, RebalStatus::Stopped); - assert_eq!(meta.pool_stats[0].info.end_time, Some(now)); - assert_eq!(meta.pool_stats[0].info.last_error, None); - - assert_eq!(meta.pool_stats[1].info.status, RebalStatus::Failed); - assert_eq!(meta.pool_stats[1].info.end_time, Some(now)); - assert_eq!(meta.pool_stats[1].info.last_error.as_deref(), Some("failed")); - } - - #[test] - fn test_stop_rebalance_state_cancels_token_and_marks_stopped_when_in_progress() { - let now = OffsetDateTime::from_unix_timestamp(10_000).unwrap(); - let cancel = CancellationToken::new(); - let cancel_clone = cancel.clone(); - let mut meta = RebalanceMeta { - cancel: Some(cancel), - pool_stats: vec![RebalanceStats { - participating: true, - info: RebalanceInfo { - status: RebalStatus::Started, - ..Default::default() - }, - ..Default::default() - }], - ..Default::default() - }; - - stop_rebalance_state(&mut meta, now); - - assert!(cancel_clone.is_cancelled()); - assert!(meta.cancel.is_none()); - assert_eq!(meta.stopped_at, Some(now)); - assert_eq!(meta.pool_stats[0].info.status, RebalStatus::Stopped); - assert_eq!(meta.pool_stats[0].info.end_time, Some(now)); - } - - #[test] - fn test_stop_rebalance_state_clears_token_without_forcing_stopped_when_not_in_progress() { - let now = OffsetDateTime::from_unix_timestamp(20_000).unwrap(); - let cancel = CancellationToken::new(); - let cancel_clone = cancel.clone(); - let mut meta = RebalanceMeta { - cancel: Some(cancel), - pool_stats: vec![RebalanceStats { - participating: true, - info: RebalanceInfo { - status: RebalStatus::Completed, - end_time: Some(now), - ..Default::default() - }, - ..Default::default() - }], - ..Default::default() - }; - - stop_rebalance_state(&mut meta, now); - - assert!(cancel_clone.is_cancelled()); - assert!(meta.cancel.is_none()); - assert_eq!(meta.stopped_at, None); - assert_eq!(meta.pool_stats[0].info.status, RebalStatus::Completed); - } - - #[test] - fn test_stop_rebalance_state_normalizes_started_pool_when_stopped_at_already_set() { - let stopped_at = OffsetDateTime::from_unix_timestamp(30_000).unwrap(); - let now = OffsetDateTime::from_unix_timestamp(40_000).unwrap(); - let cancel = CancellationToken::new(); - let cancel_clone = cancel.clone(); - let mut meta = RebalanceMeta { - cancel: Some(cancel), - stopped_at: Some(stopped_at), - pool_stats: vec![RebalanceStats { - participating: true, - info: RebalanceInfo { - status: RebalStatus::Started, - last_error: Some("stale".to_string()), - ..Default::default() - }, - ..Default::default() - }], - ..Default::default() - }; - - stop_rebalance_state(&mut meta, now); - - assert!(cancel_clone.is_cancelled()); - assert!(meta.cancel.is_none()); - assert_eq!(meta.stopped_at, Some(stopped_at)); - assert_eq!(meta.pool_stats[0].info.status, RebalStatus::Stopped); - assert_eq!(meta.pool_stats[0].info.end_time, Some(stopped_at)); - assert_eq!(meta.pool_stats[0].info.last_error, None); - } - - #[test] - fn test_stop_rebalance_meta_snapshot_returns_none_when_meta_missing() { - let now = OffsetDateTime::from_unix_timestamp(50_000).unwrap(); - assert!(stop_rebalance_meta_snapshot(None, now).is_none()); - } - - #[test] - fn test_stop_rebalance_meta_snapshot_stops_meta_and_returns_snapshot() { - let now = OffsetDateTime::from_unix_timestamp(60_000).unwrap(); - let cancel = CancellationToken::new(); - let cancel_clone = cancel.clone(); - let mut meta = RebalanceMeta { - cancel: Some(cancel), - pool_stats: vec![RebalanceStats { - participating: true, - info: RebalanceInfo { - status: RebalStatus::Started, - ..Default::default() - }, - ..Default::default() - }], - ..Default::default() - }; - - let snapshot = stop_rebalance_meta_snapshot(Some(&mut meta), now).expect("snapshot should be returned for present meta"); - - assert!(cancel_clone.is_cancelled()); - assert!(meta.cancel.is_none()); - assert_eq!(meta.stopped_at, Some(now)); - assert_eq!(meta.last_refreshed_at, Some(now)); - assert_eq!(meta.pool_stats[0].info.status, RebalStatus::Stopped); - - assert!(snapshot.cancel.is_none()); - assert_eq!(snapshot.stopped_at, Some(now)); - assert_eq!(snapshot.last_refreshed_at, Some(now)); - assert_eq!(snapshot.pool_stats[0].info.status, RebalStatus::Stopped); - assert_eq!(snapshot.pool_stats[0].info.end_time, Some(now)); - } - - #[test] - fn test_apply_rebalance_save_option_stats_keeps_pool_status_and_updates_refresh() { - let now = OffsetDateTime::from_unix_timestamp(1_000).unwrap(); - let later = OffsetDateTime::from_unix_timestamp(2_000).unwrap(); - let mut meta = RebalanceMeta { - pool_stats: vec![RebalanceStats { - participating: true, - info: RebalanceInfo { - status: RebalStatus::Started, - start_time: Some(now), - ..Default::default() - }, - buckets: vec!["bucket-a".to_string()], - ..Default::default() - }], - last_refreshed_at: Some(now), - stopped_at: None, - ..Default::default() - }; - - apply_rebalance_save_option(&mut meta, 0, RebalSaveOpt::Stats, later); - - assert_eq!(meta.last_refreshed_at, Some(later)); - assert_eq!(meta.stopped_at, None); - assert_eq!(meta.pool_stats.len(), 1); - assert!(meta.pool_stats[0].participating); - assert_eq!(meta.pool_stats[0].info.status, RebalStatus::Started); - assert_eq!(meta.pool_stats[0].info.start_time, Some(now)); - assert_eq!(meta.pool_stats[0].buckets, vec!["bucket-a".to_string()]); - } - - #[test] - fn test_apply_rebalance_save_option_stopped_at_updates_refresh_and_statuses() { - let now = OffsetDateTime::from_unix_timestamp(1_000).unwrap(); - let mut meta = RebalanceMeta { - pool_stats: vec![ - RebalanceStats { - info: RebalanceInfo { - status: RebalStatus::Started, - ..Default::default() - }, - ..Default::default() - }, - RebalanceStats { - info: RebalanceInfo { - status: RebalStatus::Failed, - last_error: Some("previous failure".to_string()), - ..Default::default() - }, - ..Default::default() - }, - ], - ..Default::default() - }; - - apply_rebalance_save_option(&mut meta, 9_000, RebalSaveOpt::StoppedAt, now); - - assert_eq!(meta.stopped_at, Some(now)); - assert_eq!(meta.last_refreshed_at, Some(now)); - assert_eq!(meta.pool_stats[0].info.status, RebalStatus::Stopped); - assert_eq!(meta.pool_stats[0].info.end_time, Some(now)); - assert!(meta.pool_stats[0].info.last_error.is_none()); - assert_eq!(meta.pool_stats[1].info.status, RebalStatus::Failed); - assert_eq!(meta.pool_stats[1].info.last_error.as_deref(), Some("previous failure")); - } - - #[test] - fn test_rebalance_stats_update_counts_and_bytes_growth() { - let mut stat = RebalanceStats { - bucket: "bucket-a".to_string(), - object: "obj-previous".to_string(), - num_objects: 3, - num_versions: 5, - bytes: 120, - ..Default::default() - }; - - let mut latest = FileInfo::new("object-1", 4, 2); - latest.name = "object-1".to_string(); - latest.is_latest = true; - latest.size = 300; - latest.deleted = false; - latest.mod_time = Some(OffsetDateTime::UNIX_EPOCH); - latest.version_id = None; - - let mut historical = FileInfo::new("object-1", 4, 2); - historical.name = "object-1".to_string(); - historical.is_latest = false; - historical.size = 128; - historical.deleted = false; - historical.mod_time = Some(OffsetDateTime::UNIX_EPOCH); - historical.version_id = None; - - let mut tombstone = FileInfo::new("object-1", 4, 2); - tombstone.name = "object-1".to_string(); - tombstone.is_latest = false; - tombstone.size = 64; - tombstone.deleted = true; - tombstone.mod_time = Some(OffsetDateTime::UNIX_EPOCH); - tombstone.version_id = None; - - stat.update("bucket-b".to_string(), &latest); - stat.update("bucket-b".to_string(), &historical); - stat.update("bucket-b".to_string(), &tombstone); - - assert_eq!(stat.bucket, "bucket-b"); - assert_eq!(stat.object, "object-1"); - assert_eq!(stat.num_objects, 4); - assert_eq!(stat.num_versions, 8); - let expected_bytes = 120_u64 - + (latest.size * (latest.erasure.data_blocks + latest.erasure.parity_blocks) as i64 - / latest.erasure.data_blocks as i64) as u64 - + (historical.size * (historical.erasure.data_blocks + historical.erasure.parity_blocks) as i64 - / historical.erasure.data_blocks as i64) as u64; - assert_eq!(stat.bytes, expected_bytes); - } - - #[test] - fn test_rebalance_stats_update_batch_matches_repeated_updates() { - let mut latest = version_normal(); - latest.is_latest = true; - latest.size = 128; - - let mut historical = version_normal(); - historical.is_latest = false; - historical.size = 64; - - let mut repeated = RebalanceStats::default(); - repeated.update("bucket-a".to_string(), &latest); - repeated.update("bucket-a".to_string(), &historical); - - let mut batched = RebalanceStats::default(); - batched.update_batch("bucket-a".to_string(), &[&latest, &historical]); - - assert_eq!(batched.bucket, repeated.bucket); - assert_eq!(batched.object, repeated.object); - assert_eq!(batched.num_objects, repeated.num_objects); - assert_eq!(batched.num_versions, repeated.num_versions); - assert_eq!(batched.bytes, repeated.bytes); - } - - #[test] - fn test_rebalance_stats_update_ignores_invalid_data_blocks() { - let mut stat = RebalanceStats { - bucket: "bucket-a".to_string(), - object: "obj-previous".to_string(), - num_objects: 1, - num_versions: 2, - bytes: 77, - ..Default::default() - }; - - let mut invalid = FileInfo::new("object-invalid", 0, 2); - invalid.name = "object-invalid".to_string(); - invalid.is_latest = true; - invalid.size = 256; - invalid.deleted = false; - invalid.mod_time = Some(OffsetDateTime::UNIX_EPOCH); - invalid.version_id = None; - - stat.update("bucket-z".to_string(), &invalid); - - assert_eq!(stat.bucket, "bucket-z"); - assert_eq!(stat.object, "object-invalid"); - assert_eq!(stat.num_objects, 2); - assert_eq!(stat.num_versions, 3); - assert_eq!(stat.bytes, 77); - } - - #[test] - fn test_rebalance_goal_reached_requires_target_ratio() { - let init_free_space = 150_u64; - let init_capacity = 800_u64; - let goal = 0.35_f64; - - assert!(!rebalance_goal_reached(init_free_space, init_capacity, 0, goal)); - assert!(!rebalance_goal_reached(init_free_space, init_capacity, 129, goal)); - assert!(rebalance_goal_reached(init_free_space, init_capacity, 130, goal)); - } -} +mod rebalance_unit_tests; diff --git a/crates/ecstore/src/rebalance/control.rs b/crates/ecstore/src/rebalance/control.rs new file mode 100644 index 000000000..d0e60f8cd --- /dev/null +++ b/crates/ecstore/src/rebalance/control.rs @@ -0,0 +1,410 @@ +use super::meta::{ + clone_first_arc, clone_rebalance_pool_stats, defer_bucket_in_rebalance_queue, ensure_valid_rebalance_pool_index, + invalid_rebalance_pool_index_error, is_rebalance_conflicting_with_decommission, mark_rebalance_bucket_done, + merge_rebalance_meta, percent_free_ratio, rebalance_metadata_not_initialized_error, record_rebalance_cleanup_warning_in_meta, + resolve_next_rebalance_bucket, should_accept_rebalance_stats_update, should_pool_participate, stop_rebalance_meta_snapshot, +}; +use super::worker::{ + rebalance_meta_lock_error, resolve_load_rebalance_stats_update_result, resolve_rebalance_meta_load_result, + resolve_rebalance_meta_save_result, +}; +use super::{ + DiskStat, EVENT_REBALANCE_BUCKET, EVENT_REBALANCE_STATE, LOG_COMPONENT_ECSTORE, LOG_SUBSYSTEM_REBALANCE, REBAL_META_NAME, + RebalStatus, RebalanceInfo, RebalanceMeta, RebalanceStats, +}; +use crate::error::{Error, Result}; +use crate::object_api::ObjectOptions; +use crate::set_disk::get_lock_acquire_timeout; +use crate::storage_api_contracts::EcstoreObjectIO; +use crate::store::ECStore; +use rustfs_filemeta::FileInfo; +use rustfs_storage_api::{NamespaceLocking as StorageNamespaceLocking, StorageAdminApi}; +use std::sync::Arc; +use time::OffsetDateTime; +use tracing::{debug, info}; +use uuid::Uuid; + +impl ECStore { + pub(super) async fn save_rebalance_meta_with_merge( + &self, + pool: Arc, + local_snapshot: &RebalanceMeta, + stage: &str, + ) -> Result<()> + where + S: EcstoreObjectIO + StorageNamespaceLocking, + { + let ns_lock = pool.new_ns_lock(crate::disk::RUSTFS_META_BUCKET, REBAL_META_NAME).await?; + let _guard = ns_lock + .get_write_lock(get_lock_acquire_timeout()) + .await + .map_err(rebalance_meta_lock_error)?; + + let opts = ObjectOptions { + no_lock: true, + ..Default::default() + }; + let mut merged = RebalanceMeta::new(); + match merged.load_with_opts(pool.clone(), opts.clone()).await { + Ok(()) => { + merge_rebalance_meta(&mut merged, local_snapshot); + } + Err(Error::ConfigNotFound) => { + merged = local_snapshot.clone(); + } + Err(err) => return Err(Error::other(format!("rebalance meta load before save failed during {stage}: {err}"))), + } + + merged.save_with_opts(pool, opts).await + } + + #[tracing::instrument(skip_all)] + pub async fn load_rebalance_meta(&self) -> Result<()> { + let mut meta = RebalanceMeta::new(); + debug!( + event = EVENT_REBALANCE_STATE, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + state = "metadata_loading", + "Loading rebalance metadata" + ); + let pool = clone_first_arc(&self.pools, "rebalanceMeta: no pools available")?; + if resolve_rebalance_meta_load_result(meta.load(pool).await)? { + { + let mut rebalance_meta = self.rebalance_meta.write().await; + + *rebalance_meta = Some(meta); + + drop(rebalance_meta); + } + + resolve_load_rebalance_stats_update_result(self.update_rebalance_stats().await)?; + debug!( + event = EVENT_REBALANCE_STATE, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + state = "metadata_loaded", + "Loaded rebalance metadata" + ); + } else { + debug!( + event = EVENT_REBALANCE_STATE, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + state = "metadata_missing", + reason = "rebalance_not_started", + "Rebalance metadata not found" + ); + } + + Ok(()) + } + + #[tracing::instrument(skip_all)] + pub async fn update_rebalance_stats(&self) -> Result<()> { + let mut ok = false; + + let pool_stats = { + let rebalance_meta = self.rebalance_meta.read().await; + clone_rebalance_pool_stats(rebalance_meta.as_ref())? + }; + + debug!( + event = EVENT_REBALANCE_STATE, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + pool_count = pool_stats.len(), + "Refreshing rebalance stats snapshot" + ); + + for i in 0..self.pools.len() { + if pool_stats.get(i).is_none() { + let mut rebalance_meta = self.rebalance_meta.write().await; + debug!( + event = EVENT_REBALANCE_STATE, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + pool_index = i, + state = "pool_stat_missing", + "Adding missing rebalance pool stats entry" + ); + if let Some(meta) = rebalance_meta.as_mut() { + meta.pool_stats.push(RebalanceStats::default()); + } + ok = true; + drop(rebalance_meta); + } + } + + if ok { + debug!( + event = EVENT_REBALANCE_STATE, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + state = "metadata_saving", + "Saving rebalance metadata after stats refresh" + ); + + let rebalance_meta = self.rebalance_meta.read().await; + if let Some(meta) = rebalance_meta.as_ref() { + let pool = clone_first_arc(&self.pools, "update_rebalance_stats: no pools available")?; + resolve_rebalance_meta_save_result( + self.save_rebalance_meta_with_merge(pool, meta, "update_rebalance_stats") + .await, + "update_rebalance_stats", + )?; + } + } + + Ok(()) + } + + #[tracing::instrument(skip(self))] + pub async fn init_rebalance_meta(&self, bucktes: Vec) -> Result { + info!( + event = EVENT_REBALANCE_STATE, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + state = "initializing", + bucket_count = bucktes.len(), + "Initializing rebalance metadata" + ); + let si = StorageAdminApi::storage_info(self).await; + + let mut disk_stats = vec![DiskStat::default(); self.pools.len()]; + + let mut total_cap = 0; + let mut total_free = 0; + for disk in si.disks.iter() { + if disk.pool_index < 0 || disk_stats.len() <= disk.pool_index as usize { + continue; + } + + total_cap += disk.total_space; + total_free += disk.available_space; + + disk_stats[disk.pool_index as usize].total_space += disk.total_space; + disk_stats[disk.pool_index as usize].available_space += disk.available_space; + } + + let percent_free_goal = percent_free_ratio(total_free, total_cap); + + let mut pool_stats = Vec::with_capacity(self.pools.len()); + + let now = OffsetDateTime::now_utc(); + + for disk_stat in disk_stats.iter() { + let mut pool_stat = RebalanceStats { + init_free_space: disk_stat.available_space, + init_capacity: disk_stat.total_space, + buckets: bucktes.clone(), + rebalanced_buckets: Vec::with_capacity(bucktes.len()), + ..Default::default() + }; + + if should_pool_participate(disk_stat.available_space, disk_stat.total_space, percent_free_goal) { + pool_stat.participating = true; + pool_stat.info = RebalanceInfo { + start_time: Some(now), + status: RebalStatus::Started, + ..Default::default() + }; + } + + pool_stats.push(pool_stat); + } + + let meta = RebalanceMeta { + id: Uuid::new_v4().to_string(), + percent_free_goal, + pool_stats, + ..Default::default() + }; + + let pool = clone_first_arc(&self.pools, "init_rebalance_meta: no pools available")?; + resolve_rebalance_meta_save_result( + self.save_rebalance_meta_with_merge(pool, &meta, "init_rebalance_meta").await, + "init_rebalance_meta", + )?; + + info!( + event = EVENT_REBALANCE_STATE, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + state = "metadata_initialized", + bucket_count = bucktes.len(), + "Rebalance metadata initialized" + ); + + let id = meta.id.clone(); + + { + let mut rebalance_meta = self.rebalance_meta.write().await; + *rebalance_meta = Some(meta); + drop(rebalance_meta); + } + + Ok(id) + } + + #[tracing::instrument(skip(self, fi))] + pub async fn update_pool_stats(&self, pool_index: usize, bucket: String, fi: &FileInfo) -> Result<()> { + self.update_pool_stats_batch(pool_index, bucket, &[fi]).await + } + + #[tracing::instrument(skip(self, versions))] + pub async fn update_pool_stats_batch(&self, pool_index: usize, bucket: String, versions: &[&FileInfo]) -> Result<()> { + if versions.is_empty() { + return Ok(()); + } + + let mut rebalance_meta = self.rebalance_meta.write().await; + if let Some(meta) = rebalance_meta.as_mut() { + if !should_accept_rebalance_stats_update(meta, pool_index) { + return Ok(()); + } + + if let Some(pool_stat) = meta.pool_stats.get_mut(pool_index) { + pool_stat.update_batch(bucket, versions); + } + } + + Ok(()) + } + + #[tracing::instrument(skip(self))] + pub async fn next_rebal_bucket(&self, pool_index: usize) -> Result> { + let rebalance_meta = self.rebalance_meta.read().await; + debug!( + event = EVENT_REBALANCE_BUCKET, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + pool_index, + has_meta = rebalance_meta.is_some(), + state = "next_bucket_lookup", + "Rebalance next bucket lookup" + ); + resolve_next_rebalance_bucket(rebalance_meta.as_ref(), pool_index) + } + + #[tracing::instrument(skip(self))] + pub async fn bucket_rebalance_done(&self, pool_index: usize, bucket: String) -> Result<()> { + let mut rebalance_meta = self.rebalance_meta.write().await; + mark_rebalance_bucket_done(rebalance_meta.as_mut(), pool_index, &bucket) + } + + pub(super) async fn record_rebalance_cleanup_warning( + &self, + pool_index: usize, + bucket: &str, + object: &str, + message: String, + ) -> Result<()> { + let mut rebalance_meta = self.rebalance_meta.write().await; + record_rebalance_cleanup_warning_in_meta( + rebalance_meta.as_mut(), + pool_index, + bucket, + object, + message, + OffsetDateTime::now_utc(), + ) + } + + pub(super) async fn defer_rebalance_bucket(&self, pool_index: usize, bucket: String, last_error: String) -> Result<()> { + let mut rebalance_meta = self.rebalance_meta.write().await; + let Some(meta) = rebalance_meta.as_mut() else { + return Err(rebalance_metadata_not_initialized_error("defer rebalance bucket")); + }; + let pool_count = meta.pool_stats.len(); + ensure_valid_rebalance_pool_index(pool_count, pool_index)?; + let Some(pool_stat) = meta.pool_stats.get_mut(pool_index) else { + return Err(invalid_rebalance_pool_index_error(pool_index, pool_count)); + }; + + defer_bucket_in_rebalance_queue(pool_stat, &bucket)?; + pool_stat.info.last_error = Some(last_error); + meta.last_refreshed_at = Some(OffsetDateTime::now_utc()); + Ok(()) + } + + pub async fn is_rebalance_started(&self) -> bool { + let rebalance_meta = self.rebalance_meta.read().await; + if let Some(meta) = rebalance_meta.as_ref() { + meta.pool_stats.iter().enumerate().for_each(|(i, v)| { + debug!( + event = EVENT_REBALANCE_STATE, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + pool_index = i, + participating = v.participating, + status = ?v.info.status, + state = "status_inspected", + "Rebalance status inspected" + ); + }); + + let started = is_rebalance_conflicting_with_decommission(meta); + if started { + debug!( + event = EVENT_REBALANCE_STATE, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + state = "running", + "Rebalance is running" + ); + return true; + } + } + + debug!( + event = EVENT_REBALANCE_STATE, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + state = "not_running", + "Rebalance is not running" + ); + false + } + + pub async fn is_rebalance_conflicting_with_decommission(&self) -> bool { + let rebalance_meta = self.rebalance_meta.read().await; + rebalance_meta + .as_ref() + .is_some_and(is_rebalance_conflicting_with_decommission) + } + + pub async fn is_pool_rebalancing(&self, pool_index: usize) -> bool { + let rebalance_meta = self.rebalance_meta.read().await; + if let Some(ref meta) = *rebalance_meta { + if meta.stopped_at.is_some() { + return false; + } + + if let Some(pool_stat) = meta.pool_stats.get(pool_index) { + return pool_stat.participating && pool_stat.info.status == RebalStatus::Started; + } + } + + false + } + + #[tracing::instrument(skip(self))] + pub async fn stop_rebalance(self: &Arc) -> Result<()> { + let meta_to_save = { + let mut rebalance_meta = self.rebalance_meta.write().await; + stop_rebalance_meta_snapshot(rebalance_meta.as_mut(), OffsetDateTime::now_utc()) + }; + + if let Some(meta_to_save) = meta_to_save { + let pool = clone_first_arc(self.pools.as_slice(), "stop_rebalance: no pools available")?; + resolve_rebalance_meta_save_result( + self.save_rebalance_meta_with_merge(pool, &meta_to_save, "stop_rebalance") + .await, + "stop_rebalance", + )?; + } + + Ok(()) + } +} diff --git a/crates/ecstore/src/rebalance/entry.rs b/crates/ecstore/src/rebalance/entry.rs new file mode 100644 index 000000000..c903fcd60 --- /dev/null +++ b/crates/ecstore/src/rebalance/entry.rs @@ -0,0 +1,513 @@ +// Copyright 2024 RustFS Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use super::meta::{ + clone_arc_by_index, ensure_valid_rebalance_pool_index, invalid_rebalance_pool_index_error, + rebalance_metadata_not_initialized_error, should_ignore_rebalance_data_usage_cache, +}; +use super::migration::migrate_entry_version; +use super::worker::{ + RebalanceEntryTask, load_rebalance_bucket_configs, rebalance_max_attempts, resolve_rebalance_bucket_error, + resolve_rebalance_entry_cleanup_delete_result, resolve_rebalance_file_info_versions_result, + resolve_rebalance_migrate_result_error, resolve_rebalance_stats_update_result, resolve_rebalance_worker_result, + run_rebalance_listing_with_retry, should_cleanup_rebalance_source_entry, should_count_rebalance_version_complete, + should_defer_rebalance_entry_failure, should_skip_rebalance_delete_marker, wait_rebalance_entry_tasks, + with_rebalance_entry_context, +}; +use super::{ + EVENT_REBALANCE_BUCKET, EVENT_REBALANCE_ENTRY, EVENT_REBALANCE_STATE, LOG_COMPONENT_ECSTORE, LOG_SUBSYSTEM_REBALANCE, + REBALANCE_DEFERRED_ENTRY_ERROR_PREFIX, RebalanceBucketConfigs, RebalanceBucketOutcome, RebalanceEntryOutcome, +}; +use crate::data_movement; +use crate::error::{Error, Result}; +use crate::object_api::{GetObjectReader, ObjectOptions}; +use crate::pools::ListCallback; +use crate::set_disk::SetDisks; +use crate::store::ECStore; +use rustfs_filemeta::MetaCacheEntry; +use rustfs_storage_api::ObjectOperations as _; +use rustfs_utils::path::encode_dir_object; +use std::sync::Arc; +use time::OffsetDateTime; +use tokio_util::sync::CancellationToken; +use tracing::{debug, error, warn}; + +impl ECStore { + #[allow(unused_assignments)] + #[tracing::instrument(skip(self, set))] + async fn rebalance_entry( + self: Arc, + bucket: String, + pool_index: usize, + entry: MetaCacheEntry, + set: Arc, + bucket_configs: Arc, + // wk: Arc, + ) -> Result { + debug!( + event = EVENT_REBALANCE_ENTRY, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + bucket = %bucket, + object = %entry.name, + pool_index, + state = "started", + "Starting rebalance entry" + ); + + // defer!(|| async { + // warn!("rebalance_entry: defer give worker start"); + // wk.give().await; + // warn!("rebalance_entry: defer give worker done"); + // }); + + if entry.is_dir() { + debug!( + event = EVENT_REBALANCE_ENTRY, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + bucket = %bucket, + object = %entry.name, + pool_index, + state = "skipped", + reason = "directory_entry", + "Skipped rebalance entry" + ); + return Ok(RebalanceEntryOutcome::Completed); + } + + if self.check_if_rebalance_done(pool_index).await { + debug!( + event = EVENT_REBALANCE_ENTRY, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + pool_index, + bucket = %bucket, + object = %entry.name, + state = "skipped", + reason = "pool_completed", + "Skipped rebalance entry" + ); + return Ok(RebalanceEntryOutcome::Completed); + } + + let mut fivs = + resolve_rebalance_file_info_versions_result(entry.file_info_versions(&bucket), bucket.as_str(), entry.name.as_str())?; + + fivs.versions + .sort_by_key(|v| (v.mod_time.is_none(), std::cmp::Reverse(v.mod_time))); + + let mut rebalanced: usize = 0; + let mut expired: usize = 0; + let mut stats_updates = Vec::with_capacity(fivs.versions.len()); + for version in fivs.versions.iter() { + if crate::pools::should_skip_lifecycle_for_data_movement( + self.clone(), + &bucket, + version, + bucket_configs.lifecycle_config.as_ref(), + bucket_configs.lock_retention.clone(), + bucket_configs.replication_config.clone(), + true, + &crate::bucket::lifecycle::bucket_lifecycle_audit::LcEventSrc::Rebal, + ) + .await + { + expired += 1; + debug!( + event = EVENT_REBALANCE_ENTRY, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + pool_index, + bucket = %bucket, + object = %version.name, + state = "skipped", + reason = "expired_by_lifecycle", + "Skipped rebalance version" + ); + continue; + } + + let remaining_versions = fivs.versions.len() - expired; + if should_skip_rebalance_delete_marker(version, remaining_versions, bucket_configs.replication_config.is_some()) { + rebalanced += 1; + debug!( + event = EVENT_REBALANCE_ENTRY, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + pool_index, + bucket = %bucket, + object = %version.name, + state = "skipped", + reason = "last_delete_marker_without_replication", + "Skipped rebalance version" + ); + continue; + } + + let version_id = version.version_id.map(|v| v.to_string()); + let mut transfer = |src_pool_idx: usize, bucket: String, rd: GetObjectReader| { + let store = self.clone(); + async move { store.rebalance_object(src_pool_idx, bucket, rd).await } + }; + let result = migrate_entry_version( + set.as_ref(), + bucket.clone(), + pool_index, + version, + version_id.clone(), + rebalance_max_attempts(), + should_ignore_rebalance_data_usage_cache(bucket.as_str()), + &mut transfer, + ) + .await; + + if result.ignored { + if should_count_rebalance_version_complete(&result) { + rebalanced += 1; + } + debug!( + event = EVENT_REBALANCE_ENTRY, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + pool_index, + bucket = %bucket, + object = %version.name, + state = "skipped", + reason = "already_deleted", + "Skipped rebalance version" + ); + continue; + } + + if result.failed { + let err = resolve_rebalance_migrate_result_error( + result.error, + pool_index, + bucket.as_str(), + version.name.as_str(), + version_id.as_deref(), + ); + error!( + "rebalance_entry {} Error rebalancing entry {}/{:?}: {:?}", + &bucket, &version.name, &version.version_id, err + ); + if should_defer_rebalance_entry_failure(&err) { + let deferred_error = format!("{REBALANCE_DEFERRED_ENTRY_ERROR_PREFIX} {err}"); + warn!( + event = EVENT_REBALANCE_ENTRY, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + pool_index, + bucket = %bucket, + object = %version.name, + state = "deferred", + error = %err, + "Deferred rebalance entry after transient migration failure" + ); + if let Err(stats_err) = self.update_rebalance_last_error(pool_index, deferred_error.clone()).await { + error!( + "rebalance_entry {} failed to record deferred transient failure for {}: {}", + &bucket, &entry.name, stats_err + ); + } + return Ok(RebalanceEntryOutcome::Deferred { + last_error: deferred_error, + }); + } + 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 + .update_pool_stats_batch(pool_index, bucket.clone(), stats_updates.as_slice()) + .await + { + error!( + "rebalance_entry {} failed to update stats before returning migration error for {}: {}", + &bucket, &entry.name, stats_err + ); + } + + return Err(entry_err); + } + + stats_updates.push(version); + if should_count_rebalance_version_complete(&result) { + rebalanced += 1; + } + } + + resolve_rebalance_stats_update_result( + self.update_pool_stats_batch(pool_index, bucket.clone(), stats_updates.as_slice()) + .await, + pool_index, + bucket.as_str(), + entry.name.as_str(), + )?; + + if should_cleanup_rebalance_source_entry(rebalanced, fivs.versions.len()) { + let cleanup_warning = resolve_rebalance_entry_cleanup_delete_result( + set.delete_object( + bucket.as_str(), + &encode_dir_object(&entry.name), + ObjectOptions { + delete_prefix: true, + delete_prefix_object: true, + + ..Default::default() + }, + ) + .await, + bucket.as_str(), + entry.name.as_str(), + )?; + if let Some(message) = cleanup_warning { + warn!( + event = EVENT_REBALANCE_ENTRY, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + pool_index, + bucket = %bucket, + object = %entry.name, + stage = "cleanup_source", + cleanup_status = "failed_ignored", + error = %message, + "Ignored rebalance source cleanup failure" + ); + if let Err(err) = self + .record_rebalance_cleanup_warning(pool_index, bucket.as_str(), entry.name.as_str(), message) + .await + { + error!( + event = EVENT_REBALANCE_ENTRY, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + pool_index, + bucket = %bucket, + object = %entry.name, + stage = "cleanup_source", + error = ?err, + "Failed to record rebalance source cleanup warning" + ); + } + } else { + debug!( + event = EVENT_REBALANCE_ENTRY, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + pool_index, + bucket = %bucket, + object = %entry.name, + state = "source_deleted", + "Deleted rebalance source entry" + ); + } + } + + Ok(RebalanceEntryOutcome::Completed) + } + + #[tracing::instrument(skip(self, rd))] + async fn rebalance_object(self: Arc, pool_idx: usize, bucket: String, rd: GetObjectReader) -> Result<()> { + data_movement::migrate_object(self, pool_idx, bucket, rd, "rebalance_object").await + } + + async fn update_rebalance_last_error(&self, pool_idx: usize, message: String) -> Result<()> { + let mut rebalance_meta = self.rebalance_meta.write().await; + let Some(meta) = rebalance_meta.as_mut() else { + return Err(rebalance_metadata_not_initialized_error("record rebalance last error")); + }; + let pool_count = meta.pool_stats.len(); + ensure_valid_rebalance_pool_index(pool_count, pool_idx)?; + let Some(pool_stat) = meta.pool_stats.get_mut(pool_idx) else { + return Err(invalid_rebalance_pool_index_error(pool_idx, pool_count)); + }; + + pool_stat.info.last_error = Some(message); + meta.last_refreshed_at = Some(OffsetDateTime::now_utc()); + Ok(()) + } + + #[tracing::instrument(skip(self, rx))] + pub(super) async fn rebalance_bucket( + self: &Arc, + rx: CancellationToken, + bucket: String, + pool_index: usize, + ) -> Result { + ensure_valid_rebalance_pool_index(self.pools.len(), pool_index)?; + + // Placeholder for actual bucket rebalance logic + debug!( + event = EVENT_REBALANCE_BUCKET, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + pool_index, + bucket = %bucket, + state = "entry_scan_started", + "Rebalance bucket entry scan started" + ); + + // TODO: other config + // if bucket != RUSTFS_META_BUCKET{ + + // } + + let pool = clone_arc_by_index(self.pools.as_slice(), pool_index, "invalid rebalance pool index")?; + let bucket_configs = Arc::new(load_rebalance_bucket_configs(&bucket).await?); + + let mut jobs = Vec::new(); + let entry_error = Arc::new(tokio::sync::Mutex::new(None::)); + let entry_workers = Arc::new(tokio::sync::Semaphore::new(pool.disk_set.len().max(1))); + + for (set_idx, set) in pool.disk_set.iter().enumerate() { + let entry_tasks = Arc::new(tokio::sync::Mutex::new(Vec::::new())); + let rebalance_entry: ListCallback = Arc::new({ + let this = Arc::clone(self); + let bucket = bucket.clone(); + let entry_error = entry_error.clone(); + let callback_rx = rx.clone(); + let set = set.clone(); + let bucket_configs = bucket_configs.clone(); + let entry_tasks = entry_tasks.clone(); + let entry_workers = entry_workers.clone(); + move |entry: MetaCacheEntry| { + let this = this.clone(); + let bucket = bucket.clone(); + let entry_error = entry_error.clone(); + let callback_rx = callback_rx.clone(); + let set = set.clone(); + let bucket_configs = bucket_configs.clone(); + let entry_tasks = entry_tasks.clone(); + let entry_workers = entry_workers.clone(); + Box::pin(async move { + if callback_rx.is_cancelled() { + return; + } + if entry_error.lock().await.is_some() { + return; + } + + let permit = tokio::select! { + _ = callback_rx.cancelled() => return, + permit = entry_workers.clone().acquire_owned() => match permit { + Ok(permit) => permit, + Err(err) => { + error!("rebalance_entry: worker semaphore closed: {err}"); + return; + } + }, + }; + + if entry_error.lock().await.is_some() { + return; + } + + let task = tokio::spawn(async move { + let _permit = permit; + debug!( + event = EVENT_REBALANCE_ENTRY, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + set_index = set_idx, + state = "task_started", + "Started rebalance entry task" + ); + let result = this.rebalance_entry(bucket, pool_index, entry, set, bucket_configs).await; + if let Err(err) = &result { + error!("rebalance_entry: rebalance entry failed: {err}"); + let mut first_err = entry_error.lock().await; + if first_err.is_none() { + *first_err = Some(err.clone()); + callback_rx.cancel(); + } + } + debug!( + event = EVENT_REBALANCE_ENTRY, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + set_index = set_idx, + state = "task_completed", + "Completed rebalance entry task" + ); + result + }); + + entry_tasks.lock().await.push(task); + }) + } + }); + + let set = set.clone(); + let rx = rx.clone(); + let bucket = bucket.clone(); + let entry_tasks = entry_tasks.clone(); + + let job = tokio::spawn(async move { + let list_result = + run_rebalance_listing_with_retry(set, rx, bucket.clone(), rebalance_entry, set_idx, rebalance_max_attempts()) + .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 { + error!("Rebalance worker {} error: {}", set_idx, err); + } else { + debug!( + event = EVENT_REBALANCE_STATE, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + set_index = set_idx, + state = "worker_completed", + "Completed rebalance worker" + ); + } + result + }); + + jobs.push((set_idx, job)); + } + + let mut worker_error: Option = None; + let mut deferred_error: Option = None; + for (set_idx, job) in jobs { + match resolve_rebalance_worker_result(set_idx, job.await) { + Ok(Some(last_error)) if deferred_error.is_none() => { + deferred_error = Some(last_error); + } + Ok(_) => {} + Err(err) if worker_error.is_none() => { + worker_error = Some(err); + } + Err(_) => {} + } + } + let entry_error = entry_error.lock().await.clone(); + resolve_rebalance_bucket_error(entry_error, worker_error)?; + if let Some(last_error) = deferred_error { + return Ok(RebalanceBucketOutcome::Deferred { last_error }); + } + + debug!( + event = EVENT_REBALANCE_BUCKET, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + pool_index, + bucket = %bucket, + state = "completed", + "Finished rebalance bucket" + ); + Ok(RebalanceBucketOutcome::Completed) + } +} diff --git a/crates/ecstore/src/rebalance/meta.rs b/crates/ecstore/src/rebalance/meta.rs index 727d497e3..a707c840e 100644 --- a/crates/ecstore/src/rebalance/meta.rs +++ b/crates/ecstore/src/rebalance/meta.rs @@ -1,13 +1,198 @@ use super::{ - EVENT_REBALANCE_BUCKET, Error, LOG_COMPONENT_ECSTORE, LOG_SUBSYSTEM_REBALANCE, REBALANCE_DEFERRED_ENTRY_ERROR_PREFIX, - RebalSaveOpt, RebalStatus, RebalanceCleanupWarnings, RebalanceMeta, RebalanceStats, Result, + EVENT_REBALANCE_BUCKET, EVENT_REBALANCE_STATE, Error, GetObjectReader, LOG_COMPONENT_ECSTORE, LOG_SUBSYSTEM_REBALANCE, + ObjectInfo, ObjectOptions, PutObjReader, REBAL_META_FMT, REBAL_META_NAME, REBAL_META_VER, + REBALANCE_DEFERRED_ENTRY_ERROR_PREFIX, RebalSaveOpt, RebalStatus, RebalanceCleanupWarnings, RebalanceMeta, RebalanceStats, + Result, }; +use crate::config::com::{read_config_with_metadata, save_config_with_opts}; use crate::error::is_err_operation_canceled; -use std::collections::HashSet; -use std::sync::Arc; +use http::HeaderMap; +use rustfs_filemeta::FileInfo; +use rustfs_storage_api::{HTTPRangeSpec, ObjectIO}; +use std::{collections::HashSet, fmt, io::Cursor, sync::Arc}; use time::OffsetDateTime; use tracing::{debug, info}; +impl RebalanceStats { + pub fn update(&mut self, bucket: String, fi: &FileInfo) { + if fi.is_latest { + self.num_objects += 1; + } + + self.num_versions += 1; + let on_disk_size = if fi.deleted || fi.erasure.data_blocks == 0 || fi.size <= 0 { + 0 + } else { + let data_blocks = fi.erasure.data_blocks as i64; + let total_blocks = fi.erasure.data_blocks.saturating_add(fi.erasure.parity_blocks) as i64; + fi.size + .saturating_mul(total_blocks) + .checked_div(data_blocks) + .unwrap_or(0) + .max(0) as u64 + }; + self.bytes = self.bytes.saturating_add(on_disk_size); + self.bucket = bucket; + self.object = fi.name.clone(); + } + + pub fn update_batch(&mut self, bucket: String, versions: &[&FileInfo]) { + for version in versions { + self.update(bucket.clone(), version); + } + } +} + +impl fmt::Display for RebalStatus { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let status = match self { + RebalStatus::None => "None", + RebalStatus::Started => "Started", + RebalStatus::Completed => "Completed", + RebalStatus::Stopped => "Stopped", + RebalStatus::Failed => "Failed", + }; + write!(f, "{status}") + } +} + +impl From for RebalStatus { + fn from(value: u8) -> Self { + match value { + 1 => RebalStatus::Started, + 2 => RebalStatus::Completed, + 3 => RebalStatus::Stopped, + 4 => RebalStatus::Failed, + _ => RebalStatus::None, + } + } +} + +impl RebalanceMeta { + pub fn new() -> Self { + Self::default() + } + pub async fn load(&mut self, store: Arc) -> Result<()> + where + S: ObjectIO< + Error = Error, + RangeSpec = HTTPRangeSpec, + HeaderMap = HeaderMap, + ObjectOptions = ObjectOptions, + ObjectInfo = ObjectInfo, + GetObjectReader = GetObjectReader, + PutObjectReader = PutObjReader, + >, + { + self.load_with_opts(store, ObjectOptions::default()).await + } + + pub async fn load_with_opts(&mut self, store: Arc, opts: ObjectOptions) -> Result<()> + where + S: ObjectIO< + Error = Error, + RangeSpec = HTTPRangeSpec, + HeaderMap = HeaderMap, + ObjectOptions = ObjectOptions, + ObjectInfo = ObjectInfo, + GetObjectReader = GetObjectReader, + PutObjectReader = PutObjReader, + >, + { + let (data, _) = read_config_with_metadata(store, REBAL_META_NAME, &opts).await?; + if data.is_empty() { + debug!( + event = EVENT_REBALANCE_STATE, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + state = "metadata_empty", + "Rebalance metadata is empty" + ); + return Ok(()); + } + if data.len() <= 4 { + return Err(rebalance_meta_load_no_data_error()); + } + + // Read header + match u16::from_le_bytes([data[0], data[1]]) { + REBAL_META_FMT => {} + fmt => return Err(rebalance_meta_load_unknown_format_error(fmt)), + } + match u16::from_le_bytes([data[2], data[3]]) { + REBAL_META_VER => {} + ver => return Err(rebalance_meta_load_unknown_version_error(ver)), + } + + let meta: Self = rmp_serde::from_read(Cursor::new(&data[4..]))?; + *self = meta; + + self.last_refreshed_at = Some(OffsetDateTime::now_utc()); + + debug!( + event = EVENT_REBALANCE_STATE, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + state = "metadata_loaded", + "Loaded rebalance metadata" + ); + Ok(()) + } + + pub async fn save(&self, store: Arc) -> Result<()> + where + S: ObjectIO< + Error = Error, + RangeSpec = HTTPRangeSpec, + HeaderMap = HeaderMap, + ObjectOptions = ObjectOptions, + ObjectInfo = ObjectInfo, + GetObjectReader = GetObjectReader, + PutObjectReader = PutObjReader, + >, + { + self.save_with_opts(store, ObjectOptions::default()).await + } + + pub async fn save_with_opts(&self, store: Arc, opts: ObjectOptions) -> Result<()> + where + S: ObjectIO< + Error = Error, + RangeSpec = HTTPRangeSpec, + HeaderMap = HeaderMap, + ObjectOptions = ObjectOptions, + ObjectInfo = ObjectInfo, + GetObjectReader = GetObjectReader, + PutObjectReader = PutObjReader, + >, + { + if self.pool_stats.is_empty() { + debug!( + event = EVENT_REBALANCE_STATE, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + state = "metadata_save_skipped", + reason = "no_pool_stats", + "Skipped rebalance metadata save" + ); + return Ok(()); + } + + let mut data = Vec::new(); + + // Initialize the header + data.extend(&REBAL_META_FMT.to_le_bytes()); + data.extend(&REBAL_META_VER.to_le_bytes()); + + let msg = rmp_serde::to_vec(self)?; + data.extend(msg); + + save_config_with_opts(store, REBAL_META_NAME, data, &opts).await?; + + Ok(()) + } +} + pub(super) fn is_rebalance_pool_started(pool_stat: &RebalanceStats) -> bool { pool_stat.participating && pool_stat.info.status == RebalStatus::Started } diff --git a/crates/ecstore/src/rebalance/rebalance_unit_tests.rs b/crates/ecstore/src/rebalance/rebalance_unit_tests.rs new file mode 100644 index 000000000..3dfbf44d6 --- /dev/null +++ b/crates/ecstore/src/rebalance/rebalance_unit_tests.rs @@ -0,0 +1,3230 @@ +// Copyright 2024 RustFS Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use super::REBALANCE_DEFERRED_ENTRY_ERROR_PREFIX; +use super::meta::{ + 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, + complete_rebalance_pools_at_goal, complete_rebalance_pools_with_empty_queue, defer_bucket_in_rebalance_queue, + ensure_rebalance_not_decommissioning, ensure_valid_rebalance_pool_index, first_rebalance_bucket, + has_deferred_rebalance_error, is_rebalance_actively_running, is_rebalance_conflicting_with_decommission, + is_rebalance_in_progress, is_rebalance_stopped_terminal_event, mark_rebalance_bucket_done, merge_rebalance_bucket_lists, + merge_rebalance_meta, next_rebal_bucket_from_stat, percent_free_ratio, rebalance_goal_reached, + rebalance_meta_load_no_data_error, rebalance_meta_load_unknown_format_error, rebalance_meta_load_unknown_version_error, + record_rebalance_cleanup_warning_in_meta, remove_rebalanced_buckets_from_queue, resolve_next_rebalance_bucket, + resolve_rebalance_participants, should_accept_rebalance_stats_update, should_ignore_rebalance_data_usage_cache, + should_pool_participate, should_preserve_rebalance_stopped_state, should_skip_start_rebalance, stop_rebalance_meta_snapshot, + stop_rebalance_state, take_bucket_from_rebalance_queue, validate_start_rebalance_state, +}; +use super::migration::{ + MigrationBackend, MigrationVersionResult, migrate_entry_version, migrate_entry_version_with_retry_wait, + rebalance_delete_marker_opts, +}; +use super::worker::{ + ensure_rebalance_listing_disks_available, is_transient_rebalance_error, load_rebalance_bucket_configs, + parse_rebalance_max_attempts, rebalance_listing_retry_delay, rebalance_migration_retry_delay, + resolve_load_rebalance_stats_update_result, resolve_rebalance_bucket_error, resolve_rebalance_bucket_result, + resolve_rebalance_entry_cleanup_delete_result, 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, should_defer_rebalance_entry_failure, + should_retry_rebalance_listing, should_skip_rebalance_delete_marker, wait_rebalance_entry_tasks, + wait_rebalance_listing_retry, with_rebalance_entry_context, +}; +use super::{ + GetObjectReader, ObjectInfo, ObjectOptions, RebalSaveOpt, RebalStatus, RebalanceBucketOutcome, RebalanceCleanupWarnings, + RebalanceEntryOutcome, RebalanceInfo, RebalanceMeta, RebalanceStats, +}; +use crate::data_movement; +use crate::data_usage::DATA_USAGE_CACHE_NAME; +use crate::disk::RUSTFS_META_BUCKET; +use crate::disk::error::DiskError; +use crate::error::{Error, Result}; +use rustfs_filemeta::FileInfo; +use rustfs_filemeta::TRANSITION_COMPLETE; +use rustfs_rio::Index; +use rustfs_storage_api::HTTPRangeSpec; +use s3s::dto::ReplicationConfiguration; +use serde::Serialize; +use std::io::Cursor; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::atomic::{AtomicUsize, Ordering}; +use time::OffsetDateTime; +use tokio::sync::mpsc; +use tokio::time::Duration; +use tokio_util::sync::CancellationToken; + +#[derive(Debug, Default, Serialize)] +struct LegacyRebalanceStats { + #[serde(rename = "ifs")] + init_free_space: u64, + #[serde(rename = "ic")] + init_capacity: u64, + #[serde(rename = "bus")] + buckets: Vec, + #[serde(rename = "rbs")] + rebalanced_buckets: Vec, + #[serde(rename = "bu")] + bucket: String, + #[serde(rename = "ob")] + object: String, + #[serde(rename = "no")] + num_objects: u64, + #[serde(rename = "nv")] + num_versions: u64, + #[serde(rename = "bs")] + bytes: u64, + #[serde(rename = "par")] + participating: bool, + #[serde(rename = "inf")] + info: RebalanceInfo, +} + +#[derive(Debug, Default, Serialize)] +struct LegacyRebalanceMeta { + #[serde(rename = "stopTs")] + stopped_at: Option, + #[serde(rename = "id")] + id: String, + #[serde(rename = "pf")] + percent_free_goal: f64, + #[serde(rename = "rss")] + pool_stats: Vec, +} + +struct MigrationBackendSpy { + get_object_reader: Mutex>>, + delete_object: Mutex>>, + move_remote: Mutex>>, + get_calls: AtomicUsize, + delete_calls: AtomicUsize, + move_remote_calls: AtomicUsize, +} + +impl MigrationBackendSpy { + fn new( + get_object_reader: Option>, + delete_object: Option>, + move_remote: Option>, + ) -> Self { + Self { + get_object_reader: Mutex::new(get_object_reader), + delete_object: Mutex::new(delete_object), + move_remote: Mutex::new(move_remote), + get_calls: AtomicUsize::new(0), + delete_calls: AtomicUsize::new(0), + move_remote_calls: AtomicUsize::new(0), + } + } + + fn get_calls(&self) -> usize { + self.get_calls.load(Ordering::SeqCst) + } + + fn delete_calls(&self) -> usize { + self.delete_calls.load(Ordering::SeqCst) + } + + fn move_remote_calls(&self) -> usize { + self.move_remote_calls.load(Ordering::SeqCst) + } + + fn make_reader() -> GetObjectReader { + GetObjectReader { + stream: Box::new(Cursor::new(vec![0_u8; 3])), + object_info: ObjectInfo::default(), + } + } +} + +#[async_trait::async_trait] +impl MigrationBackend for MigrationBackendSpy { + async fn get_object_reader_for_migration( + &self, + _bucket: &str, + _object: &str, + _range: Option, + _h: http::HeaderMap, + _opts: &ObjectOptions, + ) -> Result { + self.get_calls.fetch_add(1, Ordering::SeqCst); + if let Some(result) = self.get_object_reader.lock().unwrap().take() { + return result; + } + + Ok(Self::make_reader()) + } + + async fn delete_object_for_migration(&self, _bucket: &str, _object: &str, _opts: ObjectOptions) -> Result { + self.delete_calls.fetch_add(1, Ordering::SeqCst); + if let Some(result) = self.delete_object.lock().unwrap().take() { + return result; + } + + Ok(ObjectInfo::default()) + } + + async fn move_remote_version_for_migration( + &self, + _bucket: &str, + _object: &str, + _fi: &FileInfo, + _opts: &ObjectOptions, + ) -> Result<()> { + self.move_remote_calls.fetch_add(1, Ordering::SeqCst); + if let Some(result) = self.move_remote.lock().unwrap().take() { + return result; + } + + Ok(()) + } +} + +fn version_deleted() -> FileInfo { + let mut version = FileInfo::new("object.bin", 4, 2); + version.name = "object.bin".to_string(); + version.deleted = true; + version +} + +fn version_normal() -> FileInfo { + let mut version = FileInfo::new("object.bin", 4, 2); + version.name = "object.bin".to_string(); + version.size = 64; + version +} + +fn version_remote() -> FileInfo { + let mut version = FileInfo::new("object.bin", 4, 2); + version.name = "object.bin".to_string(); + version.transition_status = TRANSITION_COMPLETE.to_string(); + version +} + +#[test] +fn test_rebalance_delete_marker_opts_preserves_replication_state() { + let mod_time = OffsetDateTime::now_utc(); + let version = FileInfo { + mod_time: Some(mod_time), + replication_state_internal: Some(rustfs_filemeta::ReplicationState { + replica_status: rustfs_filemeta::ReplicationStatusType::Replica, + delete_marker: true, + replicate_decision_str: "existing".to_string(), + ..Default::default() + }), + ..version_deleted() + }; + + let opts = rebalance_delete_marker_opts(&version, Some("version-id".to_string()), 7); + let replication = opts.delete_replication.expect("replication state should be preserved"); + + assert!(opts.versioned); + assert!(opts.data_movement); + assert!(opts.delete_marker); + assert!(opts.skip_decommissioned); + assert_eq!(opts.src_pool_idx, 7); + assert_eq!(opts.version_id.as_deref(), Some("version-id")); + assert_eq!(opts.mod_time, Some(mod_time)); + assert_eq!(replication.replica_status, rustfs_filemeta::ReplicationStatusType::Replica); + assert!(replication.delete_marker); + assert_eq!(replication.replicate_decision_str, "existing"); +} + +#[tokio::test] +async fn test_migrate_entry_version_remote_version_is_moved_without_transfer() { + let backend = MigrationBackendSpy::new(None, Some(Ok(ObjectInfo::default())), Some(Ok(()))); + let version = version_remote(); + let transfer_count = Arc::new(AtomicUsize::new(0)); + let mut transfer = { + let transfer_count = transfer_count.clone(); + move |_, _, _| { + let transfer_count = transfer_count.clone(); + async move { + transfer_count.fetch_add(1, Ordering::SeqCst); + Ok(()) + } + } + }; + + let result = migrate_entry_version( + &backend, + "bucket".to_string(), + 0, + &version, + version.version_id.map(|v| v.to_string()), + 3, + false, + &mut transfer, + ) + .await; + + assert!(!result.ignored); + assert!(!result.cleanup_ignored); + assert!(result.moved); + assert!(!result.failed); + assert!(result.error.is_none()); + assert_eq!(transfer_count.load(Ordering::SeqCst), 0); + assert_eq!(backend.move_remote_calls(), 1); + assert_eq!(backend.get_calls(), 0); + assert_eq!(backend.delete_calls(), 0); +} + +#[tokio::test] +async fn test_migrate_entry_version_remote_not_found_is_cleanup_ignored() { + let backend = MigrationBackendSpy::new( + None, + Some(Ok(ObjectInfo::default())), + Some(Err(Error::ObjectNotFound("bucket".to_string(), "object.bin".to_string()))), + ); + let version = version_remote(); + let transfer_count = Arc::new(AtomicUsize::new(0)); + let mut transfer = { + let transfer_count = transfer_count.clone(); + move |_, _, _| { + let transfer_count = transfer_count.clone(); + async move { + transfer_count.fetch_add(1, Ordering::SeqCst); + Ok(()) + } + } + }; + + let result = migrate_entry_version( + &backend, + "bucket".to_string(), + 0, + &version, + version.version_id.map(|v| v.to_string()), + 3, + false, + &mut transfer, + ) + .await; + + assert!(result.ignored); + assert!(result.cleanup_ignored); + assert!(!result.moved); + assert!(!result.failed); + assert!(result.error.is_none()); + assert_eq!(transfer_count.load(Ordering::SeqCst), 0); + assert_eq!(backend.move_remote_calls(), 1); + assert_eq!(backend.get_calls(), 0); + assert_eq!(backend.delete_calls(), 0); +} + +#[tokio::test] +async fn test_migrate_entry_version_remote_overwrite_is_not_ignored() { + let backend = MigrationBackendSpy::new( + None, + None, + Some(Err(Error::DataMovementOverwriteErr( + "bucket".to_string(), + "object.bin".to_string(), + "vid-1".to_string(), + ))), + ); + let version = version_remote(); + let mut transfer = |_, _, _| async move { Ok(()) }; + + let result = migrate_entry_version( + &backend, + "bucket".to_string(), + 0, + &version, + Some("vid-1".to_string()), + 3, + false, + &mut transfer, + ) + .await; + + assert!(result.failed); + assert!(!result.ignored); + assert!(!result.cleanup_ignored); + assert!(!result.moved); + assert_eq!(result.stage, Some("move_remote_version")); + assert!(matches!(result.error, Some(Error::DataMovementOverwriteErr(_, _, _)))); + assert_eq!(backend.move_remote_calls(), 1); +} + +#[tokio::test] +async fn test_migrate_entry_version_remote_failure_is_reported() { + let backend = MigrationBackendSpy::new(None, Some(Ok(ObjectInfo::default())), Some(Err(Error::SlowDown))); + let version = version_remote(); + let transfer_count = Arc::new(AtomicUsize::new(0)); + let mut transfer = { + let transfer_count = transfer_count.clone(); + move |_, _, _| { + let transfer_count = transfer_count.clone(); + async move { + transfer_count.fetch_add(1, Ordering::SeqCst); + Ok(()) + } + } + }; + + let result = migrate_entry_version( + &backend, + "bucket".to_string(), + 0, + &version, + version.version_id.map(|v| v.to_string()), + 3, + false, + &mut transfer, + ) + .await; + + assert!(!result.ignored); + assert!(!result.cleanup_ignored); + assert!(!result.moved); + assert!(result.failed); + assert!(matches!(result.error, Some(Error::SlowDown))); + assert_eq!(transfer_count.load(Ordering::SeqCst), 0); + assert_eq!(backend.move_remote_calls(), 1); + assert_eq!(backend.get_calls(), 0); + assert_eq!(backend.delete_calls(), 0); +} + +#[tokio::test] +async fn test_migrate_entry_version_deleted_version_calls_delete_and_moved() { + let backend = MigrationBackendSpy::new(None, Some(Ok(ObjectInfo::default())), None); + let version = version_deleted(); + let mut transfer = |_, _, _| async move { Ok(()) }; + + let result = migrate_entry_version( + &backend, + "bucket".to_string(), + 1, + &version, + version.version_id.map(|v| v.to_string()), + 3, + false, + &mut transfer, + ) + .await; + + assert!(!result.ignored); + assert!(!result.cleanup_ignored); + assert!(result.moved); + assert!(!result.failed); + assert!(result.error.is_none()); + assert_eq!(backend.get_calls(), 0); + assert_eq!(backend.delete_calls(), 1); +} + +#[tokio::test] +async fn test_migrate_entry_version_deleted_version_not_found_is_ignored() { + let backend = MigrationBackendSpy::new( + None, + Some(Err(Error::ObjectNotFound("bucket".to_string(), "object.bin".to_string()))), + None, + ); + let version = version_deleted(); + let mut transfer = |_, _, _| async move { Ok(()) }; + + let result = migrate_entry_version( + &backend, + "bucket".to_string(), + 1, + &version, + version.version_id.map(|v| v.to_string()), + 3, + false, + &mut transfer, + ) + .await; + + assert!(result.ignored); + assert!(result.cleanup_ignored); + assert!(!result.moved); + assert!(!result.failed); + assert!(result.error.is_none()); + assert_eq!(backend.delete_calls(), 1); +} + +#[tokio::test] +async fn test_migrate_entry_version_deleted_version_overwrite_is_not_ignored() { + let backend = MigrationBackendSpy::new( + None, + Some(Err(Error::DataMovementOverwriteErr( + "bucket".to_string(), + "object.bin".to_string(), + "vid-1".to_string(), + ))), + None, + ); + let version = version_deleted(); + let mut transfer = |_, _, _| async move { Ok(()) }; + + let result = migrate_entry_version( + &backend, + "bucket".to_string(), + 1, + &version, + Some("vid-1".to_string()), + 3, + false, + &mut transfer, + ) + .await; + + assert!(result.failed); + assert!(!result.ignored); + assert!(!result.cleanup_ignored); + assert!(!result.moved); + assert_eq!(result.stage, Some("delete_marker")); + assert!(matches!(result.error, Some(Error::DataMovementOverwriteErr(_, _, _)))); + assert_eq!(backend.delete_calls(), 1); +} + +#[tokio::test] +async fn test_migrate_entry_version_reader_not_found_is_ignored() { + let backend = MigrationBackendSpy::new( + Some(Err(Error::ObjectNotFound("bucket".to_string(), "object.bin".to_string()))), + None, + None, + ); + let version = version_normal(); + let mut transfer = |_, _, _| async move { Ok(()) }; + + let result = migrate_entry_version( + &backend, + "bucket".to_string(), + 1, + &version, + version.version_id.map(|v| v.to_string()), + 3, + false, + &mut transfer, + ) + .await; + + assert!(result.ignored); + assert!(result.cleanup_ignored); + assert!(!result.moved); + assert!(!result.failed); + assert!(result.error.is_none()); + assert_eq!(backend.get_calls(), 1); + assert_eq!(backend.delete_calls(), 0); +} + +#[tokio::test] +async fn test_migrate_entry_version_reader_retries_before_success() { + let backend = MigrationBackendSpy::new(Some(Err(Error::SlowDown)), None, None); + let transfer_count = Arc::new(AtomicUsize::new(0)); + let wait_count = Arc::new(AtomicUsize::new(0)); + let mut transfer = { + let transfer_count = transfer_count.clone(); + move |_, _, _| { + let transfer_count = transfer_count.clone(); + async move { + transfer_count.fetch_add(1, Ordering::SeqCst); + Ok(()) + } + } + }; + + let version = version_normal(); + let result = migrate_entry_version_with_retry_wait( + &backend, + "bucket".to_string(), + 1, + &version, + version.version_id.map(|v| v.to_string()), + 3, + false, + &mut transfer, + { + let wait_count = wait_count.clone(); + move |_| { + let wait_count = wait_count.clone(); + async move { + wait_count.fetch_add(1, Ordering::SeqCst); + } + } + }, + ) + .await; + + assert!(!result.ignored); + assert!(!result.cleanup_ignored); + assert!(result.moved); + assert!(!result.failed); + assert!(result.error.is_none()); + assert_eq!(backend.get_calls(), 2); + assert_eq!(backend.delete_calls(), 0); + assert_eq!(transfer_count.load(Ordering::SeqCst), 1); + assert_eq!(wait_count.load(Ordering::SeqCst), 1); +} + +struct AlwaysFailGetBackend { + get_calls: AtomicUsize, +} + +impl AlwaysFailGetBackend { + fn new() -> Self { + Self { + get_calls: AtomicUsize::new(0), + } + } + + fn get_calls(&self) -> usize { + self.get_calls.load(Ordering::SeqCst) + } +} + +#[async_trait::async_trait] +impl MigrationBackend for AlwaysFailGetBackend { + async fn get_object_reader_for_migration( + &self, + _bucket: &str, + _object: &str, + _range: Option, + _h: http::HeaderMap, + _opts: &ObjectOptions, + ) -> Result { + self.get_calls.fetch_add(1, Ordering::SeqCst); + Err(Error::SlowDown) + } + + async fn delete_object_for_migration(&self, _bucket: &str, _object: &str, _opts: ObjectOptions) -> Result { + Ok(ObjectInfo::default()) + } + + async fn move_remote_version_for_migration( + &self, + _bucket: &str, + _object: &str, + _fi: &FileInfo, + _opts: &ObjectOptions, + ) -> Result<()> { + Ok(()) + } +} + +#[tokio::test] +async fn test_migrate_entry_version_reader_fails_after_retries() { + let backend = AlwaysFailGetBackend::new(); + let transfer_count = Arc::new(AtomicUsize::new(0)); + let mut transfer = { + let transfer_count = transfer_count.clone(); + move |_, _, _| { + let transfer_count = transfer_count.clone(); + async move { + transfer_count.fetch_add(1, Ordering::SeqCst); + Ok(()) + } + } + }; + + let version = version_normal(); + let result = migrate_entry_version( + &backend, + "bucket".to_string(), + 1, + &version, + version.version_id.map(|v| v.to_string()), + 3, + false, + &mut transfer, + ) + .await; + + assert!(!result.ignored); + assert!(!result.cleanup_ignored); + assert!(!result.moved); + assert!(result.failed); + assert!(matches!(result.error, Some(Error::SlowDown))); + assert_eq!(backend.get_calls(), 3); + assert_eq!(transfer_count.load(Ordering::SeqCst), 0); +} + +#[tokio::test] +async fn test_migrate_entry_version_zero_max_attempts_still_attempts_once() { + let backend = AlwaysFailGetBackend::new(); + let transfer_count = Arc::new(AtomicUsize::new(0)); + let mut transfer = { + let transfer_count = transfer_count.clone(); + move |_, _, _| { + let transfer_count = transfer_count.clone(); + async move { + transfer_count.fetch_add(1, Ordering::SeqCst); + Ok(()) + } + } + }; + + let version = version_normal(); + let result = migrate_entry_version( + &backend, + "bucket".to_string(), + 1, + &version, + version.version_id.map(|v| v.to_string()), + 0, + false, + &mut transfer, + ) + .await; + + assert!(!result.ignored); + assert!(!result.cleanup_ignored); + assert!(!result.moved); + assert!(result.failed); + assert!(matches!(result.error, Some(Error::SlowDown))); + assert_eq!(backend.get_calls(), 1); + assert_eq!(transfer_count.load(Ordering::SeqCst), 0); +} + +#[tokio::test] +async fn test_migrate_entry_version_transfer_retries_before_success() { + let backend = MigrationBackendSpy::new(Some(Ok(MigrationBackendSpy::make_reader())), None, None); + let transfer_count = Arc::new(AtomicUsize::new(0)); + let wait_count = Arc::new(AtomicUsize::new(0)); + let mut transfer = { + let transfer_count = transfer_count.clone(); + move |_, _, _| { + let transfer_count = transfer_count.clone(); + async move { + let attempt = transfer_count.fetch_add(1, Ordering::SeqCst); + if attempt == 0 { + return Err(Error::SlowDown); + } + Ok(()) + } + } + }; + + let version = version_normal(); + let result = migrate_entry_version_with_retry_wait( + &backend, + "bucket".to_string(), + 1, + &version, + version.version_id.map(|v| v.to_string()), + 3, + false, + &mut transfer, + { + let wait_count = wait_count.clone(); + move |_| { + let wait_count = wait_count.clone(); + async move { + wait_count.fetch_add(1, Ordering::SeqCst); + } + } + }, + ) + .await; + + assert!(!result.ignored); + assert!(!result.cleanup_ignored); + assert!(result.moved); + assert!(!result.failed); + assert_eq!(backend.get_calls(), 2); + assert_eq!(transfer_count.load(Ordering::SeqCst), 2); + assert_eq!(wait_count.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +async fn test_migrate_entry_version_transfer_non_transient_fails_without_retry() { + let backend = MigrationBackendSpy::new(Some(Ok(MigrationBackendSpy::make_reader())), None, None); + let transfer_count = Arc::new(AtomicUsize::new(0)); + let wait_count = Arc::new(AtomicUsize::new(0)); + let mut transfer = { + let transfer_count = transfer_count.clone(); + move |_, _, _| { + let transfer_count = transfer_count.clone(); + async move { + transfer_count.fetch_add(1, Ordering::SeqCst); + Err(Error::FileAccessDenied) + } + } + }; + + let version = version_normal(); + let result = migrate_entry_version_with_retry_wait( + &backend, + "bucket".to_string(), + 1, + &version, + version.version_id.map(|v| v.to_string()), + 3, + false, + &mut transfer, + { + let wait_count = wait_count.clone(); + move |_| { + let wait_count = wait_count.clone(); + async move { + wait_count.fetch_add(1, Ordering::SeqCst); + } + } + }, + ) + .await; + + assert!(result.failed); + assert!(!result.ignored); + assert_eq!(result.stage, Some("write_target")); + assert!(matches!(result.error, Some(Error::FileAccessDenied))); + assert_eq!(backend.get_calls(), 1); + assert_eq!(transfer_count.load(Ordering::SeqCst), 1); + assert_eq!(wait_count.load(Ordering::SeqCst), 0); +} + +#[tokio::test] +async fn test_migrate_entry_version_transfer_fails_after_retries() { + let backend = MigrationBackendSpy::new(Some(Ok(MigrationBackendSpy::make_reader())), None, None); + let transfer_count = Arc::new(AtomicUsize::new(0)); + let mut transfer = { + let transfer_count = transfer_count.clone(); + move |_, _, _| { + let transfer_count = transfer_count.clone(); + async move { + transfer_count.fetch_add(1, Ordering::SeqCst); + Err(Error::SlowDown) + } + } + }; + + let version = version_normal(); + let result = migrate_entry_version( + &backend, + "bucket".to_string(), + 1, + &version, + version.version_id.map(|v| v.to_string()), + 2, + false, + &mut transfer, + ) + .await; + + assert!(result.failed); + assert!(!result.ignored); + assert!(!result.cleanup_ignored); + assert!(!result.moved); + assert!(result.error.is_some()); + assert_eq!(backend.get_calls(), 2); + assert_eq!(transfer_count.load(Ordering::SeqCst), 2); +} + +#[tokio::test] +async fn test_migrate_entry_version_transfer_not_found_is_ignored() { + let backend = MigrationBackendSpy::new(Some(Ok(MigrationBackendSpy::make_reader())), None, None); + let transfer_count = Arc::new(AtomicUsize::new(0)); + let mut transfer = { + let transfer_count = transfer_count.clone(); + move |_, _, _| { + let transfer_count = transfer_count.clone(); + async move { + transfer_count.fetch_add(1, Ordering::SeqCst); + Err(Error::ObjectNotFound("bucket".to_string(), "object.bin".to_string())) + } + } + }; + + let version = version_normal(); + let result = migrate_entry_version( + &backend, + "bucket".to_string(), + 1, + &version, + version.version_id.map(|v| v.to_string()), + 3, + false, + &mut transfer, + ) + .await; + + assert!(result.ignored); + assert!(result.cleanup_ignored); + assert!(!result.moved); + assert!(!result.failed); + assert!(result.error.is_none()); + assert_eq!(transfer_count.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +async fn test_migrate_entry_version_transfer_overwrite_is_not_ignored() { + let backend = MigrationBackendSpy::new(Some(Ok(MigrationBackendSpy::make_reader())), None, None); + let transfer_count = Arc::new(AtomicUsize::new(0)); + let mut transfer = { + let transfer_count = transfer_count.clone(); + move |_, _, _| { + let transfer_count = transfer_count.clone(); + async move { + transfer_count.fetch_add(1, Ordering::SeqCst); + Err(Error::DataMovementOverwriteErr( + "bucket".to_string(), + "object.bin".to_string(), + "vid-1".to_string(), + )) + } + } + }; + + let version = version_normal(); + let result = migrate_entry_version( + &backend, + "bucket".to_string(), + 1, + &version, + Some("vid-1".to_string()), + 3, + false, + &mut transfer, + ) + .await; + + assert!(result.failed); + assert!(!result.ignored); + assert!(!result.cleanup_ignored); + assert!(!result.moved); + assert_eq!(result.stage, Some("write_target")); + assert!(matches!(result.error, Some(Error::DataMovementOverwriteErr(_, _, _)))); + assert_eq!(transfer_count.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +async fn test_migrate_entry_version_ignores_data_usage_cache_when_enabled() { + let backend = MigrationBackendSpy::new(Some(Ok(MigrationBackendSpy::make_reader())), None, None); + let version = { + let mut version = version_normal(); + version.name = format!("{}.{}", DATA_USAGE_CACHE_NAME, version.name); + version + }; + let transfer_count = Arc::new(AtomicUsize::new(0)); + let mut transfer = { + let transfer_count = transfer_count.clone(); + move |_, _, _| { + let transfer_count = transfer_count.clone(); + async move { + transfer_count.fetch_add(1, Ordering::SeqCst); + Ok(()) + } + } + }; + + let result = migrate_entry_version( + &backend, + RUSTFS_META_BUCKET.to_string(), + 1, + &version, + version.version_id.map(|v| v.to_string()), + 2, + true, + &mut transfer, + ) + .await; + + assert!(result.ignored); + assert!(!result.cleanup_ignored); + assert!(!result.moved); + assert!(!result.failed); + assert!(result.error.is_none()); + assert_eq!(transfer_count.load(Ordering::SeqCst), 0); + assert_eq!(backend.get_calls(), 0); + assert_eq!(backend.delete_calls(), 0); +} + +#[tokio::test] +async fn test_migrate_entry_version_data_usage_cache_moves_when_ignore_disabled() { + let backend = MigrationBackendSpy::new(Some(Ok(MigrationBackendSpy::make_reader())), None, None); + let version = { + let mut version = version_normal(); + version.name = format!("{}.{}", DATA_USAGE_CACHE_NAME, version.name); + version + }; + let transfer_count = Arc::new(AtomicUsize::new(0)); + let mut transfer = { + let transfer_count = transfer_count.clone(); + move |_, _, _| { + let transfer_count = transfer_count.clone(); + async move { + transfer_count.fetch_add(1, Ordering::SeqCst); + Ok(()) + } + } + }; + + let result = migrate_entry_version( + &backend, + RUSTFS_META_BUCKET.to_string(), + 1, + &version, + version.version_id.map(|v| v.to_string()), + 2, + false, + &mut transfer, + ) + .await; + + assert!(!result.ignored); + assert!(!result.cleanup_ignored); + assert!(result.moved); + assert!(!result.failed); + assert!(result.error.is_none()); + assert_eq!(transfer_count.load(Ordering::SeqCst), 1); + assert_eq!(backend.get_calls(), 1); + assert_eq!(backend.delete_calls(), 0); +} + +#[test] +fn test_should_ignore_rebalance_data_usage_cache_true_for_meta_bucket() { + assert!(should_ignore_rebalance_data_usage_cache(RUSTFS_META_BUCKET)); +} + +#[test] +fn test_should_ignore_rebalance_data_usage_cache_false_for_regular_bucket() { + assert!(!should_ignore_rebalance_data_usage_cache("bucket-a")); +} + +#[test] +fn test_rebalance_goal_reached_at_target() { + let init_free_space = 200_u64; + let init_capacity = 1_000_u64; + let goal = 0.45_f64; + + assert!(rebalance_goal_reached(init_free_space, init_capacity, 250, goal)); + assert!(!rebalance_goal_reached(init_free_space, init_capacity, 249, goal)); +} + +#[test] +fn test_rebalance_goal_reached_above_target() { + let init_free_space = 200_u64; + let init_capacity = 1_000_u64; + let bytes = 251_u64; + let goal = 0.45_f64; + + assert!(rebalance_goal_reached(init_free_space, init_capacity, bytes, goal)); +} + +#[test] +fn test_rebalance_goal_not_reached_outside_tolerance() { + let init_free_space = 100_u64; + let init_capacity = 1_000_u64; + let bytes = 80_u64; + let goal = 0.5_f64; + + assert!(!rebalance_goal_reached(init_free_space, init_capacity, bytes, goal)); +} + +#[test] +fn test_rebalance_goal_zero_capacity_is_false() { + assert!(!rebalance_goal_reached(100, 0, 50, 0.5)); +} + +#[test] +fn test_rebalance_goal_above_one_is_true_when_reached() { + assert!(rebalance_goal_reached(950, 1_000, 100, 1.0)); +} + +#[test] +fn test_rebalance_goal_below_zero_is_true_when_reached() { + assert!(rebalance_goal_reached(10, 1_000, 0, -0.01)); +} + +#[test] +fn test_resolve_rebalance_worker_result_passthrough() { + assert!(resolve_rebalance_worker_result(0, Ok(Ok(()))).is_ok()); + + let err = resolve_rebalance_worker_result::<()>(0, Ok(Err(Error::OperationCanceled))).unwrap_err(); + assert!(matches!(err, Error::OperationCanceled)); +} + +#[tokio::test] +async fn test_resolve_rebalance_worker_result_join_error_keeps_context() { + let join_error = tokio::spawn(async { + panic!("rebalance worker panic"); + }) + .await + .expect_err("panic task should return JoinError"); + + let err = resolve_rebalance_worker_result::<()>(7, Err(join_error)).unwrap_err(); + assert!(err.to_string().contains("rebalance worker 7 task join error")); +} + +#[tokio::test] +async fn test_wait_rebalance_entry_tasks_returns_ok_for_successful_tasks() { + let tasks = Arc::new(tokio::sync::Mutex::new(vec![tokio::spawn(async { + Ok(RebalanceEntryOutcome::Completed) + })])); + + let result = wait_rebalance_entry_tasks(1, tasks) + .await + .expect("successful entry tasks should pass"); + + assert!(result.is_none()); +} + +#[tokio::test] +async fn test_wait_rebalance_entry_tasks_returns_first_task_error() { + let tasks = Arc::new(tokio::sync::Mutex::new(vec![ + tokio::spawn(async { Ok(RebalanceEntryOutcome::Completed) }), + tokio::spawn(async { Err(Error::other("entry failed")) }), + ])); + + let err = wait_rebalance_entry_tasks(1, tasks) + .await + .expect_err("entry task failure should be returned"); + + assert!(err.to_string().contains("entry failed")); +} + +#[tokio::test] +async fn test_wait_rebalance_entry_tasks_returns_deferred_error_without_failing() { + let tasks = Arc::new(tokio::sync::Mutex::new(vec![ + tokio::spawn(async { Ok(RebalanceEntryOutcome::Completed) }), + tokio::spawn(async { + Ok(RebalanceEntryOutcome::Deferred { + last_error: "deferred transient rebalance entry failure: timeout".to_string(), + }) + }), + ])); + + let result = wait_rebalance_entry_tasks(1, tasks) + .await + .expect("deferred transient entry should not fail worker"); + + assert_eq!(result.as_deref(), Some("deferred transient rebalance entry failure: timeout")); +} + +#[test] +fn test_resolve_rebalance_save_task_result_passthrough() { + assert!(resolve_rebalance_save_task_result(0, Ok(Ok(()))).is_ok()); +} + +#[test] +fn test_resolve_rebalance_save_task_result_wraps_inner_error_context() { + let err = resolve_rebalance_save_task_result(1, Ok(Err(Error::SlowDown))) + .expect_err("inner save-task error should include pool context"); + assert!(err.to_string().contains("rebalance save_task failed for pool 1")); +} + +#[test] +fn test_resolve_rebalance_meta_save_result_passthrough() { + assert!(resolve_rebalance_meta_save_result(Ok(()), "stop_rebalance").is_ok()); +} + +#[test] +fn test_resolve_rebalance_meta_save_result_wraps_error_context() { + let err = resolve_rebalance_meta_save_result(Err(Error::SlowDown), "init_rebalance_meta") + .expect_err("meta save failure should include stage context"); + let message = err.to_string(); + assert!(message.contains("rebalance meta save failed during init_rebalance_meta")); + assert!(message.contains(Error::SlowDown.to_string().as_str())); +} + +#[test] +fn test_merge_rebalance_meta_preserves_updates_from_multiple_pools() { + let start_time = OffsetDateTime::from_unix_timestamp(1_000).unwrap(); + let warning_at = OffsetDateTime::from_unix_timestamp(1_500).unwrap(); + let mut remote = RebalanceMeta { + id: "rebal-1".to_string(), + percent_free_goal: 0.5, + pool_stats: vec![ + RebalanceStats { + buckets: vec!["bucket-a".to_string()], + participating: true, + info: RebalanceInfo { + start_time: Some(start_time), + status: RebalStatus::Started, + ..Default::default() + }, + num_versions: 4, + bytes: 400, + bucket: "bucket-a".to_string(), + object: "remote-object".to_string(), + ..Default::default() + }, + RebalanceStats { + buckets: vec!["bucket-a".to_string()], + participating: true, + info: RebalanceInfo { + start_time: Some(start_time), + status: RebalStatus::Started, + ..Default::default() + }, + ..Default::default() + }, + ], + ..Default::default() + }; + let local = RebalanceMeta { + id: "rebal-1".to_string(), + percent_free_goal: 0.5, + pool_stats: vec![ + RebalanceStats { + buckets: vec!["bucket-a".to_string()], + participating: true, + info: RebalanceInfo { + start_time: Some(start_time), + status: RebalStatus::Started, + ..Default::default() + }, + ..Default::default() + }, + RebalanceStats { + buckets: Vec::new(), + rebalanced_buckets: vec!["bucket-a".to_string()], + participating: true, + info: RebalanceInfo { + start_time: Some(start_time), + status: RebalStatus::Started, + ..Default::default() + }, + num_versions: 7, + bytes: 700, + bucket: "bucket-a".to_string(), + object: "local-object".to_string(), + cleanup_warnings: RebalanceCleanupWarnings { + count: 1, + last_message: Some("cleanup failed".to_string()), + last_bucket: Some("bucket-a".to_string()), + last_object: Some("local-object".to_string()), + last_at: Some(warning_at), + }, + ..Default::default() + }, + ], + ..Default::default() + }; + + merge_rebalance_meta(&mut remote, &local); + + assert_eq!(remote.pool_stats[0].num_versions, 4); + assert_eq!(remote.pool_stats[0].object, "remote-object"); + assert_eq!(remote.pool_stats[1].num_versions, 7); + assert_eq!(remote.pool_stats[1].object, "local-object"); + assert!(remote.pool_stats[1].buckets.is_empty()); + assert_eq!(remote.pool_stats[1].rebalanced_buckets, vec!["bucket-a"]); + assert_eq!(remote.pool_stats[1].cleanup_warnings.count, 1); + assert_eq!(remote.pool_stats[1].cleanup_warnings.last_message.as_deref(), Some("cleanup failed")); + assert_eq!(remote.pool_stats[1].cleanup_warnings.last_at, Some(warning_at)); +} + +#[test] +fn test_merge_rebalance_meta_does_not_overwrite_failed_with_started_stats() { + let now = OffsetDateTime::from_unix_timestamp(2_000).unwrap(); + let mut remote = RebalanceMeta { + id: "rebal-1".to_string(), + pool_stats: vec![RebalanceStats { + info: RebalanceInfo { + status: RebalStatus::Failed, + end_time: Some(now), + last_error: Some("timeout".to_string()), + ..Default::default() + }, + num_versions: 10, + bytes: 1_000, + ..Default::default() + }], + ..Default::default() + }; + let local = RebalanceMeta { + id: "rebal-1".to_string(), + pool_stats: vec![RebalanceStats { + participating: true, + info: RebalanceInfo { + status: RebalStatus::Started, + ..Default::default() + }, + num_versions: 8, + bytes: 800, + ..Default::default() + }], + ..Default::default() + }; + + merge_rebalance_meta(&mut remote, &local); + + assert_eq!(remote.pool_stats[0].info.status, RebalStatus::Failed); + assert_eq!(remote.pool_stats[0].info.last_error.as_deref(), Some("timeout")); + assert_eq!(remote.pool_stats[0].num_versions, 10); + assert_eq!(remote.pool_stats[0].bytes, 1_000); +} + +#[test] +fn test_merge_rebalance_meta_does_not_overwrite_stopped_with_started_stats() { + let stopped_at = OffsetDateTime::from_unix_timestamp(3_000).unwrap(); + let mut remote = RebalanceMeta { + id: "rebal-1".to_string(), + stopped_at: Some(stopped_at), + pool_stats: vec![RebalanceStats { + info: RebalanceInfo { + status: RebalStatus::Stopped, + end_time: Some(stopped_at), + ..Default::default() + }, + num_versions: 5, + ..Default::default() + }], + ..Default::default() + }; + let local = RebalanceMeta { + id: "rebal-1".to_string(), + pool_stats: vec![RebalanceStats { + info: RebalanceInfo { + status: RebalStatus::Started, + ..Default::default() + }, + num_versions: 9, + ..Default::default() + }], + ..Default::default() + }; + + merge_rebalance_meta(&mut remote, &local); + + assert_eq!(remote.stopped_at, Some(stopped_at)); + assert_eq!(remote.pool_stats[0].info.status, RebalStatus::Stopped); + assert_eq!(remote.pool_stats[0].info.end_time, Some(stopped_at)); + assert_eq!(remote.pool_stats[0].num_versions, 9); +} + +#[test] +fn test_merge_rebalance_meta_preserves_failed_status_over_stopped() { + let stopped_at = OffsetDateTime::from_unix_timestamp(3_000).unwrap(); + let failed_at = OffsetDateTime::from_unix_timestamp(4_000).unwrap(); + let mut remote = RebalanceMeta { + id: "rebal-1".to_string(), + stopped_at: Some(stopped_at), + pool_stats: vec![RebalanceStats { + info: RebalanceInfo { + status: RebalStatus::Stopped, + end_time: Some(stopped_at), + ..Default::default() + }, + ..Default::default() + }], + ..Default::default() + }; + let local = RebalanceMeta { + id: "rebal-1".to_string(), + pool_stats: vec![RebalanceStats { + info: RebalanceInfo { + status: RebalStatus::Failed, + end_time: Some(failed_at), + last_error: Some("late failure".to_string()), + ..Default::default() + }, + ..Default::default() + }], + ..Default::default() + }; + + merge_rebalance_meta(&mut remote, &local); + + assert_eq!(remote.pool_stats[0].info.status, RebalStatus::Failed); + assert_eq!(remote.pool_stats[0].info.end_time, Some(failed_at)); + assert_eq!(remote.pool_stats[0].info.last_error.as_deref(), Some("late failure")); +} + +#[test] +fn test_merge_rebalance_meta_does_not_overwrite_stopped_with_completed_status() { + let stopped_at = OffsetDateTime::from_unix_timestamp(3_000).unwrap(); + let completed_at = OffsetDateTime::from_unix_timestamp(5_000).unwrap(); + let mut remote = RebalanceMeta { + id: "rebal-1".to_string(), + stopped_at: Some(stopped_at), + pool_stats: vec![RebalanceStats { + info: RebalanceInfo { + status: RebalStatus::Stopped, + end_time: Some(stopped_at), + ..Default::default() + }, + num_versions: 5, + ..Default::default() + }], + ..Default::default() + }; + let local = RebalanceMeta { + id: "rebal-1".to_string(), + pool_stats: vec![RebalanceStats { + info: RebalanceInfo { + status: RebalStatus::Completed, + end_time: Some(completed_at), + ..Default::default() + }, + num_versions: 12, + ..Default::default() + }], + ..Default::default() + }; + + merge_rebalance_meta(&mut remote, &local); + + assert_eq!(remote.stopped_at, Some(stopped_at)); + assert_eq!(remote.pool_stats[0].info.status, RebalStatus::Stopped); + assert_eq!(remote.pool_stats[0].info.end_time, Some(stopped_at)); + assert_eq!(remote.pool_stats[0].num_versions, 12); +} + +#[test] +fn test_rebalance_meta_load_no_data_error_formats_context() { + let err = rebalance_meta_load_no_data_error(); + let rendered = err.to_string(); + + assert!(rendered.contains("rebalance metadata load failed"), "{rendered}"); + assert!(rendered.contains("payload is too short"), "{rendered}"); +} + +#[test] +fn test_rebalance_meta_load_unknown_format_error_formats_context() { + let err = rebalance_meta_load_unknown_format_error(9); + let rendered = err.to_string(); + + assert!(rendered.contains("rebalance metadata load failed"), "{rendered}"); + assert!(rendered.contains("unknown format 9"), "{rendered}"); +} + +#[test] +fn test_rebalance_meta_load_unknown_version_error_formats_context() { + let err = rebalance_meta_load_unknown_version_error(3); + let rendered = err.to_string(); + + assert!(rendered.contains("rebalance metadata load failed"), "{rendered}"); + assert!(rendered.contains("unknown version 3"), "{rendered}"); +} + +#[test] +fn test_rebalance_meta_deserializes_legacy_stats_without_cleanup_warnings() { + let legacy = LegacyRebalanceMeta { + id: "rebal-legacy".to_string(), + percent_free_goal: 0.35, + pool_stats: vec![LegacyRebalanceStats { + buckets: vec!["bucket-a".to_string()], + participating: true, + info: RebalanceInfo { + status: RebalStatus::Started, + ..Default::default() + }, + ..Default::default() + }], + ..Default::default() + }; + let data = rmp_serde::to_vec(&legacy).expect("legacy rebalance metadata should serialize"); + + let decoded: RebalanceMeta = rmp_serde::from_slice(data.as_slice()).expect("legacy rebalance metadata should deserialize"); + + assert_eq!(decoded.id, "rebal-legacy"); + assert_eq!(decoded.pool_stats.len(), 1); + assert_eq!(decoded.pool_stats[0].cleanup_warnings, RebalanceCleanupWarnings::default()); + assert_eq!(decoded.pool_stats[0].info.status, RebalStatus::Started); +} + +#[test] +fn test_resolve_rebalance_stats_update_result_passthrough() { + assert!(resolve_rebalance_stats_update_result(Ok(()), 0, "bucket", "object").is_ok()); +} + +#[test] +fn test_resolve_rebalance_stats_update_result_wraps_error_context() { + let err = resolve_rebalance_stats_update_result(Err(Error::SlowDown), 2, "bucket-a", "obj.txt") + .expect_err("stats update error should include context"); + assert!( + err.to_string() + .contains("rebalance stats update failed for pool 2 bucket bucket-a object obj.txt") + ); +} + +#[test] +fn test_resolve_load_rebalance_stats_update_result_passthrough() { + assert!(resolve_load_rebalance_stats_update_result(Ok(())).is_ok()); +} + +#[test] +fn test_resolve_load_rebalance_stats_update_result_wraps_error_context() { + let err = resolve_load_rebalance_stats_update_result(Err(Error::SlowDown)) + .expect_err("load-time stats refresh failure should include context"); + assert!(err.to_string().contains("rebalance metadata stats refresh failed after load")); +} + +#[test] +fn test_resolve_rebalance_file_info_versions_result_passthrough() { + let value = resolve_rebalance_file_info_versions_result::(Ok(7), "bucket-a", "obj.txt") + .expect("ok results should pass through"); + assert_eq!(value, 7); +} + +#[test] +fn test_resolve_rebalance_file_info_versions_result_wraps_error_context() { + let err = resolve_rebalance_file_info_versions_result::(Err(Error::SlowDown), "bucket-a", "obj.txt") + .expect_err("errors should be wrapped"); + let message = err.to_string(); + assert!(message.contains("rebalance file_info_versions failed for bucket-a/obj.txt")); +} + +#[test] +fn test_resolve_rebalance_meta_load_result_returns_true_for_loaded_meta() { + assert!(resolve_rebalance_meta_load_result(Ok(())).expect("loaded rebalance metadata should pass through")); +} + +#[test] +fn test_resolve_rebalance_meta_load_result_returns_false_for_missing_meta() { + assert!( + !resolve_rebalance_meta_load_result(Err(Error::ConfigNotFound)) + .expect("missing rebalance metadata should be treated as not started") + ); +} + +#[test] +fn test_resolve_rebalance_meta_load_result_wraps_error_context() { + let err = resolve_rebalance_meta_load_result(Err(Error::SlowDown)).expect_err("unexpected load failures should be wrapped"); + let message = err.to_string(); + assert!(message.contains("rebalance metadata load failed during load_rebalance_meta")); +} + +#[test] +fn test_resolve_rebalance_entry_cleanup_delete_result_passthrough() { + let result = resolve_rebalance_entry_cleanup_delete_result(Ok(ObjectInfo::default()), "bucket-a", "obj.txt"); + assert_eq!(result.expect("successful cleanup should pass through"), None); +} + +#[test] +fn test_resolve_rebalance_entry_cleanup_delete_result_ignores_not_found() { + let result = resolve_rebalance_entry_cleanup_delete_result( + Err(Error::ObjectNotFound("bucket-a".to_string(), "obj.txt".to_string())), + "bucket-a", + "obj.txt", + ); + assert_eq!(result.expect("missing cleanup source should be ignored"), None); +} + +#[test] +fn test_resolve_rebalance_entry_cleanup_delete_result_returns_warning_for_failures() { + let warning = resolve_rebalance_entry_cleanup_delete_result(Err(Error::SlowDown), "bucket-a", "obj.txt") + .expect("cleanup delete failures should be downgraded to warnings") + .expect("cleanup delete failure should return warning"); + let message = warning.as_str(); + assert!(message.contains("rebalance cleanup delete failed for bucket-a/obj.txt")); +} + +#[test] +fn test_resolve_rebalance_migrate_result_error_preserves_inner_error() { + let err = resolve_rebalance_migrate_result_error(Some(Error::SlowDown), 2, "bucket-a", "obj.txt", Some("vid-1")); + assert!(matches!(err, Error::SlowDown)); +} + +#[test] +fn test_resolve_rebalance_migrate_result_error_wraps_missing_error_context() { + let err = resolve_rebalance_migrate_result_error(None, 2, "bucket-a", "obj.txt", Some("vid-1")); + let message = err.to_string(); + assert!( + message.contains("rebalance migration reported failure without error for pool 2 entry bucket-a/obj.txt version vid-1"), + "{message}" + ); +} + +#[test] +fn test_resolve_rebalance_bucket_result_passthrough() { + let outcome = resolve_rebalance_bucket_result(Ok(RebalanceBucketOutcome::Completed), 2, "bucket-a") + .expect("completed bucket should pass through"); + assert_eq!(outcome, RebalanceBucketOutcome::Completed); +} + +#[test] +fn test_resolve_rebalance_bucket_result_preserves_operation_canceled() { + let err = resolve_rebalance_bucket_result(Err(Error::OperationCanceled), 2, "bucket-a") + .expect_err("operation canceled should be preserved"); + assert!(matches!(err, Error::OperationCanceled)); +} + +#[test] +fn test_resolve_rebalance_bucket_result_wraps_not_initialized_with_context() { + let err = resolve_rebalance_bucket_result(Err(Error::other("errServerNotInitialized")), 2, "bucket-a") + .expect_err("not initialized should be surfaced with context"); + let message = err.to_string(); + assert!(message.contains("rebalance bucket bucket-a failed for pool 2")); + assert!(message.contains("errServerNotInitialized")); +} + +#[test] +fn test_rebalance_listing_disks_available_rejects_empty_set() { + let err = ensure_rebalance_listing_disks_available(false, "bucket-a") + .expect_err("missing online disks should be reported with bucket context"); + assert!( + err.to_string() + .contains("failed to list objects to rebalance for bucket bucket-a: no disks available") + ); +} + +#[test] +fn test_rebalance_listing_disks_available_allows_online_disks() { + assert!(ensure_rebalance_listing_disks_available(true, "bucket-a").is_ok()); +} + +#[test] +fn test_is_transient_rebalance_error_classifies_retryable_errors() { + assert!(is_transient_rebalance_error(&Error::SlowDown)); + assert!(is_transient_rebalance_error(&Error::ErasureReadQuorum)); + assert!(is_transient_rebalance_error(&Error::ErasureWriteQuorum)); + assert!(is_transient_rebalance_error(&Error::Io(std::io::Error::new( + std::io::ErrorKind::TimedOut, + "timed out", + )))); + assert!(is_transient_rebalance_error(&Error::other("i/o timeout"))); +} + +#[test] +fn test_is_transient_rebalance_error_accepts_lock_and_rpc_timeouts() { + assert!(is_transient_rebalance_error(&Error::Io(std::io::Error::other( + "Failed to acquire read lock: ns_loc: read lock acquisition timed out on bucket/object" + )))); + assert!(is_transient_rebalance_error(&Error::other( + "Remote lock RPC timed out: RPC timed out after 50ms" + ))); + assert!(is_transient_rebalance_error(&Error::other( + "Unknown hyper error: hyper::Error(Http2, KeepAliveTimedOut)" + ))); + assert!(is_transient_rebalance_error(&Error::Lock(rustfs_lock::LockError::timeout( + "bucket/object", + Duration::from_secs(5), + )))); +} + +#[test] +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_is_transient_rebalance_error_accepts_io_timeout_message() { + assert!(is_transient_rebalance_error(&Error::Io(std::io::Error::other("timeout")))); +} + +#[test] +fn test_is_transient_rebalance_error_rejects_terminal_errors() { + assert!(!is_transient_rebalance_error(&Error::DataMovementOverwriteErr( + "bucket".to_string(), + "object".to_string(), + "version".to_string(), + ))); + assert!(!is_transient_rebalance_error(&Error::ObjectNotFound( + "bucket".to_string(), + "object".to_string(), + ))); + assert!(!is_transient_rebalance_error(&Error::FileAccessDenied)); + assert!(!is_transient_rebalance_error(&Error::Lock(rustfs_lock::LockError::already_locked( + "bucket/object", + "owner-a", + )))); + assert!(!is_transient_rebalance_error(&Error::other( + "Failed to acquire read lock: ns_loc: read lock acquisition failed on bucket/object: permission denied" + ))); +} + +#[test] +fn test_should_retry_rebalance_listing_respects_attempt_limit_and_error_type() { + assert!(should_retry_rebalance_listing(&Error::SlowDown, 0, 3)); + assert!(should_retry_rebalance_listing(&Error::SlowDown, 1, 3)); + assert!(!should_retry_rebalance_listing(&Error::SlowDown, 2, 3)); + assert!(!should_retry_rebalance_listing(&Error::FileAccessDenied, 0, 3)); +} + +#[test] +fn test_parse_rebalance_max_attempts_uses_positive_override_or_default() { + assert_eq!(parse_rebalance_max_attempts(Some("5")), 5); + assert_eq!(parse_rebalance_max_attempts(Some(" 7 ")), 7); + assert_eq!(parse_rebalance_max_attempts(Some("0")), 3); + assert_eq!(parse_rebalance_max_attempts(Some("invalid")), 3); + assert_eq!(parse_rebalance_max_attempts(None), 3); +} + +#[test] +fn test_rebalance_listing_retry_delay_scales_by_attempt() { + assert_eq!(rebalance_listing_retry_delay(0), Duration::from_millis(250)); + assert_eq!(rebalance_listing_retry_delay(1), Duration::from_millis(500)); +} + +#[test] +fn test_rebalance_migration_retry_delay_scales_for_non_lock_timeout() { + assert_eq!(rebalance_migration_retry_delay(0, &Error::SlowDown), Duration::from_millis(250)); + assert_eq!(rebalance_migration_retry_delay(1, &Error::SlowDown), Duration::from_millis(500)); +} + +#[test] +fn test_should_defer_rebalance_entry_failure_only_for_transient_errors() { + assert!(should_defer_rebalance_entry_failure(&Error::Io(std::io::Error::other( + "Failed to acquire write lock: ns_loc: write lock acquisition timed out on bucket/object" + )))); + assert!(!should_defer_rebalance_entry_failure(&Error::DataMovementOverwriteErr( + "bucket".to_string(), + "object".to_string(), + "version".to_string(), + ))); + assert!(!should_defer_rebalance_entry_failure(&Error::FileAccessDenied)); +} + +#[tokio::test] +async fn test_wait_rebalance_listing_retry_returns_canceled_without_sleeping() { + let token = CancellationToken::new(); + token.cancel(); + + let err = wait_rebalance_listing_retry(&token, Duration::from_secs(30)) + .await + .expect_err("canceled rebalance should not wait before retrying"); + + assert!(matches!(err, Error::OperationCanceled)); +} + +#[test] +fn test_with_rebalance_entry_context_formats_stage_bucket_and_object() { + let err = with_rebalance_entry_context("migrate", "bucket-a", "obj.txt", Error::SlowDown); + let message = err.to_string(); + assert!(message.contains("rebalance entry migrate failed for bucket-a/obj.txt")); + assert!(message.contains("Please reduce your request rate")); +} + +#[test] +fn test_with_rebalance_entry_context_formats_precise_stage() { + let err = with_rebalance_entry_context("write_target", "bucket-a", "obj.txt", Error::SlowDown); + let message = err.to_string(); + assert!(message.contains("rebalance entry write_target failed for bucket-a/obj.txt")); + assert!(message.contains("Please reduce your request rate")); +} + +#[tokio::test] +async fn test_migrate_entry_version_transfer_failure_reports_write_target_stage() { + let backend = MigrationBackendSpy::new(Some(Ok(MigrationBackendSpy::make_reader())), None, None); + let mut transfer = |_, _, _| async { Err(Error::SlowDown) }; + let version = version_normal(); + + let result = migrate_entry_version( + &backend, + "bucket".to_string(), + 1, + &version, + version.version_id.map(|v| v.to_string()), + 1, + false, + &mut transfer, + ) + .await; + + assert!(result.failed); + assert_eq!(result.stage, Some("write_target")); + assert!(matches!(result.error, Some(Error::SlowDown))); +} + +#[tokio::test] +async fn test_migrate_entry_version_reader_failure_reports_read_source_stage() { + let backend = AlwaysFailGetBackend::new(); + let mut transfer = |_, _, _| async { Ok(()) }; + let version = version_normal(); + + let result = migrate_entry_version( + &backend, + "bucket".to_string(), + 1, + &version, + version.version_id.map(|v| v.to_string()), + 1, + false, + &mut transfer, + ) + .await; + + assert!(result.failed); + assert_eq!(result.stage, Some("read_source")); + assert!(matches!(result.error, Some(Error::SlowDown))); +} + +#[test] +fn test_should_count_rebalance_version_complete_for_cleanup_safe_ignored_result() { + let result = MigrationVersionResult { + ignored: true, + cleanup_ignored: true, + ..Default::default() + }; + assert!(should_count_rebalance_version_complete(&result)); +} + +#[test] +fn test_should_count_rebalance_version_complete_rejects_skip_only_ignored_result() { + let result = MigrationVersionResult { + ignored: true, + cleanup_ignored: false, + ..Default::default() + }; + assert!(!should_count_rebalance_version_complete(&result)); +} + +#[test] +fn test_should_count_rebalance_version_complete_for_moved_result() { + let result = MigrationVersionResult { + moved: true, + ..Default::default() + }; + assert!(should_count_rebalance_version_complete(&result)); +} + +#[test] +fn test_should_count_rebalance_version_complete_rejects_failed_result() { + let result = MigrationVersionResult { + moved: true, + failed: true, + ..Default::default() + }; + assert!(!should_count_rebalance_version_complete(&result)); +} + +#[test] +fn test_should_count_rebalance_version_complete_rejects_incomplete_result() { + assert!(!should_count_rebalance_version_complete(&MigrationVersionResult::default())); +} + +#[test] +fn test_should_accept_rebalance_stats_update_only_for_started_pool() { + let meta = RebalanceMeta { + pool_stats: vec![RebalanceStats { + info: RebalanceInfo { + status: RebalStatus::Started, + ..Default::default() + }, + ..Default::default() + }], + ..Default::default() + }; + + assert!(should_accept_rebalance_stats_update(&meta, 0)); +} + +#[test] +fn test_should_accept_rebalance_stats_update_rejects_stopped_meta() { + let meta = RebalanceMeta { + stopped_at: Some(OffsetDateTime::UNIX_EPOCH), + pool_stats: vec![RebalanceStats { + info: RebalanceInfo { + status: RebalStatus::Started, + ..Default::default() + }, + ..Default::default() + }], + ..Default::default() + }; + + assert!(!should_accept_rebalance_stats_update(&meta, 0)); +} + +#[test] +fn test_should_accept_rebalance_stats_update_rejects_terminal_or_invalid_pool() { + for status in [RebalStatus::Completed, RebalStatus::Stopped, RebalStatus::Failed] { + let meta = RebalanceMeta { + pool_stats: vec![RebalanceStats { + info: RebalanceInfo { + status, + ..Default::default() + }, + ..Default::default() + }], + ..Default::default() + }; + + assert!(!should_accept_rebalance_stats_update(&meta, 0)); + assert!(!should_accept_rebalance_stats_update(&meta, 1)); + } +} + +#[test] +fn test_should_skip_rebalance_delete_marker_when_last_remaining_without_replication() { + assert!(should_skip_rebalance_delete_marker(&version_deleted(), 1, false)); +} + +#[test] +fn test_should_skip_rebalance_delete_marker_rejects_configured_replication() { + assert!(!should_skip_rebalance_delete_marker(&version_deleted(), 1, true)); +} + +#[test] +fn test_should_skip_rebalance_delete_marker_rejects_non_deleted_versions() { + assert!(!should_skip_rebalance_delete_marker(&version_normal(), 1, false)); +} + +#[test] +fn test_should_skip_rebalance_delete_marker_rejects_multiple_remaining_versions() { + assert!(!should_skip_rebalance_delete_marker(&version_deleted(), 2, false)); +} + +#[test] +fn test_should_cleanup_rebalance_source_entry_accepts_all_versions_completed() { + assert!(should_cleanup_rebalance_source_entry(3, 3)); +} + +#[test] +fn test_should_cleanup_rebalance_source_entry_rejects_versions_only_expired_by_lifecycle() { + assert!(!should_cleanup_rebalance_source_entry(2, 3)); +} + +#[test] +fn test_resolve_rebalance_optional_bucket_config_result_passthrough() { + let result = resolve_rebalance_optional_bucket_config_result( + "bucket-a", + "replication", + Ok((ReplicationConfiguration::default(), OffsetDateTime::UNIX_EPOCH)), + ) + .expect("bucket config should pass through"); + assert!(result.is_some()); +} + +#[test] +fn test_resolve_rebalance_optional_bucket_config_result_returns_none_for_missing_config() { + let result = resolve_rebalance_optional_bucket_config_result::<()>("bucket-a", "versioning", Err(Error::ConfigNotFound)) + .expect("missing bucket config should map to None"); + assert!(result.is_none()); +} + +#[test] +fn test_resolve_rebalance_optional_bucket_config_result_wraps_other_errors() { + let err = resolve_rebalance_optional_bucket_config_result::<()>("bucket-a", "replication", Err(Error::SlowDown)) + .expect_err("unexpected bucket config errors should be wrapped with context"); + assert!( + err.to_string() + .contains("rebalance replication config load failed for bucket bucket-a") + ); +} + +#[tokio::test] +async fn test_load_rebalance_bucket_configs_skips_meta_bucket_lookup() { + let configs = load_rebalance_bucket_configs(RUSTFS_META_BUCKET) + .await + .expect("meta bucket config loading should short-circuit"); + assert!(configs.lifecycle_config.is_none()); + assert!(configs.lock_retention.is_none()); + assert!(configs.replication_config.is_none()); +} + +#[tokio::test] +async fn test_resolve_rebalance_save_task_result_join_error_keeps_context() { + let join_error = tokio::spawn(async { + panic!("rebalance save task panic"); + }) + .await + .expect_err("panic task should return JoinError"); + + let err = resolve_rebalance_save_task_result(3, Err(join_error)).unwrap_err(); + assert!(err.to_string().contains("rebalance save_task for pool 3 join error")); +} + +#[tokio::test] +async fn test_send_rebalance_done_signal_sends_message() { + let (tx, mut rx) = mpsc::channel(1); + + send_rebalance_done_signal(&tx, Ok(()), 2) + .await + .expect("send should succeed when receiver is active"); + + let received = rx.recv().await.expect("receiver should get signal"); + assert!(received.is_ok()); +} + +#[tokio::test] +async fn test_send_rebalance_done_signal_reports_closed_channel() { + let (tx, rx) = mpsc::channel(1); + drop(rx); + + let err = send_rebalance_done_signal(&tx, Ok(()), 5) + .await + .expect_err("send should fail when receiver is closed"); + assert!(err.to_string().contains("rebalance done signal send failed for pool 5")); +} + +#[test] +fn test_resolve_rebalance_terminal_error_keeps_primary_when_signal_ok() { + let err = resolve_rebalance_terminal_error(Error::SlowDown, Ok(())); + assert!(matches!(err, Error::SlowDown)); +} + +#[test] +fn test_resolve_rebalance_terminal_error_wraps_signal_failure_context() { + let err = resolve_rebalance_terminal_error(Error::SlowDown, Err(Error::OperationCanceled)); + assert!(err.to_string().contains("rebalance terminal signal failed after error")); +} + +#[test] +fn test_resolve_rebalance_bucket_error_prefers_entry_error() { + let err = resolve_rebalance_bucket_error(Some(Error::OperationCanceled), Some(Error::SlowDown)).unwrap_err(); + assert!(matches!(err, Error::OperationCanceled)); +} + +#[test] +fn test_resolve_rebalance_bucket_error_uses_worker_error_when_entry_ok() { + let err = resolve_rebalance_bucket_error(None, Some(Error::SlowDown)).unwrap_err(); + assert!(matches!(err, Error::SlowDown)); +} + +#[test] +fn test_resolve_rebalance_bucket_error_is_ok_when_no_errors() { + assert!(resolve_rebalance_bucket_error(None, None).is_ok()); +} + +#[test] +fn test_ensure_valid_rebalance_pool_index_allows_in_range() { + assert!(ensure_valid_rebalance_pool_index(3, 2).is_ok()); +} + +#[test] +fn test_ensure_valid_rebalance_pool_index_rejects_out_of_range() { + let err = ensure_valid_rebalance_pool_index(2, 2).expect_err("out of range index should fail"); + assert!(err.to_string().contains("invalid rebalance pool index")); +} + +#[test] +fn test_clone_first_arc_returns_first_value() { + let values = vec![Arc::new(7_u8), Arc::new(9_u8)]; + let first = clone_first_arc(values.as_slice(), "empty values").expect("first value should be returned"); + assert_eq!(*first, 7_u8); +} + +#[test] +fn test_clone_first_arc_rejects_empty_values() { + let values: Vec> = Vec::new(); + let err = clone_first_arc(values.as_slice(), "empty values").expect_err("empty values should fail"); + assert!(err.to_string().contains("empty values")); +} + +#[test] +fn test_clone_arc_by_index_returns_value() { + let values = vec![Arc::new(7_u8), Arc::new(9_u8)]; + let value = + clone_arc_by_index(values.as_slice(), 1, "invalid rebalance pool index").expect("index within bounds should work"); + assert_eq!(*value, 9_u8); +} + +#[test] +fn test_clone_arc_by_index_rejects_out_of_range() { + let values = vec![Arc::new(7_u8)]; + let err = + clone_arc_by_index(values.as_slice(), 2, "invalid rebalance pool index").expect_err("out of range index should fail"); + assert!(err.to_string().contains("invalid rebalance pool index: 2")); +} + +#[test] +fn test_classify_rebalance_terminal_event_completed() { + let now = OffsetDateTime::now_utc(); + match classify_rebalance_terminal_event(Some(Ok(())), now) { + RebalanceTerminalEvent::Completed { msg } => assert!(msg.contains("Rebalance completed")), + _ => panic!("expected completed terminal event"), + } +} + +#[test] +fn test_classify_rebalance_terminal_event_stopped() { + let now = OffsetDateTime::now_utc(); + match classify_rebalance_terminal_event(Some(Err(Error::OperationCanceled)), now) { + RebalanceTerminalEvent::Stopped { msg } => assert!(msg.contains("Rebalance stopped")), + _ => panic!("expected stopped terminal event"), + } +} + +#[test] +fn test_classify_rebalance_terminal_event_failed() { + let now = OffsetDateTime::now_utc(); + match classify_rebalance_terminal_event(Some(Err(Error::SlowDown)), now) { + RebalanceTerminalEvent::Failed { msg, last_error } => { + assert!(msg.contains("Rebalance failed")); + assert!(msg.contains("with err")); + assert_eq!(last_error, Error::SlowDown.to_string()); + } + _ => panic!("expected failed terminal event"), + } +} + +#[test] +fn test_classify_rebalance_terminal_event_channel_closed() { + let now = OffsetDateTime::now_utc(); + match classify_rebalance_terminal_event(None, now) { + RebalanceTerminalEvent::ChannelClosed { msg, last_error } => { + assert!(msg.contains("channel closed")); + assert!(last_error.contains("before terminal event")); + assert!(msg.contains("at")); + assert!(last_error.contains("at")); + } + _ => panic!("expected channel closed terminal event"), + } +} + +#[test] +fn test_apply_rebalance_terminal_event_channel_closed_marks_failed() { + let now = OffsetDateTime::now_utc(); + let mut status = RebalStatus::Started; + let mut end_time = None; + let mut last_error = None; + + apply_rebalance_terminal_event( + &mut status, + &mut end_time, + &mut last_error, + RebalanceTerminalEvent::ChannelClosed { + msg: "channel closed".to_string(), + last_error: "rebalance save channel closed before terminal event".to_string(), + }, + now, + ); + + assert_eq!(status, RebalStatus::Failed); + assert_eq!(end_time, Some(now)); + assert_eq!(last_error.as_deref(), Some("rebalance save channel closed before terminal event")); +} + +#[test] +fn test_apply_rebalance_terminal_event_stopped_clears_error() { + let now = OffsetDateTime::now_utc(); + let mut status = RebalStatus::Started; + let mut end_time = None; + let mut last_error = Some("old-error".to_string()); + + apply_rebalance_terminal_event( + &mut status, + &mut end_time, + &mut last_error, + RebalanceTerminalEvent::Stopped { + msg: "rebalance stopped".to_string(), + }, + now, + ); + + assert_eq!(status, RebalStatus::Stopped); + assert_eq!(end_time, Some(now)); + assert_eq!(last_error, None); +} + +#[test] +fn test_is_rebalance_stopped_terminal_event_only_matches_stopped_variant() { + let stopped = RebalanceTerminalEvent::Stopped { + msg: "stopped".to_string(), + }; + let completed = RebalanceTerminalEvent::Completed { + msg: "completed".to_string(), + }; + + assert!(is_rebalance_stopped_terminal_event(&stopped)); + assert!(!is_rebalance_stopped_terminal_event(&completed)); +} + +#[test] +fn test_should_preserve_rebalance_stopped_state_when_meta_marked_stopped() { + let event = RebalanceTerminalEvent::Completed { + msg: "completed".to_string(), + }; + + assert!(should_preserve_rebalance_stopped_state(true, RebalStatus::Started, &event)); +} + +#[test] +fn test_should_preserve_rebalance_stopped_state_when_pool_already_stopped() { + let event = RebalanceTerminalEvent::Failed { + msg: "failed".to_string(), + last_error: "boom".to_string(), + }; + + assert!(should_preserve_rebalance_stopped_state(false, RebalStatus::Stopped, &event)); +} + +#[test] +fn test_should_preserve_rebalance_stopped_state_allows_stopped_terminal_update() { + let event = RebalanceTerminalEvent::Stopped { + msg: "stopped".to_string(), + }; + + assert!(!should_preserve_rebalance_stopped_state(true, RebalStatus::Started, &event)); +} + +#[test] +fn test_ensure_rebalance_not_decommissioning_rejects_running_decommission() { + assert!(!ensure_rebalance_not_decommissioning(true)); +} + +#[test] +fn test_ensure_rebalance_not_decommissioning_allows_idle_decommission() { + assert!(ensure_rebalance_not_decommissioning(false)); +} + +#[test] +fn test_validate_start_rebalance_state_rejects_running_decommission() { + let err = validate_start_rebalance_state(true, true).expect_err("running decommission should block rebalance start"); + assert!(matches!(err, Error::DecommissionAlreadyRunning)); +} + +#[test] +fn test_validate_start_rebalance_state_rejects_missing_meta() { + let err = validate_start_rebalance_state(false, false).expect_err("missing rebalance meta should fail"); + assert!(matches!(err, Error::ConfigNotFound)); +} + +#[test] +fn test_validate_start_rebalance_state_allows_loaded_meta() { + validate_start_rebalance_state(false, true).expect("loaded rebalance meta should allow start"); +} + +#[test] +fn test_percent_free_ratio_zero_capacity_is_zero() { + assert_eq!(percent_free_ratio(100, 0), 0.0); +} + +#[test] +fn test_percent_free_ratio_normal_case() { + assert_eq!(percent_free_ratio(250, 1_000), 0.25); +} + +#[test] +fn test_should_pool_participate_false_when_capacity_zero() { + assert!(!should_pool_participate(0, 0, 0.2)); +} + +#[test] +fn test_should_pool_participate_true_when_ratio_below_goal() { + assert!(should_pool_participate(200, 1_000, 0.3)); +} + +#[test] +fn test_should_pool_participate_false_when_ratio_meets_goal() { + assert!(!should_pool_participate(300, 1_000, 0.3)); +} + +#[test] +fn test_rebalance_goal_not_reached_for_issue_3137_initial_imbalance() { + let pool0_capacity = 10_000_u64; + let pool0_free = 9_135_u64; + let pool1_capacity = 10_000_u64; + let pool1_free = 9_800_u64; + let goal = percent_free_ratio(pool0_free + pool1_free, pool0_capacity + pool1_capacity); + + assert!(should_pool_participate(pool0_free, pool0_capacity, goal)); + assert!(!rebalance_goal_reached(pool0_free, pool0_capacity, 0, goal)); +} + +#[test] +fn test_complete_rebalance_pools_at_goal_marks_started_participants_completed() { + let now = OffsetDateTime::from_unix_timestamp(1_000).unwrap(); + let mut meta = RebalanceMeta { + 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() + }, + RebalanceStats { + participating: true, + init_free_space: 100, + init_capacity: 1_000, + bytes: 0, + info: RebalanceInfo { + status: RebalStatus::Started, + ..Default::default() + }, + ..Default::default() + }, + ], + ..Default::default() + }; + + assert!(complete_rebalance_pools_at_goal(&mut meta, now)); + assert_eq!(meta.pool_stats[0].info.status, RebalStatus::Completed); + assert_eq!(meta.pool_stats[0].info.end_time, Some(now)); + assert_eq!(meta.pool_stats[1].info.status, RebalStatus::Started); +} + +#[test] +fn test_complete_rebalance_pools_at_goal_skips_deferred_transient_error() { + let now = OffsetDateTime::from_unix_timestamp(1_000).unwrap(); + let mut meta = RebalanceMeta { + 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, + last_error: Some(format!("{REBALANCE_DEFERRED_ENTRY_ERROR_PREFIX} timeout")), + ..Default::default() + }, + ..Default::default() + }], + ..Default::default() + }; + + assert!(!complete_rebalance_pools_at_goal(&mut meta, now)); + assert_eq!(meta.pool_stats[0].info.status, RebalStatus::Started); + assert!(has_deferred_rebalance_error(&meta.pool_stats[0])); +} + +#[test] +fn test_complete_rebalance_pools_with_empty_queue_marks_started_participants_completed() { + let now = OffsetDateTime::from_unix_timestamp(1_000).unwrap(); + let mut meta = RebalanceMeta { + pool_stats: vec![ + RebalanceStats { + participating: true, + buckets: Vec::new(), + info: RebalanceInfo { + status: RebalStatus::Started, + last_error: Some("stale error".to_string()), + ..Default::default() + }, + ..Default::default() + }, + RebalanceStats { + participating: true, + buckets: vec!["bucket-a".to_string()], + info: RebalanceInfo { + status: RebalStatus::Started, + ..Default::default() + }, + ..Default::default() + }, + ], + ..Default::default() + }; + + assert!(complete_rebalance_pools_with_empty_queue(&mut meta, now)); + assert_eq!(meta.pool_stats[0].info.status, RebalStatus::Completed); + assert_eq!(meta.pool_stats[0].info.end_time, Some(now)); + assert!(meta.pool_stats[0].info.last_error.is_none()); + assert_eq!(meta.pool_stats[1].info.status, RebalStatus::Started); +} + +#[test] +fn test_should_skip_start_rebalance_only_when_running_and_cancel_attached() { + assert!(should_skip_start_rebalance(true, true)); + assert!(!should_skip_start_rebalance(true, false)); + assert!(!should_skip_start_rebalance(false, true)); + assert!(!should_skip_start_rebalance(false, false)); +} + +#[test] +fn test_new_multipart_abort_flag_defaults_to_abort_enabled() { + let flag = data_movement::new_multipart_abort_flag(); + assert!(data_movement::should_abort_multipart_upload(&flag)); +} + +#[test] +fn test_mark_multipart_upload_completed_disables_abort_cleanup() { + let flag = data_movement::new_multipart_abort_flag(); + data_movement::mark_multipart_upload_completed(&flag); + assert!(!data_movement::should_abort_multipart_upload(&flag)); +} + +#[test] +fn test_decode_part_index_returns_none_when_absent() { + assert!(data_movement::decode_part_index(None).is_none()); +} + +#[test] +fn test_decode_part_index_returns_none_for_invalid_payload() { + let invalid = bytes::Bytes::from_static(b"not-a-valid-index"); + assert!(data_movement::decode_part_index(Some(&invalid)).is_none()); +} + +#[test] +fn test_decode_part_index_returns_some_for_valid_payload() { + let mut index = Index::new(); + index.add(0, 0).expect("first index entry should be accepted"); + index + .add(2_097_152, 2_097_152) + .expect("second index entry should advance totals"); + + let encoded = index.into_vec(); + let decoded = data_movement::decode_part_index(Some(&encoded)).expect("valid index payload should decode"); + + assert_eq!(decoded.total_uncompressed, 2_097_152); + assert_eq!(decoded.total_compressed, 2_097_152); +} + +#[test] +fn test_merge_rebalance_bucket_lists_keeps_existing_order_and_skips_duplicates() { + let mut remote = vec!["bucket-a".to_string(), "bucket-b".to_string()]; + let local = vec![ + "bucket-b".to_string(), + "bucket-c".to_string(), + "bucket-c".to_string(), + "bucket-d".to_string(), + ]; + + merge_rebalance_bucket_lists(&mut remote, &local); + + assert_eq!( + remote, + vec![ + "bucket-a".to_string(), + "bucket-b".to_string(), + "bucket-c".to_string(), + "bucket-d".to_string() + ] + ); +} + +#[test] +fn test_remove_rebalanced_buckets_from_queue_filters_membership_linearly() { + let mut pool_stat = RebalanceStats { + buckets: vec![ + "bucket-a".to_string(), + "bucket-b".to_string(), + "bucket-c".to_string(), + "bucket-b".to_string(), + ], + rebalanced_buckets: vec!["bucket-b".to_string(), "bucket-d".to_string()], + ..Default::default() + }; + + remove_rebalanced_buckets_from_queue(&mut pool_stat); + + assert_eq!(pool_stat.buckets, vec!["bucket-a".to_string(), "bucket-c".to_string()]); +} + +#[test] +fn test_resolve_rebalance_participants_respects_runtime_pool_count() { + let now = OffsetDateTime::now_utc(); + let stats = vec![ + RebalanceStats { + participating: true, + info: RebalanceInfo { + status: RebalStatus::Started, + start_time: Some(now), + ..Default::default() + }, + ..Default::default() + }, + RebalanceStats { + participating: false, + info: RebalanceInfo { + status: RebalStatus::Started, + start_time: Some(now), + ..Default::default() + }, + ..Default::default() + }, + RebalanceStats { + participating: true, + info: RebalanceInfo { + status: RebalStatus::Started, + start_time: Some(now), + ..Default::default() + }, + ..Default::default() + }, + ]; + + let participants = resolve_rebalance_participants(stats.as_slice(), 2); + assert_eq!(participants, vec![true, false]); +} + +#[test] +fn test_resolve_rebalance_participants_requires_started_status() { + let now = OffsetDateTime::now_utc(); + let stats = vec![ + RebalanceStats { + participating: true, + info: RebalanceInfo { + status: RebalStatus::Completed, + start_time: Some(now), + ..Default::default() + }, + ..Default::default() + }, + RebalanceStats { + participating: true, + info: RebalanceInfo { + status: RebalStatus::Started, + start_time: Some(now), + ..Default::default() + }, + ..Default::default() + }, + ]; + + let participants = resolve_rebalance_participants(stats.as_slice(), 2); + assert_eq!(participants, vec![false, true]); +} + +#[test] +fn test_is_rebalance_actively_running_requires_cancel_and_started_state() { + let now = OffsetDateTime::now_utc(); + let mut meta = RebalanceMeta { + pool_stats: vec![RebalanceStats { + participating: true, + info: RebalanceInfo { + start_time: Some(now), + status: RebalStatus::Started, + ..Default::default() + }, + ..Default::default() + }], + ..Default::default() + }; + + assert!(!is_rebalance_actively_running(&meta)); + meta.cancel = Some(CancellationToken::new()); + assert!(is_rebalance_actively_running(&meta)); +} + +#[test] +fn test_is_rebalance_in_progress_only_started_participants() { + let now = OffsetDateTime::now_utc(); + let meta = RebalanceMeta { + stopped_at: None, + pool_stats: vec![ + RebalanceStats { + participating: true, + info: RebalanceInfo { + start_time: Some(now), + status: RebalStatus::Completed, + ..Default::default() + }, + ..Default::default() + }, + RebalanceStats { + participating: true, + info: RebalanceInfo { + start_time: Some(now), + status: RebalStatus::Started, + ..Default::default() + }, + ..Default::default() + }, + ], + ..Default::default() + }; + + assert!(is_rebalance_in_progress(&meta)); +} + +#[test] +fn test_is_rebalance_conflicting_with_decommission_true_when_in_progress() { + let now = OffsetDateTime::now_utc(); + let meta = RebalanceMeta { + stopped_at: None, + pool_stats: vec![RebalanceStats { + participating: true, + info: RebalanceInfo { + start_time: Some(now), + status: RebalStatus::Started, + ..Default::default() + }, + ..Default::default() + }], + ..Default::default() + }; + + assert!(is_rebalance_conflicting_with_decommission(&meta)); +} + +#[test] +fn test_is_rebalance_conflicting_with_decommission_false_when_stopped() { + let now = OffsetDateTime::now_utc(); + let meta = RebalanceMeta { + stopped_at: Some(now), + pool_stats: vec![RebalanceStats { + participating: true, + info: RebalanceInfo { + start_time: Some(now), + status: RebalStatus::Started, + ..Default::default() + }, + ..Default::default() + }], + ..Default::default() + }; + + assert!(!is_rebalance_conflicting_with_decommission(&meta)); +} + +#[test] +fn test_is_rebalance_in_progress_stopped_takes_precedence() { + let now = OffsetDateTime::now_utc(); + let meta = RebalanceMeta { + stopped_at: Some(now), + pool_stats: vec![RebalanceStats { + participating: true, + info: RebalanceInfo { + start_time: Some(now), + status: RebalStatus::Started, + ..Default::default() + }, + ..Default::default() + }], + ..Default::default() + }; + + assert!(!is_rebalance_in_progress(&meta)); +} + +#[test] +fn test_first_rebalance_bucket_returns_first_name() { + let pool_stat = RebalanceStats { + buckets: vec!["bucket-a".to_string(), "bucket-b".to_string()], + ..Default::default() + }; + + assert_eq!(first_rebalance_bucket(&pool_stat), Some("bucket-a".to_string())); +} + +#[test] +fn test_first_rebalance_bucket_returns_none_when_empty() { + let pool_stat = RebalanceStats::default(); + + assert_eq!(first_rebalance_bucket(&pool_stat), None); +} + +#[test] +fn test_next_rebal_bucket_from_stat_respects_empty_queue() { + let pool_stat = RebalanceStats { + buckets: vec![], + ..Default::default() + }; + + assert_eq!(next_rebal_bucket_from_stat(&pool_stat), None); +} + +#[test] +fn test_next_rebal_bucket_from_stat_returns_first_bucket() { + let now = OffsetDateTime::now_utc(); + let pool_stat = RebalanceStats { + participating: true, + info: RebalanceInfo { + status: RebalStatus::Started, + start_time: Some(now), + ..Default::default() + }, + buckets: vec!["bucket-a".to_string(), "bucket-b".to_string()], + ..Default::default() + }; + + assert_eq!(next_rebal_bucket_from_stat(&pool_stat), Some("bucket-a".to_string())); +} + +#[test] +fn test_clone_rebalance_pool_stats_rejects_missing_meta() { + let err = clone_rebalance_pool_stats(None).expect_err("missing rebalance meta should fail"); + assert!( + err.to_string() + .contains("failed to clone rebalance pool stats: rebalance metadata not initialized") + ); +} + +#[test] +fn test_clone_rebalance_pool_stats_clones_entries() { + let meta = RebalanceMeta { + pool_stats: vec![RebalanceStats::default()], + ..Default::default() + }; + + let stats = clone_rebalance_pool_stats(Some(&meta)).expect("metadata should clone pool stats"); + assert_eq!(stats.len(), 1); +} + +#[test] +fn test_resolve_next_rebalance_bucket_rejects_missing_meta() { + let err = resolve_next_rebalance_bucket(None, 0).expect_err("missing meta should fail"); + assert!( + err.to_string() + .contains("failed to resolve next rebalance bucket: rebalance metadata not initialized") + ); +} + +#[test] +fn test_resolve_next_rebalance_bucket_rejects_invalid_pool_index() { + let meta = RebalanceMeta { + pool_stats: vec![RebalanceStats::default()], + ..Default::default() + }; + + let err = resolve_next_rebalance_bucket(Some(&meta), 3).expect_err("invalid pool index should fail"); + assert!(err.to_string().contains("invalid rebalance pool index 3 for 1 pools")); +} + +#[test] +fn test_resolve_next_rebalance_bucket_returns_none_for_completed_pool() { + let meta = RebalanceMeta { + pool_stats: vec![RebalanceStats { + participating: true, + info: RebalanceInfo { + status: RebalStatus::Completed, + ..Default::default() + }, + buckets: vec!["bucket-a".to_string()], + ..Default::default() + }], + ..Default::default() + }; + + let next = resolve_next_rebalance_bucket(Some(&meta), 0).expect("completed pool should return none"); + assert!(next.is_none()); +} + +#[test] +fn test_resolve_next_rebalance_bucket_returns_first_bucket_for_active_pool() { + let now = OffsetDateTime::now_utc(); + let meta = RebalanceMeta { + pool_stats: vec![RebalanceStats { + participating: true, + info: RebalanceInfo { + status: RebalStatus::Started, + start_time: Some(now), + ..Default::default() + }, + buckets: vec!["bucket-a".to_string(), "bucket-b".to_string()], + ..Default::default() + }], + ..Default::default() + }; + + let next = resolve_next_rebalance_bucket(Some(&meta), 0).expect("active pool should return first bucket"); + assert_eq!(next.as_deref(), Some("bucket-a")); +} + +#[test] +fn test_take_bucket_from_rebalance_queue_moves_bucket_and_keeps_remaining() { + let mut pool_stat = RebalanceStats { + buckets: vec!["bucket-a".to_string(), "bucket-b".to_string(), "bucket-a".to_string()], + rebalanced_buckets: Vec::new(), + ..Default::default() + }; + + assert!(take_bucket_from_rebalance_queue(&mut pool_stat, "bucket-a")); + assert_eq!(pool_stat.buckets, vec!["bucket-b".to_string()]); + assert_eq!(pool_stat.rebalanced_buckets, vec!["bucket-a".to_string(), "bucket-a".to_string()]); +} + +#[test] +fn test_take_bucket_from_rebalance_queue_no_match_keeps_queue() { + let mut pool_stat = RebalanceStats { + buckets: vec!["bucket-a".to_string(), "bucket-b".to_string()], + rebalanced_buckets: Vec::new(), + ..Default::default() + }; + + assert!(!take_bucket_from_rebalance_queue(&mut pool_stat, "bucket-c")); + assert_eq!(pool_stat.buckets, vec!["bucket-a".to_string(), "bucket-b".to_string()]); + assert!(pool_stat.rebalanced_buckets.is_empty()); +} + +#[test] +fn test_defer_bucket_in_rebalance_queue_moves_bucket_to_back() { + let mut pool_stat = RebalanceStats { + buckets: vec!["bucket-a".to_string(), "bucket-b".to_string(), "bucket-c".to_string()], + rebalanced_buckets: Vec::new(), + ..Default::default() + }; + + defer_bucket_in_rebalance_queue(&mut pool_stat, "bucket-a").expect("queued bucket should be deferred"); + + assert_eq!( + pool_stat.buckets, + vec!["bucket-b".to_string(), "bucket-c".to_string(), "bucket-a".to_string()] + ); + assert!(pool_stat.rebalanced_buckets.is_empty()); +} + +#[test] +fn test_defer_bucket_in_rebalance_queue_rejects_missing_bucket() { + let mut pool_stat = RebalanceStats { + buckets: vec!["bucket-a".to_string(), "bucket-b".to_string()], + rebalanced_buckets: Vec::new(), + ..Default::default() + }; + + let err = defer_bucket_in_rebalance_queue(&mut pool_stat, "bucket-c").expect_err("missing bucket should not be deferred"); + + assert!(err.to_string().contains("failed to defer rebalance bucket bucket-c")); + assert_eq!(pool_stat.buckets, vec!["bucket-a".to_string(), "bucket-b".to_string()]); + assert!(pool_stat.rebalanced_buckets.is_empty()); +} + +#[test] +fn test_mark_rebalance_bucket_done_rejects_missing_meta() { + let err = mark_rebalance_bucket_done(None, 0, "bucket-a").expect_err("missing meta should fail"); + assert!( + err.to_string() + .contains("failed to mark rebalance bucket done: rebalance metadata not initialized") + ); +} + +#[test] +fn test_mark_rebalance_bucket_done_rejects_invalid_pool_index() { + let mut meta = RebalanceMeta { + pool_stats: vec![RebalanceStats::default()], + ..Default::default() + }; + + let err = mark_rebalance_bucket_done(Some(&mut meta), 3, "bucket-a").expect_err("invalid pool index should fail"); + assert!(err.to_string().contains("invalid rebalance pool index 3 for 1 pools")); +} + +#[test] +fn test_mark_rebalance_bucket_done_rejects_missing_bucket() { + let mut meta = RebalanceMeta { + pool_stats: vec![RebalanceStats { + buckets: vec!["bucket-a".to_string()], + ..Default::default() + }], + ..Default::default() + }; + + let err = mark_rebalance_bucket_done(Some(&mut meta), 0, "bucket-x").expect_err("missing bucket should fail"); + assert!( + err.to_string() + .contains("failed to mark rebalance bucket done: bucket bucket-x was not queued for pool 0") + ); +} + +#[test] +fn test_mark_rebalance_bucket_done_marks_bucket_as_rebalanced() { + let mut meta = RebalanceMeta { + pool_stats: vec![RebalanceStats { + buckets: vec!["bucket-a".to_string(), "bucket-b".to_string()], + rebalanced_buckets: Vec::new(), + ..Default::default() + }], + ..Default::default() + }; + + mark_rebalance_bucket_done(Some(&mut meta), 0, "bucket-a").expect("bucket in queue should be marked done"); + assert_eq!(meta.pool_stats[0].buckets, vec!["bucket-b".to_string()]); + assert_eq!(meta.pool_stats[0].rebalanced_buckets, vec!["bucket-a".to_string()]); +} + +#[test] +fn test_mark_rebalance_bucket_done_clears_deferred_transient_error() { + let mut meta = RebalanceMeta { + pool_stats: vec![RebalanceStats { + buckets: vec!["bucket-a".to_string()], + info: RebalanceInfo { + last_error: Some(format!("{REBALANCE_DEFERRED_ENTRY_ERROR_PREFIX} timeout")), + ..Default::default() + }, + ..Default::default() + }], + ..Default::default() + }; + + mark_rebalance_bucket_done(Some(&mut meta), 0, "bucket-a") + .expect("bucket in queue should be marked done after deferred retry succeeds"); + + assert!(meta.pool_stats[0].info.last_error.is_none()); +} + +#[test] +fn test_record_rebalance_cleanup_warning_in_meta_preserves_last_error() { + let now = OffsetDateTime::from_unix_timestamp(10_000).unwrap(); + let mut meta = RebalanceMeta { + pool_stats: vec![RebalanceStats { + info: RebalanceInfo { + last_error: Some("old-error".to_string()), + ..Default::default() + }, + ..Default::default() + }], + ..Default::default() + }; + + record_rebalance_cleanup_warning_in_meta(Some(&mut meta), 0, "bucket-a", "obj.txt", "cleanup failed".to_string(), now) + .expect("cleanup warning should be recorded"); + + assert_eq!(meta.pool_stats[0].info.last_error.as_deref(), Some("old-error")); + assert_eq!(meta.pool_stats[0].cleanup_warnings.count, 1); + assert_eq!(meta.pool_stats[0].cleanup_warnings.last_message.as_deref(), Some("cleanup failed")); + assert_eq!(meta.pool_stats[0].cleanup_warnings.last_bucket.as_deref(), Some("bucket-a")); + assert_eq!(meta.pool_stats[0].cleanup_warnings.last_object.as_deref(), Some("obj.txt")); + assert_eq!(meta.pool_stats[0].cleanup_warnings.last_at, Some(now)); + assert_eq!(meta.last_refreshed_at, Some(now)); +} + +#[test] +fn test_complete_rebalance_pools_with_empty_queue_preserves_cleanup_warnings() { + let warning_at = OffsetDateTime::from_unix_timestamp(9_000).unwrap(); + let completed_at = OffsetDateTime::from_unix_timestamp(10_000).unwrap(); + let mut meta = RebalanceMeta { + pool_stats: vec![RebalanceStats { + participating: true, + info: RebalanceInfo { + status: RebalStatus::Started, + ..Default::default() + }, + cleanup_warnings: RebalanceCleanupWarnings { + count: 1, + last_message: Some("cleanup failed".to_string()), + last_bucket: Some("bucket-a".to_string()), + last_object: Some("obj.txt".to_string()), + last_at: Some(warning_at), + }, + ..Default::default() + }], + ..Default::default() + }; + + assert!(complete_rebalance_pools_with_empty_queue(&mut meta, completed_at)); + + assert_eq!(meta.pool_stats[0].info.status, RebalStatus::Completed); + assert!(meta.pool_stats[0].info.last_error.is_none()); + assert_eq!(meta.pool_stats[0].cleanup_warnings.count, 1); + assert_eq!(meta.pool_stats[0].cleanup_warnings.last_message.as_deref(), Some("cleanup failed")); + assert_eq!(meta.pool_stats[0].cleanup_warnings.last_at, Some(warning_at)); +} + +#[test] +fn test_apply_stopped_at_transitions_started_pools_only() { + let now = OffsetDateTime::now_utc(); + let mut meta = RebalanceMeta { + pool_stats: vec![ + RebalanceStats { + info: RebalanceInfo { + status: RebalStatus::Started, + end_time: None, + last_error: Some("old".to_string()), + ..Default::default() + }, + ..Default::default() + }, + RebalanceStats { + info: RebalanceInfo { + status: RebalStatus::Failed, + end_time: Some(now), + last_error: Some("failed".to_string()), + ..Default::default() + }, + ..Default::default() + }, + ], + ..Default::default() + }; + + apply_stopped_at(&mut meta, now); + + assert_eq!(meta.stopped_at, Some(now)); + assert_eq!(meta.pool_stats[0].info.status, RebalStatus::Stopped); + assert_eq!(meta.pool_stats[0].info.end_time, Some(now)); + assert_eq!(meta.pool_stats[0].info.last_error, None); + + assert_eq!(meta.pool_stats[1].info.status, RebalStatus::Failed); + assert_eq!(meta.pool_stats[1].info.end_time, Some(now)); + assert_eq!(meta.pool_stats[1].info.last_error.as_deref(), Some("failed")); +} + +#[test] +fn test_stop_rebalance_state_cancels_token_and_marks_stopped_when_in_progress() { + let now = OffsetDateTime::from_unix_timestamp(10_000).unwrap(); + let cancel = CancellationToken::new(); + let cancel_clone = cancel.clone(); + let mut meta = RebalanceMeta { + cancel: Some(cancel), + pool_stats: vec![RebalanceStats { + participating: true, + info: RebalanceInfo { + status: RebalStatus::Started, + ..Default::default() + }, + ..Default::default() + }], + ..Default::default() + }; + + stop_rebalance_state(&mut meta, now); + + assert!(cancel_clone.is_cancelled()); + assert!(meta.cancel.is_none()); + assert_eq!(meta.stopped_at, Some(now)); + assert_eq!(meta.pool_stats[0].info.status, RebalStatus::Stopped); + assert_eq!(meta.pool_stats[0].info.end_time, Some(now)); +} + +#[test] +fn test_stop_rebalance_state_clears_token_without_forcing_stopped_when_not_in_progress() { + let now = OffsetDateTime::from_unix_timestamp(20_000).unwrap(); + let cancel = CancellationToken::new(); + let cancel_clone = cancel.clone(); + let mut meta = RebalanceMeta { + cancel: Some(cancel), + pool_stats: vec![RebalanceStats { + participating: true, + info: RebalanceInfo { + status: RebalStatus::Completed, + end_time: Some(now), + ..Default::default() + }, + ..Default::default() + }], + ..Default::default() + }; + + stop_rebalance_state(&mut meta, now); + + assert!(cancel_clone.is_cancelled()); + assert!(meta.cancel.is_none()); + assert_eq!(meta.stopped_at, None); + assert_eq!(meta.pool_stats[0].info.status, RebalStatus::Completed); +} + +#[test] +fn test_stop_rebalance_state_normalizes_started_pool_when_stopped_at_already_set() { + let stopped_at = OffsetDateTime::from_unix_timestamp(30_000).unwrap(); + let now = OffsetDateTime::from_unix_timestamp(40_000).unwrap(); + let cancel = CancellationToken::new(); + let cancel_clone = cancel.clone(); + let mut meta = RebalanceMeta { + cancel: Some(cancel), + stopped_at: Some(stopped_at), + pool_stats: vec![RebalanceStats { + participating: true, + info: RebalanceInfo { + status: RebalStatus::Started, + last_error: Some("stale".to_string()), + ..Default::default() + }, + ..Default::default() + }], + ..Default::default() + }; + + stop_rebalance_state(&mut meta, now); + + assert!(cancel_clone.is_cancelled()); + assert!(meta.cancel.is_none()); + assert_eq!(meta.stopped_at, Some(stopped_at)); + assert_eq!(meta.pool_stats[0].info.status, RebalStatus::Stopped); + assert_eq!(meta.pool_stats[0].info.end_time, Some(stopped_at)); + assert_eq!(meta.pool_stats[0].info.last_error, None); +} + +#[test] +fn test_stop_rebalance_meta_snapshot_returns_none_when_meta_missing() { + let now = OffsetDateTime::from_unix_timestamp(50_000).unwrap(); + assert!(stop_rebalance_meta_snapshot(None, now).is_none()); +} + +#[test] +fn test_stop_rebalance_meta_snapshot_stops_meta_and_returns_snapshot() { + let now = OffsetDateTime::from_unix_timestamp(60_000).unwrap(); + let cancel = CancellationToken::new(); + let cancel_clone = cancel.clone(); + let mut meta = RebalanceMeta { + cancel: Some(cancel), + pool_stats: vec![RebalanceStats { + participating: true, + info: RebalanceInfo { + status: RebalStatus::Started, + ..Default::default() + }, + ..Default::default() + }], + ..Default::default() + }; + + let snapshot = stop_rebalance_meta_snapshot(Some(&mut meta), now).expect("snapshot should be returned for present meta"); + + assert!(cancel_clone.is_cancelled()); + assert!(meta.cancel.is_none()); + assert_eq!(meta.stopped_at, Some(now)); + assert_eq!(meta.last_refreshed_at, Some(now)); + assert_eq!(meta.pool_stats[0].info.status, RebalStatus::Stopped); + + assert!(snapshot.cancel.is_none()); + assert_eq!(snapshot.stopped_at, Some(now)); + assert_eq!(snapshot.last_refreshed_at, Some(now)); + assert_eq!(snapshot.pool_stats[0].info.status, RebalStatus::Stopped); + assert_eq!(snapshot.pool_stats[0].info.end_time, Some(now)); +} + +#[test] +fn test_apply_rebalance_save_option_stats_keeps_pool_status_and_updates_refresh() { + let now = OffsetDateTime::from_unix_timestamp(1_000).unwrap(); + let later = OffsetDateTime::from_unix_timestamp(2_000).unwrap(); + let mut meta = RebalanceMeta { + pool_stats: vec![RebalanceStats { + participating: true, + info: RebalanceInfo { + status: RebalStatus::Started, + start_time: Some(now), + ..Default::default() + }, + buckets: vec!["bucket-a".to_string()], + ..Default::default() + }], + last_refreshed_at: Some(now), + stopped_at: None, + ..Default::default() + }; + + apply_rebalance_save_option(&mut meta, 0, RebalSaveOpt::Stats, later); + + assert_eq!(meta.last_refreshed_at, Some(later)); + assert_eq!(meta.stopped_at, None); + assert_eq!(meta.pool_stats.len(), 1); + assert!(meta.pool_stats[0].participating); + assert_eq!(meta.pool_stats[0].info.status, RebalStatus::Started); + assert_eq!(meta.pool_stats[0].info.start_time, Some(now)); + assert_eq!(meta.pool_stats[0].buckets, vec!["bucket-a".to_string()]); +} + +#[test] +fn test_apply_rebalance_save_option_stopped_at_updates_refresh_and_statuses() { + let now = OffsetDateTime::from_unix_timestamp(1_000).unwrap(); + let mut meta = RebalanceMeta { + pool_stats: vec![ + RebalanceStats { + info: RebalanceInfo { + status: RebalStatus::Started, + ..Default::default() + }, + ..Default::default() + }, + RebalanceStats { + info: RebalanceInfo { + status: RebalStatus::Failed, + last_error: Some("previous failure".to_string()), + ..Default::default() + }, + ..Default::default() + }, + ], + ..Default::default() + }; + + apply_rebalance_save_option(&mut meta, 9_000, RebalSaveOpt::StoppedAt, now); + + assert_eq!(meta.stopped_at, Some(now)); + assert_eq!(meta.last_refreshed_at, Some(now)); + assert_eq!(meta.pool_stats[0].info.status, RebalStatus::Stopped); + assert_eq!(meta.pool_stats[0].info.end_time, Some(now)); + assert!(meta.pool_stats[0].info.last_error.is_none()); + assert_eq!(meta.pool_stats[1].info.status, RebalStatus::Failed); + assert_eq!(meta.pool_stats[1].info.last_error.as_deref(), Some("previous failure")); +} + +#[test] +fn test_rebalance_stats_update_counts_and_bytes_growth() { + let mut stat = RebalanceStats { + bucket: "bucket-a".to_string(), + object: "obj-previous".to_string(), + num_objects: 3, + num_versions: 5, + bytes: 120, + ..Default::default() + }; + + let mut latest = FileInfo::new("object-1", 4, 2); + latest.name = "object-1".to_string(); + latest.is_latest = true; + latest.size = 300; + latest.deleted = false; + latest.mod_time = Some(OffsetDateTime::UNIX_EPOCH); + latest.version_id = None; + + let mut historical = FileInfo::new("object-1", 4, 2); + historical.name = "object-1".to_string(); + historical.is_latest = false; + historical.size = 128; + historical.deleted = false; + historical.mod_time = Some(OffsetDateTime::UNIX_EPOCH); + historical.version_id = None; + + let mut tombstone = FileInfo::new("object-1", 4, 2); + tombstone.name = "object-1".to_string(); + tombstone.is_latest = false; + tombstone.size = 64; + tombstone.deleted = true; + tombstone.mod_time = Some(OffsetDateTime::UNIX_EPOCH); + tombstone.version_id = None; + + stat.update("bucket-b".to_string(), &latest); + stat.update("bucket-b".to_string(), &historical); + stat.update("bucket-b".to_string(), &tombstone); + + assert_eq!(stat.bucket, "bucket-b"); + assert_eq!(stat.object, "object-1"); + assert_eq!(stat.num_objects, 4); + assert_eq!(stat.num_versions, 8); + let expected_bytes = 120_u64 + + (latest.size * (latest.erasure.data_blocks + latest.erasure.parity_blocks) as i64 / latest.erasure.data_blocks as i64) + as u64 + + (historical.size * (historical.erasure.data_blocks + historical.erasure.parity_blocks) as i64 + / historical.erasure.data_blocks as i64) as u64; + assert_eq!(stat.bytes, expected_bytes); +} + +#[test] +fn test_rebalance_stats_update_batch_matches_repeated_updates() { + let mut latest = version_normal(); + latest.is_latest = true; + latest.size = 128; + + let mut historical = version_normal(); + historical.is_latest = false; + historical.size = 64; + + let mut repeated = RebalanceStats::default(); + repeated.update("bucket-a".to_string(), &latest); + repeated.update("bucket-a".to_string(), &historical); + + let mut batched = RebalanceStats::default(); + batched.update_batch("bucket-a".to_string(), &[&latest, &historical]); + + assert_eq!(batched.bucket, repeated.bucket); + assert_eq!(batched.object, repeated.object); + assert_eq!(batched.num_objects, repeated.num_objects); + assert_eq!(batched.num_versions, repeated.num_versions); + assert_eq!(batched.bytes, repeated.bytes); +} + +#[test] +fn test_rebalance_stats_update_ignores_invalid_data_blocks() { + let mut stat = RebalanceStats { + bucket: "bucket-a".to_string(), + object: "obj-previous".to_string(), + num_objects: 1, + num_versions: 2, + bytes: 77, + ..Default::default() + }; + + let mut invalid = FileInfo::new("object-invalid", 0, 2); + invalid.name = "object-invalid".to_string(); + invalid.is_latest = true; + invalid.size = 256; + invalid.deleted = false; + invalid.mod_time = Some(OffsetDateTime::UNIX_EPOCH); + invalid.version_id = None; + + stat.update("bucket-z".to_string(), &invalid); + + assert_eq!(stat.bucket, "bucket-z"); + assert_eq!(stat.object, "object-invalid"); + assert_eq!(stat.num_objects, 2); + assert_eq!(stat.num_versions, 3); + assert_eq!(stat.bytes, 77); +} + +#[test] +fn test_rebalance_goal_reached_requires_target_ratio() { + let init_free_space = 150_u64; + let init_capacity = 800_u64; + let goal = 0.35_f64; + + assert!(!rebalance_goal_reached(init_free_space, init_capacity, 0, goal)); + assert!(!rebalance_goal_reached(init_free_space, init_capacity, 129, goal)); + assert!(rebalance_goal_reached(init_free_space, init_capacity, 130, goal)); +} diff --git a/crates/ecstore/src/rebalance/runtime.rs b/crates/ecstore/src/rebalance/runtime.rs new file mode 100644 index 000000000..391087712 --- /dev/null +++ b/crates/ecstore/src/rebalance/runtime.rs @@ -0,0 +1,574 @@ +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, + has_deferred_rebalance_error, is_rebalance_in_progress, rebalance_goal_reached, resolve_rebalance_participants, + should_preserve_rebalance_stopped_state, should_skip_start_rebalance, validate_start_rebalance_state, +}; +use super::worker::{ + resolve_rebalance_bucket_result, resolve_rebalance_meta_save_result, resolve_rebalance_save_task_result, + resolve_rebalance_terminal_error, send_rebalance_done_signal, +}; +use super::{ + EVENT_REBALANCE_BUCKET, EVENT_REBALANCE_STATE, LOG_COMPONENT_ECSTORE, LOG_SUBSYSTEM_REBALANCE, RebalSaveOpt, RebalStatus, + RebalanceBucketOutcome, +}; +use crate::error::{Error, Result}; +use crate::global::get_global_endpoints; +use crate::store::ECStore; +use std::collections::HashSet; +use std::sync::Arc; +use time::OffsetDateTime; +use tokio::time::{Duration, Instant}; +use tokio_util::sync::CancellationToken; +use tracing::{debug, error, info, warn}; + +impl ECStore { + #[tracing::instrument(skip_all)] + pub async fn start_rebalance(self: &Arc) -> Result<()> { + info!( + event = EVENT_REBALANCE_STATE, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + state = "starting", + "Starting rebalance" + ); + let decommission_running = self.is_decommission_running().await; + // let rebalance_meta = self.rebalance_meta.read().await; + + let cancel_tx = CancellationToken::new(); + let rx = cancel_tx.clone(); + let mut meta_to_save = None; + + { + let mut rebalance_meta = self.rebalance_meta.write().await; + validate_start_rebalance_state(decommission_running, rebalance_meta.is_some())?; + + let Some(meta) = rebalance_meta.as_mut() else { + return Err(Error::ConfigNotFound); + }; + if should_skip_start_rebalance(meta.cancel.is_some(), is_rebalance_in_progress(meta)) { + debug!( + event = EVENT_REBALANCE_STATE, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + state = "start_skipped", + reason = "already_in_progress", + "Skipped duplicate rebalance start" + ); + return Ok(()); + } + let now = OffsetDateTime::now_utc(); + if complete_rebalance_pools_at_goal(meta, now) { + meta_to_save = Some(meta.clone()); + } + if complete_rebalance_pools_with_empty_queue(meta, now) { + meta_to_save = Some(meta.clone()); + } + meta.cancel = Some(cancel_tx); + + drop(rebalance_meta); + } + + if let Some(meta) = meta_to_save { + let pool = clone_first_arc(self.pools.as_slice(), "start_rebalance: no pools available")?; + 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", + )?; + } + + let participants = if let Some(ref meta) = *self.rebalance_meta.read().await { + resolve_rebalance_participants(meta.pool_stats.as_slice(), self.pools.len()) + } else { + debug!( + event = EVENT_REBALANCE_STATE, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + state = "start_skipped", + reason = "metadata_missing", + "Skipped rebalance start because metadata is unavailable" + ); + Vec::new() + }; + + if !participants.iter().any(|participating| *participating) { + debug!( + event = EVENT_REBALANCE_STATE, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + state = "start_skipped", + reason = "no_participants", + "Skipped rebalance start because no pools are participating" + ); + return Ok(()); + } + + let mut workers_started = 0usize; + for (idx, participating) in participants.iter().enumerate() { + if !*participating { + debug!( + event = EVENT_REBALANCE_STATE, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + pool_index = idx, + state = "pool_skipped", + reason = "not_participating", + "Skipped rebalance pool" + ); + continue; + } + + if !get_global_endpoints() + .as_ref() + .get(idx) + .is_some_and(|v| v.endpoints.as_ref().first().is_some_and(|e| e.is_local)) + { + debug!( + event = EVENT_REBALANCE_STATE, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + pool_index = idx, + state = "pool_skipped", + reason = "not_local", + "Skipped non-local rebalance pool" + ); + continue; + } + + let pool_idx = idx; + let store = self.clone(); + let rx_clone = rx.clone(); + workers_started += 1; + tokio::spawn(async move { + if let Err(err) = store.rebalance_buckets(rx_clone, pool_idx).await { + error!( + event = EVENT_REBALANCE_STATE, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + pool_index = pool_idx, + state = "pool_failed", + error = %err, + "Rebalance pool failed" + ); + } else { + debug!( + event = EVENT_REBALANCE_STATE, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + pool_index = pool_idx, + state = "completed", + "Rebalance pool completed" + ); + } + }); + } + + if workers_started == 0 { + debug!( + event = EVENT_REBALANCE_STATE, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + state = "start_skipped", + reason = "no_local_participants", + "Skipped rebalance start because no local pools are participating" + ); + return Ok(()); + } + + info!( + event = EVENT_REBALANCE_STATE, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + state = "started", + worker_count = workers_started, + "Rebalance started" + ); + Ok(()) + } + + #[tracing::instrument(skip(self, rx))] + async fn rebalance_buckets(self: &Arc, rx: CancellationToken, pool_index: usize) -> Result<()> { + ensure_valid_rebalance_pool_index(self.pools.len(), pool_index)?; + + let (done_tx, mut done_rx) = tokio::sync::mpsc::channel::>(1); + + // Save rebalance metadata periodically + let store = self.clone(); + let save_task = tokio::spawn(async move { + let mut timer = tokio::time::interval_at(Instant::now() + Duration::from_secs(30), Duration::from_secs(10)); + let mut msg: String; + let mut quit = false; + + loop { + tokio::select! { + result = done_rx.recv() => { + quit = true; + let now = OffsetDateTime::now_utc(); + let terminal_event = classify_rebalance_terminal_event(result, now); + msg = terminal_event.message().to_string(); + let mut rebalance_meta = store.rebalance_meta.write().await; + if let Some(meta) = rebalance_meta.as_mut() { + let meta_stopped = meta.stopped_at.is_some(); + if let Some(pool_stat) = meta.pool_stats.get_mut(pool_index) { + if should_preserve_rebalance_stopped_state( + meta_stopped, + pool_stat.info.status, + &terminal_event, + ) { + debug!( + event = EVENT_REBALANCE_STATE, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + pool_index, + state = "stopped_preserved", + "Preserved stopped rebalance status" + ); + } else { + apply_rebalance_terminal_event( + &mut pool_stat.info.status, + &mut pool_stat.info.end_time, + &mut pool_stat.info.last_error, + terminal_event, + now, + ); + } + } + } + } + _ = timer.tick() => { + let now = OffsetDateTime::now_utc(); + msg = format!("Saving rebalance metadata at {now:?}"); + } + } + + if let Err(err) = store.save_rebalance_stats(pool_index, RebalSaveOpt::Stats).await { + let wrapped = Error::other(format!("rebalance save_task stats save failed for pool {pool_index}: {err}")); + error!("{} err: {:?}", msg, wrapped); + if quit { + return Err(wrapped); + } + } else { + debug!( + event = EVENT_REBALANCE_STATE, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + pool_index, + state = "metadata_saved", + message = %msg, + "Saved rebalance metadata" + ); + } + + if quit { + debug!( + event = EVENT_REBALANCE_STATE, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + pool_index, + state = "save_task_exiting", + message = %msg, + "Exiting rebalance save task" + ); + return Ok(()); + } + + timer.reset(); + } + }); + + info!( + event = EVENT_REBALANCE_STATE, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + pool_index, + state = "pool_started", + "Rebalance worker started" + ); + let mut final_result: Result<()> = Ok(()); + let mut deferred_buckets = HashSet::new(); + + loop { + if rx.is_cancelled() { + info!( + event = EVENT_REBALANCE_STATE, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + pool_index, + state = "pool_stopped", + reason = "cancelled", + "Stopped rebalance worker" + ); + let err = Error::OperationCanceled; + final_result = Err(resolve_rebalance_terminal_error( + err.clone(), + send_rebalance_done_signal(&done_tx, Err(err.clone()), pool_index).await, + )); + break; + } + + let next_bucket = match self.next_rebal_bucket(pool_index).await { + Ok(bucket) => bucket, + Err(err) => { + error!( + event = EVENT_REBALANCE_BUCKET, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + pool_index, + state = "next_bucket_failed", + error = ?err, + "Rebalance next bucket lookup failed" + ); + final_result = Err(resolve_rebalance_terminal_error( + err.clone(), + send_rebalance_done_signal(&done_tx, Err(err.clone()), pool_index).await, + )); + break; + } + }; + + if let Some(bucket) = next_bucket { + debug!( + event = EVENT_REBALANCE_BUCKET, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + pool_index, + bucket = %bucket, + state = "started", + "Starting rebalance bucket" + ); + + let outcome = match resolve_rebalance_bucket_result( + self.rebalance_bucket(rx.clone(), bucket.clone(), pool_index).await, + pool_index, + &bucket, + ) { + Ok(outcome) => outcome, + Err(err) => { + error!( + event = EVENT_REBALANCE_BUCKET, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + pool_index, + bucket = %bucket, + state = "bucket_failed", + error = ?err, + "Rebalance bucket failed" + ); + final_result = Err(resolve_rebalance_terminal_error( + err.clone(), + send_rebalance_done_signal(&done_tx, Err(err.clone()), pool_index).await, + )); + break; + } + }; + + if let RebalanceBucketOutcome::Deferred { last_error } = outcome { + if !deferred_buckets.insert(bucket.clone()) { + let err = Error::other(format!( + "rebalance bucket {bucket} deferred repeatedly due to transient object failures: {last_error}" + )); + error!( + event = EVENT_REBALANCE_BUCKET, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + pool_index, + bucket = %bucket, + state = "bucket_deferred_repeatedly", + error = ?err, + "Rebalance bucket failed after repeated deferral" + ); + final_result = Err(resolve_rebalance_terminal_error( + err.clone(), + send_rebalance_done_signal(&done_tx, Err(err.clone()), pool_index).await, + )); + break; + } + + warn!( + event = EVENT_REBALANCE_BUCKET, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + pool_index, + bucket = %bucket, + state = "deferred", + error = %last_error, + "Deferred rebalance bucket after transient object failures" + ); + if let Err(err) = self.defer_rebalance_bucket(pool_index, bucket.clone(), last_error).await { + error!( + event = EVENT_REBALANCE_BUCKET, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + pool_index, + bucket = %bucket, + state = "defer_failed", + error = ?err, + "Rebalance bucket defer failed" + ); + final_result = Err(resolve_rebalance_terminal_error( + err.clone(), + send_rebalance_done_signal(&done_tx, Err(err.clone()), pool_index).await, + )); + break; + } + continue; + } + + debug!( + event = EVENT_REBALANCE_BUCKET, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + pool_index, + bucket = %bucket, + state = "completed", + "Completed rebalance bucket" + ); + if let Err(err) = self.bucket_rebalance_done(pool_index, bucket).await { + error!( + event = EVENT_REBALANCE_BUCKET, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + pool_index, + state = "bucket_done_mark_failed", + error = ?err, + "Rebalance bucket completion mark failed" + ); + final_result = Err(resolve_rebalance_terminal_error( + err.clone(), + send_rebalance_done_signal(&done_tx, Err(err.clone()), pool_index).await, + )); + break; + } + } else { + debug!( + event = EVENT_REBALANCE_BUCKET, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + pool_index, + state = "idle", + reason = "no_bucket_to_rebalance", + "No rebalance bucket available" + ); + break; + } + } + + info!( + event = EVENT_REBALANCE_STATE, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + pool_index, + state = "pool_done", + "Rebalance worker finished" + ); + + if final_result.is_ok() + && let Err(err) = send_rebalance_done_signal(&done_tx, Ok(()), pool_index).await + { + final_result = Err(err); + } + drop(done_tx); + if let Err(err) = resolve_rebalance_save_task_result(pool_index, save_task.await) + && final_result.is_ok() + { + final_result = Err(err); + } + debug!( + event = EVENT_REBALANCE_STATE, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + pool_index, + state = "pool_result_returned", + "Rebalance worker result returned" + ); + final_result + } + + pub(super) async fn check_if_rebalance_done(&self, pool_index: usize) -> bool { + let mut rebalance_meta = self.rebalance_meta.write().await; + + if let Some(meta) = rebalance_meta.as_mut() + && let Some(pool_stat) = meta.pool_stats.get_mut(pool_index) + { + // Check if the pool's rebalance status is already completed + if pool_stat.info.status == RebalStatus::Completed { + debug!( + event = EVENT_REBALANCE_STATE, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + pool_index, + state = "already_completed", + "Rebalance pool is already completed" + ); + return true; + } + + // Mark pool rebalance as done only after it reaches the PercentFreeGoal. + let pfi = if pool_stat.init_capacity == 0 { + 0.0 + } else { + (pool_stat.init_free_space + pool_stat.bytes) as f64 / pool_stat.init_capacity as f64 + }; + + if !has_deferred_rebalance_error(pool_stat) + && rebalance_goal_reached( + pool_stat.init_free_space, + pool_stat.init_capacity, + pool_stat.bytes, + meta.percent_free_goal, + ) + { + pool_stat.info.status = RebalStatus::Completed; + pool_stat.info.end_time = Some(OffsetDateTime::now_utc()); + info!( + event = EVENT_REBALANCE_STATE, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + pool_index, + state = "completed", + percent_free = pfi, + "Marked rebalance pool completed" + ); + return true; + } + } + + false + } +} + +impl ECStore { + #[tracing::instrument(skip(self))] + pub async fn save_rebalance_stats(&self, pool_idx: usize, opt: RebalSaveOpt) -> Result<()> { + let meta_to_save = { + let mut rebalance_meta = self.rebalance_meta.write().await; + let Some(meta) = rebalance_meta.as_mut() else { + return Ok(()); + }; + + let now = OffsetDateTime::now_utc(); + apply_rebalance_save_option(meta, pool_idx, opt, now); + meta.clone() + }; + + let pool = clone_first_arc(&self.pools, "save_rebalance_stats: no pools available")?; + + debug!( + event = EVENT_REBALANCE_STATE, + component = LOG_COMPONENT_ECSTORE, + subsystem = LOG_SUBSYSTEM_REBALANCE, + pool_index = pool_idx, + save_opt = ?opt, + state = "metadata_save_requested", + "Rebalance metadata save requested" + ); + let stage = format!("save_rebalance_stats for pool {pool_idx} opt {opt:?}"); + resolve_rebalance_meta_save_result( + self.save_rebalance_meta_with_merge(pool, &meta_to_save, stage.as_str()).await, + stage.as_str(), + )?; + + Ok(()) + } +} diff --git a/crates/ecstore/src/rebalance/types.rs b/crates/ecstore/src/rebalance/types.rs new file mode 100644 index 000000000..e9efbe4df --- /dev/null +++ b/crates/ecstore/src/rebalance/types.rs @@ -0,0 +1,119 @@ +use serde::{Deserialize, Serialize}; +use std::sync::Arc; +use time::OffsetDateTime; +use tokio_util::sync::CancellationToken; + +#[derive(Debug, Default, Clone, Serialize, Deserialize)] +pub struct RebalanceStats { + #[serde(rename = "ifs")] + pub init_free_space: u64, // Pool free space at the start of rebalance + #[serde(rename = "ic")] + pub init_capacity: u64, // Pool capacity at the start of rebalance + #[serde(rename = "bus")] + pub buckets: Vec, // Buckets being rebalanced or to be rebalanced + #[serde(rename = "rbs")] + pub rebalanced_buckets: Vec, // Buckets rebalanced + #[serde(rename = "bu")] + pub bucket: String, // Last rebalanced bucket + #[serde(rename = "ob")] + pub object: String, // Last rebalanced object + #[serde(rename = "no")] + pub num_objects: u64, // Number of objects rebalanced + #[serde(rename = "nv")] + pub num_versions: u64, // Number of versions rebalanced + #[serde(rename = "bs")] + pub bytes: u64, // Number of bytes rebalanced + #[serde(rename = "par")] + pub participating: bool, // Whether the pool is participating in rebalance + #[serde(rename = "inf")] + pub info: RebalanceInfo, // Rebalance operation info + #[serde(rename = "cw", default)] + pub cleanup_warnings: RebalanceCleanupWarnings, +} + +pub type RStats = Vec>; + +#[derive(Debug, Default)] +pub(super) struct RebalanceBucketConfigs { + pub(super) lifecycle_config: Option, + pub(super) lock_retention: Option, + pub(super) replication_config: Option<(s3s::dto::ReplicationConfiguration, OffsetDateTime)>, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(super) enum RebalanceBucketOutcome { + Completed, + Deferred { last_error: String }, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(super) enum RebalanceEntryOutcome { + Completed, + Deferred { last_error: String }, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)] +pub enum RebalStatus { + #[default] + None, + Started, + Completed, + Stopped, + Failed, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)] +pub enum RebalSaveOpt { + #[default] + Stats, + StoppedAt, +} + +#[derive(Debug, Default, Clone, Serialize, Deserialize)] +pub struct RebalanceInfo { + #[serde(rename = "startTs")] + pub start_time: Option, // Time at which rebalance-start was issued + #[serde(rename = "stopTs")] + pub end_time: Option, // Time at which rebalance operation completed or rebalance-stop was called + #[serde(rename = "err")] + pub last_error: Option, // Last rebalance error message + #[serde(rename = "status")] + pub status: RebalStatus, // Current state of rebalance operation +} + +#[derive(Debug, Default, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct RebalanceCleanupWarnings { + #[serde(rename = "count", default)] + pub count: u64, + #[serde(rename = "lastMsg", default)] + pub last_message: Option, + #[serde(rename = "lastBucket", default)] + pub last_bucket: Option, + #[serde(rename = "lastObject", default)] + pub last_object: Option, + #[serde(rename = "lastAt", default)] + pub last_at: Option, +} + +#[allow(dead_code)] +#[derive(Debug, Clone, Default)] +pub struct DiskStat { + pub total_space: u64, + pub available_space: u64, +} + +#[derive(Debug, Default, Serialize, Deserialize, Clone)] +pub struct RebalanceMeta { + #[serde(skip)] + pub cancel: Option, // To be invoked on rebalance-stop + #[serde(skip)] + pub last_refreshed_at: Option, + #[serde(rename = "stopTs")] + pub stopped_at: Option, // Time when rebalance-stop was issued + #[serde(rename = "id")] + pub id: String, // ID of the ongoing rebalance operation + #[serde(rename = "pf")] + pub percent_free_goal: f64, // Computed from total free space and capacity at the start of rebalance + #[serde(rename = "rss")] + pub pool_stats: Vec, // Per-pool rebalance stats keyed by pool index +} diff --git a/crates/utils/src/os/windows.rs b/crates/utils/src/os/windows.rs index 5d017d78d..c1e94c3d3 100644 --- a/crates/utils/src/os/windows.rs +++ b/crates/utils/src/os/windows.rs @@ -170,7 +170,7 @@ pub fn get_physical_device_ids(disk: &str) -> std::io::Result> { /// This is the Windows equivalent of Linux's `st_dev` major:minor pair and is /// useful for diagnostic purposes when validating physical disk independence. // SAFETY: `path_wide` is a valid null-terminated UTF-16 string and all output -// pointers target initialized stack variables with documented MAX_PATH sizing. +// SAFETY: pointers target initialized stack variables with documented MAX_PATH sizing. #[allow(unsafe_code)] pub fn get_volume_serial_number(path: &str) -> std::io::Result { let path_wide = to_wide_path(Path::new(path)); diff --git a/docs/architecture/migration-progress.md b/docs/architecture/migration-progress.md index c9f6d4950..483fecb27 100644 --- a/docs/architecture/migration-progress.md +++ b/docs/architecture/migration-progress.md @@ -5,19 +5,22 @@ Status values: `[ ]` not started, `[~]` in progress, `[x]` complete, `[!]` block ## Current Context - Issue: [`rustfs/backlog#660`](https://github.com/rustfs/backlog/issues/660) -- Branch: `overtrue/arch-ecstore-rebalance-migration-helpers` -- Baseline: completed `E-009/E-REBALANCE-003`. +- Branch: `overtrue/arch-ecstore-rebalance-entry-flow` +- Baseline: completed `E-015/E-REBALANCE-009`. - Stacked on: merged ECStore layout foundation, format layout ownership, and endpoint layout move slices plus the set-format heal, pool-space layout helper, rebalance support helper, pool-space builder, rebalance metadata - helper, and rebalance worker helper slices. + helper, rebalance worker helper, rebalance migration helper, rebalance state + impl, rebalance control impl, rebalance runtime loop, rebalance entry flow, + rebalance unit-test split, and rebalance type-owner slices. - PR type for this branch: `pure-move` - Runtime behavior changes: none. -- Rust code changes: move ECStore rebalance migration backend, version-result, - delete-marker option, remote-tier option, and version migration retry helpers - into `rebalance::migration` while preserving ECStore orchestration. +- Rust code changes: move ECStore rebalance control, runtime, entry flow, + inline tests, and type contracts into dedicated rebalance submodules while + preserving public root paths. - CI/script changes: none. -- Docs changes: record the ECStore rebalance migration helper slice. +- Docs changes: record the combined ECStore rebalance control/runtime/entry + flow, test-module split, and type-owner split slices. ## Phase 0 Tasks @@ -2363,20 +2366,109 @@ Status values: `[ ]` not started, `[~]` in progress, `[x]` complete, `[!]` block checks, migration/layer guards, formatting, diff hygiene, Rust risk scan, branch freshness check, pre-commit quality gate, and three-expert review. +- [x] `E-011/E-REBALANCE-005` Move ECStore rebalance state impls. + - Do: move `RebalanceStats` update helpers, `RebalStatus` conversions, and + `RebalanceMeta` load/save impls into `rebalance::meta` while leaving public + wire structs in `rebalance.rs`. + - Acceptance: `rebalance::meta` owns metadata/state behavior, `rebalance.rs` + keeps data contracts and ECStore orchestration, and focused rebalance tests + keep covering moved behavior. + - Must preserve: serialized rebalance metadata header format/version, + empty/short/unknown metadata handling, last refresh timestamps, save-skip + behavior for empty pool stats, object/version/byte accounting, batch update + behavior, status display labels, and legacy status byte mapping. + - Verification: focused ECStore rebalance tests, ECStore/RustFS/Heal compile + checks, migration/layer guards, formatting, diff hygiene, Rust risk scan, + branch freshness check, pre-commit quality gate, and three-expert review. + +- [x] `E-012/E-REBALANCE-006` Move ECStore rebalance control impls. + - Do: move ECStore rebalance metadata save/load/update/init/status/stop + control methods into `rebalance::control` while leaving the worker loop and + entry migration orchestration in `rebalance.rs`. + - Acceptance: `rebalance::control` owns metadata/control methods, + `rebalance.rs` keeps public data contracts and worker orchestration, and + focused rebalance tests keep covering moved behavior. + - Must preserve: metadata merge locking, load/save error wrapping, pool stats + refresh and extension, init free-space goal, pool stat update behavior, + bucket queue done/defer behavior, cleanup warning recording, start/stop + status checks, decommission conflict checks, and stop snapshot persistence. + - Verification: focused ECStore rebalance tests, ECStore/RustFS/Heal compile + checks, migration/layer guards, formatting, diff hygiene, Rust risk scan, + branch freshness check, pre-commit quality gate, and three-expert review. + +- [x] `E-013/E-REBALANCE-007` Move ECStore rebalance runtime loop. + - Do: move `start_rebalance`, the pool rebalance worker loop, completion + check, and periodic stats save loop into `rebalance::runtime` while leaving + entry/object/bucket migration orchestration in `rebalance.rs`. + - Acceptance: `rebalance::runtime` owns start and pool runtime orchestration, + `rebalance.rs` keeps public data contracts and entry/object/bucket + migration flow, and focused rebalance tests keep covering moved behavior. + - Must preserve: decommission/start validation, duplicate-start skipping, + pool-at-goal and empty-queue completion persistence, participant/local + endpoint filtering, cancellation handling, deferred-bucket repeated failure + guard, bucket done/defer behavior, terminal event application, save-task + error precedence, goal completion math, and save option persistence. + - Verification: focused ECStore rebalance tests, ECStore/RustFS/Heal compile + checks, migration/layer guards, formatting, diff hygiene, Rust risk scan, + branch freshness check, pre-commit quality gate, and three-expert review. + +- [x] `E-014/E-REBALANCE-008` Move ECStore rebalance entry flow. + - Do: move the remaining entry, object-transfer, deferred-error, and bucket + entry-scan migration flow into `rebalance::entry` while leaving public data + contracts in `rebalance.rs`. + - Acceptance: `rebalance::entry` owns bucket/entry migration flow, + `rebalance::runtime` keeps pool-level orchestration, and focused rebalance + tests keep covering moved behavior. + - Must preserve: directory and completed-pool skips, lifecycle-expired + filtering, delete-marker skip semantics, data-movement retry flow, deferred + transient failure recording, batch stats updates, source cleanup warning + recording, entry worker semaphore limits, cancellation handling, listing + retry flow, and bucket outcome precedence. + - Verification: focused ECStore rebalance tests, ECStore/RustFS/Heal compile + checks, migration/layer guards, formatting, diff hygiene, Rust risk scan, + branch freshness check, pre-commit quality gate, and three-expert review. + +- [x] `E-015/E-REBALANCE-009` Split ECStore rebalance unit tests. + - Do: move the large inline `rebalance_unit_tests` module out of + `rebalance.rs` into `rebalance/rebalance_unit_tests.rs` while preserving + the module name and test filter path. + - Acceptance: `rebalance.rs` is reduced to public rebalance data contracts + plus submodule wiring, rebalance unit tests remain under + `rebalance::rebalance_unit_tests`, and focused rebalance tests keep covering + moved behavior. + - Must preserve: test coverage, helper visibility, legacy metadata + serialization coverage, migration backend spies, panic-context tests, and + every existing rebalance unit-test filter path. + - Verification: focused ECStore rebalance tests, migration/layer guards, + formatting, diff hygiene, Rust risk scan, branch freshness check, + pre-commit quality gate, and three-expert review. + +- [x] `E-016/E-REBALANCE-010` Move ECStore rebalance type contracts. + - Do: move rebalance stats, status, info, metadata DTOs, and internal + bucket/entry outcomes into `rebalance::types` while preserving root + re-exports. + - Acceptance: public `crate::rebalance::*` paths remain stable, internal + submodules keep `super::...` access, and `rebalance.rs` only wires shared + constants, modules, and re-exports. + - Must preserve: serde field names/defaults, rebalance metadata wire shape, + status/save-option defaults, cancellation/refresh metadata fields, and + internal bucket/entry outcome semantics. + - Verification: focused ECStore rebalance tests, migration/layer guards, + formatting, diff hygiene, Rust risk scan, branch freshness check, + pre-commit quality gate, and three-expert review. + ## Next PRs -1. `pure-move`: continue moving ECStore rebalance/layout support helpers once - E-010/E-REBALANCE-004 lands. -2. `pure-move`: continue pruning residual embedded startup-only orchestration +1. `pure-move`: continue pruning residual embedded startup-only orchestration once the lifecycle helpers are merged. ## Pre-Push Review Log | Expert | Status | Notes | |---|---|---| -| Quality/architecture | passed | E-010/E-REBALANCE-004 moves migration helpers into a child module while leaving ECStore orchestration in place. | -| Migration preservation | passed | Remote-tier movement, delete-marker options, retry/backoff classification, not-found handling, stage labels, and cleanup accounting remain preserved. | -| Testing/verification | passed | Focused rebalance tests, compile checks, guards, formatting, diff hygiene, Rust risk scan, and full pre-commit passed. | +| Quality/architecture | passed | E-012 through E-016 move control/runtime/entry flow, tests, and type contracts into rebalance submodules while preserving root exports. | +| Migration preservation | passed | Lifecycle filtering, delete-marker skips, deferred transient failures, source cleanup warnings, cancellation, bucket outcome precedence, serde contracts, and test coverage remain preserved. | +| Testing/verification | passed | Focused rebalance tests, compile checks, guards, formatting, diff hygiene, Rust risk scan, and pre-commit gate passed. | ## Verification Notes @@ -2536,6 +2628,67 @@ Passed before push: - `make pre-commit`: passed. - Three-expert review: passed. +- Issue #660 E-012/E-REBALANCE-006 current slice: + - `cargo test -p rustfs-ecstore rebalance::rebalance_unit_tests -- --nocapture`: passed. + - `cargo check -p rustfs-ecstore -p rustfs -p rustfs-heal`: passed. + - `./scripts/check_architecture_migration_rules.sh`: passed. + - `./scripts/check_layer_dependencies.sh`: passed. + - `cargo fmt --all --check`: passed. + - `git diff --check`: passed. + - Rust risk scan on changed Rust files: passed; added casts are moved + pool-index accounting from the existing implementation and remain guarded. + - `make pre-commit`: passed. + - Three-expert review: passed. + +- Issue #660 E-013/E-REBALANCE-007 current slice: + - `cargo test -p rustfs-ecstore rebalance::rebalance_unit_tests -- --nocapture`: passed. + - `cargo check -p rustfs-ecstore -p rustfs -p rustfs-heal`: passed. + - `./scripts/check_architecture_migration_rules.sh`: passed. + - `./scripts/check_layer_dependencies.sh`: passed. + - `cargo fmt --all --check`: passed. + - `git diff --check`: passed. + - Rust risk scan on changed Rust files: passed; moved casts are existing + pool completion math and remain guarded. + - `make pre-commit`: passed. + - Three-expert review: passed. + +- Issue #660 E-014/E-REBALANCE-008 current slice: + - `cargo test -p rustfs-ecstore rebalance::rebalance_unit_tests -- --nocapture`: passed. + - `cargo check -p rustfs-ecstore -p rustfs -p rustfs-heal`: passed. + - `./scripts/check_architecture_migration_rules.sh`: passed. + - `./scripts/check_layer_dependencies.sh`: passed. + - `cargo fmt --all --check`: passed. + - `git diff --check`: passed. + - Rust risk scan on changed Rust files: passed; moved casts and unwraps are + existing test or migration-flow code and remain guarded. + - `make pre-commit`: passed. + - Three-expert review: passed. + +- Issue #660 E-015/E-REBALANCE-009 current slice: + - `cargo test -p rustfs-ecstore rebalance::rebalance_unit_tests -- --nocapture`: passed. + - `./scripts/check_unsafe_code_allowances.sh`: passed. + - `./scripts/check_architecture_migration_rules.sh`: passed. + - `./scripts/check_layer_dependencies.sh`: passed. + - `cargo fmt --all --check`: passed. + - `git diff --check`: passed. + - Rust risk scan on changed Rust files: passed; the runtime diff is a test + module move plus a SAFETY-comment proximity fix required by the guard. + - `make pre-commit`: passed. + - Three-expert review: passed. + +- Issue #660 E-016/E-REBALANCE-010 current slice: + - `cargo test -p rustfs-ecstore rebalance::rebalance_unit_tests -- --nocapture`: passed. + - `cargo check -p rustfs-ecstore -p rustfs -p rustfs-heal`: passed. + - `./scripts/check_unsafe_code_allowances.sh`: passed. + - `./scripts/check_architecture_migration_rules.sh`: passed. + - `./scripts/check_layer_dependencies.sh`: passed. + - `cargo fmt --all --check`: passed. + - `git diff --check`: passed. + - Rust risk scan on changed Rust files: passed; production changes are a + type-contract move and existing Windows FFI casts remain unchanged. + - `make pre-commit`: passed. + - Three-expert review: passed. + - Issue #660 X-012 current slice: - `cargo test -p rustfs-extension-schema`: passed. - `cargo check -p rustfs-extension-schema`: passed.