From 68a35fc82fe46a896e83a760ff08b03c16a8bfc9 Mon Sep 17 00:00:00 2001 From: overtrue Date: Wed, 28 Oct 2026 09:52:18 +0800 Subject: [PATCH] test(heal): cover same-id scheduler ownership --- crates/heal/src/heal/manager/scheduler.rs | 17 ++- crates/heal/src/heal/manager/tests.rs | 160 ++++++++++++++++++++++ 2 files changed, 175 insertions(+), 2 deletions(-) diff --git a/crates/heal/src/heal/manager/scheduler.rs b/crates/heal/src/heal/manager/scheduler.rs index 24b5e063c..e2b44bd7f 100644 --- a/crates/heal/src/heal/manager/scheduler.rs +++ b/crates/heal/src/heal/manager/scheduler.rs @@ -254,7 +254,9 @@ impl HealManager { let config_for_spawn = config.clone(); // start heal task - tokio::spawn(async move { + #[cfg(test)] + let worker_completion = tests::take_worker_completion_hook(task.as_ref()); + let worker = async move { debug!( target: "rustfs::heal::manager", event = EVENT_HEAL_SCHEDULER_STATE, @@ -718,7 +720,18 @@ impl HealManager { }); } notify_clone.notify_one(); - }); + }; + #[cfg(test)] + if let Some(completed) = worker_completion { + tokio::spawn(async move { + worker.await; + let _ = completed.send(()); + }); + } else { + tokio::spawn(worker); + } + #[cfg(not(test))] + tokio::spawn(worker); tasks_started += 1; } else { delayed_by_mainline_throttle = mainline_pressure.is_some(); diff --git a/crates/heal/src/heal/manager/tests.rs b/crates/heal/src/heal/manager/tests.rs index 1b168726e..53f6eaa8f 100644 --- a/crates/heal/src/heal/manager/tests.rs +++ b/crates/heal/src/heal/manager/tests.rs @@ -354,6 +354,19 @@ struct CompletedRetentionHook { static COMPLETED_RETENTION_HOOKS: LazyLock>>> = LazyLock::new(|| Mutex::new(HashMap::new())); +static WORKER_COMPLETION_HOOKS: LazyLock>>> = + LazyLock::new(|| StdMutex::new(HashMap::new())); + +pub(super) fn take_worker_completion_hook(task: &HealTask) -> Option> { + let HealType::Object { bucket, .. } = &task.heal_type else { + return None; + }; + WORKER_COMPLETION_HOOKS + .lock() + .expect("worker completion hooks") + .remove(bucket) +} + pub(super) async fn pause_completed_retention_handoff(task_id: &str) { let hook = COMPLETED_RETENTION_HOOKS.lock().await.get(task_id).cloned(); if let Some(hook) = hook { @@ -477,6 +490,153 @@ async fn completed_retention_cancel_wins_over_a_prepared_retry_snapshot() { .retain(|key, _| key != bucket && key != &task_id); } +#[tokio::test] +async fn cancelled_worker_does_not_remove_a_same_id_successor() { + let old_bucket = format!("scheduler-old-{}", Uuid::new_v4()); + let new_bucket = format!("scheduler-new-{}", Uuid::new_v4()); + let manager = HealManager::new_without_root_recovery_for_test( + Arc::new(MockStorage), + Some(HealConfig { + queue_size: 32, + max_concurrent_heals: 1, + task_timeout: Duration::from_secs(60), + event_driven_scheduler_enable: false, + mainline_throttle_enable: false, + ..Default::default() + }), + ); + let old_request = HealRequest::object(old_bucket.clone(), "object".to_string(), None); + let task_id = old_request.id.clone(); + let old_hook = Arc::new(CompletedRetentionHook { + pause_before_publish: true, + ..Default::default() + }); + let new_hook = Arc::new(CompletedRetentionHook::default()); + let (worker_completion, worker_finished) = tokio::sync::oneshot::channel(); + WORKER_COMPLETION_HOOKS + .lock() + .expect("install worker completion hook") + .insert(old_bucket.clone(), worker_completion); + { + let mut hooks = COMPLETED_RETENTION_HOOKS.lock().await; + hooks.insert(old_bucket.clone(), Arc::clone(&old_hook)); + hooks.insert(new_bucket.clone(), Arc::clone(&new_hook)); + hooks.insert(task_id.clone(), Arc::clone(&old_hook)); + } + + assert_eq!( + manager.submit_heal_request(old_request).await.expect("admit original owner"), + HealAdmissionResult::Accepted + ); + process_manager_queue_once(&manager).await; + tokio::time::timeout(Duration::from_secs(5), old_hook.started.notified()) + .await + .expect("original owner enters mock storage"); + let old_task = manager + .active_heals + .lock() + .await + .get(&task_id) + .cloned() + .expect("original owner is active"); + old_hook.execute.notify_one(); + tokio::time::timeout(Duration::from_secs(5), old_hook.before_publish.notified()) + .await + .expect("original owner pauses after execution and before publication"); + assert!(matches!(*old_hook.prepared_status.lock().await, Some(HealTaskStatus::Completed))); + manager.cancel_task(&task_id).await.expect("cancel original owner"); + assert!(!manager.active_heals.lock().await.contains_key(&task_id)); + + let mut new_request = HealRequest::object(new_bucket.clone(), "object".to_string(), None); + new_request.id.clone_from(&task_id); + new_request.source = HealRequestSource::Admin; + new_request.force_start = true; + let receipt = manager + .submit_heal_request_with_receipt(new_request) + .await + .expect("successor admission decision"); + if receipt.result == HealAdmissionResult::Dropped(HealAdmissionDropReason::AlreadyRunning) { + assert!(!manager.active_heals.lock().await.contains_key(&task_id)); + assert!( + !manager + .heal_queue + .lock() + .await + .requests() + .any(|request| request.id == task_id) + ); + old_hook.publish.notify_one(); + old_hook.finish.notify_one(); + tokio::time::timeout(Duration::from_secs(5), worker_finished) + .await + .expect("cancelled original worker exits") + .expect("cancelled original worker returns normally"); + let status = manager + .get_task_status(&task_id) + .await + .expect("cancelled original remains queryable"); + let mut hooks = COMPLETED_RETENTION_HOOKS.lock().await; + hooks.remove(&old_bucket); + hooks.remove(&new_bucket); + hooks.remove(&task_id); + assert_eq!(status, HealTaskStatus::Cancelled); + return; + } + assert_eq!( + receipt.result, + HealAdmissionResult::Accepted, + "successor must either be admitted or safely rejected" + ); + assert_eq!(receipt.task_id, task_id); + process_manager_queue_once(&manager).await; + tokio::time::timeout(Duration::from_secs(5), new_hook.started.notified()) + .await + .expect("successor enters mock storage"); + let new_task = manager + .active_heals + .lock() + .await + .get(&task_id) + .cloned() + .expect("successor is active"); + assert!(!Arc::ptr_eq(&old_task, &new_task), "successor is a different task"); + assert_eq!( + manager + .get_task_status_for_path(&format!("{new_bucket}/object"), &task_id) + .await + .expect("successor status before old handoff"), + HealTaskStatus::Running + ); + + old_hook.publish.notify_one(); + old_hook.finish.notify_one(); + tokio::time::timeout(Duration::from_secs(5), worker_finished) + .await + .expect("original worker exits") + .expect("original worker returns normally"); + let successor_is_active = manager + .active_heals + .lock() + .await + .get(&task_id) + .is_some_and(|current| Arc::ptr_eq(current, &new_task)); + let successor_status = manager + .get_task_status_for_path(&format!("{new_bucket}/object"), &task_id) + .await; + println!( + "same_id_successor admission={:?} successor_is_active={} successor_status={:?}", + receipt.result, successor_is_active, successor_status + ); + { + let mut hooks = COMPLETED_RETENTION_HOOKS.lock().await; + hooks.remove(&old_bucket); + hooks.remove(&new_bucket); + hooks.remove(&task_id); + } + assert!(successor_is_active, "the old worker must not remove the active successor"); + assert!(matches!(successor_status, Ok(HealTaskStatus::Running))); +} + #[tokio::test] async fn completed_retention_scheduler_preserves_progress_aliases_and_atomic_handoff() { for outcome in ["success", "failed", "cancelled"] {