Add REST API endpoints for WAGED Rebalancer

helix-rest/src/main/java/org/apache/helix/rest/server/resources/AbstractResource.java

@@ -66,7 +66,15 @@ public enum Command {
     rebalance,
     reset,
     resetPartitions,
-    removeInstanceTag
+    removeInstanceTag,
+    addResource,
+    addWagedResource,
+    getResource,
+    validateWeight,
+    enableWagedRebalance,
+    enableWagedRebalanceForAllResources,
+    getInstance,
+    getAllInstances
   }
 
   @Context

helix-rest/src/main/java/org/apache/helix/rest/server/resources/helix/ClusterAccessor.java

@@ -240,7 +240,15 @@ public Response updateCluster(@PathParam("clusterId") String clusterId,
       helixAdmin.manuallyEnableMaintenanceMode(clusterId, command == Command.enableMaintenanceMode,
           content, customFieldsMap);
       break;
-
+    case enableWagedRebalanceForAllResources:
+      // Enable WAGED rebalance for all resources in the cluster
+      List<String> resources = helixAdmin.getResourcesInCluster(clusterId);
+      try {
+        helixAdmin.enableWagedRebalance(clusterId, resources);
+      } catch (HelixException e) {
+        return badRequest(e.getMessage());
+      }
+      break;
     default:
       return badRequest("Unsupported command " + command);
     }

helix-rest/src/main/java/org/apache/helix/rest/server/resources/helix/InstancesAccessor.java

@@ -1,5 +1,24 @@
 package org.apache.helix.rest.server.resources.helix;
 
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collections;
@@ -8,6 +27,7 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.TreeSet;
+import javax.ws.rs.DefaultValue;
 import javax.ws.rs.GET;
 import javax.ws.rs.POST;
 import javax.ws.rs.Path;
@@ -60,41 +80,65 @@ public enum InstanceHealthSelectionBase {
   }
 
   @GET
