fix(heal): respect scoped object repair limits (#5855)

crates/ecstore/src/core/sets.rs

@@ -286,6 +286,23 @@ impl Sets {
         self.get_disks(self.get_hashed_set_index(key))
     }
 
+    fn get_disks_for_heal_object(&self, key: &str, opts: &HealOpts) -> Result<Arc<SetDisks>> {
+        match opts.set {
+            Some(set_idx) => self.disk_set.get(set_idx).cloned().ok_or_else(|| {
+                StorageError::InvalidArgument(
+                    "heal".to_string(),
+                    "set".to_string(),
+                    format!(
+                        "invalid heal set index {set_idx} for pool {} with {} sets",
+                        self.pool_idx,
+                        self.disk_set.len()
+                    ),
+                )
+            }),
+            None => Ok(self.get_disks_by_key(key)),
+        }
+    }
+
     pub(crate) async fn storage_info_snapshot(&self) -> rustfs_madmin::StorageInfo {
         let mut futures = Vec::with_capacity(self.disk_set.len());
 
@@ -1101,7 +1118,7 @@ impl crate::storage_api_contracts::heal::HealOperations for Sets {
         version_id: &str,
         opts: &HealOpts,
     ) -> Result<(HealResultItem, Option<Error>)> {
-        self.get_disks_by_key(object)
+        self.get_disks_for_heal_object(object, opts)?
             .heal_object(bucket, object, version_id, opts)
             .await
     }
@@ -1431,6 +1448,53 @@ mod tests {
         (temp_dirs, sets)
     }
 
+    #[tokio::test]
+    async fn heal_object_uses_explicit_set_scope() {
+        let (_temp_dirs, sets) = two_set_test_sets().await;
+        let selected = sets
+            .get_disks_for_heal_object(
+                "object",
+                &HealOpts {
+                    set: Some(1),
+                    ..Default::default()
+                },
+            )
+            .expect("requested set should be selected");
+
+        assert!(Arc::ptr_eq(&selected, &sets.disk_set[1]));
+    }
+
+    #[tokio::test]
+    async fn heal_object_without_set_scope_keeps_hash_routing() {
+        let (_temp_dirs, sets) = two_set_test_sets().await;
+        let object = "object";
+        let selected = sets
+            .get_disks_for_heal_object(object, &HealOpts::default())
+            .expect("hash-routed set should be selected");
+
+        assert!(Arc::ptr_eq(&selected, &sets.get_disks_by_key(object)));
+    }
+
+    #[tokio::test]
+    async fn heal_object_rejects_invalid_set_scope() {
+        let (_temp_dirs, sets) = two_set_test_sets().await;
+        let err = sets
+            .get_disks_for_heal_object(
+                "object",
+                &HealOpts {
+                    set: Some(2),
+                    ..Default::default()
+                },
+            )
+            .expect_err("out-of-range set scope must fail closed");
+
+        assert!(
+            matches!(err, StorageError::InvalidArgument(_, ref field, ref reason)
+                if field == "set" && reason.contains("invalid heal set index 2 for pool 0 with 2 sets")),
+            "unexpected invalid set error: {err:?}"
+        );
+    }
+
     #[tokio::test]
     async fn delete_prefix_surfaces_a_hard_error_from_any_set() {
         let (_temp_dirs, sets) = two_set_test_sets().await;

crates/ecstore/src/store/heal.rs

@@ -21,7 +21,27 @@ const LOG_SUBSYSTEM_HEAL: &str = "heal";
 const EVENT_HEAL_FORMAT_COMPLETED: &str = "heal_format_completed";
 const EVENT_HEAL_OBJECT_STARTED: &str = "heal_object_started";
 
+fn invalid_heal_pool_index(pool_idx: usize, pool_count: usize) -> Error {
+    StorageError::InvalidArgument(
+        "heal".to_string(),
+        "pool".to_string(),
+        format!("invalid heal pool index {pool_idx} for {pool_count} pools"),
+    )
+}
+
 impl ECStore {
+    fn get_pools_for_heal_object(&self, opts: &HealOpts) -> Result<Vec<Arc<Sets>>> {
+        match opts.pool {
+            Some(pool_idx) => Ok(vec![
+                self.pools
+                    .get(pool_idx)
+                    .cloned()
+                    .ok_or_else(|| invalid_heal_pool_index(pool_idx, self.pools.len()))?,
+            ]),
+            None => Ok(self.pools.clone()),
+        }
+    }
+
     #[instrument(skip(self))]
     pub(super) async fn handle_heal_format(&self, dry_run: bool) -> Result<(HealResultItem, Option<Error>)> {
         let mut r = HealResultItem {
@@ -105,8 +125,10 @@ impl ECStore {
         );
         let object = encode_dir_object(object);
 
-        let mut futures = Vec::with_capacity(self.pools.len());
+        let pools = self.get_pools_for_heal_object(opts)?;
-        for pool in self.pools.iter() {
+
+        let mut futures = Vec::with_capacity(pools.len());
+        for pool in pools.iter() {
             if self.is_suspended(pool.pool_idx).await {
                 continue;
             }
@@ -178,6 +200,82 @@ mod tests {
     use crate::layout::endpoints::{Endpoints, PoolEndpoints};
     use crate::store::init_format::{load_format_erasure, save_format_file};
 
+    async fn minimal_heal_pool(pool_idx: usize) -> Arc<Sets> {
+        let format = FormatV3::new(1, 1);
+        let endpoint_url = format!("http://127.0.0.1:{}/data", 19000 + pool_idx);
+        let mut endpoint = Endpoint::try_from(endpoint_url.as_str()).expect("endpoint should parse");
+        endpoint.set_pool_index(pool_idx);
+        endpoint.set_set_index(0);
+        endpoint.set_disk_index(0);
+
+        Sets::new(
+            vec![None],
+            &PoolEndpoints {
+                legacy: false,
+                set_count: 1,
+                drives_per_set: 1,
+                endpoints: Endpoints::from(vec![endpoint]),
+                cmd_line: String::new(),
+                platform: String::new(),
+            },
+            &format,
+            pool_idx,
+            0,
+        )
+        .await
+        .expect("minimal pool should build")
+    }
+
+    async fn minimal_heal_store() -> ECStore {
+        ECStore {
+            id: Uuid::new_v4(),
+            disk_map: HashMap::new(),
+            pools: vec![minimal_heal_pool(0).await, minimal_heal_pool(1).await],
+            peer_sys: S3PeerSys {
+                clients: Vec::new(),
+                pools_count: 2,
+            },
+            pool_meta: RwLock::new(PoolMeta::default()),
+            rebalance_meta: RwLock::new(None),
+            decommission_cancelers: RwLock::new(Vec::new()),
+            start_gate: Mutex::new(()),
+            pool_meta_save_gate: Mutex::new(()),
+            ctx: crate::runtime::instance::bootstrap_ctx(),
+            bucket_fence_registry: std::sync::Arc::default(),
+        }
+    }
+
+    #[tokio::test]
+    async fn heal_object_pool_scope_selects_only_requested_pool() {
+        let store = minimal_heal_store().await;
+        let pools = store
+            .get_pools_for_heal_object(&HealOpts {
+                pool: Some(1),
+                ..Default::default()
+            })
+            .expect("requested pool should be selected");
+
+        assert_eq!(pools.len(), 1);
+        assert!(Arc::ptr_eq(&pools[0], &store.pools[1]));
+    }
+
+    #[tokio::test]
+    async fn heal_object_pool_scope_rejects_invalid_pool() {
+        let store = minimal_heal_store().await;
+        let err = store
+            .get_pools_for_heal_object(&HealOpts {
+                pool: Some(2),
+                ..Default::default()
+            })
+            .expect_err("out-of-range pool scope must fail closed");
+
+        assert!(
+            matches!(err, StorageError::InvalidArgument(_, ref field, ref reason)
+                if field == "pool" && reason.contains("invalid heal pool index 2 for 2 pools")),
+            "unexpected invalid pool error: {err:?}"
+        );
+    }
+
     #[tokio::test]
     async fn handle_heal_format_continues_after_a_pool_error() {
         let canonical_format = FormatV3::new(1, 3);

crates/heal/src/heal/manager.rs

@@ -2475,7 +2475,7 @@ impl HealManager {
             return;
         }
 
-        let mut running_per_set = running_erasure_set_counts(&active_heals_guard);
+        let mut running_per_set = running_heal_set_counts(&active_heals_guard);
         let mut tasks_started = 0usize;
         let mut delayed_by_mainline_throttle = false;
 
@@ -2856,6 +2856,14 @@ impl std::fmt::Debug for HealManager {
 fn heal_request_set_key(request: &HealRequest) -> Option<String> {
     match &request.heal_type {
         HealType::ErasureSet { set_disk_id, .. } => Some(set_disk_id.clone()),
+        HealType::Object { .. } => heal_options_set_key(&request.options),
+        _ => None,
+    }
+}
+
+fn heal_options_set_key(options: &HealOptions) -> Option<String> {
+    match (options.pool_index, options.set_index) {
+        (Some(pool), Some(set)) => Some(format!("pool_{pool}_set_{set}")),
         _ => None,
     }
 }
@@ -2905,16 +2913,24 @@ fn update_task_running_metric_for_task(active_heals: &HashMap<String, Arc<HealTa
     .set(count as f64);
 }
 
-fn running_erasure_set_counts(active_heals: &HashMap<String, Arc<HealTask>>) -> HashMap<String, usize> {
+fn running_heal_set_counts(active_heals: &HashMap<String, Arc<HealTask>>) -> HashMap<String, usize> {
     let mut running = HashMap::new();
     for task in active_heals.values() {
-        if let HealType::ErasureSet { set_disk_id, .. } = &task.heal_type {
+        if let Some(set_key) = heal_request_set_key_for_task(task) {
-            *running.entry(set_disk_id.clone()).or_insert(0) += 1;
+            *running.entry(set_key).or_insert(0) += 1;
         }
     }
     running
 }
 
+fn heal_request_set_key_for_task(task: &HealTask) -> Option<String> {
+    match &task.heal_type {
+        HealType::ErasureSet { set_disk_id, .. } => Some(set_disk_id.clone()),
+        HealType::Object { .. } => heal_options_set_key(&task.options),
+        _ => None,
+    }
+}
+
 fn prune_completed_heal_statuses(completed_heals: &mut HashMap<String, CompletedHealStatus>) {
     let Ok(now) = SystemTime::now().duration_since(SystemTime::UNIX_EPOCH) else {
         return;
@@ -3523,6 +3539,30 @@ mod tests {
         assert!(can_schedule_request(&request, &running, 2));
     }
 
+    #[test]
+    fn test_can_schedule_scoped_object_request_respects_per_set_limit() {
+        let options = HealOptions {
+            pool_index: Some(0),
+            set_index: Some(1),
+            ..Default::default()
+        };
+        let request = HealRequest::new(
+            HealType::Object {
+                bucket: "bucket".to_string(),
+                object: "object".to_string(),
+                version_id: None,
+            },
+            options,
+            HealPriority::Normal,
+        );
+
+        let mut running = HashMap::new();
+        running.insert("pool_0_set_1".to_string(), 1);
+
+        assert!(!can_schedule_request(&request, &running, 1));
+        assert!(can_schedule_request(&request, &running, 2));
+    }
+
     #[tokio::test]
     async fn test_submit_heal_request_returns_merged_for_duplicate() {
         let storage: Arc<dyn HealStorageAPI> = Arc::new(MockStorage);
@@ -5136,7 +5176,7 @@ mod tests {
     }
 
     #[test]
-    fn test_running_erasure_set_counts_groups_only_erasure_tasks() {
+    fn test_running_heal_set_counts_groups_set_scoped_tasks() {
         let storage: Arc<dyn HealStorageAPI> = Arc::new(MockStorage);
         let erasure_task = Arc::new(HealTask::from_request(
             HealRequest::new(
@@ -5149,6 +5189,23 @@ mod tests {
             ),
             storage.clone(),
         ));
+        let scoped_options = HealOptions {
+            pool_index: Some(0),
+            set_index: Some(1),
+            ..Default::default()
+        };
+        let scoped_object_task = Arc::new(HealTask::from_request(
+            HealRequest::new(
+                HealType::Object {
+                    bucket: "bucket".to_string(),
+                    object: "scoped-object".to_string(),
+                    version_id: None,
+                },
+                scoped_options,
+                HealPriority::Normal,
+            ),
+            storage.clone(),
+        ));
         let object_task = Arc::new(HealTask::from_request(
             HealRequest::new(
                 HealType::Object {
@@ -5164,10 +5221,11 @@ mod tests {
 
         let mut active = HashMap::new();
         active.insert(erasure_task.id.clone(), erasure_task);
+        active.insert(scoped_object_task.id.clone(), scoped_object_task);
         active.insert(object_task.id.clone(), object_task);
 
-        let counts = running_erasure_set_counts(&active);
+        let counts = running_heal_set_counts(&active);
-        assert_eq!(counts.get("pool_0_set_1"), Some(&1));
+        assert_eq!(counts.get("pool_0_set_1"), Some(&2));
         assert_eq!(counts.len(), 1);
     }