diff --git a/crates/heal/src/heal/manager.rs b/crates/heal/src/heal/manager.rs index d36c9f9a7..971590fdb 100644 --- a/crates/heal/src/heal/manager.rs +++ b/crates/heal/src/heal/manager.rs @@ -45,6 +45,11 @@ use tracing::{debug, error, info, warn}; use super::{DiskError, Endpoint, HealDiskExt as _, local_disk_map_read}; const KEEP_HEAL_TASK_STATUS_DURATION: Duration = Duration::from_secs(10 * 60); +// Each cache includes alias tokens in its count and byte budget. Eviction +// removes every token sharing a snapshot; neither cache retains repair state. +const MAX_COMPLETED_HEAL_TOKENS: usize = 1024; +const MAX_COMPLETED_HEAL_BYTES: usize = 64 * 1024 * 1024; +const MAX_COMPLETED_HEAL_RESULT_BYTES: usize = 1024 * 1024; const DISPLACED_HEAL_REASON: &str = "reason=displaced; retry_hint=submit_again"; const LOG_COMPONENT_HEAL: &str = "heal"; const LOG_SUBSYSTEM_DISK_SCANNER: &str = "disk_scanner"; @@ -180,6 +185,8 @@ fn record_displaced_terminal( request: &HealRequest, ) -> Arc { let terminal = Arc::new(CompletedHealStatus { + progress: None, + retained_bytes: std::sync::OnceLock::new(), heal_type: request.heal_type.clone(), status: HealTaskStatus::Failed { error: format!("heal task displaced by a higher-priority request ({DISPLACED_HEAL_REASON})"), @@ -193,6 +200,7 @@ fn record_displaced_terminal( let mut terminals = lock_displaced_terminals(registry); prune_completed_heal_statuses(&mut terminals); terminals.insert(request.id.clone(), Arc::clone(&terminal)); + prune_completed_heal_statuses(&mut terminals); terminal } @@ -209,9 +217,15 @@ async fn remove_displaced_task_aliases( .collect::>(); let mut displaced_terminals = lock_displaced_terminals(terminals); prune_completed_heal_statuses(&mut displaced_terminals); - for alias_id in alias_ids { - displaced_terminals.insert(alias_id, Arc::clone(terminal)); + if displaced_terminals + .get(task_id) + .is_some_and(|current| Arc::ptr_eq(current, terminal)) + { + for alias_id in alias_ids { + displaced_terminals.insert(alias_id, Arc::clone(terminal)); + } } + prune_completed_heal_statuses(&mut displaced_terminals); aliases.retain(|alias_id, alias| alias_id != task_id && alias.task_id != task_id); } @@ -222,6 +236,36 @@ async fn remove_task_aliases_for_task(registry: &Arc +// retrying (when needed) -> aliases -> completed; queries release aliases +// before looking up active state. Publishing aliases before removing their +// mapping keeps both an already-resolved token and a new lookup valid. +async fn publish_completed_heal( + completed_heals: &Mutex>>, + task_aliases: &Mutex>, + task_id: &str, + completed: CompletedHealStatus, + terminal: bool, +) { + let completed = Arc::new(completed); + completed.retained_bytes(); + let mut aliases = task_aliases.lock().await; + let mut retained = completed_heals.lock().await; + if let Some(previous) = retained.get(task_id).cloned() { + for entry in retained.values_mut().filter(|entry| Arc::ptr_eq(entry, &previous)) { + *entry = Arc::clone(&completed); + } + } + retained.insert(task_id.to_owned(), Arc::clone(&completed)); + if terminal { + for (alias_id, _) in aliases.iter().filter(|(_, alias)| alias.task_id == task_id) { + retained.insert(alias_id.clone(), Arc::clone(&completed)); + } + aliases.retain(|alias_id, alias| alias_id != task_id && alias.task_id != task_id); + } + prune_completed_heal_statuses(&mut retained); +} + #[derive(Debug, Clone)] pub struct HealTaskReport { pub status: HealTaskStatus, @@ -268,7 +312,7 @@ fn completed_task_report(completed: &CompletedHealStatus, since: Option) -> let result_items = match since { None => completed.seqed_items.iter().map(|(_, item)| item.clone()).collect(), Some(cursor) => { - if cursor + 1 < completed.min_seq { + if cursor.saturating_add(1) < completed.min_seq { lagged = true; } completed @@ -283,7 +327,7 @@ fn completed_task_report(completed: &CompletedHealStatus, since: Option) -> status: completed.status.clone(), result_items, result_items_truncated: completed.result_items_truncated || lagged, - progress: None, + progress: completed.progress.clone(), next_seq: completed.next_seq, min_seq: completed.min_seq, } @@ -1847,14 +1891,14 @@ impl HealManager { pub async fn get_task_progress(&self, task_id: &str) -> Result { let canonical_task_id = self.canonical_task_id(task_id).await; - let active_heals = self.active_heals.lock().await; - if let Some(task) = active_heals.get(&canonical_task_id) { - Ok(task.get_progress().await) - } else { - Err(Error::TaskNotFound { - task_id: task_id.to_string(), - }) - } + let progress = match self.lookup_task_state(&canonical_task_id, None).await { + TaskStateLookup::Active(task) => Some(task.get_progress().await), + TaskStateLookup::Completed(completed) => completed.progress.clone(), + _ => None, + }; + progress.ok_or_else(|| Error::TaskNotFound { + task_id: task_id.to_string(), + }) } /// Cancel task @@ -1864,6 +1908,8 @@ impl HealManager { let mut active_heals = self.active_heals.lock().await; if let Some(task) = active_heals.get(&canonical_task_id) { task.cancel().await?; + let completed = CompletedHealStatus::snapshot(task, HealTaskStatus::Cancelled).await; + publish_completed_heal(&self.completed_heals, &self.task_aliases, &canonical_task_id, completed, true).await; active_heals.remove(&canonical_task_id); publish_active_heal_count(&active_heals); info!( @@ -1940,6 +1986,8 @@ impl HealManager { for task_id in &task_ids { if let Some(task) = active_heals.get(task_id) { task.cancel().await?; + let completed = CompletedHealStatus::snapshot(task, HealTaskStatus::Cancelled).await; + publish_completed_heal(&self.completed_heals, &self.task_aliases, task_id, completed, true).await; } active_heals.remove(task_id); cancelled += 1; diff --git a/crates/heal/src/heal/manager/queue.rs b/crates/heal/src/heal/manager/queue.rs index b4608f27b..aceabe42a 100644 --- a/crates/heal/src/heal/manager/queue.rs +++ b/crates/heal/src/heal/manager/queue.rs @@ -82,6 +82,8 @@ pub(super) enum QueuePushOutcome { pub(super) struct CompletedHealStatus { pub(super) heal_type: HealType, pub(super) status: HealTaskStatus, + pub(super) progress: Option, + pub(super) retained_bytes: std::sync::OnceLock, pub(super) result_items_truncated: bool, pub(super) completed_at: SystemTime, /// Sequence-stamped retained window, archived with the completion so @@ -92,6 +94,133 @@ pub(super) struct CompletedHealStatus { pub(super) min_seq: u64, } +impl CompletedHealStatus { + // Account for owned capacities, including nested drive arrays. Aliases + // conservatively charge the shared allocation again, keeping both token + // count and retained payload bounded without a second ownership index. + pub(super) fn retained_bytes(&self) -> usize { + *self.retained_bytes.get_or_init(|| self.measure_retained_bytes()) + } + + fn measure_retained_bytes(&self) -> usize { + let mut bytes = size_of::(); + let mut add = |amount: usize| bytes = bytes.saturating_add(amount); + match &self.heal_type { + HealType::Cluster => {} + HealType::Bucket { bucket } => add(bucket.capacity()), + HealType::Object { + bucket, + object, + version_id, + } + | HealType::ECDecode { + bucket, + object, + version_id, + } => { + add(bucket.capacity()); + add(object.capacity()); + add(version_id.as_ref().map_or(0, String::capacity)); + } + HealType::Prefix { bucket, prefix } => { + add(bucket.capacity()); + add(prefix.capacity()); + } + HealType::Metadata { bucket, object } => { + add(bucket.capacity()); + add(object.capacity()); + } + HealType::ErasureSet { buckets, set_disk_id } => { + add(buckets.capacity().saturating_mul(size_of::())); + for bucket in buckets { + add(bucket.capacity()); + } + add(set_disk_id.capacity()); + } + } + if let HealTaskStatus::Failed { error } | HealTaskStatus::Retrying { error, .. } = &self.status { + add(error.capacity()); + } + add(self + .progress + .as_ref() + .and_then(|progress| progress.current_object.as_ref()) + .map_or(0, String::capacity)); + add(self.seqed_items.capacity().saturating_mul(size_of::<(u64, HealResultItem)>())); + for (_, item) in &self.seqed_items { + add(Self::result_item_heap_bytes(item)); + } + bytes + } + + fn result_item_heap_bytes(item: &HealResultItem) -> usize { + let mut bytes = 0usize; + let mut add = |amount: usize| bytes = bytes.saturating_add(amount); + for value in [ + &item.heal_item_type, + &item.bucket, + &item.object, + &item.version_id, + &item.detail, + ] { + add(value.capacity()); + } + for infos in [&item.before, &item.after] { + add(infos + .drives + .capacity() + .saturating_mul(size_of::())); + for drive in &infos.drives { + add(drive.uuid.capacity()); + add(drive.endpoint.capacity()); + add(drive.state.capacity()); + } + } + bytes + } + + pub(super) fn bound_result_window(&mut self) { + let mut bytes = 0usize; + let retained = self + .seqed_items + .iter() + .rev() + .take_while(|(_, item)| { + bytes = bytes + .saturating_add(size_of::<(u64, HealResultItem)>()) + .saturating_add(Self::result_item_heap_bytes(item)); + bytes <= MAX_COMPLETED_HEAL_RESULT_BYTES + }) + .count(); + let truncated = retained < self.seqed_items.len(); + if truncated { + self.seqed_items.drain(..self.seqed_items.len() - retained); + self.seqed_items.shrink_to_fit(); + self.min_seq = self.seqed_items.first().map_or(self.next_seq, |(seq, _)| *seq); + self.result_items_truncated = true; + self.retained_bytes.take(); + } + } + + pub(super) async fn snapshot(task: &HealTask, status: HealTaskStatus) -> Self { + let seqed_items = task.get_seqed_result_items().await; + let (next_seq, min_seq) = task.result_seq_cursors(); + let mut snapshot = Self { + heal_type: task.heal_type.clone(), + status, + progress: Some(task.get_progress().await), + retained_bytes: std::sync::OnceLock::new(), + result_items_truncated: task.result_items_truncated(), + completed_at: SystemTime::now(), + seqed_items, + next_seq, + min_seq, + }; + snapshot.bound_result_window(); + snapshot + } +} + #[derive(Debug, Clone)] pub(super) struct HealTaskAlias { pub(super) task_id: String, diff --git a/crates/heal/src/heal/manager/scheduler.rs b/crates/heal/src/heal/manager/scheduler.rs index c0deae524..feb0c1231 100644 --- a/crates/heal/src/heal/manager/scheduler.rs +++ b/crates/heal/src/heal/manager/scheduler.rs @@ -264,7 +264,7 @@ impl HealManager { error: error.clone(), retry_attempt: request.retry_attempts, }); - let retry_request_for_queue = retry_request; + let mut retry_request_for_queue = retry_request; let retry_cancel_token = retry_request_for_queue.as_ref().map(|_| CancellationToken::new()); if retry_request_for_queue.is_none() { replacement_recovery_anchors_clone @@ -272,7 +272,35 @@ impl HealManager { .unwrap_or_else(|poisoned| poisoned.into_inner()) .remove(&task_id); } + let mut completed_status = match retry_request_for_status { + Some(status) => status, + None => task.get_status().await, + }; + let mut completed_status_entry = CompletedHealStatus::snapshot(&task, completed_status.clone()).await; + let completed_progress = task.get_progress().await; + #[cfg(test)] + tests::pause_completed_retention_before_publish(&task_id, &completed_status).await; let mut active_heals_guard = active_heals_clone.lock().await; + let owns_completion = active_heals_guard.contains_key(&task_id); + let cancelled_completion = if owns_completion { + false + } else { + // Cancellation can win while a finished worker waits + // for active ownership. It must not resurrect a retry + // or replace an acknowledged cancellation with success. + retry_request_for_queue = None; + completed_heals_clone + .lock() + .await + .get(&task_id) + .is_some_and(|completed| completed.status == HealTaskStatus::Cancelled) + }; + if cancelled_completion { + completed_status = HealTaskStatus::Cancelled; + completed_status_entry.status = HealTaskStatus::Cancelled; + } + let terminal_completion = !matches!(completed_status, HealTaskStatus::Retrying { .. }); + let successful_completion = matches!(completed_status, HealTaskStatus::Completed); // Keep retry ownership continuous: status snapshots acquire // these locks in the same active -> retrying order. let mut retrying_heals_guard = if let (Some((request, _, error)), Some(cancel_token)) = @@ -295,6 +323,16 @@ impl HealManager { } else { None }; + if owns_completion || cancelled_completion { + publish_completed_heal( + &completed_heals_clone, + &task_aliases_clone, + &task_id, + completed_status_entry, + terminal_completion, + ) + .await; + } let completed_task = active_heals_guard.remove(&task_id); if let Some(completed_task) = completed_task.as_ref() { publish_active_heal_count(&active_heals_guard); @@ -304,33 +342,10 @@ impl HealManager { drop(retrying_heals_guard.take()); drop(active_heals_guard); - if let Some(completed_task) = completed_task { - let completed_status = if let Some(status) = retry_request_for_status { - status - } else { - completed_task.get_status().await - }; - let terminal_completion = !matches!(completed_status, HealTaskStatus::Retrying { .. }); - let successful_completion = matches!(completed_status, HealTaskStatus::Completed); - let completed_progress = completed_task.get_progress().await; - // Single snapshot of the retained window: the task is - // finished and already off the active map, so there is - // no concurrent writer to race with. - let seqed_items = completed_task.get_seqed_result_items().await; - let (next_seq, min_seq) = completed_task.result_seq_cursors(); - let completed_status_entry = CompletedHealStatus { - heal_type: completed_task.heal_type.clone(), - status: completed_status.clone(), - result_items_truncated: completed_task.result_items_truncated(), - completed_at: SystemTime::now(), - seqed_items, - next_seq, - min_seq, - }; - let mut completed_heals_guard = completed_heals_clone.lock().await; - prune_completed_heal_statuses(&mut completed_heals_guard); - completed_heals_guard.insert(task_id.clone(), Arc::new(completed_status_entry)); - drop(completed_heals_guard); + #[cfg(test)] + tests::pause_completed_retention_handoff(&task_id).await; + + if completed_task.is_some() { // update statistics let mut stats = statistics_clone.write().await; match completed_status { @@ -352,10 +367,6 @@ impl HealManager { } else { release_mrf_repair_notice_targets(notice_targets); } - task_aliases_clone - .lock() - .await - .retain(|alias_id, alias| alias_id != &task_id && alias.task_id != task_id); } } @@ -718,17 +729,42 @@ pub(super) fn heal_request_set_key_for_task(task: &HealTask) -> Option { } pub(super) fn prune_completed_heal_statuses(completed_heals: &mut HashMap>) { - let Ok(now) = SystemTime::now().duration_since(SystemTime::UNIX_EPOCH) else { - return; - }; + prune_completed_heal_statuses_at(completed_heals, SystemTime::now()); +} +pub(super) fn prune_completed_heal_statuses_at(completed_heals: &mut HashMap>, now: SystemTime) { completed_heals.retain(|_, completed| { - completed - .completed_at - .duration_since(SystemTime::UNIX_EPOCH) - .map(|completed_at| now.saturating_sub(completed_at) <= KEEP_HEAL_TASK_STATUS_DURATION) + now.duration_since(completed.completed_at) + .map(|age| age <= KEEP_HEAL_TASK_STATUS_DURATION) .unwrap_or(false) }); + let entry_bytes = |key: &String, value: &Arc| { + key.capacity() + .saturating_add(size_of::<(String, Arc)>()) + .saturating_add(value.retained_bytes()) + }; + let mut bytes = completed_heals + .iter() + .fold(0usize, |total, (key, value)| total.saturating_add(entry_bytes(key, value))); + while completed_heals.len() > MAX_COMPLETED_HEAL_TOKENS || bytes > MAX_COMPLETED_HEAL_BYTES { + let Some(oldest) = completed_heals + .iter() + .min_by(|(left_id, left), (right_id, right)| { + left.completed_at.cmp(&right.completed_at).then_with(|| left_id.cmp(right_id)) + }) + .map(|(_, value)| Arc::clone(value)) + else { + break; + }; + completed_heals.retain(|key, value| { + if Arc::ptr_eq(value, &oldest) { + bytes = bytes.saturating_sub(entry_bytes(key, value)); + false + } else { + true + } + }); + } } pub(super) fn can_schedule_request( diff --git a/crates/heal/src/heal/manager/tests.rs b/crates/heal/src/heal/manager/tests.rs index 37a100f4b..aa1c1293f 100644 --- a/crates/heal/src/heal/manager/tests.rs +++ b/crates/heal/src/heal/manager/tests.rs @@ -101,6 +101,326 @@ async fn process_manager_queue_once(manager: &HealManager) { struct MockStorage; +fn completed_retention_fixture(completed_at: SystemTime) -> CompletedHealStatus { + CompletedHealStatus { + heal_type: HealType::Cluster, + status: HealTaskStatus::Completed, + progress: Some(HealProgress { + objects_scanned: 9, + objects_healed: 8, + objects_failed: 1, + ..Default::default() + }), + retained_bytes: std::sync::OnceLock::new(), + result_items_truncated: false, + completed_at, + seqed_items: vec![(3, HealResultItem::default()), (4, HealResultItem::default())], + next_seq: 5, + min_seq: 3, + } +} + +#[test] +fn completed_retention_cursor_boundaries_preserve_progress() { + let completed = completed_retention_fixture(SystemTime::now()); + for (cursor, count, lagged) in [ + (0, 2, true), + (1, 2, true), + (2, 2, false), + (3, 1, false), + (4, 0, false), + (5, 0, false), + (u64::MAX, 0, false), + ] { + let report = completed_task_report(&completed, Some(cursor)); + assert_eq!(report.result_items.len(), count, "cursor={cursor}"); + assert_eq!(report.result_items_truncated, lagged, "cursor={cursor}"); + assert_eq!(report.progress, completed.progress); + assert_eq!((report.next_seq, report.min_seq), (5, 3)); + } + assert_eq!(completed_task_report(&completed, None).result_items.len(), 2); +} + +#[tokio::test] +async fn completed_retention_displaced_alias_does_not_resurrect_evicted_snapshot() { + let manager = HealManager::new(Arc::new(MockStorage), None); + let request = HealRequest::bucket("bucket".to_string()); + manager.insert_task_alias("alias", &request.id).await; + let terminal = record_displaced_terminal(&manager.displaced_terminals, &request); + lock_displaced_terminals(&manager.displaced_terminals).remove(&request.id); + remove_displaced_task_aliases(&manager.task_aliases, &manager.displaced_terminals, &request.id, &terminal).await; + for token in [&request.id, &"alias".to_string()] { + assert!(matches!(manager.get_task_report(token).await, Err(Error::TaskNotFound { .. }))); + } + assert!(manager.task_aliases.lock().await.is_empty()); + assert!(lock_displaced_terminals(&manager.displaced_terminals).is_empty()); +} + +#[test] +fn completed_retention_count_ttl_and_alias_eviction_are_bounded() { + let now = SystemTime::now(); + let mut entries = HashMap::new(); + let oldest = Arc::new(completed_retention_fixture(now - KEEP_HEAL_TASK_STATUS_DURATION)); + entries.insert("oldest".to_string(), Arc::clone(&oldest)); + entries.insert("oldest-alias".to_string(), Arc::clone(&oldest)); + for index in 2..MAX_COMPLETED_HEAL_TOKENS { + entries.insert(format!("task-{index}"), Arc::new(completed_retention_fixture(now))); + } + prune_completed_heal_statuses_at(&mut entries, now); + assert_eq!(entries.len(), MAX_COMPLETED_HEAL_TOKENS); + entries.insert("cap-plus-one".to_string(), Arc::new(completed_retention_fixture(now))); + prune_completed_heal_statuses_at(&mut entries, now); + assert_eq!(entries.len(), MAX_COMPLETED_HEAL_TOKENS - 1); + assert!(!entries.contains_key("oldest")); + assert!(!entries.contains_key("oldest-alias")); + entries.clear(); + entries.insert("ttl-boundary".to_string(), oldest); + entries.insert( + "expired".to_string(), + Arc::new(completed_retention_fixture( + now - KEEP_HEAL_TASK_STATUS_DURATION - Duration::from_nanos(1), + )), + ); + entries.insert("future".to_string(), Arc::new(completed_retention_fixture(now + Duration::from_nanos(1)))); + prune_completed_heal_statuses_at(&mut entries, now); + assert_eq!(entries.len(), 1); + assert!(entries.contains_key("ttl-boundary")); + prune_completed_heal_statuses_at(&mut entries, now + Duration::from_nanos(1)); + assert!(entries.is_empty()); +} + +#[test] +fn completed_retention_total_byte_cap_and_cap_plus_one() { + let now = SystemTime::now(); + let key = "large".to_string(); + let mut entry = completed_retention_fixture(now); + let base_bytes = entry.retained_bytes() + key.capacity() + size_of::<(String, Arc)>(); + entry.retained_bytes.take(); + entry.status = HealTaskStatus::Failed { + error: "x".repeat(MAX_COMPLETED_HEAL_BYTES - base_bytes), + }; + assert_eq!( + entry.retained_bytes() + key.capacity() + size_of::<(String, Arc)>(), + MAX_COMPLETED_HEAL_BYTES + ); + let mut entries = HashMap::from([(key, Arc::new(entry))]); + prune_completed_heal_statuses_at(&mut entries, now); + assert_eq!(entries.len(), 1, "exact byte cap remains retained"); + let mut over = Arc::try_unwrap(entries.remove("large").expect("entry retained")).expect("entry not shared"); + over.retained_bytes.take(); + if let HealTaskStatus::Failed { error } = &mut over.status { + *error = "x".repeat(error.len() + 1); + } + entries.insert("large".to_string(), Arc::new(over)); + prune_completed_heal_statuses_at(&mut entries, now); + assert!(entries.is_empty(), "oversized metadata cannot escape total byte bound"); +} + +#[tokio::test] +async fn completed_retention_large_window_keeps_cursors_and_progress() { + let task = HealTask::from_request(HealRequest::bucket("bucket".to_string()), Arc::new(MockStorage)); + let mut snapshot = completed_retention_fixture(SystemTime::now()); + snapshot.seqed_items[0].1.detail = "x".repeat(MAX_COMPLETED_HEAL_RESULT_BYTES); + snapshot.bound_result_window(); + assert_eq!(snapshot.seqed_items.len(), 1); + assert_eq!((snapshot.min_seq, snapshot.next_seq), (4, 5)); + assert!(snapshot.result_items_truncated); + assert!(snapshot.retained_bytes() < MAX_COMPLETED_HEAL_RESULT_BYTES); + let report = completed_task_report(&snapshot, Some(0)); + assert_eq!(report.progress.expect("progress retained").objects_scanned, 9); + assert!(report.result_items_truncated); + let active_max = task.get_result_items_since(Some(u64::MAX)).await; + assert!(active_max.items.is_empty()); + assert!(!active_max.lagged); +} + +#[test] +fn completed_retention_result_byte_cap_and_cap_plus_one() { + for extra in [0, 1] { + let mut snapshot = completed_retention_fixture(SystemTime::now()); + snapshot.seqed_items = vec![( + 4, + HealResultItem { + detail: "x".repeat(MAX_COMPLETED_HEAL_RESULT_BYTES - size_of::<(u64, HealResultItem)>() + extra), + ..Default::default() + }, + )]; + snapshot.min_seq = 4; + snapshot.bound_result_window(); + assert_eq!(snapshot.seqed_items.len(), 1 - extra); + assert_eq!(snapshot.result_items_truncated, extra == 1); + assert_eq!(snapshot.min_seq, if extra == 0 { 4 } else { 5 }); + assert_eq!(snapshot.next_seq, 5); + assert_eq!(snapshot.progress.as_ref().expect("progress retained").objects_scanned, 9); + } +} + +#[derive(Default)] +struct CompletedRetentionHook { + started: Notify, + execute: Notify, + handoff: Notify, + finish: Notify, + pause_before_publish: bool, + before_publish: Notify, + publish: Notify, + prepared_status: Mutex>, +} + +static COMPLETED_RETENTION_HOOKS: LazyLock>>> = + LazyLock::new(|| Mutex::new(HashMap::new())); + +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 { + hook.handoff.notify_one(); + hook.finish.notified().await; + } +} + +pub(super) async fn pause_completed_retention_before_publish(task_id: &str, status: &HealTaskStatus) { + let hook = COMPLETED_RETENTION_HOOKS.lock().await.get(task_id).cloned(); + if let Some(hook) = hook.filter(|hook| hook.pause_before_publish) { + *hook.prepared_status.lock().await = Some(status.clone()); + hook.before_publish.notify_one(); + hook.publish.notified().await; + } +} + +#[tokio::test] +async fn completed_retention_cancel_wins_over_a_prepared_retry_snapshot() { + let bucket = "completed-retention-retry-cancel"; + let manager = HealManager::new(Arc::new(MockStorage), None); + let request = HealRequest::object(bucket.to_string(), "object".to_string(), None); + let task_id = request.id.clone(); + let duplicate = HealRequest::object(bucket.to_string(), "object".to_string(), None); + let alias = duplicate.id.clone(); + let hook = Arc::new(CompletedRetentionHook { + pause_before_publish: true, + ..Default::default() + }); + { + let mut hooks = COMPLETED_RETENTION_HOOKS.lock().await; + hooks.insert(bucket.to_string(), Arc::clone(&hook)); + hooks.insert(task_id.clone(), Arc::clone(&hook)); + } + manager.submit_heal_request(request).await.expect("admit original"); + manager.submit_heal_request(duplicate).await.expect("admit alias"); + process_manager_queue_once(&manager).await; + tokio::time::timeout(Duration::from_secs(5), hook.started.notified()) + .await + .expect("scheduler starts"); + let task = manager.active_heals.lock().await.get(&task_id).cloned().expect("active task"); + task.progress.write().await.update_object_progress(1, 1, 0, 0, 4096); + hook.execute.notify_one(); + tokio::time::timeout(Duration::from_secs(5), hook.before_publish.notified()) + .await + .expect("retry snapshot prepared"); + manager.cancel_task(&alias).await.expect("cancel wins active ownership"); + assert!(matches!(*hook.prepared_status.lock().await, Some(HealTaskStatus::Retrying { .. }))); + hook.publish.notify_one(); + tokio::time::timeout(Duration::from_secs(5), hook.handoff.notified()) + .await + .expect("scheduler finishes handoff"); + for token in [&task_id, &alias] { + let report = manager.get_task_report(token).await.expect("cancelled token retained"); + assert_eq!(report.status, HealTaskStatus::Cancelled); + assert_eq!(report.progress.expect("frozen progress").objects_scanned, 1); + } + assert!(!manager.retrying_heals.lock().await.contains_key(&task_id)); + assert!(!manager.heal_queue.lock().await.contains_request_id(&task_id)); + hook.finish.notify_one(); + COMPLETED_RETENTION_HOOKS + .lock() + .await + .retain(|key, _| key != bucket && key != &task_id); +} + +#[tokio::test] +async fn completed_retention_scheduler_preserves_progress_aliases_and_atomic_handoff() { + for outcome in ["success", "failed", "cancelled"] { + let bucket = format!("completed-retention-{outcome}"); + let hook = Arc::new(CompletedRetentionHook::default()); + let manager = Arc::new(HealManager::new(Arc::new(MockStorage), None)); + let request = HealRequest::object(bucket.clone(), "object".to_string(), None); + let task_id = request.id.clone(); + let duplicate = HealRequest::object(bucket.clone(), "object".to_string(), None); + let alias = duplicate.id.clone(); + { + let mut hooks = COMPLETED_RETENTION_HOOKS.lock().await; + hooks.insert(bucket.clone(), Arc::clone(&hook)); + hooks.insert(task_id.clone(), Arc::clone(&hook)); + } + manager.submit_heal_request(request).await.expect("admit original"); + manager.submit_heal_request(duplicate).await.expect("admit alias"); + process_manager_queue_once(&manager).await; + tokio::time::timeout(Duration::from_secs(5), hook.started.notified()) + .await + .expect("scheduler reaches storage"); + let task = manager + .active_heals + .lock() + .await + .get(&task_id) + .cloned() + .expect("task is active"); + task.progress.write().await.update_object_progress(1, 1, 0, 0, 4096); + let before = manager.get_task_report(&alias).await.expect("alias resolves active progress"); + assert_eq!(before.progress.as_ref().expect("active progress").objects_scanned, 1); + let poll_manager = Arc::clone(&manager); + let poll_alias = alias.clone(); + let stop = CancellationToken::new(); + let poll_stop = stop.clone(); + let polling = tokio::spawn(async move { + while !poll_stop.is_cancelled() { + let report = poll_manager + .get_task_report(&poll_alias) + .await + .expect("handoff must never return NotFound"); + assert!(report.progress.expect("progress never disappears").objects_scanned >= 1); + tokio::task::yield_now().await; + } + }); + if outcome == "cancelled" { + manager.cancel_task(&alias).await.expect("cancel active task by alias"); + } else { + hook.execute.notify_one(); + } + tokio::time::timeout(Duration::from_secs(5), hook.handoff.notified()) + .await + .expect("scheduler archives terminal"); + assert!(!manager.active_heals.lock().await.contains_key(&task_id)); + let expected = task.get_progress().await; + for token in [&task_id, &alias] { + assert_eq!(manager.get_task_progress(token).await.expect("terminal progress query"), expected); + let report = manager + .get_task_report_for_path_since(&format!("{bucket}/object"), token, Some(u64::MAX)) + .await + .expect("terminal token remains queryable at handoff"); + assert_eq!(report.progress.as_ref(), Some(&expected)); + assert!(report.result_items.is_empty()); + match outcome { + "success" => assert_eq!(report.status, HealTaskStatus::Completed), + "failed" => assert!(matches!(report.status, HealTaskStatus::Failed { .. })), + _ => assert_eq!(report.status, HealTaskStatus::Cancelled), + } + } + let retained = manager.completed_heals.lock().await; + assert!(Arc::ptr_eq(&retained[&task_id], &retained[&alias])); + drop(retained); + stop.cancel(); + polling.await.expect("concurrent polling succeeds"); + // Archived progress must not alias a mutable live progress object. + task.progress.write().await.objects_scanned = 999; + assert_eq!(manager.get_task_report(&alias).await.expect("frozen report").progress, Some(expected)); + hook.finish.notify_one(); + COMPLETED_RETENTION_HOOKS + .lock() + .await + .retain(|key, _| key != &bucket && key != &task_id); + } +} + #[async_trait::async_trait] impl HealStorageAPI for MockStorage { async fn get_object_meta(&self, _bucket: &str, _object: &str) -> Result> { @@ -123,6 +443,12 @@ impl HealStorageAPI for MockStorage { } async fn object_exists(&self, bucket: &str, _object: &str) -> Result { + let hook = COMPLETED_RETENTION_HOOKS.lock().await.get(bucket).cloned(); + if let Some(hook) = hook { + hook.started.notify_one(); + hook.execute.notified().await; + return Ok(true); + } Ok(bucket == "retry-transition") } @@ -133,13 +459,18 @@ impl HealStorageAPI for MockStorage { _version_id: Option<&str>, _opts: &HealOpts, ) -> Result<(HealResultItem, Option)> { + if bucket == "completed-retention-failed" { + return Err(Error::TaskExecutionFailed { + message: "retention fixture failure".to_string(), + }); + } if let Some(hook) = manager_recovery_test_hook() { *hook .heal_object_calls .lock() .expect("manager recovery object call lock should not poison") += 1; } - if bucket == "retry-transition" { + if matches!(bucket, "retry-transition" | "completed-retention-retry-cancel") { return Ok(( HealResultItem::default(), Some(Error::Storage(EcstoreError::InsufficientReadQuorum( @@ -1145,7 +1476,13 @@ async fn test_active_duplicate_token_can_query_and_cancel_original_task() { .expect("duplicate token should cancel merged active task"); assert!(manager.active_heals.lock().await.get(&active_task_id).is_none()); - assert!(matches!(manager.get_task_status(&active_task_id).await, Err(Error::TaskNotFound { .. }))); + assert_eq!( + manager + .get_task_status(&active_task_id) + .await + .expect("cancelled task remains queryable"), + HealTaskStatus::Cancelled + ); } #[tokio::test] @@ -1638,6 +1975,8 @@ async fn insert_retrying_request(manager: &HealManager, request: HealRequest) -> manager.completed_heals.lock().await.insert( task_id, Arc::new(CompletedHealStatus { + progress: None, + retained_bytes: std::sync::OnceLock::new(), heal_type: request.heal_type, status: HealTaskStatus::Retrying { error: "Lock acquisition timeout".to_string(), @@ -2053,7 +2392,7 @@ async fn admin_force_start_cancels_overlapping_active_task_first() { "the overlapping admin task must be cancelled (removed from the active table) before the new one starts" ); assert!( - matches!(manager.get_task_status(&old_id).await, Err(Error::TaskNotFound { .. })), + matches!(manager.get_task_status(&old_id).await, Ok(HealTaskStatus::Cancelled)), "a cancelled task must no longer resolve as an active heal" ); } @@ -2360,6 +2699,8 @@ async fn test_retrying_completion_outranks_the_queue_for_the_same_id() { manager.completed_heals.lock().await.insert( task_id.clone(), Arc::new(CompletedHealStatus { + progress: None, + retained_bytes: std::sync::OnceLock::new(), heal_type: request.heal_type.clone(), status: HealTaskStatus::Retrying { error: "transient disk failure".to_string(), @@ -2395,6 +2736,8 @@ async fn test_get_task_status_reads_recent_completed_status() { manager.completed_heals.lock().await.insert( "completed-token".to_string(), Arc::new(CompletedHealStatus { + progress: None, + retained_bytes: std::sync::OnceLock::new(), heal_type: HealType::Bucket { bucket: "bucket".to_string(), }, @@ -2424,6 +2767,8 @@ async fn test_get_task_report_for_path_reads_completed_items() { manager.completed_heals.lock().await.insert( "completed-token".to_string(), Arc::new(CompletedHealStatus { + progress: None, + retained_bytes: std::sync::OnceLock::new(), heal_type: HealType::Object { bucket: "bucket".to_string(), object: "object".to_string(), diff --git a/crates/heal/src/heal/task.rs b/crates/heal/src/heal/task.rs index c95770834..c4c1d902b 100644 --- a/crates/heal/src/heal/task.rs +++ b/crates/heal/src/heal/task.rs @@ -999,7 +999,7 @@ impl HealTask { let items = match since { None => result_items.iter().map(|(_, item)| item.clone()).collect::>(), Some(cursor) => { - if cursor + 1 < min_seq { + if cursor.saturating_add(1) < min_seq { lagged = true; } result_items