diff --git a/crates/heal/src/heal/manager.rs b/crates/heal/src/heal/manager.rs index 2f03b30d4..2db72b651 100644 --- a/crates/heal/src/heal/manager.rs +++ b/crates/heal/src/heal/manager.rs @@ -227,11 +227,11 @@ impl ForegroundPressure { struct CompletedHealStatus { heal_type: HealType, status: HealTaskStatus, - result_items: Vec, result_items_truncated: bool, completed_at: SystemTime, /// Sequence-stamped retained window, archived with the completion so /// incremental consumers keep their cursor across the transition (HS-06). + /// The un-stamped legacy view is derived from it on demand. seqed_items: Vec<(u64, HealResultItem)>, next_seq: u64, min_seq: u64, @@ -293,7 +293,7 @@ fn empty_task_report(status: HealTaskStatus) -> HealTaskReport { fn completed_task_report(completed: &CompletedHealStatus, since: Option) -> HealTaskReport { let mut lagged = false; let result_items = match since { - None => completed.result_items.clone(), + None => completed.seqed_items.iter().map(|(_, item)| item.clone()).collect(), Some(cursor) => { if cursor + 1 < completed.min_seq { lagged = true; @@ -1027,8 +1027,10 @@ pub struct HealManager { active_heals: Arc>>>, /// Heal queue (priority-based) heal_queue: Arc>, - /// Recently completed heal statuses retained for status queries. - completed_heals: Arc>>, + /// Recently completed heal statuses retained for status queries. Values + /// are shared so the lookup helper can hand a completed entry to a + /// caller without cloning the retained result window. + completed_heals: Arc>>>, /// Client tokens merged into an existing task id. task_aliases: Arc>>, /// Heal tasks waiting for a retry backoff to expire. @@ -1051,10 +1053,21 @@ pub struct HealManager { workload_provider: Option, } +/// Where a task-id lookup resolved. The variants carry the resolved state +/// so both the status and the report adapters can consume one shared +/// cascade without re-locking. +enum TaskStateLookup { + Active(Arc), + Retrying(HealTaskStatus), + Completed(Arc), + Queued, + NotFound, +} + struct HealQueueContext<'a> { heal_queue: &'a Arc>, active_heals: &'a Arc>>>, - completed_heals: &'a Arc>>, + completed_heals: &'a Arc>>>, retrying_heals: &'a Arc>>, replacement_recovery_anchors: &'a Arc>>, config: &'a Arc>, @@ -2160,47 +2173,79 @@ impl HealManager { } /// Get task status - pub async fn get_task_status(&self, task_id: &str) -> Result { - let canonical_task_id = self.canonical_task_id(task_id).await; + /// Ordered task-state lookup shared by every status/report query. The + /// map precedence mirrors the historical per-method cascades exactly: + /// active, then retrying, then completed — where a completed entry in a + /// retrying state outranks the queue so a retrying task reports + /// Retrying, never Pending — then the queue, and finally a terminal + /// completed entry. `heal_path` additionally constrains the map matches + /// the way the `*_for_path` variants always have. + async fn lookup_task_state(&self, canonical_task_id: &str, heal_path: Option<&str>) -> TaskStateLookup { + let matches_path = |heal_type: &HealType| heal_path.is_none_or(|path| heal_type_matches_path(heal_type, path)); + { let active_heals = self.active_heals.lock().await; - if let Some(task) = active_heals.get(&canonical_task_id) { - return Ok(task.get_status().await); + if let Some(task) = active_heals + .get(canonical_task_id) + .filter(|task| matches_path(&task.heal_type)) + { + return TaskStateLookup::Active(Arc::clone(task)); } } { let retrying_heals = self.retrying_heals.lock().await; - if let Some(retrying) = retrying_heals.get(&canonical_task_id) { - return Ok(retrying.status()); + if let Some(retrying) = retrying_heals + .get(canonical_task_id) + .filter(|retrying| matches_path(&retrying.request.heal_type)) + { + return TaskStateLookup::Retrying(retrying.status()); + } + } + + // One completed-map pass (single lock + prune): a retrying completion + // returns immediately; a terminal completion is held back until the + // queue has been checked, so queued work outranks it. + let mut terminal_completed: Option> = None; + { + let mut completed_heals = self.completed_heals.lock().await; + prune_completed_heal_statuses(&mut completed_heals); + if let Some(completed) = completed_heals.get(canonical_task_id).filter(|c| matches_path(&c.heal_type)) { + if completed_status_is_retrying(&completed.status) { + return TaskStateLookup::Completed(Arc::clone(completed)); + } + terminal_completed = Some(Arc::clone(completed)); } } { - let mut completed_heals = self.completed_heals.lock().await; - prune_completed_heal_statuses(&mut completed_heals); - if let Some(completed) = completed_heals.get(&canonical_task_id) - && completed_status_is_retrying(&completed.status) - { - return Ok(completed.status.clone()); + let queue = self.heal_queue.lock().await; + let queued = match heal_path { + Some(path) => queue.contains_request_id_matching_path(canonical_task_id, path), + None => queue.contains_request_id(canonical_task_id), + }; + if queued { + return TaskStateLookup::Queued; } } - let queue = self.heal_queue.lock().await; - if queue.contains_request_id(&canonical_task_id) { - return Ok(HealTaskStatus::Pending); + match terminal_completed { + Some(completed) => TaskStateLookup::Completed(completed), + None => TaskStateLookup::NotFound, } - drop(queue); + } - let mut completed_heals = self.completed_heals.lock().await; - prune_completed_heal_statuses(&mut completed_heals); - if let Some(completed) = completed_heals.get(&canonical_task_id) { - return Ok(completed.status.clone()); + pub async fn get_task_status(&self, task_id: &str) -> Result { + let canonical_task_id = self.canonical_task_id(task_id).await; + match self.lookup_task_state(&canonical_task_id, None).await { + TaskStateLookup::Active(task) => Ok(task.get_status().await), + TaskStateLookup::Retrying(status) => Ok(status), + TaskStateLookup::Completed(completed) => Ok(completed.status.clone()), + TaskStateLookup::Queued => Ok(HealTaskStatus::Pending), + TaskStateLookup::NotFound => Err(Error::TaskNotFound { + task_id: task_id.to_string(), + }), } - - Err(Error::TaskNotFound { - task_id: task_id.to_string(), - }) } pub async fn get_task_report(&self, task_id: &str) -> Result { @@ -2212,46 +2257,15 @@ impl HealManager { /// full-snapshot semantics. pub async fn get_task_report_since(&self, task_id: &str, since: Option) -> 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) { - return Ok(active_task_report(task, since).await); - } + match self.lookup_task_state(&canonical_task_id, None).await { + TaskStateLookup::Active(task) => Ok(active_task_report(&task, since).await), + TaskStateLookup::Retrying(status) => Ok(empty_task_report(status)), + TaskStateLookup::Completed(completed) => Ok(completed_task_report(&completed, since)), + TaskStateLookup::Queued => Ok(empty_task_report(HealTaskStatus::Pending)), + TaskStateLookup::NotFound => Err(Error::TaskNotFound { + task_id: task_id.to_string(), + }), } - - { - let retrying_heals = self.retrying_heals.lock().await; - if let Some(retrying) = retrying_heals.get(&canonical_task_id) { - return Ok(empty_task_report(retrying.status())); - } - } - - { - let mut completed_heals = self.completed_heals.lock().await; - prune_completed_heal_statuses(&mut completed_heals); - if let Some(completed) = completed_heals.get(&canonical_task_id) - && completed_status_is_retrying(&completed.status) - { - return Ok(completed_task_report(completed, since)); - } - } - - { - let queue = self.heal_queue.lock().await; - if queue.contains_request_id(&canonical_task_id) { - return Ok(empty_task_report(HealTaskStatus::Pending)); - } - } - - let mut completed_heals = self.completed_heals.lock().await; - prune_completed_heal_statuses(&mut completed_heals); - if let Some(completed) = completed_heals.get(&canonical_task_id) { - return Ok(completed_task_report(completed, since)); - } - - Err(Error::TaskNotFound { - task_id: task_id.to_string(), - }) } pub async fn get_task_report_for_path(&self, heal_path: &str, task_id: &str) -> Result { @@ -2266,59 +2280,20 @@ impl HealManager { since: Option, ) -> 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) - && heal_type_matches_path(&task.heal_type, heal_path) - { - return Ok(active_task_report(task, since).await); + match self.lookup_task_state(&canonical_task_id, Some(heal_path)).await { + TaskStateLookup::Active(task) => Ok(active_task_report(&task, since).await), + TaskStateLookup::Retrying(status) => Ok(empty_task_report(status)), + TaskStateLookup::Completed(completed) => Ok(completed_task_report(&completed, since)), + TaskStateLookup::Queued => Ok(empty_task_report(HealTaskStatus::Pending)), + TaskStateLookup::NotFound => { + if self.path_has_task(heal_path).await { + return Err(Error::InvalidClientToken); + } + Err(Error::TaskNotFound { + task_id: task_id.to_string(), + }) } } - - { - let retrying_heals = self.retrying_heals.lock().await; - if let Some(retrying) = retrying_heals.get(&canonical_task_id) - && heal_type_matches_path(&retrying.request.heal_type, heal_path) - { - return Ok(empty_task_report(retrying.status())); - } - } - - { - let mut completed_heals = self.completed_heals.lock().await; - prune_completed_heal_statuses(&mut completed_heals); - if let Some(completed) = completed_heals.get(&canonical_task_id) - && heal_type_matches_path(&completed.heal_type, heal_path) - && completed_status_is_retrying(&completed.status) - { - return Ok(completed_task_report(completed, since)); - } - } - - { - let queue = self.heal_queue.lock().await; - if queue.contains_request_id_matching_path(&canonical_task_id, heal_path) { - return Ok(empty_task_report(HealTaskStatus::Pending)); - } - } - - { - let mut completed_heals = self.completed_heals.lock().await; - prune_completed_heal_statuses(&mut completed_heals); - if let Some(completed) = completed_heals.get(&canonical_task_id) - && heal_type_matches_path(&completed.heal_type, heal_path) - { - return Ok(completed_task_report(completed, since)); - } - } - - if self.path_has_task(heal_path).await { - return Err(Error::InvalidClientToken); - } - - Err(Error::TaskNotFound { - task_id: task_id.to_string(), - }) } /// Get task status for a path-bound client token. @@ -2328,59 +2303,20 @@ impl HealManager { /// recently completed task, a different token is invalid for that path. pub async fn get_task_status_for_path(&self, heal_path: &str, 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) - && heal_type_matches_path(&task.heal_type, heal_path) - { - return Ok(task.get_status().await); + match self.lookup_task_state(&canonical_task_id, Some(heal_path)).await { + TaskStateLookup::Active(task) => Ok(task.get_status().await), + TaskStateLookup::Retrying(status) => Ok(status), + TaskStateLookup::Completed(completed) => Ok(completed.status.clone()), + TaskStateLookup::Queued => Ok(HealTaskStatus::Pending), + TaskStateLookup::NotFound => { + if self.path_has_task(heal_path).await { + return Err(Error::InvalidClientToken); + } + Err(Error::TaskNotFound { + task_id: task_id.to_string(), + }) } } - - { - let retrying_heals = self.retrying_heals.lock().await; - if let Some(retrying) = retrying_heals.get(&canonical_task_id) - && heal_type_matches_path(&retrying.request.heal_type, heal_path) - { - return Ok(retrying.status()); - } - } - - { - let mut completed_heals = self.completed_heals.lock().await; - prune_completed_heal_statuses(&mut completed_heals); - if let Some(completed) = completed_heals.get(&canonical_task_id) - && heal_type_matches_path(&completed.heal_type, heal_path) - && completed_status_is_retrying(&completed.status) - { - return Ok(completed.status.clone()); - } - } - - { - let queue = self.heal_queue.lock().await; - if queue.contains_request_id_matching_path(&canonical_task_id, heal_path) { - return Ok(HealTaskStatus::Pending); - } - } - - { - let mut completed_heals = self.completed_heals.lock().await; - prune_completed_heal_statuses(&mut completed_heals); - if let Some(completed) = completed_heals.get(&canonical_task_id) - && heal_type_matches_path(&completed.heal_type, heal_path) - { - return Ok(completed.status.clone()); - } - } - - if self.path_has_task(heal_path).await { - return Err(Error::InvalidClientToken); - } - - Err(Error::TaskNotFound { - task_id: task_id.to_string(), - }) } async fn path_has_task(&self, heal_path: &str) -> bool { @@ -3503,20 +3439,23 @@ impl HealManager { completed_task.get_status().await }; let completed_progress = completed_task.get_progress().await; - let final_window = completed_task.get_result_items_since(None).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: final_window.items.clone(), result_items_truncated: completed_task.result_items_truncated(), completed_at: SystemTime::now(), - seqed_items: completed_task.get_seqed_result_items().await, - next_seq: final_window.next_seq, - min_seq: final_window.min_seq, + 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(), completed_status_entry); + completed_heals_guard.insert(task_id.clone(), Arc::new(completed_status_entry)); // update statistics let mut stats = statistics_clone.write().await; match completed_status { @@ -3808,7 +3747,7 @@ fn heal_request_set_key_for_task(task: &HealTask) -> Option { } } -fn prune_completed_heal_statuses(completed_heals: &mut HashMap) { +fn prune_completed_heal_statuses(completed_heals: &mut HashMap>) { let Ok(now) = SystemTime::now().duration_since(SystemTime::UNIX_EPOCH) else { return; }; @@ -5275,19 +5214,18 @@ mod tests { ); manager.completed_heals.lock().await.insert( task_id, - CompletedHealStatus { + Arc::new(CompletedHealStatus { heal_type: request.heal_type, status: HealTaskStatus::Retrying { error: "Lock acquisition timeout".to_string(), retry_attempt: request.retry_attempts, }, - result_items: Vec::new(), result_items_truncated: false, seqed_items: Vec::new(), next_seq: 0, min_seq: 0, completed_at: SystemTime::now(), - }, + }), ); cancel_token } @@ -5985,6 +5923,47 @@ mod tests { assert!(report.result_items.is_empty()); } + #[tokio::test] + async fn test_retrying_completion_outranks_the_queue_for_the_same_id() { + let storage: Arc = Arc::new(MockStorage); + let manager = HealManager::new(storage, None); + + // A completed entry recorded in a Retrying state for a task whose + // request is also (still) queued under the same id: the retrying + // completion must win the lookup, or the task would read back as + // Pending while it is actually waiting out a retry backoff. + let request = HealRequest::object("bucket".to_string(), "object".to_string(), None); + let task_id = request.id.clone(); + manager.completed_heals.lock().await.insert( + task_id.clone(), + Arc::new(CompletedHealStatus { + heal_type: request.heal_type.clone(), + status: HealTaskStatus::Retrying { + error: "transient disk failure".to_string(), + retry_attempt: 1, + }, + result_items_truncated: false, + seqed_items: Vec::new(), + next_seq: 0, + min_seq: 0, + completed_at: SystemTime::now(), + }), + ); + manager.heal_queue.lock().await.push(HealRequest { + id: task_id.clone(), + heal_type: request.heal_type, + ..request + }); + + assert_eq!( + manager.get_task_status(&task_id).await.expect("task must resolve"), + HealTaskStatus::Retrying { + error: "transient disk failure".to_string(), + retry_attempt: 1 + } + ); + } + #[tokio::test] async fn test_get_task_status_reads_recent_completed_status() { let storage: Arc = Arc::new(MockStorage); @@ -5992,18 +5971,17 @@ mod tests { manager.completed_heals.lock().await.insert( "completed-token".to_string(), - CompletedHealStatus { + Arc::new(CompletedHealStatus { heal_type: HealType::Bucket { bucket: "bucket".to_string(), }, status: HealTaskStatus::Completed, - result_items: Vec::new(), result_items_truncated: false, seqed_items: Vec::new(), next_seq: 0, min_seq: 0, completed_at: SystemTime::now(), - }, + }), ); assert_eq!( @@ -6022,25 +6000,27 @@ mod tests { manager.completed_heals.lock().await.insert( "completed-token".to_string(), - CompletedHealStatus { + Arc::new(CompletedHealStatus { heal_type: HealType::Object { bucket: "bucket".to_string(), object: "object".to_string(), version_id: None, }, status: HealTaskStatus::Completed, - result_items: vec![HealResultItem { - bucket: "bucket".to_string(), - object: "object".to_string(), - object_size: 1024, - ..Default::default() - }], result_items_truncated: true, - seqed_items: Vec::new(), - next_seq: 0, - min_seq: 0, + seqed_items: vec![( + 1, + HealResultItem { + bucket: "bucket".to_string(), + object: "object".to_string(), + object_size: 1024, + ..Default::default() + }, + )], + next_seq: 2, + min_seq: 1, completed_at: SystemTime::now(), - }, + }), ); let report = manager @@ -6052,6 +6032,10 @@ mod tests { assert_eq!(report.status, HealTaskStatus::Completed); assert_eq!(report.result_items.len(), 1); assert_eq!(report.result_items[0].object_size, 1024); + // The archived cursors pass through to the report so an incremental + // consumer can resume against the next expected sequence. + assert_eq!(report.next_seq, 2); + assert_eq!(report.min_seq, 1); } #[tokio::test] diff --git a/crates/heal/src/heal/task.rs b/crates/heal/src/heal/task.rs index 4db62f271..5949cd1a3 100644 --- a/crates/heal/src/heal/task.rs +++ b/crates/heal/src/heal/task.rs @@ -29,6 +29,7 @@ use rustfs_madmin::heal_commands::HealResultItem; use rustfs_utils::path::SLASH_SEPARATOR; use serde::{Deserialize, Serialize}; use std::{ + collections::VecDeque, future::Future, sync::{ Arc, @@ -384,7 +385,7 @@ pub struct HealTask { /// monotonically increasing sequence number for incremental consumption /// (the client passes the last seen seq back and receives only newer /// items; see `get_result_items_since`). - pub result_items: Arc>>, + pub result_items: Arc>>, /// Next sequence number to assign; starts at 1. next_item_seq: Arc, /// Sequence number of the oldest item still inside the retention window; @@ -440,7 +441,7 @@ impl HealTask { replacement_resume_endpoint: None, status: Arc::new(RwLock::new(HealTaskStatus::Pending)), progress: Arc::new(RwLock::new(HealProgress::new())), - result_items: Arc::new(RwLock::new(Vec::new())), + result_items: Arc::new(RwLock::new(VecDeque::with_capacity(MAX_RETAINED_HEAL_RESULT_ITEMS))), next_item_seq: Arc::new(AtomicU64::new(1)), min_available_seq: Arc::new(AtomicU64::new(1)), result_items_truncated: Arc::new(AtomicBool::new(false)), @@ -931,7 +932,14 @@ impl HealTask { /// Sequence-stamped retained window, used when archiving a completed /// task so incremental cursors survive the transition (HS-06). pub async fn get_seqed_result_items(&self) -> Vec<(u64, HealResultItem)> { - self.result_items.read().await.clone() + self.result_items.read().await.iter().cloned().collect::>() + } + + /// Sequence cursors of the retained window (next to assign, oldest + /// retained) — the same pair `get_result_items_since` reports, without + /// copying the items. Used when archiving a finished task. + pub fn result_seq_cursors(&self) -> (u64, u64) { + (self.next_item_seq.load(Ordering::Relaxed), self.min_available_seq.load(Ordering::Relaxed)) } /// Incremental result window (HS-06): `since = None` returns the full @@ -974,14 +982,14 @@ impl HealTask { let seq = self.next_item_seq.fetch_add(1, Ordering::Relaxed); let mut result_items = self.result_items.write().await; if result_items.len() < MAX_RETAINED_HEAL_RESULT_ITEMS { - result_items.push((seq, result)); + result_items.push_back((seq, result)); } else { // Slide the window: the oldest item leaves and the cursor for the // oldest still-available item moves forward with it. - result_items.remove(0); + result_items.pop_front(); self.min_available_seq - .store(result_items.first().map_or(seq, |(oldest, _)| *oldest), Ordering::Relaxed); - result_items.push((seq, result)); + .store(result_items.front().map_or(seq, |(oldest, _)| *oldest), Ordering::Relaxed); + result_items.push_back((seq, result)); self.result_items_truncated.store(true, Ordering::Relaxed); } }