From 98362ebacdbbca0b6cebc647bc2cc549574a7a6b Mon Sep 17 00:00:00 2001
From: Xiaxuan Gao <gaox8499@gmail.com>
Date: Wed, 18 Oct 2023 11:03:13 -0700
Subject: [PATCH] Add support for random single-zone stoppable check

---
 .../common/StoppableInstancesSelector.java    | 42 +++++++++++++------
 .../resources/helix/InstancesAccessor.java    | 10 +++--
 2 files changed, 36 insertions(+), 16 deletions(-)

diff --git a/helix-rest/src/main/java/org/apache/helix/rest/common/StoppableInstancesSelector.java b/helix-rest/src/main/java/org/apache/helix/rest/common/StoppableInstancesSelector.java
index 11b9b11924..63c3943457 100644
--- a/helix-rest/src/main/java/org/apache/helix/rest/common/StoppableInstancesSelector.java
+++ b/helix-rest/src/main/java/org/apache/helix/rest/common/StoppableInstancesSelector.java
@@ -73,7 +73,7 @@ public StoppableInstancesSelector(String clusterId, List<String> orderOfZone,
    */
   public void getStoppableInstancesInSingleZone(List<String> instances) throws IOException {
     List<String> zoneBasedInstance =
-        getZoneBasedInstances(instances, _orderOfZone, _clusterTopology.toZoneMapping());
+        getZoneBasedInstances(instances, _clusterTopology.toZoneMapping());
     Map<String, StoppableCheck> instancesStoppableChecks =
         _maintenanceService.batchGetInstancesStoppableChecks(_clusterId, zoneBasedInstance,
             _customizedInput);
@@ -110,6 +110,30 @@ public void getStoppableInstancesCrossZones() {
     throw new NotImplementedException("Not Implemented");
   }
 
+  /**
+   * Determines the order of zones. If an order is provided by the user, it will be used directly.
+   * Otherwise, zones will be ordered by their associated instance count in descending order.
+   *
+   * If `isRandom` is true, the order of zones will be randomized regardless of any previous order.
+   *
+   * @param isRandom Indicates whether to randomize the order of zones.
+   */
+  public void calculateOrderOfZone(boolean isRandom) {
+    // If the orderedZones is not specified, we will order all zones by their instances count in descending order.
+    if (_orderOfZone == null) {
+      _orderOfZone =
+          new ArrayList<>(getOrderedZoneToInstancesMap(_clusterTopology.toZoneMapping()).keySet());
+    }
+
+    if (_orderOfZone.isEmpty()) {
+      return;
+    }
+
+    if (isRandom) {
+      Collections.shuffle(_orderOfZone);
+    }
+  }
+
   /**
    * Get instances belongs to the first zone. If the zone is already empty, Helix will iterate zones
    * by order until find the zone contains instances.
@@ -118,23 +142,17 @@ public void getStoppableInstancesCrossZones() {
    * zones by the number of associated instances in descending order.
    *
    * @param instances
-   * @param orderedZones
+   * @param zoneMapping
    * @return
    */
-  private List<String> getZoneBasedInstances(List<String> instances, List<String> orderedZones,
+  private List<String> getZoneBasedInstances(List<String> instances,
       Map<String, Set<String>> zoneMapping) {
-
-    // If the orderedZones is not specified, we will order all zones by their instances count in descending order.
-    if (orderedZones == null) {
-      orderedZones = new ArrayList<>(getOrderedZoneToInstancesMap(zoneMapping).keySet());
-    }
-
-    if (orderedZones.isEmpty()) {
-      return orderedZones;
+    if (_orderOfZone.isEmpty()) {
+      return _orderOfZone;
     }
 
     Set<String> instanceSet = null;
-    for (String zone : orderedZones) {
+    for (String zone : _orderOfZone) {
       instanceSet = new TreeSet<>(instances);
       Set<String> currentZoneInstanceSet = new HashSet<>(zoneMapping.get(zone));
       instanceSet.retainAll(currentZoneInstanceSet);
diff --git a/helix-rest/src/main/java/org/apache/helix/rest/server/resources/helix/InstancesAccessor.java b/helix-rest/src/main/java/org/apache/helix/rest/server/resources/helix/InstancesAccessor.java
index 514b244a11..0901d30d31 100644
--- a/helix-rest/src/main/java/org/apache/helix/rest/server/resources/helix/InstancesAccessor.java
+++ b/helix-rest/src/main/java/org/apache/helix/rest/server/resources/helix/InstancesAccessor.java
@@ -148,7 +148,8 @@ public Response instancesOperations(@PathParam("clusterId") String clusterId,
       @QueryParam("command") String command,
       @QueryParam("continueOnFailures") boolean continueOnFailures,
       @QueryParam("skipZKRead") boolean skipZKRead,
-      @QueryParam("skipHealthCheckCategories") String skipHealthCheckCategories, String content) {
+      @QueryParam("skipHealthCheckCategories") String skipHealthCheckCategories,
+      @QueryParam("random") boolean random, String content) {
     Command cmd;
     try {
       cmd = Command.valueOf(command);
@@ -193,7 +194,7 @@ public Response instancesOperations(@PathParam("clusterId") String clusterId,
           break;
         case stoppable:
           return batchGetStoppableInstances(clusterId, node, skipZKRead, continueOnFailures,
-              skipHealthCheckCategorySet);
+              skipHealthCheckCategorySet, random);
         default:
           _logger.error("Unsupported command :" + command);
           return badRequest("Unsupported command :" + command);
@@ -210,8 +211,8 @@ public Response instancesOperations(@PathParam("clusterId") String clusterId,
   }
 
   private Response batchGetStoppableInstances(String clusterId, JsonNode node, boolean skipZKRead,
-      boolean continueOnFailures, Set<StoppableCheck.Category> skipHealthCheckCategories)
-      throws IOException {
+      boolean continueOnFailures, Set<StoppableCheck.Category> skipHealthCheckCategories,
+      boolean random) throws IOException {
     try {
       // TODO: Process input data from the content
       InstancesAccessor.InstanceHealthSelectionBase selectionBase =
@@ -258,6 +259,7 @@ private Response batchGetStoppableInstances(String clusterId, JsonNode node, boo
               .setMaintenanceService(maintenanceService)
               .setClusterTopology(clusterTopology)
               .build();
+      stoppableInstancesSelector.calculateOrderOfZone(random);
       switch (selectionBase) {
         case zone_based:
           stoppableInstancesSelector.getStoppableInstancesInSingleZone(instances);