Add on-demand rebalance API (#2601)

Add on-demand rebalance API
apache · Aug 25, 2023 · 0237ef2 · 0237ef2
1 parent 0b5b54f
commit 0237ef2
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 1 deletion.
diff --git a/helix-rest/src/main/java/org/apache/helix/rest/server/resources/AbstractResource.java b/helix-rest/src/main/java/org/apache/helix/rest/server/resources/AbstractResource.java
@@ -85,7 +85,8 @@ public enum Command {
     purgeOfflineParticipants,
     getInstance,
     getAllInstances,
-    setInstanceOperation // TODO: Name is just a place holder, may change in future
+    setInstanceOperation, // TODO: Name is just a place holder, may change in future
+    onDemandRebalance
   }
 
   @Context

diff --git a/helix-rest/src/main/java/org/apache/helix/rest/server/resources/helix/ClusterAccessor.java b/helix-rest/src/main/java/org/apache/helix/rest/server/resources/helix/ClusterAccessor.java
@@ -347,6 +347,15 @@ public Response updateCluster(@PathParam("clusterId") String clusterId,
           helixAdmin.purgeOfflineInstances(clusterId, duration);
         }
         break;
+      case onDemandRebalance:
+        try {
+          helixAdmin.onDemandRebalance(clusterId);
+        } catch (Exception ex) {
+          LOG.error(
+              "Cannot start on-demand rebalance for cluster: {}, Exception: {}", clusterId, ex);
+          return serverError(ex);
+        }
+        break;
       default:
         return badRequest("Unsupported command {}." + command);
     }

diff --git a/helix-rest/src/test/java/org/apache/helix/rest/server/TestClusterAccessor.java b/helix-rest/src/test/java/org/apache/helix/rest/server/TestClusterAccessor.java
@@ -1414,6 +1414,26 @@ public void testUpdateCustomizedConfig() throws IOException {
     System.out.println("End test :" + TestHelper.getTestMethodName());
   }
 
+  @Test(dependsOnMethods = "testUpdateCustomizedConfig")
+  public void testOnDemandRebalance() throws IOException {
+    System.out.println("Start test :" + TestHelper.getTestMethodName());
+    long currentTime = System.currentTimeMillis();
+    String cluster = "TestCluster_1";
+    new JerseyUriRequestBuilder("clusters/{}?command=onDemandRebalance").format(cluster)
+        .post(this, Entity.entity("{}", MediaType.APPLICATION_JSON_TYPE));
+
+    ClusterConfig config = _configAccessor.getClusterConfig(cluster);
+    long lastOnDemandRebalanceTime = config.getLastOnDemandRebalanceTimestamp();
+    Assert.assertFalse(lastOnDemandRebalanceTime == -1L,
+        "The last on-demand rebalance timestamp is not found.");
+    Assert.assertTrue(lastOnDemandRebalanceTime > currentTime, String.format(
+        "The last on-demand rebalance timestamp {} is stale. Expect a timestamp that is larger than {}.",
+        lastOnDemandRebalanceTime, currentTime));
+    // restore the state
+    config.setLastOnDemandRebalanceTimestamp(-1L);
+    System.out.println("End test :" + TestHelper.getTestMethodName());
+  }
+
   @Test
   public void testClusterFreezeMode() throws Exception {
     String cluster = "TestCluster_0";