gsadmin/rustfs
1
0
Fork 0
mirror of https://github.com/rustfs/rustfs.git synced 2026-08-12 16:16:55 +00:00
Code Issues Packages Projects Releases Wiki Activity
Files
24fa03e04b84ab23759e89a865b8effd6293bcf1
rustfs/crates/ecstore/src/rebalance.rs
T
安正超 53337a6f71 refactor: move HTTP range helper contracts (#3533)
2026-06-18 00:47:23 +08:00

6377 lines
230 KiB
Rust
Raw Blame History

// 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 crate::cache_value::metacache_set::{ListPathRawOptions, list_path_raw};
use crate::config::com::{read_config_with_metadata, save_config_with_opts};
use crate::data_movement;
use crate::data_usage::DATA_USAGE_CACHE_NAME;
use crate::disk::error::DiskError;
use crate::error::{Error, Result};
use crate::error::{is_err_object_not_found, is_err_operation_canceled, is_err_version_not_found, is_network_or_host_down};
use crate::global::get_global_endpoints;
use crate::pools::ListCallback;
use crate::set_disk::{SetDisks, get_lock_acquire_timeout};
use crate::store::ECStore;
use crate::store_api::{GetObjectReader, NamespaceLocking, ObjectIO, ObjectInfo, ObjectOperations, ObjectOptions};
use http::HeaderMap;
use rand::RngExt as _;
use rustfs_filemeta::{FileInfo, MetaCacheEntries, MetaCacheEntry, MetadataResolutionParams};
use rustfs_storage_api::{HTTPRangeSpec, StorageAdminApi};
use rustfs_utils::path::encode_dir_object;
use serde::{Deserialize, Serialize};
use std::collections::HashSet;
use std::fmt;
use std::future::Future;
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};
const LOG_COMPONENT_ECSTORE: &str = "ecstore";
const LOG_SUBSYSTEM_REBALANCE: &str = "rebalance";
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
const REBAL_META_NAME: &str = "rebalance.bin";
const DEFAULT_REBALANCE_MAX_ATTEMPTS: usize = 3;
const REBALANCE_MAX_ATTEMPTS_ENV: &str = "RUSTFS_REBALANCE_MAX_ATTEMPTS";
const REBALANCE_LISTING_RETRY_BASE_DELAY: Duration = Duration::from_millis(250);
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:";
#[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<String>, // Buckets being rebalanced or to be rebalanced
#[serde(rename = "rbs")]
pub rebalanced_buckets: Vec<String>, // 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<Arc<RebalanceStats>>;
#[derive(Debug, Default)]
struct RebalanceBucketConfigs {
lifecycle_config: Option<s3s::dto::BucketLifecycleConfiguration>,
lock_retention: Option<s3s::dto::DefaultRetention>,
replication_config: Option<(s3s::dto::ReplicationConfiguration, OffsetDateTime)>,
}
#[derive(Debug, Default, Clone)]
pub(crate) struct MigrationVersionResult {
pub moved: bool,
pub ignored: bool,
pub cleanup_ignored: bool,
pub failed: bool,
pub stage: Option<&'static str>,
pub error: Option<Error>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
enum RebalanceBucketOutcome {
Completed,
Deferred { last_error: String },
}
#[derive(Debug, Clone, PartialEq, Eq)]
enum RebalanceEntryOutcome {
Completed,
Deferred { last_error: String },
}
fn rebalance_delete_marker_opts(version: &FileInfo, version_id: Option<String>, src_pool_idx: usize) -> ObjectOptions {
ObjectOptions {
versioned: true,
version_id,
mod_time: version.mod_time,
src_pool_idx,
data_movement: true,
delete_marker: true,
skip_decommissioned: true,
delete_replication: version.replication_state_internal.clone(),
..Default::default()
}
}
fn rebalance_remote_tiered_opts(version: &FileInfo, version_id: Option<String>, src_pool_idx: usize) -> ObjectOptions {
ObjectOptions {
versioned: version_id.is_some(),
version_id,
mod_time: version.mod_time,
user_defined: version.metadata.clone(),
src_pool_idx,
data_movement: true,
..Default::default()
}
}
#[async_trait::async_trait]
pub(crate) trait MigrationBackend: Send + Sync {
async fn get_object_reader_for_migration(
&self,
bucket: &str,
object: &str,
range: Option<HTTPRangeSpec>,
h: HeaderMap,
opts: &ObjectOptions,
) -> Result<GetObjectReader>;
async fn delete_object_for_migration(&self, bucket: &str, object: &str, opts: ObjectOptions) -> Result<ObjectInfo>;
async fn move_remote_version_for_migration(
&self,
bucket: &str,
object: &str,
fi: &FileInfo,
opts: &ObjectOptions,
) -> Result<()>;
}
#[async_trait::async_trait]
impl MigrationBackend for SetDisks {
async fn get_object_reader_for_migration(
&self,
bucket: &str,
object: &str,
range: Option<HTTPRangeSpec>,
h: HeaderMap,
opts: &ObjectOptions,
) -> Result<GetObjectReader> {
self.get_object_reader(bucket, object, range, h, opts).await
}
async fn delete_object_for_migration(&self, bucket: &str, object: &str, opts: ObjectOptions) -> Result<ObjectInfo> {
self.delete_object(bucket, object, opts).await
}
async fn move_remote_version_for_migration(
&self,
bucket: &str,
object: &str,
fi: &FileInfo,
opts: &ObjectOptions,
) -> Result<()> {
self.decommission_tiered_object(bucket, object, fi, opts).await
}
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn migrate_entry_version<Backend, F, Fut>(
set: &Backend,
bucket: String,
pool_index: usize,
version: &FileInfo,
version_id: Option<String>,
max_attempts: usize,
ignore_data_usage_cache: bool,
transfer: F,
) -> MigrationVersionResult
where
Backend: MigrationBackend + ?Sized,
F: FnMut(usize, String, GetObjectReader) -> Fut + Send,
Fut: Future<Output = Result<()>> + Send,
{
migrate_entry_version_with_retry_wait(
set,
bucket,
pool_index,
version,
version_id,
max_attempts,
ignore_data_usage_cache,
transfer,
sleep_rebalance_migration_retry,
)
.await
}
#[allow(clippy::too_many_arguments)]
async fn migrate_entry_version_with_retry_wait<Backend, F, Fut, W, WFut>(
set: &Backend,
bucket: String,
pool_index: usize,
version: &FileInfo,
version_id: Option<String>,
max_attempts: usize,
ignore_data_usage_cache: bool,
mut transfer: F,
mut wait_retry: W,
) -> MigrationVersionResult
where
Backend: MigrationBackend + ?Sized,
F: FnMut(usize, String, GetObjectReader) -> Fut + Send,
Fut: Future<Output = Result<()>> + Send,
W: FnMut(Duration) -> WFut + Send,
WFut: Future<Output = ()> + Send,
{
let max_attempts = max_attempts.max(1);
if ignore_data_usage_cache && bucket == crate::disk::RUSTFS_META_BUCKET && version.name.contains(DATA_USAGE_CACHE_NAME) {
return MigrationVersionResult {
moved: false,
ignored: true,
cleanup_ignored: false,
failed: false,
stage: None,
error: None,
};
}
if version.is_remote() {
if let Err(err) = set
.move_remote_version_for_migration(
&bucket,
&version.name,
version,
&rebalance_remote_tiered_opts(version, version_id, pool_index),
)
.await
{
if is_err_object_not_found(&err) || is_err_version_not_found(&err) {
return MigrationVersionResult {
moved: false,
ignored: true,
cleanup_ignored: true,
failed: false,
stage: Some("move_remote_version"),
error: None,
};
}
return MigrationVersionResult {
moved: false,
ignored: false,
cleanup_ignored: false,
failed: true,
stage: Some("move_remote_version"),
error: Some(err),
};
}
return MigrationVersionResult {
moved: true,
ignored: false,
cleanup_ignored: false,
failed: false,
stage: None,
error: None,
};
}
if version.deleted {
if let Err(err) = set
.delete_object_for_migration(&bucket, &version.name, rebalance_delete_marker_opts(version, version_id, pool_index))
.await
{
if is_err_object_not_found(&err) || is_err_version_not_found(&err) {
return MigrationVersionResult {
moved: false,
ignored: true,
cleanup_ignored: true,
failed: false,
stage: Some("delete_marker"),
error: None,
};
}
return MigrationVersionResult {
moved: false,
ignored: false,
cleanup_ignored: false,
failed: true,
stage: Some("delete_marker"),
error: Some(err),
};
}
return MigrationVersionResult {
moved: true,
ignored: false,
cleanup_ignored: false,
failed: false,
stage: None,
error: None,
};
}
let mut last_error: Option<Error> = None;
for attempt in 0..max_attempts {
let rd = match set
.get_object_reader_for_migration(
&bucket,
&encode_dir_object(&version.name),
None,
HeaderMap::new(),
&ObjectOptions {
version_id: version_id.clone(),
no_lock: true,
..Default::default()
},
)
.await
{
Ok(rd) => rd,
Err(err) => {
if is_err_object_not_found(&err) || is_err_version_not_found(&err) {
return MigrationVersionResult {
moved: false,
ignored: true,
cleanup_ignored: true,
failed: false,
stage: Some("read_source"),
error: None,
};
}
last_error = Some(err);
let Some(err) = last_error.as_ref() else {
continue;
};
if attempt + 1 >= max_attempts || !is_transient_rebalance_error(err) {
return MigrationVersionResult {
moved: false,
ignored: false,
cleanup_ignored: false,
failed: true,
stage: Some("read_source"),
error: last_error,
};
}
wait_retry(rebalance_migration_retry_delay(attempt, err)).await;
continue;
}
};
if let Err(err) = transfer(pool_index, bucket.clone(), rd).await {
if is_err_object_not_found(&err) || is_err_version_not_found(&err) {
return MigrationVersionResult {
moved: false,
ignored: true,
cleanup_ignored: true,
failed: false,
stage: Some("write_target"),
error: None,
};
}
last_error = Some(err);
let Some(err) = last_error.as_ref() else {
continue;
};
if attempt + 1 >= max_attempts || !is_transient_rebalance_error(err) {
return MigrationVersionResult {
moved: false,
ignored: false,
cleanup_ignored: false,
failed: true,
stage: Some("write_target"),
error: last_error,
};
}
wait_retry(rebalance_migration_retry_delay(attempt, err)).await;
continue;
}
return MigrationVersionResult {
moved: true,
ignored: false,
cleanup_ignored: false,
failed: false,
stage: None,
error: None,
};
}
MigrationVersionResult {
moved: false,
ignored: false,
cleanup_ignored: false,
failed: true,
stage: Some("migrate"),
error: last_error,
}
}
#[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<u8> 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<OffsetDateTime>, // Time at which rebalance-start was issued
#[serde(rename = "stopTs")]
pub end_time: Option<OffsetDateTime>, // Time at which rebalance operation completed or rebalance-stop was called
#[serde(rename = "err")]
pub last_error: Option<String>, // 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<String>,
#[serde(rename = "lastBucket", default)]
pub last_bucket: Option<String>,
#[serde(rename = "lastObject", default)]
pub last_object: Option<String>,
#[serde(rename = "lastAt", default)]
pub last_at: Option<OffsetDateTime>,
}
#[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<CancellationToken>, // To be invoked on rebalance-stop
#[serde(skip)]
pub last_refreshed_at: Option<OffsetDateTime>,
#[serde(rename = "stopTs")]
pub stopped_at: Option<OffsetDateTime>, // 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<RebalanceStats>, // Per-pool rebalance stats keyed by pool index
}
fn is_rebalance_pool_started(pool_stat: &RebalanceStats) -> bool {
pool_stat.participating && pool_stat.info.status == RebalStatus::Started
}
fn is_rebalance_in_progress(meta: &RebalanceMeta) -> bool {
if meta.stopped_at.is_some() {
return false;
}
meta.pool_stats.iter().any(is_rebalance_pool_started)
}
fn is_rebalance_conflicting_with_decommission(meta: &RebalanceMeta) -> bool {
is_rebalance_in_progress(meta)
}
fn first_rebalance_bucket(pool_stat: &RebalanceStats) -> Option<String> {
pool_stat.buckets.first().cloned()
}
fn rebalance_meta_load_no_data_error() -> Error {
Error::other("rebalance metadata load failed: metadata payload is too short")
}
fn rebalance_meta_load_unknown_format_error(fmt: u16) -> Error {
Error::other(format!("rebalance metadata load failed: unknown format {fmt}"))
}
fn rebalance_meta_load_unknown_version_error(ver: u16) -> Error {
Error::other(format!("rebalance metadata load failed: unknown version {ver}"))
}
impl RebalanceMeta {
pub fn new() -> Self {
Self::default()
}
pub async fn load<S: ObjectIO>(&mut self, store: Arc<S>) -> Result<()> {
self.load_with_opts(store, ObjectOptions::default()).await
}
pub async fn load_with_opts<S: ObjectIO>(&mut self, store: Arc<S>, opts: ObjectOptions) -> Result<()> {
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<S: ObjectIO>(&self, store: Arc<S>) -> Result<()> {
self.save_with_opts(store, ObjectOptions::default()).await
}
pub async fn save_with_opts<S: ObjectIO>(&self, store: Arc<S>, opts: ObjectOptions) -> Result<()> {
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<S: ObjectIO + NamespaceLocking>(
&self,
pool: Arc<S>,
local_snapshot: &RebalanceMeta,
stage: &str,
) -> Result<()> {
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<usize> {
// 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<String>) -> Result<String> {
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<Option<String>> {
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<Self>) -> 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<Self>) -> 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<Self>, 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::<Result<()>>(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
}
}
fn rebalance_goal_reached(init_free_space: u64, init_capacity: u64, bytes: u64, percent_free_goal: f64) -> bool {
if init_capacity == 0 {
return false;
}
let pfi = (init_free_space + bytes) as f64 / init_capacity as f64;
pfi + f64::EPSILON >= percent_free_goal
}
fn percent_free_ratio(total_free: u64, total_cap: u64) -> f64 {
if total_cap == 0 {
return 0.0;
}
total_free as f64 / total_cap as f64
}
fn next_rebal_bucket_from_stat(pool_stat: &RebalanceStats) -> Option<String> {
if pool_stat.buckets.is_empty() {
return None;
}
first_rebalance_bucket(pool_stat)
}
fn rebalance_metadata_not_initialized_error(operation: &str) -> Error {
Error::other(format!("failed to {operation}: rebalance metadata not initialized"))
}
fn invalid_rebalance_pool_index_error(pool_index: usize, pool_count: usize) -> Error {
Error::other(format!("invalid rebalance pool index {pool_index} for {pool_count} pools"))
}
fn clone_rebalance_pool_stats(meta: Option<&RebalanceMeta>) -> Result<Vec<RebalanceStats>> {
let Some(meta) = meta else {
return Err(rebalance_metadata_not_initialized_error("clone rebalance pool stats"));
};
Ok(meta.pool_stats.clone())
}
fn should_accept_rebalance_stats_update(meta: &RebalanceMeta, pool_index: usize) -> bool {
if meta.stopped_at.is_some() {
return false;
}
meta.pool_stats
.get(pool_index)
.is_some_and(|pool_stat| pool_stat.info.status == RebalStatus::Started)
}
fn resolve_next_rebalance_bucket(meta: Option<&RebalanceMeta>, pool_index: usize) -> Result<Option<String>> {
let Some(meta) = meta else {
return Err(rebalance_metadata_not_initialized_error("resolve next rebalance bucket"));
};
ensure_valid_rebalance_pool_index(meta.pool_stats.len(), pool_index)?;
let Some(pool_stat) = meta.pool_stats.get(pool_index) else {
return Err(invalid_rebalance_pool_index_error(pool_index, meta.pool_stats.len()));
};
if pool_stat.info.status == RebalStatus::Completed || !pool_stat.participating {
debug!(
event = EVENT_REBALANCE_BUCKET,
component = LOG_COMPONENT_ECSTORE,
subsystem = LOG_SUBSYSTEM_REBALANCE,
pool_index,
state = "unavailable",
reason = "completed_or_not_participating",
"No rebalance bucket available"
);
return Ok(None);
}
if pool_stat.buckets.is_empty() {
debug!(
event = EVENT_REBALANCE_BUCKET,
component = LOG_COMPONENT_ECSTORE,
subsystem = LOG_SUBSYSTEM_REBALANCE,
pool_index,
state = "unavailable",
reason = "bucket_queue_empty",
"No rebalance bucket available"
);
return Ok(None);
}
if let Some(bucket) = next_rebal_bucket_from_stat(pool_stat) {
debug!(
event = EVENT_REBALANCE_BUCKET,
component = LOG_COMPONENT_ECSTORE,
subsystem = LOG_SUBSYSTEM_REBALANCE,
pool_index,
bucket = %bucket,
state = "selected",
"Selected rebalance bucket"
);
return Ok(Some(bucket));
}
debug!(
event = EVENT_REBALANCE_BUCKET,
component = LOG_COMPONENT_ECSTORE,
subsystem = LOG_SUBSYSTEM_REBALANCE,
pool_index,
state = "unavailable",
reason = "selection_returned_none",
"No rebalance bucket available"
);
Ok(None)
}
fn mark_rebalance_bucket_done(meta: Option<&mut RebalanceMeta>, pool_index: usize, bucket: &str) -> Result<()> {
let Some(meta) = meta else {
return Err(rebalance_metadata_not_initialized_error("mark rebalance bucket done"));
};
ensure_valid_rebalance_pool_index(meta.pool_stats.len(), pool_index)?;
let Some(pool_stat) = meta.pool_stats.get_mut(pool_index) else {
return Err(invalid_rebalance_pool_index_error(pool_index, meta.pool_stats.len()));
};
if take_bucket_from_rebalance_queue(pool_stat, bucket) {
debug!(
event = EVENT_REBALANCE_BUCKET,
component = LOG_COMPONENT_ECSTORE,
subsystem = LOG_SUBSYSTEM_REBALANCE,
pool_index,
bucket = %bucket,
state = "queue_removed",
remaining_bucket_count = pool_stat.buckets.len(),
"Removed bucket from rebalance queue"
);
if has_deferred_rebalance_error(pool_stat) {
pool_stat.info.last_error = None;
}
Ok(())
} else {
Err(Error::other(format!(
"failed to mark rebalance bucket done: bucket {bucket} was not queued for pool {pool_index}"
)))
}
}
fn record_rebalance_cleanup_warning_in_meta(
meta: Option<&mut RebalanceMeta>,
pool_index: usize,
bucket: &str,
object: &str,
message: String,
now: OffsetDateTime,
) -> Result<()> {
let Some(meta) = meta else {
return Err(rebalance_metadata_not_initialized_error("record rebalance cleanup warning"));
};
ensure_valid_rebalance_pool_index(meta.pool_stats.len(), pool_index)?;
let Some(pool_stat) = meta.pool_stats.get_mut(pool_index) else {
return Err(invalid_rebalance_pool_index_error(pool_index, meta.pool_stats.len()));
};
pool_stat.cleanup_warnings.count = pool_stat.cleanup_warnings.count.saturating_add(1);
pool_stat.cleanup_warnings.last_message = Some(message);
pool_stat.cleanup_warnings.last_bucket = Some(bucket.to_string());
pool_stat.cleanup_warnings.last_object = Some(object.to_string());
pool_stat.cleanup_warnings.last_at = Some(now);
meta.last_refreshed_at = Some(now);
Ok(())
}
fn take_bucket_from_rebalance_queue(pool_stat: &mut RebalanceStats, bucket: &str) -> bool {
let mut found = false;
pool_stat.buckets.retain(|name| {
if name == bucket {
found = true;
pool_stat.rebalanced_buckets.push(name.clone());
false
} else {
true
}
});
found
}
fn defer_bucket_in_rebalance_queue(pool_stat: &mut RebalanceStats, bucket: &str) -> Result<()> {
let Some(pos) = pool_stat.buckets.iter().position(|name| name == bucket) else {
return Err(Error::other(format!("failed to defer rebalance bucket {bucket}: bucket was not queued")));
};
let bucket = pool_stat.buckets.remove(pos);
pool_stat.buckets.push(bucket);
Ok(())
}
fn should_pool_participate(init_free_space: u64, init_capacity: u64, percent_free_goal: f64) -> bool {
init_capacity > 0 && percent_free_ratio(init_free_space, init_capacity) < percent_free_goal
}
fn complete_rebalance_pools_at_goal(meta: &mut RebalanceMeta, now: OffsetDateTime) -> bool {
let mut changed = false;
for pool_stat in meta.pool_stats.iter_mut() {
if !is_rebalance_pool_started(pool_stat) || has_deferred_rebalance_error(pool_stat) {
continue;
}
if 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(now);
pool_stat.info.last_error = None;
changed = true;
}
}
changed
}
fn complete_rebalance_pools_with_empty_queue(meta: &mut RebalanceMeta, now: OffsetDateTime) -> bool {
let mut changed = false;
for pool_stat in meta.pool_stats.iter_mut() {
if !is_rebalance_pool_started(pool_stat) || !pool_stat.buckets.is_empty() {
continue;
}
pool_stat.info.status = RebalStatus::Completed;
pool_stat.info.end_time = Some(now);
pool_stat.info.last_error = None;
changed = true;
}
changed
}
fn has_deferred_rebalance_error(pool_stat: &RebalanceStats) -> bool {
pool_stat
.info
.last_error
.as_deref()
.is_some_and(|last_error| last_error.starts_with(REBALANCE_DEFERRED_ENTRY_ERROR_PREFIX))
}
fn resolve_rebalance_worker_result<T>(
set_idx: usize,
worker_result: std::result::Result<Result<T>, tokio::task::JoinError>,
) -> Result<T> {
match worker_result {
Ok(result) => result,
Err(err) => Err(Error::other(format!("rebalance worker {set_idx} task join error: {err}"))),
}
}
type RebalanceEntryTask = tokio::task::JoinHandle<Result<RebalanceEntryOutcome>>;
async fn wait_rebalance_entry_tasks(
set_idx: usize,
tasks: Arc<tokio::sync::Mutex<Vec<RebalanceEntryTask>>>,
) -> Result<Option<String>> {
let tasks = {
let mut tasks = tasks.lock().await;
std::mem::take(&mut *tasks)
};
let mut first_error = None;
let mut first_deferred = None;
for task in tasks {
match task.await {
Ok(Ok(RebalanceEntryOutcome::Completed)) => {}
Ok(Ok(RebalanceEntryOutcome::Deferred { last_error })) => {
if first_deferred.is_none() {
first_deferred = Some(last_error);
}
}
Ok(Err(err)) => {
error!("rebalance entry task failed for set {}: {}", set_idx, err);
if first_error.is_none() {
first_error = Some(err);
}
}
Err(err) => {
let err = Error::other(format!("rebalance entry task join error for set {set_idx}: {err}"));
error!("{}", err);
if first_error.is_none() {
first_error = Some(err);
}
}
}
}
if let Some(err) = first_error {
Err(err)
} else {
Ok(first_deferred)
}
}
fn resolve_rebalance_save_task_result(
pool_idx: usize,
save_task_result: std::result::Result<Result<()>, tokio::task::JoinError>,
) -> Result<()> {
match save_task_result {
Ok(result) => result.map_err(|err| Error::other(format!("rebalance save_task failed for pool {pool_idx}: {err}"))),
Err(err) => Err(Error::other(format!("rebalance save_task for pool {pool_idx} join error: {err}"))),
}
}
fn resolve_rebalance_meta_save_result(result: Result<()>, stage: &str) -> Result<()> {
result.map_err(|err| Error::other(format!("rebalance meta save failed during {stage}: {err}")))
}
fn rebalance_meta_lock_error(err: rustfs_lock::LockError) -> Error {
match err {
rustfs_lock::LockError::QuorumNotReached { required, achieved } => Error::NamespaceLockQuorumUnavailable {
mode: "write",
bucket: crate::disk::RUSTFS_META_BUCKET.to_string(),
object: REBAL_META_NAME.to_string(),
required,
achieved,
},
other => Error::other(format!(
"failed to acquire rebalance metadata write lock on {}/{}: {other}",
crate::disk::RUSTFS_META_BUCKET,
REBAL_META_NAME
)),
}
}
fn resolve_rebalance_meta_load_result(result: Result<()>) -> Result<bool> {
match result {
Ok(()) => Ok(true),
Err(Error::ConfigNotFound) => Ok(false),
Err(err) => {
error!("rebalanceMeta: load rebalance meta err {:?}", &err);
Err(Error::other(format!("rebalance metadata load failed during load_rebalance_meta: {err}")))
}
}
}
fn resolve_rebalance_stats_update_result(result: Result<()>, pool_idx: usize, bucket: &str, object_name: &str) -> Result<()> {
result.map_err(|err| {
Error::other(format!(
"rebalance stats update failed for pool {pool_idx} bucket {bucket} object {object_name}: {err}"
))
})
}
fn resolve_rebalance_file_info_versions_result<T, E>(
result: std::result::Result<T, E>,
bucket: &str,
object_name: &str,
) -> Result<T>
where
E: std::fmt::Display,
{
result.map_err(|err| Error::other(format!("rebalance file_info_versions failed for {bucket}/{object_name}: {err}")))
}
fn resolve_rebalance_entry_cleanup_delete_result(
result: Result<ObjectInfo>,
bucket: &str,
object_name: &str,
) -> Result<Option<String>> {
match result {
Ok(_) => Ok(None),
Err(err) if is_err_object_not_found(&err) || is_err_version_not_found(&err) => Ok(None),
Err(err) => Ok(Some(format!("rebalance cleanup delete failed for {bucket}/{object_name}: {err}"))),
}
}
fn resolve_rebalance_migrate_result_error(
err: Option<Error>,
pool_idx: usize,
bucket: &str,
object_name: &str,
version_id: Option<&str>,
) -> Error {
err.unwrap_or_else(|| {
Error::other(format!(
"rebalance migration reported failure without error for pool {pool_idx} entry {bucket}/{object_name} version {}",
version_id.unwrap_or("none")
))
})
}
fn should_defer_rebalance_entry_failure(err: &Error) -> bool {
is_transient_rebalance_error(err)
}
fn resolve_load_rebalance_stats_update_result(result: Result<()>) -> Result<()> {
result.map_err(|err| Error::other(format!("rebalance metadata stats refresh failed after load: {err}")))
}
async fn send_rebalance_done_signal(
done_tx: &tokio::sync::mpsc::Sender<Result<()>>,
signal: Result<()>,
pool_idx: usize,
) -> Result<()> {
done_tx
.send(signal)
.await
.map_err(|err| Error::other(format!("rebalance done signal send failed for pool {pool_idx}: {err}")))
}
fn resolve_rebalance_terminal_error(primary_err: Error, signal_result: Result<()>) -> Error {
match signal_result {
Ok(()) => primary_err,
Err(signal_err) => Error::other(format!("rebalance terminal signal failed after error {primary_err}: {signal_err}")),
}
}
fn resolve_rebalance_bucket_error(entry_error: Option<Error>, worker_error: Option<Error>) -> Result<()> {
if let Some(err) = entry_error {
return Err(err);
}
if let Some(err) = worker_error {
return Err(err);
}
Ok(())
}
fn resolve_rebalance_bucket_result(
result: Result<RebalanceBucketOutcome>,
pool_idx: usize,
bucket: &str,
) -> Result<RebalanceBucketOutcome> {
match result {
Ok(outcome) => Ok(outcome),
Err(err) if is_err_operation_canceled(&err) => Err(err),
Err(err) => Err(Error::other(format!("rebalance bucket {bucket} failed for pool {pool_idx}: {err}"))),
}
}
fn is_transient_rebalance_error(err: &Error) -> bool {
match err {
Error::SlowDown
| Error::ErasureReadQuorum
| Error::ErasureWriteQuorum
| Error::InsufficientReadQuorum(_, _)
| Error::InsufficientWriteQuorum(_, _) => true,
Error::Lock(lock_err) => is_rebalance_transient_lock_error(lock_err),
Error::Io(io_err) => is_rebalance_transient_io_error(io_err) || is_rebalance_transient_message(&io_err.to_string()),
_ => is_rebalance_transient_message(&err.to_string()) || is_network_or_host_down(&err.to_string(), true),
}
}
fn is_rebalance_transient_lock_error(err: &rustfs_lock::LockError) -> bool {
match err {
rustfs_lock::LockError::Timeout { .. } | rustfs_lock::LockError::Network { .. } => true,
rustfs_lock::LockError::Internal { message } => is_rebalance_transient_message(message),
_ => false,
}
}
fn is_rebalance_transient_io_error(err: &std::io::Error) -> bool {
if err.kind() == std::io::ErrorKind::TimedOut {
return true;
}
if let Some(disk_err) = err.get_ref().and_then(|err| err.downcast_ref::<DiskError>())
&& *disk_err == DiskError::Timeout
{
return true;
}
let message = err.to_string();
message.eq_ignore_ascii_case("timeout") || is_rebalance_transient_message(&message)
}
fn is_rebalance_transient_message(message: &str) -> bool {
let message = message.to_ascii_lowercase();
message.contains("lock acquisition timed out")
|| message.contains("remote lock rpc timed out")
|| message.contains("keepalivetimedout")
|| message.contains("i/o timeout")
|| message.contains("operation timed out")
}
fn should_retry_rebalance_listing(err: &Error, attempt: usize, max_attempts: usize) -> bool {
attempt + 1 < max_attempts && is_transient_rebalance_error(err)
}
fn parse_rebalance_max_attempts(value: Option<&str>) -> usize {
value
.and_then(|value| value.trim().parse::<usize>().ok())
.filter(|attempts| *attempts > 0)
.unwrap_or(DEFAULT_REBALANCE_MAX_ATTEMPTS)
}
fn rebalance_max_attempts() -> usize {
parse_rebalance_max_attempts(std::env::var(REBALANCE_MAX_ATTEMPTS_ENV).ok().as_deref())
}
fn rebalance_listing_retry_delay(attempt: usize) -> Duration {
let multiplier = u32::try_from(attempt.saturating_add(1)).unwrap_or(u32::MAX);
REBALANCE_LISTING_RETRY_BASE_DELAY.saturating_mul(multiplier)
}
fn is_rebalance_lock_or_rpc_timeout(err: &Error) -> bool {
match err {
Error::Lock(rustfs_lock::LockError::Timeout { .. }) | Error::Lock(rustfs_lock::LockError::Network { .. }) => true,
Error::Io(io_err) => is_rebalance_lock_or_rpc_timeout_message(&io_err.to_string()),
_ => is_rebalance_lock_or_rpc_timeout_message(&err.to_string()),
}
}
fn is_rebalance_lock_or_rpc_timeout_message(message: &str) -> bool {
let message = message.to_ascii_lowercase();
message.contains("lock acquisition timed out")
|| message.contains("remote lock rpc timed out")
|| message.contains("keepalivetimedout")
}
fn rebalance_migration_retry_delay(attempt: usize, err: &Error) -> Duration {
if is_rebalance_lock_or_rpc_timeout(err) {
return rebalance_lock_retry_delay(attempt);
}
let multiplier = u32::try_from(attempt.saturating_add(1)).unwrap_or(u32::MAX);
REBALANCE_MIGRATION_RETRY_BASE_DELAY.saturating_mul(multiplier)
}
fn rebalance_lock_retry_delay(attempt: usize) -> Duration {
let lock_timeout = get_lock_acquire_timeout();
let attempt_shift = u32::try_from(attempt.min(4)).unwrap_or(4);
let multiplier = 1_u32.checked_shl(attempt_shift).unwrap_or(u32::MAX);
let cap = lock_timeout
.saturating_mul(multiplier)
.min(REBALANCE_MIGRATION_LOCK_RETRY_CAP)
.max(REBALANCE_MIGRATION_RETRY_BASE_DELAY);
let max_millis = u64::try_from(cap.as_millis()).unwrap_or(u64::MAX).max(1);
let jitter_millis = rand::rng().random_range(1..=max_millis);
Duration::from_millis(jitter_millis)
}
async fn sleep_rebalance_migration_retry(delay: Duration) {
tokio::time::sleep(delay).await;
}
async fn wait_rebalance_listing_retry(rx: &CancellationToken, delay: Duration) -> Result<()> {
tokio::select! {
_ = rx.cancelled() => Err(Error::OperationCanceled),
_ = tokio::time::sleep(delay) => Ok(()),
}
}
fn ensure_rebalance_listing_disks_available(has_disks: bool, bucket: &str) -> Result<()> {
if !has_disks {
return Err(Error::other(format!(
"failed to list objects to rebalance for bucket {bucket}: no disks available"
)));
}
Ok(())
}
fn with_rebalance_entry_context(stage: &str, bucket: &str, object_name: &str, err: Error) -> Error {
Error::other(format!("rebalance entry {stage} failed for {bucket}/{object_name}: {err}"))
}
fn should_count_rebalance_version_complete(result: &MigrationVersionResult) -> bool {
result.cleanup_ignored || (result.moved && !result.failed)
}
fn should_cleanup_rebalance_source_entry(rebalanced: usize, total_versions: usize) -> bool {
rebalanced == total_versions
}
fn should_skip_rebalance_delete_marker(version: &FileInfo, remaining_versions: usize, replication_configured: bool) -> bool {
version.deleted && remaining_versions == 1 && !replication_configured
}
fn resolve_rebalance_optional_bucket_config_result<T>(bucket: &str, stage: &str, result: Result<T>) -> Result<Option<T>> {
match result {
Ok(config) => Ok(Some(config)),
Err(Error::ConfigNotFound) => Ok(None),
Err(err) => Err(Error::other(format!("rebalance {stage} config load failed for bucket {bucket}: {err}"))),
}
}
async fn load_rebalance_bucket_configs(bucket: &str) -> Result<RebalanceBucketConfigs> {
if bucket == crate::disk::RUSTFS_META_BUCKET {
return Ok(RebalanceBucketConfigs::default());
}
let _ = resolve_rebalance_optional_bucket_config_result(
bucket,
"versioning",
crate::bucket::versioning_sys::BucketVersioningSys::get(bucket).await,
)?;
Ok(RebalanceBucketConfigs {
lifecycle_config: crate::global::GLOBAL_LifecycleSys.get(bucket).await,
lock_retention: crate::bucket::object_lock::objectlock_sys::BucketObjectLockSys::get(bucket).await,
replication_config: resolve_rebalance_optional_bucket_config_result(
bucket,
"replication",
crate::bucket::metadata_sys::get_replication_config(bucket).await,
)?,
})
}
fn clone_first_arc<T>(values: &[Arc<T>], err_msg: &str) -> Result<Arc<T>> {
values.first().cloned().ok_or_else(|| Error::other(err_msg))
}
fn clone_arc_by_index<T>(values: &[Arc<T>], idx: usize, err_prefix: &str) -> Result<Arc<T>> {
values
.get(idx)
.cloned()
.ok_or_else(|| Error::other(format!("{err_prefix}: {idx}")))
}
fn ensure_valid_rebalance_pool_index(pool_count: usize, idx: usize) -> Result<()> {
if idx >= pool_count {
return Err(invalid_rebalance_pool_index_error(idx, pool_count));
}
Ok(())
}
enum RebalanceTerminalEvent {
Completed { msg: String },
Stopped { msg: String },
Failed { msg: String, last_error: String },
ChannelClosed { msg: String, last_error: String },
}
impl RebalanceTerminalEvent {
fn message(&self) -> &str {
match self {
RebalanceTerminalEvent::Completed { msg }
| RebalanceTerminalEvent::Stopped { msg }
| RebalanceTerminalEvent::Failed { msg, .. }
| RebalanceTerminalEvent::ChannelClosed { msg, .. } => msg,
}
}
}
fn apply_rebalance_terminal_event(
status: &mut RebalStatus,
end_time: &mut Option<OffsetDateTime>,
last_error: &mut Option<String>,
terminal_event: RebalanceTerminalEvent,
now: OffsetDateTime,
) {
match terminal_event {
RebalanceTerminalEvent::Completed { .. } => {
*status = RebalStatus::Completed;
*end_time = Some(now);
*last_error = None;
}
RebalanceTerminalEvent::Stopped { .. } => {
*status = RebalStatus::Stopped;
*end_time = Some(now);
*last_error = None;
}
RebalanceTerminalEvent::Failed { last_error: err, .. }
| RebalanceTerminalEvent::ChannelClosed { last_error: err, .. } => {
*status = RebalStatus::Failed;
*end_time = Some(now);
*last_error = Some(err);
}
}
}
fn classify_rebalance_terminal_event(signal: Option<Result<()>>, now: OffsetDateTime) -> RebalanceTerminalEvent {
match signal {
Some(Ok(())) => RebalanceTerminalEvent::Completed {
msg: format!("Rebalance completed at {now:?}"),
},
Some(Err(err)) => {
if is_err_operation_canceled(&err) {
RebalanceTerminalEvent::Stopped {
msg: format!("Rebalance stopped at {now:?}"),
}
} else {
RebalanceTerminalEvent::Failed {
msg: format!("Rebalance failed at {now:?} with err {err:?}"),
last_error: err.to_string(),
}
}
}
None => RebalanceTerminalEvent::ChannelClosed {
msg: format!("Rebalance save task channel closed unexpectedly at {now:?}"),
last_error: format!("rebalance save channel closed before terminal event at {now:?}"),
},
}
}
fn ensure_rebalance_not_decommissioning(decommission_running: bool) -> bool {
!decommission_running
}
fn validate_start_rebalance_state(decommission_running: bool, meta_loaded: bool) -> Result<()> {
if !ensure_rebalance_not_decommissioning(decommission_running) {
return Err(Error::DecommissionAlreadyRunning);
}
if !meta_loaded {
return Err(Error::ConfigNotFound);
}
Ok(())
}
fn should_skip_start_rebalance(cancel_attached: bool, in_progress: bool) -> bool {
cancel_attached && in_progress
}
fn is_rebalance_stopped_terminal_event(terminal_event: &RebalanceTerminalEvent) -> bool {
matches!(terminal_event, RebalanceTerminalEvent::Stopped { .. })
}
fn should_preserve_rebalance_stopped_state(
meta_stopped: bool,
status: RebalStatus,
terminal_event: &RebalanceTerminalEvent,
) -> bool {
(meta_stopped || status == RebalStatus::Stopped) && !is_rebalance_stopped_terminal_event(terminal_event)
}
fn resolve_rebalance_participants(pool_stats: &[RebalanceStats], pool_count: usize) -> Vec<bool> {
let mut participants = vec![false; pool_count];
for (idx, pool_stat) in pool_stats.iter().enumerate() {
if idx >= participants.len() {
break;
}
if pool_stat.info.status == RebalStatus::Started {
participants[idx] = pool_stat.participating;
}
}
participants
}
fn is_rebalance_actively_running(meta: &RebalanceMeta) -> bool {
meta.cancel.is_some() && is_rebalance_in_progress(meta)
}
fn should_ignore_rebalance_data_usage_cache(bucket: &str) -> bool {
bucket == crate::disk::RUSTFS_META_BUCKET
}
fn apply_rebalance_save_option(meta: &mut RebalanceMeta, pool_idx: usize, opt: RebalSaveOpt, now: OffsetDateTime) {
match opt {
RebalSaveOpt::Stats => {
if pool_idx >= meta.pool_stats.len() {
info!("save_rebalance_stats: pool_idx {pool_idx} out of range for pool_stats");
}
}
RebalSaveOpt::StoppedAt => {
apply_stopped_at(meta, now);
}
}
meta.last_refreshed_at = Some(now);
}
fn is_rebalance_terminal_status(status: RebalStatus) -> bool {
matches!(status, RebalStatus::Completed | RebalStatus::Stopped | RebalStatus::Failed)
}
fn merge_rebalance_bucket_lists(remote: &mut Vec<String>, local: &[String]) {
let mut existing: HashSet<String> = remote.iter().cloned().collect();
for bucket in local {
if existing.insert(bucket.clone()) {
remote.push(bucket.clone());
}
}
}
fn remove_rebalanced_buckets_from_queue(pool_stat: &mut RebalanceStats) {
let rebalanced_buckets: HashSet<String> = pool_stat.rebalanced_buckets.iter().cloned().collect();
pool_stat.buckets.retain(|bucket| !rebalanced_buckets.contains(bucket));
}
fn merge_rebalance_cleanup_warnings(remote: &mut RebalanceCleanupWarnings, local: &RebalanceCleanupWarnings) {
remote.count = remote.count.max(local.count);
if should_replace_rebalance_cleanup_warning(remote.last_at, local.last_at) {
remote.last_message = local.last_message.clone();
remote.last_bucket = local.last_bucket.clone();
remote.last_object = local.last_object.clone();
remote.last_at = local.last_at;
}
}
fn should_replace_rebalance_cleanup_warning(remote_at: Option<OffsetDateTime>, local_at: Option<OffsetDateTime>) -> bool {
match (remote_at, local_at) {
(_, None) => false,
(None, Some(_)) => true,
(Some(remote), Some(local)) => local >= remote,
}
}
fn merge_rebalance_pool_stats(remote: &mut RebalanceStats, local: &RebalanceStats) {
remote.init_free_space = remote.init_free_space.max(local.init_free_space);
remote.init_capacity = remote.init_capacity.max(local.init_capacity);
remote.participating |= local.participating;
merge_rebalance_bucket_lists(&mut remote.buckets, &local.buckets);
merge_rebalance_bucket_lists(&mut remote.rebalanced_buckets, &local.rebalanced_buckets);
remove_rebalanced_buckets_from_queue(remote);
let local_is_newer = local.num_versions >= remote.num_versions;
remote.num_objects = remote.num_objects.max(local.num_objects);
remote.num_versions = remote.num_versions.max(local.num_versions);
remote.bytes = remote.bytes.max(local.bytes);
merge_rebalance_cleanup_warnings(&mut remote.cleanup_warnings, &local.cleanup_warnings);
if local_is_newer {
remote.bucket = local.bucket.clone();
remote.object = local.object.clone();
}
if remote.info.start_time.is_none() {
remote.info.start_time = local.info.start_time;
}
if is_rebalance_terminal_status(remote.info.status) && local.info.status == RebalStatus::Started {
return;
}
if remote.info.status == RebalStatus::Stopped && matches!(local.info.status, RebalStatus::Started | RebalStatus::Completed) {
return;
}
match local.info.status {
RebalStatus::Failed => {
remote.info.status = RebalStatus::Failed;
remote.info.end_time = local.info.end_time.or(remote.info.end_time);
remote.info.last_error = local.info.last_error.clone().or_else(|| remote.info.last_error.clone());
}
RebalStatus::Stopped => {
if remote.info.status != RebalStatus::Failed {
remote.info.status = RebalStatus::Stopped;
remote.info.end_time = local.info.end_time.or(remote.info.end_time);
remote.info.last_error = None;
}
}
RebalStatus::Completed => {
if !matches!(remote.info.status, RebalStatus::Failed | RebalStatus::Stopped) {
remote.info.status = RebalStatus::Completed;
remote.info.end_time = local.info.end_time.or(remote.info.end_time);
remote.info.last_error = None;
}
}
RebalStatus::Started => {
if !is_rebalance_terminal_status(remote.info.status) {
remote.info.status = RebalStatus::Started;
remote.info.last_error = local.info.last_error.clone();
}
}
RebalStatus::None => {}
}
}
fn merge_rebalance_meta(remote: &mut RebalanceMeta, local: &RebalanceMeta) {
if remote.id.is_empty() {
*remote = local.clone();
return;
}
if !local.id.is_empty() && remote.id != local.id {
*remote = local.clone();
return;
}
remote.percent_free_goal = local.percent_free_goal;
remote.last_refreshed_at = Some(OffsetDateTime::now_utc());
if remote.stopped_at.is_none() {
remote.stopped_at = local.stopped_at;
}
if remote.pool_stats.len() < local.pool_stats.len() {
remote.pool_stats.resize_with(local.pool_stats.len(), RebalanceStats::default);
}
for (idx, local_pool_stat) in local.pool_stats.iter().enumerate() {
if let Some(remote_pool_stat) = remote.pool_stats.get_mut(idx) {
merge_rebalance_pool_stats(remote_pool_stat, local_pool_stat);
}
}
}
fn mark_started_rebalance_pools_stopped(meta: &mut RebalanceMeta, stop_time: OffsetDateTime) {
for pool_stat in meta.pool_stats.iter_mut() {
if pool_stat.info.status == RebalStatus::Started {
pool_stat.info.status = RebalStatus::Stopped;
pool_stat.info.end_time.get_or_insert(stop_time);
pool_stat.info.last_error = None;
}
}
}
fn apply_stopped_at(meta: &mut RebalanceMeta, now: OffsetDateTime) {
meta.stopped_at = Some(now);
mark_started_rebalance_pools_stopped(meta, now);
}
fn stop_rebalance_state(meta: &mut RebalanceMeta, now: OffsetDateTime) {
if let Some(cancel_tx) = meta.cancel.take() {
cancel_tx.cancel();
}
let stop_time = meta.stopped_at.unwrap_or(now);
if meta.stopped_at.is_none() && is_rebalance_in_progress(meta) {
meta.stopped_at = Some(stop_time);
}
if meta.stopped_at.is_some() {
mark_started_rebalance_pools_stopped(meta, stop_time);
}
}
fn stop_rebalance_meta_snapshot(meta: Option<&mut RebalanceMeta>, now: OffsetDateTime) -> Option<RebalanceMeta> {
meta.map(|meta| {
stop_rebalance_state(meta, now);
meta.last_refreshed_at = Some(now);
meta.clone()
})
}
impl ECStore {
#[allow(unused_assignments)]
#[tracing::instrument(skip(self, set))]
async fn rebalance_entry(
self: Arc<Self>,
bucket: String,
pool_index: usize,
entry: MetaCacheEntry,
set: Arc<SetDisks>,
bucket_configs: Arc<RebalanceBucketConfigs>,
// wk: Arc<Workers>,
) -> Result<RebalanceEntryOutcome> {
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<Self>, 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<Self>,
rx: CancellationToken,
bucket: String,
pool_index: usize,
) -> Result<RebalanceBucketOutcome> {
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::<Error>));
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::<RebalanceEntryTask>::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<Error> = None;
let mut deferred_error: Option<String> = 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(())
}
}
async fn run_rebalance_listing_with_retry(
set: Arc<SetDisks>,
rx: CancellationToken,
bucket: String,
cb: ListCallback,
set_idx: usize,
max_attempts: usize,
) -> Result<()> {
let max_attempts = max_attempts.max(1);
let mut last_error = None;
for attempt in 0..max_attempts {
match set.list_objects_to_rebalance(rx.clone(), bucket.clone(), cb.clone()).await {
Ok(()) => return Ok(()),
Err(err) if should_retry_rebalance_listing(&err, attempt, max_attempts) => {
let next_attempt = attempt + 2;
let delay = rebalance_listing_retry_delay(attempt);
error!(
"rebalance listing failed for bucket {} set {} attempt {}/{}: {}; retrying in {:?}",
bucket,
set_idx,
attempt + 1,
max_attempts,
err,
delay
);
last_error = Some(err);
wait_rebalance_listing_retry(&rx, delay).await?;
info!(
"rebalance listing retrying bucket {} set {} attempt {}/{}",
bucket, set_idx, next_attempt, max_attempts
);
}
Err(err) => {
return Err(Error::other(format!(
"rebalance listing failed for bucket {bucket} set {set_idx} attempt {}/{}: {err}",
attempt + 1,
max_attempts
)));
}
}
}
Err(Error::other(format!(
"rebalance listing failed for bucket {bucket} set {set_idx} after {max_attempts} attempts: {}",
last_error
.map(|err| err.to_string())
.unwrap_or_else(|| "unknown listing failure".to_string())
)))
}
impl SetDisks {
#[tracing::instrument(skip(self, rx, cb))]
pub async fn list_objects_to_rebalance(
self: &Arc<Self>,
rx: CancellationToken,
bucket: String,
cb: ListCallback,
) -> Result<()> {
debug!(
event = EVENT_REBALANCE_LISTING,
component = LOG_COMPONENT_ECSTORE,
subsystem = LOG_SUBSYSTEM_REBALANCE,
bucket = %bucket,
state = "started",
"Rebalance listing started"
);
// Placeholder for actual object listing logic
let (disks, _) = self.get_online_disks_with_healing(false).await;
ensure_rebalance_listing_disks_available(!disks.is_empty(), &bucket)?;
debug!(
event = EVENT_REBALANCE_LISTING,
component = LOG_COMPONENT_ECSTORE,
subsystem = LOG_SUBSYSTEM_REBALANCE,
bucket = %bucket,
disk_count = disks.len(),
state = "disks_resolved",
"Rebalance listing disks resolved"
);
let listing_quorum = self.set_drive_count.div_ceil(2);
let resolver = MetadataResolutionParams {
dir_quorum: listing_quorum,
obj_quorum: listing_quorum,
bucket: bucket.clone(),
..Default::default()
};
let cb1 = cb.clone();
list_path_raw(
rx,
ListPathRawOptions {
disks: disks.iter().cloned().map(Some).collect(),
bucket: bucket.clone(),
recursive: true,
min_disks: listing_quorum,
skip_walkdir_total_timeout: true,
agreed: Some(Box::new(move |entry: MetaCacheEntry| {
debug!(
event = EVENT_REBALANCE_LISTING,
component = LOG_COMPONENT_ECSTORE,
subsystem = LOG_SUBSYSTEM_REBALANCE,
entry = %entry.name,
state = "agreed_entry",
"Rebalance listing agreed entry"
);
Box::pin(cb1(entry))
})),
partial: Some(Box::new(move |entries: MetaCacheEntries, _: &[Option<DiskError>]| {
// let cb = cb.clone();
let resolver = resolver.clone();
let cb = cb.clone();
match entries.resolve(resolver) {
Some(entry) => {
debug!(
event = EVENT_REBALANCE_LISTING,
component = LOG_COMPONENT_ECSTORE,
subsystem = LOG_SUBSYSTEM_REBALANCE,
entry = %entry.name,
state = "resolved_partial_entry",
"Rebalance listing resolved partial entry"
);
Box::pin(async move { cb(entry).await })
}
None => {
debug!(
event = EVENT_REBALANCE_LISTING,
component = LOG_COMPONENT_ECSTORE,
subsystem = LOG_SUBSYSTEM_REBALANCE,
state = "partial_entry_missing",
"Rebalance listing partial entry missing"
);
Box::pin(async {})
}
}
})),
..Default::default()
},
)
.await?;
debug!(
event = EVENT_REBALANCE_LISTING,
component = LOG_COMPONENT_ECSTORE,
subsystem = LOG_SUBSYSTEM_REBALANCE,
bucket = %bucket,
state = "completed",
"Rebalance listing completed"
);
Ok(())
}
}
#[cfg(test)]
mod rebalance_unit_tests {
use super::REBALANCE_DEFERRED_ENTRY_ERROR_PREFIX;
use super::first_rebalance_bucket;
use super::is_rebalance_actively_running;
use super::is_rebalance_conflicting_with_decommission;
use super::is_rebalance_in_progress;
use super::percent_free_ratio;
use super::rebalance_goal_reached;
use super::{
DiskError, GetObjectReader, MigrationBackend, MigrationVersionResult, ObjectInfo, ObjectOptions, RebalSaveOpt,
RebalStatus, RebalanceBucketOutcome, RebalanceCleanupWarnings, RebalanceEntryOutcome, RebalanceInfo, RebalanceMeta,
RebalanceStats, 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_listing_disks_available, ensure_rebalance_not_decommissioning, ensure_valid_rebalance_pool_index,
has_deferred_rebalance_error, is_rebalance_stopped_terminal_event, is_transient_rebalance_error,
load_rebalance_bucket_configs, mark_rebalance_bucket_done, merge_rebalance_meta, migrate_entry_version,
migrate_entry_version_with_retry_wait, next_rebal_bucket_from_stat, parse_rebalance_max_attempts,
rebalance_delete_marker_opts, rebalance_listing_retry_delay, rebalance_meta_load_no_data_error,
rebalance_meta_load_unknown_format_error, rebalance_meta_load_unknown_version_error, rebalance_migration_retry_delay,
record_rebalance_cleanup_warning_in_meta, resolve_load_rebalance_stats_update_result, resolve_next_rebalance_bucket,
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_participants,
resolve_rebalance_save_task_result, resolve_rebalance_stats_update_result, resolve_rebalance_terminal_error,
resolve_rebalance_worker_result, send_rebalance_done_signal, should_accept_rebalance_stats_update,
should_cleanup_rebalance_source_entry, should_count_rebalance_version_complete, should_defer_rebalance_entry_failure,
should_ignore_rebalance_data_usage_cache, should_pool_participate, should_preserve_rebalance_stopped_state,
should_retry_rebalance_listing, should_skip_rebalance_delete_marker, should_skip_start_rebalance,
stop_rebalance_meta_snapshot, stop_rebalance_state, take_bucket_from_rebalance_queue, validate_start_rebalance_state,
wait_rebalance_listing_retry, with_rebalance_entry_context,
};
use crate::data_movement;
use crate::data_usage::DATA_USAGE_CACHE_NAME;
use crate::disk::RUSTFS_META_BUCKET;
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<String>,
#[serde(rename = "rbs")]
rebalanced_buckets: Vec<String>,
#[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<OffsetDateTime>,
#[serde(rename = "id")]
id: String,
#[serde(rename = "pf")]
percent_free_goal: f64,
#[serde(rename = "rss")]
pool_stats: Vec<LegacyRebalanceStats>,
}
struct MigrationBackendSpy {
get_object_reader: Mutex<Option<core::result::Result<GetObjectReader, Error>>>,
delete_object: Mutex<Option<core::result::Result<ObjectInfo, Error>>>,
move_remote: Mutex<Option<core::result::Result<(), Error>>>,
get_calls: AtomicUsize,
delete_calls: AtomicUsize,
move_remote_calls: AtomicUsize,
}
impl MigrationBackendSpy {
fn new(
get_object_reader: Option<core::result::Result<GetObjectReader, Error>>,
delete_object: Option<core::result::Result<ObjectInfo, Error>>,
move_remote: Option<core::result::Result<(), Error>>,
) -> 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<HTTPRangeSpec>,
_h: http::HeaderMap,
_opts: &ObjectOptions,
) -> Result<GetObjectReader> {
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<ObjectInfo> {
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<HTTPRangeSpec>,
_h: http::HeaderMap,
_opts: &ObjectOptions,
) -> Result<GetObjectReader> {
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<ObjectInfo> {
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::<usize, Error>(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::<usize, Error>(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<Arc<u8>> = 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(),
];
super::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()
};
super::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));
}
}
Reference in New Issue View Git Blame Copy Permalink