Add preserveOrder check for CustomChecks

LZD-PratyushBhatt · LZD-PratyushBhatt · commit e7ec03b84f61 · 2025-05-12T11:58:46.000+05:30
diff --git a/helix-rest/src/main/java/org/apache/helix/rest/clusterMaintenanceService/MaintenanceManagementService.java b/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;
   }
 
@@ -502,12 +502,12 @@ private List<String> batchHelixInstanceStoppableCheck(String clusterId,
     addMinActiveReplicaChecks(clusterId, helixInstanceChecks, toBeStoppedInstances);
 
     // finalStoppableChecks contains instances that does not pass this health check
-    return filterInstancesForNextCheck(helixInstanceChecks, finalStoppableChecks);
+    return filterInstancesForNextCheck(helixInstanceChecks, finalStoppableChecks, preserveOrder);
   }
 
   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,11 +553,15 @@ 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);
+          filterInstancesForNextCheck(customInstanceLevelChecks, finalStoppableChecks, preserveOrder);
     }
 
     if (!instancesForCustomPartitionLevelChecks.isEmpty() && !_skipHealthCheckCategories.contains(
@@ -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!");
       }
@@ -675,18 +679,31 @@ private void addStoppableCheck(Map<String, StoppableCheck> stoppableChecks, Stri
   private List<String> filterInstancesForNextCheck(
       Map<String, Future<StoppableCheck>> futureStoppableCheckByInstance,
       Map<String, StoppableCheck> finalStoppableCheckByInstance) {
-    List<String> instancesForNextCheck = new ArrayList<>();
-    for (Map.Entry<String, Future<StoppableCheck>> entry : futureStoppableCheckByInstance
-        .entrySet()) {
+    return filterInstancesForNextCheck(futureStoppableCheckByInstance, finalStoppableCheckByInstance, false);
+  }
+
+  /**
+   * Filters instances that have passed the stoppable check for the next check stage.
+   * When preserveOrder is true, the original order of instances is maintained.
+   *
+   * @param futureStoppableCheckByInstance Map of instance name to future stoppable check
+   * @param finalStoppableCheckByInstance Map to collect the final stoppable check results
+   * @param preserveOrder Whether to preserve the original order of instances
+   * @return List of instances that passed the stoppable check
+   */
+  private List<String> filterInstancesForNextCheck(
+      Map<String, Future<StoppableCheck>> futureStoppableCheckByInstance,
+      Map<String, StoppableCheck> finalStoppableCheckByInstance,
+      boolean preserveOrder) {
+    // Use LinkedHashMap when preserveOrder is true to maintain the original order of instances
+    Map<String, Boolean> instanceStoppableMap = preserveOrder ? new LinkedHashMap<>() : new HashMap<>();
+
+    for (Map.Entry<String, Future<StoppableCheck>> entry : futureStoppableCheckByInstance.entrySet()) {
       String instance = entry.getKey();
       try {
         StoppableCheck stoppableCheck = entry.getValue().get();
         addStoppableCheck(finalStoppableCheckByInstance, instance, stoppableCheck);
-        if (stoppableCheck.isStoppable() || isNonBlockingCheck(stoppableCheck)) {
-          // instance passed this around of check or mandatory all checks
-          // will be checked in the next round
-          instancesForNextCheck.add(instance);
-        }
+        instanceStoppableMap.put(instance, stoppableCheck.isStoppable() || isNonBlockingCheck(stoppableCheck));
       } catch (Exception e) {
         String errorMessage =
             String.format("Failed to get StoppableChecks in parallel. Instance: %s", instance);
@@ -695,6 +712,16 @@ private List<String> filterInstancesForNextCheck(
       }
     }
 
+    // Filter instances that passed the check
+    List<String> instancesForNextCheck = new ArrayList<>();
+    for (Map.Entry<String, Boolean> entry : instanceStoppableMap.entrySet()) {
+      if (entry.getValue()) {
+        // Instance passed this round of check or mandatory all checks
+        // will be checked in the next round
+        instancesForNextCheck.add(entry.getKey());
+      }
+    }
+
     return instancesForNextCheck;
   }
 
@@ -785,8 +812,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 +828,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;
   }
@@ -912,7 +945,7 @@ protected Map<String, Boolean> getInstanceHealthStatus(String clusterId, String
     return healthStatus;
   }
 
-  // Adds the result of the min_active replica check for each stoppable check passed in futureStoppableCheckByInstance
+
   private void addMinActiveReplicaChecks(String clusterId, Map<String, Future<StoppableCheck>> futureStoppableCheckByInstance,
       Set<String> toBeStoppedInstances) {
     // Do not perform check if in the skip list
diff --git a/helix-rest/src/main/java/org/apache/helix/rest/clusterMaintenanceService/StoppableInstancesSelector.java b/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