-  public Response getAllInstances(@PathParam("clusterId") String clusterId) {
+  public Response getAllInstances(@PathParam("clusterId") String clusterId,
+      @DefaultValue("getAllInstances") @QueryParam("command") String command) {
+    // Get the command. If not provided, the default would be "getAllInstances"
+    Command cmd;
+    try {
+      cmd = Command.valueOf(command);
+    } catch (Exception e) {
+      return badRequest("Invalid command : " + command);
+    }
+
     HelixDataAccessor accessor = getDataAccssor(clusterId);
     List<String> instances = accessor.getChildNames(accessor.keyBuilder().instanceConfigs());
-
     if (instances == null) {
       return notFound();
     }
 
-    ObjectNode root = JsonNodeFactory.instance.objectNode();
-    root.put(Properties.id.name(), JsonNodeFactory.instance.textNode(clusterId));
+    switch (cmd) {
+    case getAllInstances:
+      ObjectNode root = JsonNodeFactory.instance.objectNode();
+      root.put(Properties.id.name(), JsonNodeFactory.instance.textNode(clusterId));
 
-    ArrayNode instancesNode = root.putArray(InstancesAccessor.InstancesProperties.instances.name());
-    instancesNode.addAll((ArrayNode) OBJECT_MAPPER.valueToTree(instances));
-    ArrayNode onlineNode = root.putArray(InstancesAccessor.InstancesProperties.online.name());
-    ArrayNode disabledNode = root.putArray(InstancesAccessor.InstancesProperties.disabled.name());
+      ArrayNode instancesNode =
+          root.putArray(InstancesAccessor.InstancesProperties.instances.name());
+      instancesNode.addAll((ArrayNode) OBJECT_MAPPER.valueToTree(instances));
+      ArrayNode onlineNode = root.putArray(InstancesAccessor.InstancesProperties.online.name());
+      ArrayNode disabledNode = root.putArray(InstancesAccessor.InstancesProperties.disabled.name());
 
-    List<String> liveInstances = accessor.getChildNames(accessor.keyBuilder().liveInstances());
-    ClusterConfig clusterConfig = accessor.getProperty(accessor.keyBuilder().clusterConfig());
+      List<String> liveInstances = accessor.getChildNames(accessor.keyBuilder().liveInstances());
+      ClusterConfig clusterConfig = accessor.getProperty(accessor.keyBuilder().clusterConfig());
 
-    for (String instanceName : instances) {
-      InstanceConfig instanceConfig =
-          accessor.getProperty(accessor.keyBuilder().instanceConfig(instanceName));
-      if (instanceConfig != null) {
-        if (!instanceConfig.getInstanceEnabled() || (clusterConfig.getDisabledInstances() != null
-            && clusterConfig.getDisabledInstances().containsKey(instanceName))) {
-          disabledNode.add(JsonNodeFactory.instance.textNode(instanceName));
-        }
+      for (String instanceName : instances) {
+        InstanceConfig instanceConfig =
+            accessor.getProperty(accessor.keyBuilder().instanceConfig(instanceName));
+        if (instanceConfig != null) {
+          if (!instanceConfig.getInstanceEnabled() || (clusterConfig.getDisabledInstances() != null
+              && clusterConfig.getDisabledInstances().containsKey(instanceName))) {
+            disabledNode.add(JsonNodeFactory.instance.textNode(instanceName));
+          }
 
-        if (liveInstances.contains(instanceName)){
-          onlineNode.add(JsonNodeFactory.instance.textNode(instanceName));
+          if (liveInstances.contains(instanceName)) {
+            onlineNode.add(JsonNodeFactory.instance.textNode(instanceName));
+          }
         }
       }
+      return JSONRepresentation(root);
+    case validateWeight:
+      // Validate all instances for WAGED rebalance
+      HelixAdmin admin = getHelixAdmin();
+      Map<String, Boolean> validationResultMap;
+      try {
+        validationResultMap = admin.validateInstancesForWagedRebalance(clusterId, instances);
+      } catch (HelixException e) {
+        return badRequest(e.getMessage());
+      }
+      return JSONRepresentation(validationResultMap);
+    default:
+      _logger.error("Unsupported command :" + command);
+      return badRequest("Unsupported command :" + command);
     }
-
-    return JSONRepresentation(root);
   }
 
   @POST

helix-rest/src/main/java/org/apache/helix/rest/server/resources/helix/PerInstanceAccessor.java

@@ -20,7 +20,9 @@
  */
 
 import java.io.IOException;
+import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 
 import javax.ws.rs.*;
 import javax.ws.rs.core.MediaType;
@@ -67,16 +69,41 @@ public enum PerInstanceProperties {
 
   @GET
   public Response getInstanceById(@PathParam("clusterId") String clusterId,
-      @PathParam("instanceName") String instanceName) throws IOException {
-    ObjectMapper objectMapper = new ObjectMapper();
-    HelixDataAccessor dataAccessor = getDataAccssor(clusterId);
-    // TODO reduce GC by dependency injection
-    InstanceService instanceService =
-        new InstanceServiceImpl(new HelixDataAccessorWrapper((ZKHelixDataAccessor) dataAccessor), getConfigAccessor());
-    InstanceInfo instanceInfo = instanceService.getInstanceInfo(clusterId, instanceName,
-        InstanceService.HealthCheck.STARTED_AND_HEALTH_CHECK_LIST);
+      @PathParam("instanceName") String instanceName,
+      @DefaultValue("getInstance") @QueryParam("command") String command) throws IOException {
+    // Get the command. If not provided, the default would be "getInstance"
+    Command cmd;
+    try {
+      cmd = Command.valueOf(command);
+    } catch (Exception e) {
+      return badRequest("Invalid command : " + command);
+    }
 
-    return OK(objectMapper.writeValueAsString(instanceInfo));
+    switch (cmd) {
+    case getInstance:
+      ObjectMapper objectMapper = new ObjectMapper();
+      HelixDataAccessor dataAccessor = getDataAccssor(clusterId);
+      // TODO reduce GC by dependency injection
+      InstanceService instanceService = new InstanceServiceImpl(
+          new HelixDataAccessorWrapper((ZKHelixDataAccessor) dataAccessor), getConfigAccessor());
+      InstanceInfo instanceInfo = instanceService.getInstanceInfo(clusterId, instanceName,
+          InstanceService.HealthCheck.STARTED_AND_HEALTH_CHECK_LIST);
+      return OK(objectMapper.writeValueAsString(instanceInfo));
+    case validateWeight:
+      // Validates instanceConfig for WAGED rebalance
+      HelixAdmin admin = getHelixAdmin();
+      Map<String, Boolean> validationResultMap;
+      try {
+        validationResultMap = admin.validateInstancesForWagedRebalance(clusterId,
+            Collections.singletonList(instanceName));
+      } catch (HelixException e) {
+        return badRequest(e.getMessage());
+      }
+      return JSONRepresentation(validationResultMap);
+    default:
+      LOG.error("Unsupported command :" + command);
+      return badRequest("Unsupported command :" + command);
+    }
   }
 
   @POST
@@ -331,7 +358,8 @@ public Response getResourceOnInstance(@PathParam("clusterId") String clusterId,
     return notFound();
   }
 
-  @GET @Path("errors")
+  @GET
+  @Path("errors")
   public Response getErrorsOnInstance(@PathParam("clusterId") String clusterId,
       @PathParam("instanceName") String instanceName) throws IOException {
     HelixDataAccessor accessor = getDataAccssor(clusterId);