Add preserveOrder check for CustomChecks

Author: LZD-PratyushBhatt

helix-rest/src/main/java/org/apache/helix/rest/clusterMaintenanceService/MaintenanceManagementService.java

@@ -394,7 +394,7 @@ public Map<String, StoppableCheck> batchGetInstancesStoppableChecks(String clust
             toBeStoppedInstances, preserveOrder);
     // custom check, includes partition check.
     batchCustomInstanceStoppableCheck(clusterId, instancesForCustomInstanceLevelChecks,
-        toBeStoppedInstances, finalStoppableChecks, getMapFromJsonPayload(jsonContent));
+        toBeStoppedInstances, finalStoppableChecks, getMapFromJsonPayload(jsonContent), preserveOrder);
     return finalStoppableChecks;
   }
 
@@ -507,7 +507,7 @@ private List<String> batchHelixInstanceStoppableCheck(String clusterId,
 
   private List<String> batchCustomInstanceStoppableCheck(String clusterId, List<String> instances,
       Set<String> toBeStoppedInstances, Map<String, StoppableCheck> finalStoppableChecks,
-      Map<String, String> customPayLoads) {
+      Map<String, String> customPayLoads, boolean preserveOrder) {
     if (instances.isEmpty()) {
       // if all instances failed at previous checks, then all following checks are not required.
       return instances;
@@ -536,7 +536,7 @@ private List<String> batchCustomInstanceStoppableCheck(String clusterId, List<St
       Map<String, StoppableCheck> clusterLevelCustomCheckResult =
           performAggregatedCustomCheck(clusterId, instanceIdsForCustomCheck,
               restConfig.getCompleteConfiguredHealthUrl().get(), customPayLoads,
-              toBeStoppedInstances);
+              toBeStoppedInstances, preserveOrder);
       List<String> instancesForNextCheck = new ArrayList<>();
       clusterLevelCustomCheckResult.forEach((instance, stoppableCheck) -> {
         addStoppableCheck(finalStoppableChecks, instance, stoppableCheck);
@@ -553,9 +553,13 @@ private List<String> batchCustomInstanceStoppableCheck(String clusterId, List<St
     List<String> instancesForCustomPartitionLevelChecks = instanceIdsForCustomCheck;
     if (!_skipHealthCheckCategories.contains(StoppableCheck.Category.CUSTOM_INSTANCE_CHECK)) {
       Map<String, Future<StoppableCheck>> customInstanceLevelChecks = instances.stream().collect(
-          Collectors.toMap(Function.identity(), instance -> POOL.submit(
-              () -> performCustomInstanceCheck(clusterId, instance, restConfig.getBaseUrl(instance),
-                  customPayLoads))));
+          Collectors.toMap(
+              Function.identity(),
+              instance -> POOL.submit(() -> performCustomInstanceCheck(clusterId, instance, restConfig.getBaseUrl(instance), customPayLoads)),
+              (existing, replacement) -> existing,
+              // Use LinkedHashMap when preserveOrder is true to maintain the original order of instances
+              preserveOrder ? LinkedHashMap::new : HashMap::new
+          ));
       instancesForCustomPartitionLevelChecks =
           filterInstancesForNextCheck(customInstanceLevelChecks, finalStoppableChecks);
     }
@@ -638,7 +642,7 @@ private Map<String, MaintenanceManagementInstanceInfo> batchInstanceHealthCheck(
         // custom check, includes custom Instance check and partition check.
         instancesForNext =
             batchCustomInstanceStoppableCheck(clusterId, instancesForNext, Collections.emptySet(),
-                finalStoppableChecks, healthCheckConfig);
+                finalStoppableChecks, healthCheckConfig, false);
       } else {
         throw new UnsupportedOperationException(healthCheck + " is not supported yet!");
       }
@@ -785,8 +789,10 @@ private Map<String, StoppableCheck> performPartitionsCheck(List<String> instance
 
   private Map<String, StoppableCheck> performAggregatedCustomCheck(String clusterId,
       List<String> instances, String url, Map<String, String> customPayLoads,
-      Set<String> toBeStoppedInstances) {
-    Map<String, StoppableCheck> aggregatedStoppableChecks = new HashMap<>();
+      Set<String> toBeStoppedInstances, boolean preserveOrder) {
+    // Use LinkedHashMap when preserveOrder is true to maintain the original order of instances
+    Map<String, StoppableCheck> aggregatedStoppableChecks = preserveOrder ?
+        new LinkedHashMap<>() : new HashMap<>();
     try {
       Map<String, List<String>> customCheckResult =
           _customRestClient.getAggregatedStoppableCheck(url, instances, toBeStoppedInstances,
@@ -799,9 +805,13 @@ private Map<String, StoppableCheck> performAggregatedCustomCheck(String clusterI
       }
     } catch (IOException ex) {
       LOG.error("Custom client side aggregated health check for {} failed.", clusterId, ex);
-      return instances.stream().collect(Collectors.toMap(Function.identity(),
+      return instances.stream().collect(Collectors.toMap(
+          Function.identity(),
           instance -> new StoppableCheck(false, Collections.singletonList(instance),
-              StoppableCheck.Category.CUSTOM_AGGREGATED_CHECK)));
+              StoppableCheck.Category.CUSTOM_AGGREGATED_CHECK),
+          (existing, replacement) -> existing,
+          // Use LinkedHashMap when preserveOrder is true to maintain the original order of instances
+          preserveOrder ? LinkedHashMap::new : HashMap::new));
     }
     return aggregatedStoppableChecks;
   }

helix-rest/src/main/java/org/apache/helix/rest/clusterMaintenanceService/StoppableInstancesSelector.java

@@ -75,6 +75,25 @@ private StoppableInstancesSelector(String clusterId, List<String> orderOfZone,
     _dataAccessor = dataAccessor;
   }
 
+  /**
+   * Evaluates and collects stoppable instances within a specified or determined zone based on the order of zones.
+   * If _orderOfZone is specified, the method targets the first non-empty zone; otherwise, it targets the zone with
+   * the highest instance count. The method iterates through instances, performing stoppable checks, and records
+   * reasons for non-stoppability.
+   *
+   * @param instances A list of instance to be evaluated.
+   * @param toBeStoppedInstances A list of instances presumed to be already stopped
+   * @return An ObjectNode containing:
+   *         - 'stoppableNode': List of instances that can be stopped.
+   *         - 'instance_not_stoppable_with_reasons': A map with the instance name as the key and
+   *         a list of reasons for non-stoppability as the value.
+   * @throws IOException
+   */
+  public ObjectNode getStoppableInstancesInSingleZone(List<String> instances,
+      List<String> toBeStoppedInstances) throws IOException {
+    return getStoppableInstancesInSingleZone(instances, toBeStoppedInstances, false);
+  }
+
   /**
    * Evaluates and collects stoppable instances within a specified or determined zone based on the order of zones.
    * If _orderOfZone is specified, the method targets the first non-empty zone; otherwise, it targets the zone with
@@ -115,7 +134,6 @@ public ObjectNode getStoppableInstancesInSingleZone(List<String> instances,
    *
    * @param instances A list of instance to be evaluated.
    * @param toBeStoppedInstances A list of instances presumed to be already stopped
-   * @param preserveOrder Indicates whether to preserve the original order of instances
    * @return An ObjectNode containing:
    *         - 'stoppableNode': List of instances that can be stopped.
    *         - 'instance_not_stoppable_with_reasons': A map with the instance name as the key and
@@ -153,7 +171,6 @@ public ObjectNode getStoppableInstancesCrossZones(List<String> instances,
    *
    * @param instances A list of instance to be evaluated.
    * @param toBeStoppedInstances A list of instances presumed to be already stopped
-   * @param preserveOrder Indicates whether to preserve the original order of instances
    * @return An ObjectNode containing:
    *         - 'stoppableNode': List of instances that can be stopped.
    *         - 'instance_not_stoppable_with_reasons': A map with the instance name as the key and