From 1ac0841f6f2a08a8526ce5a9138caf4eb23b4e15 Mon Sep 17 00:00:00 2001 From: cxymds Date: Sun, 19 Jul 2026 22:20:37 +0800 Subject: [PATCH] fix(heal): initialize the runtime atomically (#4989) * fix(rpc): bind internode auth to exact targets * fix(heal): initialize the runtime atomically --- crates/heal/src/lib.rs | 341 ++++++++++++++++++++++++++++++++++++----- 1 file changed, 300 insertions(+), 41 deletions(-) diff --git a/crates/heal/src/lib.rs b/crates/heal/src/lib.rs index 73ef02988..f39192b7f 100644 --- a/crates/heal/src/lib.rs +++ b/crates/heal/src/lib.rs @@ -21,8 +21,10 @@ pub use heal::{ channel::HealChannelProcessor, progress::HealProgress, }; use rustfs_concurrency::WorkloadAdmissionSnapshotProvider; +use std::future::Future; use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{Arc, OnceLock}; +use tokio::sync::{Mutex, OnceCell}; use tokio_util::sync::CancellationToken; use tracing::{error, info}; @@ -59,14 +61,61 @@ pub fn shutdown_ahm_services() { } } -/// Global heal manager instance -static GLOBAL_HEAL_MANAGER: OnceLock> = OnceLock::new(); +struct HealRuntime { + manager: Arc, + channel_processor: Arc>, +} -/// Global heal channel processor instance -static GLOBAL_HEAL_CHANNEL_PROCESSOR: OnceLock>> = OnceLock::new(); +/// Process-wide heal runtime. Async single-flight initialization prevents two +/// callers from starting independent managers or channel processors. +static GLOBAL_HEAL_RUNTIME: OnceCell = OnceCell::const_new(); +static GLOBAL_HEAL_RUNTIME_INIT: Mutex<()> = Mutex::const_new(()); static GLOBAL_HEAL_ACTIVE_TASKS: AtomicU64 = AtomicU64::new(0); static GLOBAL_HEAL_QUEUE_LENGTH: AtomicU64 = AtomicU64::new(0); +#[cfg(test)] +#[derive(Default)] +struct HealRuntimeInitTestHook { + starts: AtomicU64, + stops: AtomicU64, + fail_channel_init: std::sync::atomic::AtomicBool, + pause_after_start: std::sync::atomic::AtomicBool, + start_paused: tokio::sync::Notify, + release_start: tokio::sync::Notify, +} + +#[cfg(test)] +static HEAL_RUNTIME_INIT_TEST_HOOK: OnceLock> = OnceLock::new(); + +#[cfg(test)] +async fn test_hook_after_manager_start() { + if let Some(hook) = HEAL_RUNTIME_INIT_TEST_HOOK.get() { + hook.starts.fetch_add(1, Ordering::SeqCst); + if hook.pause_after_start.load(Ordering::SeqCst) { + hook.start_paused.notify_one(); + hook.release_start.notified().await; + } + } +} + +async fn stop_initializing_manager(manager: &HealManager) -> Result<()> { + #[cfg(test)] + if let Some(hook) = HEAL_RUNTIME_INIT_TEST_HOOK.get() { + hook.stops.fetch_add(1, Ordering::SeqCst); + } + manager.stop().await +} + +async fn run_owned_initialization(initialization: F) -> Result +where + T: Send + 'static, + F: Future> + Send + 'static, +{ + tokio::spawn(initialization) + .await + .map_err(|err| Error::Other(format!("Heal runtime initialization task failed: {err}")))? +} + /// Initialize and start heal manager with channel processor pub async fn init_heal_manager( storage: Arc, @@ -81,33 +130,51 @@ pub async fn init_heal_manager_with_workload_provider( config: Option, workload_provider: Option>, ) -> Result> { - // Create heal manager - let heal_manager = Arc::new(HealManager::new_with_workload_provider(storage, config, workload_provider)); + // Run initialization in an owned task so cancelling an HTTP/startup caller + // cannot abandon a manager after its scheduler has been spawned. + run_owned_initialization(async move { + let _init_guard = GLOBAL_HEAL_RUNTIME_INIT.lock().await; + if GLOBAL_HEAL_RUNTIME.get().is_some() { + return Err(Error::Config("Heal manager already initialized".to_string())); + } - // Start heal manager - heal_manager.start().await?; + let heal_manager = Arc::new(HealManager::new_with_workload_provider(storage, config, workload_provider)); + if let Err(err) = heal_manager.start().await { + let _ = stop_initializing_manager(&heal_manager).await; + return Err(err); + } - // Store global instance - GLOBAL_HEAL_MANAGER - .set(heal_manager.clone()) - .map_err(|_| Error::Config("Heal manager already initialized".to_string()))?; + #[cfg(test)] + test_hook_after_manager_start().await; - // Initialize heal channel - let channel_receiver = rustfs_common::heal_channel::init_heal_channel().map_err(|err| Error::Config(err.to_string()))?; + #[cfg(test)] + let force_channel_failure = HEAL_RUNTIME_INIT_TEST_HOOK + .get() + .is_some_and(|hook| hook.fail_channel_init.load(Ordering::SeqCst)); + #[cfg(not(test))] + let force_channel_failure = false; + let channel_receiver = if force_channel_failure { + Err("forced heal channel initialization failure") + } else { + rustfs_common::heal_channel::init_heal_channel() + }; + let receiver = match channel_receiver { + Ok(receiver) => receiver, + Err(err) => { + stop_initializing_manager(&heal_manager).await?; + return Err(Error::Config(err.to_string())); + } + }; + let channel_processor = Arc::new(Mutex::new(HealChannelProcessor::new(heal_manager.clone()))); + GLOBAL_HEAL_RUNTIME + .set(HealRuntime { + manager: heal_manager.clone(), + channel_processor: channel_processor.clone(), + }) + .map_err(|_| Error::Config("Heal manager already initialized".to_string()))?; - // Create channel processor - let channel_processor = HealChannelProcessor::new(heal_manager.clone()); - - // Store channel processor instance first - GLOBAL_HEAL_CHANNEL_PROCESSOR - .set(Arc::new(tokio::sync::Mutex::new(channel_processor))) - .map_err(|_| Error::Config("Heal channel processor already initialized".to_string()))?; - - // Start channel processor in background - let receiver = channel_receiver; - tokio::spawn(async move { - if let Some(processor_guard) = GLOBAL_HEAL_CHANNEL_PROCESSOR.get() { - let mut processor = processor_guard.lock().await; + tokio::spawn(async move { + let mut processor = channel_processor.lock().await; if let Err(e) = processor.start(receiver).await { error!( target: "rustfs::heal", @@ -119,28 +186,29 @@ pub async fn init_heal_manager_with_workload_provider( "Heal runtime channel processor failed" ); } - } - }); + }); - info!( - target: "rustfs::heal", - event = EVENT_HEAL_RUNTIME_STATE, - component = LOG_COMPONENT_HEAL, - subsystem = LOG_SUBSYSTEM_RUNTIME, - state = "initialized", - "Heal runtime initialized" - ); - Ok(heal_manager) + info!( + target: "rustfs::heal", + event = EVENT_HEAL_RUNTIME_STATE, + component = LOG_COMPONENT_HEAL, + subsystem = LOG_SUBSYSTEM_RUNTIME, + state = "initialized", + "Heal runtime initialized" + ); + Ok(heal_manager) + }) + .await } /// Get global heal manager instance pub fn get_heal_manager() -> Option<&'static Arc> { - GLOBAL_HEAL_MANAGER.get() + GLOBAL_HEAL_RUNTIME.get().map(|runtime| &runtime.manager) } /// Get global heal channel processor instance -pub fn get_heal_channel_processor() -> Option<&'static Arc>> { - GLOBAL_HEAL_CHANNEL_PROCESSOR.get() +pub fn get_heal_channel_processor() -> Option<&'static Arc>> { + GLOBAL_HEAL_RUNTIME.get().map(|runtime| &runtime.channel_processor) } pub fn heal_runtime_initialized() -> bool { @@ -186,3 +254,194 @@ pub(crate) fn set_heal_active_tasks(count: usize) { pub(crate) fn set_heal_queue_length(count: usize) { GLOBAL_HEAL_QUEUE_LENGTH.store(usize_to_u64_saturated(count), Ordering::Relaxed); } + +#[cfg(test)] +mod tests { + use super::{ + Error, HEAL_RUNTIME_INIT_TEST_HOOK, HealRuntimeInitTestHook, get_heal_channel_processor, get_heal_manager, + heal::DiskStore, heal::Endpoint, heal::manager::HealConfig, heal::storage::DiskStatus, heal::storage::HealListItem, + heal::storage::HealObjectInfo, heal::storage::HealStorageAPI, init_heal_manager, run_owned_initialization, + }; + use crate::heal::storage_api::status::BucketInfo; + use rustfs_common::heal_channel::HealOpts; + use rustfs_madmin::heal_commands::HealResultItem; + use std::sync::Arc; + use std::sync::atomic::{AtomicBool, Ordering}; + use tokio::sync::{Notify, oneshot}; + + struct MockStorage; + + #[async_trait::async_trait] + impl HealStorageAPI for MockStorage { + async fn get_object_meta(&self, _bucket: &str, _object: &str) -> Result, Error> { + Ok(None) + } + + async fn get_object_data(&self, _bucket: &str, _object: &str) -> Result>, Error> { + Ok(None) + } + + async fn put_object_data(&self, _bucket: &str, _object: &str, _data: &[u8]) -> Result<(), Error> { + Ok(()) + } + + async fn delete_object(&self, _bucket: &str, _object: &str) -> Result<(), Error> { + Ok(()) + } + + async fn verify_object_integrity(&self, _bucket: &str, _object: &str) -> Result { + Ok(true) + } + + async fn ec_decode_rebuild(&self, _bucket: &str, _object: &str) -> Result, Error> { + Ok(Vec::new()) + } + + async fn get_disk_status(&self, _endpoint: &Endpoint) -> Result { + Ok(DiskStatus::Ok) + } + + async fn format_disk(&self, _endpoint: &Endpoint) -> Result<(), Error> { + Ok(()) + } + + async fn get_bucket_info(&self, _bucket: &str) -> Result, Error> { + Ok(None) + } + + async fn heal_bucket_metadata(&self, _bucket: &str) -> Result<(), Error> { + Ok(()) + } + + async fn list_buckets(&self) -> Result, Error> { + Ok(Vec::new()) + } + + async fn object_exists(&self, _bucket: &str, _object: &str) -> Result { + Ok(false) + } + + async fn get_object_size(&self, _bucket: &str, _object: &str) -> Result, Error> { + Ok(None) + } + + async fn get_object_checksum(&self, _bucket: &str, _object: &str) -> Result, Error> { + Ok(None) + } + + async fn heal_object( + &self, + _bucket: &str, + _object: &str, + _version_id: Option<&str>, + _opts: &HealOpts, + ) -> Result<(HealResultItem, Option), Error> { + Ok((HealResultItem::default(), None)) + } + + async fn heal_bucket(&self, _bucket: &str, _opts: &HealOpts) -> Result { + Ok(HealResultItem::default()) + } + + async fn heal_format(&self, _dry_run: bool) -> Result<(HealResultItem, Option), Error> { + Ok((HealResultItem::default(), None)) + } + + async fn list_objects_for_heal(&self, _bucket: &str, _prefix: &str) -> Result, Error> { + Ok(Vec::new()) + } + + async fn list_objects_for_heal_page( + &self, + _bucket: &str, + _prefix: &str, + _continuation_token: Option<&str>, + ) -> Result<(Vec, Option, bool), Error> { + Ok((Vec::new(), None, false)) + } + + async fn get_disk_for_resume(&self, _set_disk_id: &str) -> Result { + Err(Error::other("not implemented in tests")) + } + } + + #[tokio::test] + async fn owned_initialization_survives_waiter_cancellation() { + let completed = Arc::new(AtomicBool::new(false)); + let release = Arc::new(Notify::new()); + let (started_tx, started_rx) = oneshot::channel(); + let task_completed = completed.clone(); + let task_release = release.clone(); + let waiter = tokio::spawn(async move { + run_owned_initialization(async move { + let _ = started_tx.send(()); + task_release.notified().await; + task_completed.store(true, Ordering::Release); + Ok(()) + }) + .await + }); + + started_rx.await.expect("owned initializer should start"); + waiter.abort(); + release.notify_one(); + tokio::time::timeout(std::time::Duration::from_secs(1), async { + while !completed.load(Ordering::Acquire) { + tokio::task::yield_now().await; + } + }) + .await + .expect("owned initializer should outlive its cancelled waiter"); + assert!(completed.load(Ordering::Acquire)); + } + + #[tokio::test] + async fn runtime_initialization_is_atomic_cancellation_safe_and_retryable() { + let hook = Arc::new(HealRuntimeInitTestHook::default()); + assert!( + HEAL_RUNTIME_INIT_TEST_HOOK.set(hook.clone()).is_ok(), + "runtime initialization hook should be installed once" + ); + let config = HealConfig { + enable_auto_heal: false, + ..HealConfig::default() + }; + + hook.fail_channel_init.store(true, Ordering::SeqCst); + let failed = init_heal_manager(Arc::new(MockStorage), Some(config.clone())).await; + assert!(matches!(failed, Err(Error::Config(message)) if message == "forced heal channel initialization failure")); + assert_eq!(hook.starts.load(Ordering::SeqCst), 1); + assert_eq!(hook.stops.load(Ordering::SeqCst), 1); + assert!(get_heal_manager().is_none()); + assert!(get_heal_channel_processor().is_none()); + + hook.fail_channel_init.store(false, Ordering::SeqCst); + hook.pause_after_start.store(true, Ordering::SeqCst); + let first_config = config.clone(); + let first = tokio::spawn(async move { init_heal_manager(Arc::new(MockStorage), Some(first_config)).await }); + hook.start_paused.notified().await; + let second = tokio::spawn(async move { init_heal_manager(Arc::new(MockStorage), Some(config)).await }); + tokio::task::yield_now().await; + + assert_eq!(hook.starts.load(Ordering::SeqCst), 2, "concurrent loser must not start a manager"); + assert!(get_heal_manager().is_none(), "manager must not publish before the runtime is complete"); + assert!( + get_heal_channel_processor().is_none(), + "channel processor must publish atomically with the manager" + ); + + first.abort(); + hook.release_start.notify_one(); + let conflict = second + .await + .expect("second caller should complete") + .expect_err("second caller should observe the initialized runtime"); + assert!(matches!(conflict, Error::Config(message) if message == "Heal manager already initialized")); + let manager = get_heal_manager().expect("cancelled caller's owned initializer should publish the manager"); + assert!(get_heal_channel_processor().is_some()); + assert_eq!(hook.starts.load(Ordering::SeqCst), 2); + assert_eq!(hook.stops.load(Ordering::SeqCst), 1); + + manager.stop().await.expect("manager should stop cleanly"); + } +}