You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@helix.apache.org by ka...@apache.org on 2013/10/28 23:29:37 UTC

git commit: [HELIX-209] Fix Rebalancer and model package conflicts

Updated Branches:
  refs/heads/helix-logical-model a9aa77638 -> 842035a4b


[HELIX-209] Fix Rebalancer and model package conflicts


Project: http://git-wip-us.apache.org/repos/asf/incubator-helix/repo
Commit: http://git-wip-us.apache.org/repos/asf/incubator-helix/commit/842035a4
Tree: http://git-wip-us.apache.org/repos/asf/incubator-helix/tree/842035a4
Diff: http://git-wip-us.apache.org/repos/asf/incubator-helix/diff/842035a4

Branch: refs/heads/helix-logical-model
Commit: 842035a4b3edd1cc2f822e1561764c8ecd527467
Parents: a9aa776
Author: Kanak Biscuitwala <ka...@apache.org>
Authored: Mon Oct 28 14:19:53 2013 -0700
Committer: Kanak Biscuitwala <ka...@apache.org>
Committed: Mon Oct 28 14:19:53 2013 -0700

----------------------------------------------------------------------
 .../controller/rebalancer/CustomRebalancer.java |  2 +-
 .../rebalancer/FullAutoRebalancer.java          |  2 +-
 .../controller/rebalancer/HelixRebalancer.java  | 64 ++++++++++++++++++++
 .../helix/controller/rebalancer/Rebalancer.java | 39 ++++++++----
 .../controller/rebalancer/RebalancerRef.java    |  8 +--
 .../rebalancer/SemiAutoRebalancer.java          |  2 +-
 .../rebalancer/context/RebalancerConfig.java    |  8 +--
 .../stages/BestPossibleStateCalcStage.java      |  4 +-
 .../stages/ExternalViewComputeStage.java        |  4 +-
 .../stages/MessageSelectionStage.java           |  2 +-
 .../controller/rebalancer/Rebalancer.java       | 56 -----------------
 .../java/org/apache/helix/model/IdealState.java | 21 ++++++-
 .../java/org/apache/helix/model/Message.java    |  2 +-
 .../apache/helix/model/ResourceAssignment.java  | 34 +----------
 .../helix/model/StateModelDefinition.java       | 28 ++++++---
 .../helix/model/builder/AutoModeISBuilder.java  | 62 ++++++++++++++++++-
 .../builder/AutoRebalanceModeISBuilder.java     | 32 +++++++++-
 .../builder/ClusterConstraintsBuilder.java      |  8 ++-
 .../model/builder/CurrentStateBuilder.java      | 11 ++--
 .../model/builder/CustomModeISBuilder.java      | 61 ++++++++++++++++---
 .../model/builder/ExternalViewBuilder.java      | 24 --------
 .../TestCustomizedIdealStateRebalancer.java     |  4 +-
 22 files changed, 308 insertions(+), 170 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/incubator-helix/blob/842035a4/helix-core/src/main/java/org/apache/helix/controller/rebalancer/CustomRebalancer.java
----------------------------------------------------------------------
diff --git a/helix-core/src/main/java/org/apache/helix/controller/rebalancer/CustomRebalancer.java b/helix-core/src/main/java/org/apache/helix/controller/rebalancer/CustomRebalancer.java
index ac4d328..69c379a 100644
--- a/helix-core/src/main/java/org/apache/helix/controller/rebalancer/CustomRebalancer.java
+++ b/helix-core/src/main/java/org/apache/helix/controller/rebalancer/CustomRebalancer.java
@@ -37,7 +37,7 @@ import org.apache.log4j.Logger;
  * under the License.
  */
 
-public class CustomRebalancer implements Rebalancer {
+public class CustomRebalancer implements HelixRebalancer {
 
   private static final Logger LOG = Logger.getLogger(CustomRebalancer.class);
 

http://git-wip-us.apache.org/repos/asf/incubator-helix/blob/842035a4/helix-core/src/main/java/org/apache/helix/controller/rebalancer/FullAutoRebalancer.java
----------------------------------------------------------------------
diff --git a/helix-core/src/main/java/org/apache/helix/controller/rebalancer/FullAutoRebalancer.java b/helix-core/src/main/java/org/apache/helix/controller/rebalancer/FullAutoRebalancer.java
index b0c11d4..d0a96cc 100644
--- a/helix-core/src/main/java/org/apache/helix/controller/rebalancer/FullAutoRebalancer.java
+++ b/helix-core/src/main/java/org/apache/helix/controller/rebalancer/FullAutoRebalancer.java
@@ -49,7 +49,7 @@ import com.google.common.collect.Lists;
  * under the License.
  */
 
-public class FullAutoRebalancer implements Rebalancer {
+public class FullAutoRebalancer implements HelixRebalancer {
   // These should be final, but are initialized in init rather than a constructor
   private AutoRebalanceStrategy _algorithm;
 

http://git-wip-us.apache.org/repos/asf/incubator-helix/blob/842035a4/helix-core/src/main/java/org/apache/helix/controller/rebalancer/HelixRebalancer.java
----------------------------------------------------------------------
diff --git a/helix-core/src/main/java/org/apache/helix/controller/rebalancer/HelixRebalancer.java b/helix-core/src/main/java/org/apache/helix/controller/rebalancer/HelixRebalancer.java
new file mode 100644
index 0000000..7fcbba5
--- /dev/null
+++ b/helix-core/src/main/java/org/apache/helix/controller/rebalancer/HelixRebalancer.java
@@ -0,0 +1,64 @@
+package org.apache.helix.controller.rebalancer;
+
+import org.apache.helix.HelixManager;
+import org.apache.helix.api.Cluster;
+import org.apache.helix.controller.rebalancer.context.RebalancerConfig;
+import org.apache.helix.controller.stages.ResourceCurrentState;
+import org.apache.helix.model.ResourceAssignment;
+
+/*
+ * 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.
+ */
+
+/**
+ * Allows one to come up with custom implementation of a rebalancer.<br/>
+ * This will be invoked on all changes that happen in the cluster.<br/>
+ * Simply return the resource assignment for a resource in this method.<br/>
+ */
+public interface HelixRebalancer {
+  /**
+   * Initialize the rebalancer with a HelixManager if necessary
+   * @param manager
+   */
+  public void init(HelixManager helixManager);
+
+  /**
+   * Given an ideal state for a resource and liveness of participants, compute a assignment of
+   * instances and states to each partition of a resource. This method provides all the relevant
+   * information needed to rebalance a resource. If you need additional information use
+   * manager.getAccessor to read and write the cluster data. This allows one to compute the
+   * ResourceAssignment according to app-specific requirements.<br/>
+   * <br/>
+   * Say that you have:<br/>
+   * 
+   * <pre>
+   * class MyRebalancerContext implements RebalancerContext
+   * </pre>
+   * 
+   * as your rebalancer context. To extract it from a RebalancerConfig, do the following:<br/>
+   * 
+   * <pre>
+   * MyRebalancerContext context = rebalancerConfig.getRebalancerContext(MyRebalancerContext.class);
+   * </pre>
+   * @param rebalancerConfig the properties of the resource for which a mapping will be computed
+   * @param cluster complete snapshot of the cluster
+   * @param currentState the current states of all partitions
+   */
+  public ResourceAssignment computeResourceMapping(RebalancerConfig rebalancerConfig,
+      Cluster cluster, ResourceCurrentState currentState);
+}

http://git-wip-us.apache.org/repos/asf/incubator-helix/blob/842035a4/helix-core/src/main/java/org/apache/helix/controller/rebalancer/Rebalancer.java
----------------------------------------------------------------------
diff --git a/helix-core/src/main/java/org/apache/helix/controller/rebalancer/Rebalancer.java b/helix-core/src/main/java/org/apache/helix/controller/rebalancer/Rebalancer.java
index 5a6a24e..0e6cb10 100644
--- a/helix-core/src/main/java/org/apache/helix/controller/rebalancer/Rebalancer.java
+++ b/helix-core/src/main/java/org/apache/helix/controller/rebalancer/Rebalancer.java
@@ -1,11 +1,5 @@
 package org.apache.helix.controller.rebalancer;
 
-import org.apache.helix.HelixManager;
-import org.apache.helix.api.Cluster;
-import org.apache.helix.controller.rebalancer.context.RebalancerConfig;
-import org.apache.helix.controller.stages.ResourceCurrentState;
-import org.apache.helix.model.ResourceAssignment;
-
 /*
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements.  See the NOTICE file
@@ -25,15 +19,38 @@ import org.apache.helix.model.ResourceAssignment;
  * under the License.
  */
 
+import org.apache.helix.HelixManager;
+import org.apache.helix.controller.stages.ClusterDataCache;
+import org.apache.helix.controller.stages.CurrentStateOutput;
+import org.apache.helix.model.IdealState;
+
 /**
  * Allows one to come up with custom implementation of a rebalancer.<br/>
  * This will be invoked on all changes that happen in the cluster.<br/>
- * Simply return the resource assignment for a resource in this method.<br/>
+ * Simply return the newIdealState for a resource in this method.<br/>
+ * <br/>
+ * Deprecated. Use {@link HelixRebalancer} instead.
  */
+@Deprecated
 public interface Rebalancer {
+  /**
+   * Initialize the rebalancer with a HelixManager if necessary
+   * @param manager
+   */
+  public void init(HelixManager manager);
 
-  public void init(HelixManager helixManager);
-
-  public ResourceAssignment computeResourceMapping(RebalancerConfig rebalancerConfig, Cluster cluster,
-      ResourceCurrentState currentState);
+  /**
+   * Given an ideal state for a resource and liveness of instances, compute a assignment of
+   * instances and states to each partition of a resource. This method provides all the relevant
+   * information needed to rebalance a resource. If you need additional information use
+   * manager.getAccessor to read the cluster data. This allows one to compute the newIdealState
+   * according to app specific requirements.
+   * @param resourceName the resource for which a mapping will be computed
+   * @param currentIdealState the IdealState that corresponds to this resource
+   * @param currentStateOutput the current states of all partitions
+   * @param clusterData cache of the cluster state
+   */
+  public IdealState computeResourceMapping(final String resourceName,
+      final IdealState currentIdealState, final CurrentStateOutput currentStateOutput,
+      final ClusterDataCache clusterData);
 }

http://git-wip-us.apache.org/repos/asf/incubator-helix/blob/842035a4/helix-core/src/main/java/org/apache/helix/controller/rebalancer/RebalancerRef.java
----------------------------------------------------------------------
diff --git a/helix-core/src/main/java/org/apache/helix/controller/rebalancer/RebalancerRef.java b/helix-core/src/main/java/org/apache/helix/controller/rebalancer/RebalancerRef.java
index 79e4ba0..974222d 100644
--- a/helix-core/src/main/java/org/apache/helix/controller/rebalancer/RebalancerRef.java
+++ b/helix-core/src/main/java/org/apache/helix/controller/rebalancer/RebalancerRef.java
@@ -26,7 +26,7 @@ import org.codehaus.jackson.annotate.JsonIgnore;
 import org.codehaus.jackson.annotate.JsonProperty;
 
 /**
- * Reference to a class that extends {@link Rebalancer}. It loads the class automatically.
+ * Reference to a class that extends {@link HelixRebalancer}. It loads the class automatically.
  */
 public class RebalancerRef {
   private static final Logger LOG = Logger.getLogger(RebalancerRef.class);
@@ -44,9 +44,9 @@ public class RebalancerRef {
    * @return Rebalancer or null if instantiation failed
    */
   @JsonIgnore
-  public Rebalancer getRebalancer() {
+  public HelixRebalancer getRebalancer() {
     try {
-      return (Rebalancer) (HelixUtil.loadClass(getClass(), _rebalancerClassName).newInstance());
+      return (HelixRebalancer) (HelixUtil.loadClass(getClass(), _rebalancerClassName).newInstance());
     } catch (Exception e) {
       LOG.warn("Exception while invoking custom rebalancer class:" + _rebalancerClassName, e);
     }
@@ -85,7 +85,7 @@ public class RebalancerRef {
    * @param rebalancerClass class that implements Rebalancer
    * @return RebalancerRef
    */
-  public static RebalancerRef from(Class<? extends Rebalancer> rebalancerClass) {
+  public static RebalancerRef from(Class<? extends HelixRebalancer> rebalancerClass) {
     if (rebalancerClass == null) {
       return null;
     }

http://git-wip-us.apache.org/repos/asf/incubator-helix/blob/842035a4/helix-core/src/main/java/org/apache/helix/controller/rebalancer/SemiAutoRebalancer.java
----------------------------------------------------------------------
diff --git a/helix-core/src/main/java/org/apache/helix/controller/rebalancer/SemiAutoRebalancer.java b/helix-core/src/main/java/org/apache/helix/controller/rebalancer/SemiAutoRebalancer.java
index 96e3d4b..c5a7f22 100644
--- a/helix-core/src/main/java/org/apache/helix/controller/rebalancer/SemiAutoRebalancer.java
+++ b/helix-core/src/main/java/org/apache/helix/controller/rebalancer/SemiAutoRebalancer.java
@@ -40,7 +40,7 @@ import org.apache.log4j.Logger;
  * Rebalancer for the SEMI_AUTO mode. It expects a RebalancerConfig that understands the preferred
  * locations of each partition replica
  */
-public class SemiAutoRebalancer implements Rebalancer {
+public class SemiAutoRebalancer implements HelixRebalancer {
   private static final Logger LOG = Logger.getLogger(SemiAutoRebalancer.class);
 
   @Override

http://git-wip-us.apache.org/repos/asf/incubator-helix/blob/842035a4/helix-core/src/main/java/org/apache/helix/controller/rebalancer/context/RebalancerConfig.java
----------------------------------------------------------------------
diff --git a/helix-core/src/main/java/org/apache/helix/controller/rebalancer/context/RebalancerConfig.java b/helix-core/src/main/java/org/apache/helix/controller/rebalancer/context/RebalancerConfig.java
index 846fd01..64f70c3 100644
--- a/helix-core/src/main/java/org/apache/helix/controller/rebalancer/context/RebalancerConfig.java
+++ b/helix-core/src/main/java/org/apache/helix/controller/rebalancer/context/RebalancerConfig.java
@@ -2,7 +2,7 @@ package org.apache.helix.controller.rebalancer.context;
 
 import org.apache.helix.api.Scope;
 import org.apache.helix.api.config.NamespacedConfig;
-import org.apache.helix.controller.rebalancer.Rebalancer;
+import org.apache.helix.controller.rebalancer.HelixRebalancer;
 import org.apache.helix.model.ResourceConfiguration;
 import org.apache.helix.util.HelixUtil;
 import org.apache.log4j.Logger;
@@ -39,7 +39,7 @@ public final class RebalancerConfig {
 
   private static final Logger LOG = Logger.getLogger(RebalancerConfig.class);
   private ContextSerializer _serializer;
-  private Rebalancer _rebalancer;
+  private HelixRebalancer _rebalancer;
   private final RebalancerContext _context;
   private final NamespacedConfig _config;
 
@@ -117,7 +117,7 @@ public final class RebalancerConfig {
    * Get a rebalancer class instance
    * @return Rebalancer
    */
-  public Rebalancer getRebalancer() {
+  public HelixRebalancer getRebalancer() {
     // cache the rebalancer to avoid loading and instantiating it excessively
     if (_rebalancer == null) {
       if (_context == null || _context.getRebalancerRef() == null) {
@@ -137,7 +137,7 @@ public final class RebalancerConfig {
     try {
       return contextClass.cast(_context);
     } catch (ClassCastException e) {
-      LOG.info(contextClass + " is incompatible with context class: " + _context.getClass());
+      LOG.warn(contextClass + " is incompatible with context class: " + _context.getClass());
     }
     return null;
   }

http://git-wip-us.apache.org/repos/asf/incubator-helix/blob/842035a4/helix-core/src/main/java/org/apache/helix/controller/stages/BestPossibleStateCalcStage.java
----------------------------------------------------------------------
diff --git a/helix-core/src/main/java/org/apache/helix/controller/stages/BestPossibleStateCalcStage.java b/helix-core/src/main/java/org/apache/helix/controller/stages/BestPossibleStateCalcStage.java
index aae33b4..051a2f3 100644
--- a/helix-core/src/main/java/org/apache/helix/controller/stages/BestPossibleStateCalcStage.java
+++ b/helix-core/src/main/java/org/apache/helix/controller/stages/BestPossibleStateCalcStage.java
@@ -32,7 +32,7 @@ import org.apache.helix.api.id.ResourceId;
 import org.apache.helix.api.id.StateModelDefId;
 import org.apache.helix.controller.pipeline.AbstractBaseStage;
 import org.apache.helix.controller.pipeline.StageException;
-import org.apache.helix.controller.rebalancer.Rebalancer;
+import org.apache.helix.controller.rebalancer.HelixRebalancer;
 import org.apache.helix.controller.rebalancer.context.RebalancerConfig;
 import org.apache.helix.controller.rebalancer.context.RebalancerContext;
 import org.apache.helix.controller.rebalancer.util.ConstraintBasedAssignment;
@@ -119,7 +119,7 @@ public class BestPossibleStateCalcStage extends AbstractBaseStage {
       RebalancerConfig rebalancerConfig = resourceConfig.getRebalancerConfig();
       ResourceAssignment resourceAssignment = null;
       if (rebalancerConfig != null) {
-        Rebalancer rebalancer = rebalancerConfig.getRebalancer();
+        HelixRebalancer rebalancer = rebalancerConfig.getRebalancer();
         if (rebalancer != null) {
           HelixManager manager = event.getAttribute("helixmanager");
           rebalancer.init(manager);

http://git-wip-us.apache.org/repos/asf/incubator-helix/blob/842035a4/helix-core/src/main/java/org/apache/helix/controller/stages/ExternalViewComputeStage.java
----------------------------------------------------------------------
diff --git a/helix-core/src/main/java/org/apache/helix/controller/stages/ExternalViewComputeStage.java b/helix-core/src/main/java/org/apache/helix/controller/stages/ExternalViewComputeStage.java
index 1d287fa..55a5e54 100644
--- a/helix-core/src/main/java/org/apache/helix/controller/stages/ExternalViewComputeStage.java
+++ b/helix-core/src/main/java/org/apache/helix/controller/stages/ExternalViewComputeStage.java
@@ -196,7 +196,7 @@ public class ExternalViewComputeStage extends AbstractBaseStage {
           Message innerMessage =
               schedulerTaskConfig.getInnerMessage(PartitionId.from(taskPartitionName));
           if (innerMessage != null) {
-            String controllerMsgId = innerMessage.getControllerMessagId();
+            String controllerMsgId = innerMessage.getControllerMessageId();
             if (controllerMsgId != null) {
               LOG.info(taskPartitionName + " finished with controllerMsg " + controllerMsgId);
               if (!controllerMsgUpdates.containsKey(controllerMsgId)) {
@@ -211,7 +211,7 @@ public class ExternalViewComputeStage extends AbstractBaseStage {
     // fill the controllerMsgIdCountMap
     for (PartitionId taskId : schedulerTaskConfig.getPartitionSet()) {
       Message innerMessage = schedulerTaskConfig.getInnerMessage(taskId);
-      String controllerMsgId = innerMessage.getControllerMessagId();
+      String controllerMsgId = innerMessage.getControllerMessageId();
 
       if (controllerMsgId != null) {
         Integer curCnt = controllerMsgIdCountMap.get(controllerMsgId);

http://git-wip-us.apache.org/repos/asf/incubator-helix/blob/842035a4/helix-core/src/main/java/org/apache/helix/controller/stages/MessageSelectionStage.java
----------------------------------------------------------------------
diff --git a/helix-core/src/main/java/org/apache/helix/controller/stages/MessageSelectionStage.java b/helix-core/src/main/java/org/apache/helix/controller/stages/MessageSelectionStage.java
index 5c24314..4a3fe28 100644
--- a/helix-core/src/main/java/org/apache/helix/controller/stages/MessageSelectionStage.java
+++ b/helix-core/src/main/java/org/apache/helix/controller/stages/MessageSelectionStage.java
@@ -307,7 +307,7 @@ public class MessageSelectionStage extends AbstractBaseStage {
   // so that behavior is consistent
   private Map<String, Integer> getStateTransitionPriorityMap(StateModelDefinition stateModelDef) {
     Map<String, Integer> stateTransitionPriorities = new HashMap<String, Integer>();
-    List<String> stateTransitionPriorityList = stateModelDef.getStateTransitionPriorityStringList();
+    List<String> stateTransitionPriorityList = stateModelDef.getStateTransitionPriorityList();
     for (int i = 0; i < stateTransitionPriorityList.size(); i++) {
       stateTransitionPriorities.put(stateTransitionPriorityList.get(i), i);
     }

http://git-wip-us.apache.org/repos/asf/incubator-helix/blob/842035a4/helix-core/src/main/java/org/apache/helix/deprecated/controller/rebalancer/Rebalancer.java
----------------------------------------------------------------------
diff --git a/helix-core/src/main/java/org/apache/helix/deprecated/controller/rebalancer/Rebalancer.java b/helix-core/src/main/java/org/apache/helix/deprecated/controller/rebalancer/Rebalancer.java
deleted file mode 100644
index e4c2f26..0000000
--- a/helix-core/src/main/java/org/apache/helix/deprecated/controller/rebalancer/Rebalancer.java
+++ /dev/null
@@ -1,56 +0,0 @@
-package org.apache.helix.deprecated.controller.rebalancer;
-
-/*
- * 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 org.apache.helix.HelixManager;
-import org.apache.helix.controller.stages.ClusterDataCache;
-import org.apache.helix.controller.stages.CurrentStateOutput;
-import org.apache.helix.model.IdealState;
-
-/**
- * Allows one to come up with custom implementation of a rebalancer.<br/>
- * This will be invoked on all changes that happen in the cluster.<br/>
- * Simply return the newIdealState for a resource in this method.<br/>
- * <br/>
- * Deprecated. Use {@link org.apache.helix.controller.rebalancer.Rebalancer} instead.
- */
-@Deprecated
-public interface Rebalancer {
-  /**
-   * Initialize the rebalancer with a HelixManager if necessary
-   * @param manager
-   */
-  public void init(HelixManager manager);
-
-  /**
-   * Given an ideal state for a resource and liveness of instances, compute a assignment of
-   * instances and states to each partition of a resource. This method provides all the relevant
-   * information needed to rebalance a resource. If you need additional information use
-   * manager.getAccessor to read the cluster data. This allows one to compute the newIdealState
-   * according to app specific requirements.
-   * @param resourceName the resource for which a mapping will be computed
-   * @param currentIdealState the IdealState that corresponds to this resource
-   * @param currentStateOutput the current states of all partitions
-   * @param clusterData cache of the cluster state
-   */
-  public IdealState computeResourceMapping(final String resourceName,
-      final IdealState currentIdealState, final CurrentStateOutput currentStateOutput,
-      final ClusterDataCache clusterData);
-}

http://git-wip-us.apache.org/repos/asf/incubator-helix/blob/842035a4/helix-core/src/main/java/org/apache/helix/model/IdealState.java
----------------------------------------------------------------------
diff --git a/helix-core/src/main/java/org/apache/helix/model/IdealState.java b/helix-core/src/main/java/org/apache/helix/model/IdealState.java
index df4043a..09aed50 100644
--- a/helix-core/src/main/java/org/apache/helix/model/IdealState.java
+++ b/helix-core/src/main/java/org/apache/helix/model/IdealState.java
@@ -38,6 +38,7 @@ import org.apache.helix.api.id.PartitionId;
 import org.apache.helix.api.id.ResourceId;
 import org.apache.helix.api.id.StateModelDefId;
 import org.apache.helix.api.id.StateModelFactoryId;
+import org.apache.helix.controller.rebalancer.HelixRebalancer;
 import org.apache.helix.controller.rebalancer.RebalancerRef;
 import org.apache.log4j.Logger;
 
@@ -169,7 +170,7 @@ public class IdealState extends HelixProperty {
   }
 
   /**
-   * Define a custom rebalancer that implements {@link Rebalancer}
+   * Define a custom rebalancer that implements {@link HelixRebalancer}
    * @param rebalancerClassName the name of the custom rebalancing class
    */
   public void setRebalancerClassName(String rebalancerClassName) {
@@ -310,6 +311,15 @@ public class IdealState extends HelixProperty {
 
   /**
    * Set the current mapping of a partition
+   * @param partitionName the partition to set
+   * @param instanceStateMap (participant name, state) pairs
+   */
+  public void setInstanceStateMap(String partitionName, Map<String, String> instanceStateMap) {
+    _record.setMapField(partitionName, instanceStateMap);
+  }
+
+  /**
+   * Set the current mapping of a partition
    * @param partitionId the partition to set
    * @param participantStateMap (participant id, state) pairs
    */
@@ -385,6 +395,15 @@ public class IdealState extends HelixProperty {
 
   /**
    * Set the preference list of a partition
+   * @param partitionName the name of the partition to set
+   * @param preferenceList a list of participants that can serve replicas of the partition
+   */
+  public void setPreferenceList(String partitionName, List<String> preferenceList) {
+    _record.setListField(partitionName, preferenceList);
+  }
+
+  /**
+   * Set the preference list of a partition
    * @param partitionId the id of the partition to set
    * @param preferenceList a list of participants that can serve replicas of the partition
    */

http://git-wip-us.apache.org/repos/asf/incubator-helix/blob/842035a4/helix-core/src/main/java/org/apache/helix/model/Message.java
----------------------------------------------------------------------
diff --git a/helix-core/src/main/java/org/apache/helix/model/Message.java b/helix-core/src/main/java/org/apache/helix/model/Message.java
index 2bd313a..d465a80 100644
--- a/helix-core/src/main/java/org/apache/helix/model/Message.java
+++ b/helix-core/src/main/java/org/apache/helix/model/Message.java
@@ -910,7 +910,7 @@ public class Message extends HelixProperty {
    * Get controller message id, used for scheduler-task-queue state model only
    * @return controller message id
    */
-  public String getControllerMessagId() {
+  public String getControllerMessageId() {
     return _record.getSimpleField(DefaultSchedulerMessageHandlerFactory.CONTROLLER_MSG_ID);
   }
 

http://git-wip-us.apache.org/repos/asf/incubator-helix/blob/842035a4/helix-core/src/main/java/org/apache/helix/model/ResourceAssignment.java
----------------------------------------------------------------------
diff --git a/helix-core/src/main/java/org/apache/helix/model/ResourceAssignment.java b/helix-core/src/main/java/org/apache/helix/model/ResourceAssignment.java
index 96d0ca7..8672e7e 100644
--- a/helix-core/src/main/java/org/apache/helix/model/ResourceAssignment.java
+++ b/helix-core/src/main/java/org/apache/helix/model/ResourceAssignment.java
@@ -33,7 +33,6 @@ import org.apache.helix.api.id.ResourceId;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
 
 /**
@@ -48,10 +47,9 @@ import com.google.common.collect.Maps;
  * can be in s1.
  */
 public class ResourceAssignment extends HelixProperty {
-
   /**
    * Initialize an empty mapping
-   * @param resourceName the resource being mapped
+   * @param resourceId the resource being mapped
    */
   public ResourceAssignment(ResourceId resourceId) {
     super(resourceId.stringify());
@@ -86,14 +84,6 @@ public class ResourceAssignment extends HelixProperty {
   }
 
   /**
-   * Get the currently mapped partitions
-   * @return list of Partition objects (immutable)
-   */
-  public List<String> getMappedPartitions() {
-    return Lists.newArrayList(_record.getMapFields().keySet());
-  }
-
-  /**
    * Get the entire map of a resource
    * @return map of partition to participant to state
    */
@@ -121,19 +111,6 @@ public class ResourceAssignment extends HelixProperty {
   }
 
   /**
-   * Get the participant, state pairs for a partition
-   * @param partition the Partition to look up
-   * @return map of (participant id, state)
-   */
-  public Map<String, String> getReplicaMap(String partitionId) {
-    Map<String, String> rawReplicaMap = _record.getMapField(partitionId);
-    if (rawReplicaMap == null) {
-      return Collections.emptyMap();
-    }
-    return rawReplicaMap;
-  }
-
-  /**
    * Add participant, state pairs for a partition
    * @param partitionId the partition to set
    * @param replicaMap map of (participant name, state)
@@ -147,15 +124,6 @@ public class ResourceAssignment extends HelixProperty {
   }
 
   /**
-   * Add participant, state pairs for a partition
-   * @param partitionId the partition to set
-   * @param replicaMap map of (participant name, state)
-   */
-  public void addReplicaMap(String partitionId, Map<String, String> replicaMap) {
-    _record.setMapField(partitionId, replicaMap);
-  }
-
-  /**
    * Helper for converting a map of strings to a concrete replica map
    * @param rawMap map of participant name to state name
    * @return map of participant id to state

http://git-wip-us.apache.org/repos/asf/incubator-helix/blob/842035a4/helix-core/src/main/java/org/apache/helix/model/StateModelDefinition.java
----------------------------------------------------------------------
diff --git a/helix-core/src/main/java/org/apache/helix/model/StateModelDefinition.java b/helix-core/src/main/java/org/apache/helix/model/StateModelDefinition.java
index a9a6e49..2e759e6 100644
--- a/helix-core/src/main/java/org/apache/helix/model/StateModelDefinition.java
+++ b/helix-core/src/main/java/org/apache/helix/model/StateModelDefinition.java
@@ -153,7 +153,7 @@ public class StateModelDefinition extends HelixProperty {
    * Get an ordered priority list of transitions
    * @return transitions in the form SRC-DEST, the first of which is highest priority
    */
-  public List<String> getStateTransitionPriorityStringList() {
+  public List<String> getStateTransitionPriorityList() {
     return _stateTransitionPriorityList;
   }
 
@@ -161,9 +161,9 @@ public class StateModelDefinition extends HelixProperty {
    * Get an ordered priority list of transitions
    * @return Transition objects, the first of which is highest priority (immutable)
    */
-  public List<Transition> getStateTransitionPriorityList() {
+  public List<Transition> getTypedStateTransitionPriorityList() {
     ImmutableList.Builder<Transition> builder = new ImmutableList.Builder<Transition>();
-    for (String transition : getStateTransitionPriorityStringList()) {
+    for (String transition : getStateTransitionPriorityList()) {
       String fromState = transition.substring(0, transition.indexOf('-'));
       String toState = transition.substring(transition.indexOf('-') + 1);
       builder.add(Transition.from(State.from(fromState), State.from(toState)));
@@ -283,8 +283,8 @@ public class StateModelDefinition extends HelixProperty {
     Map<String, String> stateConstraintMap;
 
     /**
-     * Start building a state model with a name
-     * @param name state model name
+     * Start building a state model with a id
+     * @param stateModelDefId state model id
      */
     public Builder(StateModelDefId stateModelDefId) {
       this._statemodelName = stateModelDefId.stringify();
@@ -294,9 +294,17 @@ public class StateModelDefinition extends HelixProperty {
     }
 
     /**
+     * Start building a state model with a name
+     * @param stateModelDefId state model name
+     */
+    public Builder(String stateModelName) {
+      this(StateModelDefId.from(stateModelName));
+    }
+
+    /**
      * initial state of a replica when it starts, most commonly used initial
      * state is OFFLINE
-     * @param state
+     * @param initialState
      */
     public Builder initialState(State initialState) {
       return initialState(initialState.toString());
@@ -305,7 +313,7 @@ public class StateModelDefinition extends HelixProperty {
     /**
      * initial state of a replica when it starts, most commonly used initial
      * state is OFFLINE
-     * @param state
+     * @param initialState
      */
     public Builder initialState(String initialState) {
       this.initialState = initialState;
@@ -318,7 +326,8 @@ public class StateModelDefinition extends HelixProperty {
      * STATE2 has a constraint of 3 but only one node is up then Helix will uses
      * the priority to see STATE constraint has to be given higher preference <br/>
      * Use -1 to indicates states with no constraints, like OFFLINE
-     * @param states
+     * @param state the state to add
+     * @param priority the state priority, lower number is higher priority
      */
     public Builder addState(State state, int priority) {
       return addState(state.toString(), priority);
@@ -330,7 +339,8 @@ public class StateModelDefinition extends HelixProperty {
      * STATE2 has a constraint of 3 but only one node is up then Helix will uses
      * the priority to see STATE constraint has to be given higher preference <br/>
      * Use -1 to indicates states with no constraints, like OFFLINE
-     * @param states
+     * @param state the state to add
+     * @param priority the state priority, lower number is higher priority
      */
     public Builder addState(String state, int priority) {
       statesMap.put(state, priority);

http://git-wip-us.apache.org/repos/asf/incubator-helix/blob/842035a4/helix-core/src/main/java/org/apache/helix/model/builder/AutoModeISBuilder.java
----------------------------------------------------------------------
diff --git a/helix-core/src/main/java/org/apache/helix/model/builder/AutoModeISBuilder.java b/helix-core/src/main/java/org/apache/helix/model/builder/AutoModeISBuilder.java
index 72b2bc9..cda2f9e 100644
--- a/helix-core/src/main/java/org/apache/helix/model/builder/AutoModeISBuilder.java
+++ b/helix-core/src/main/java/org/apache/helix/model/builder/AutoModeISBuilder.java
@@ -22,24 +22,84 @@ package org.apache.helix.model.builder;
 import java.util.ArrayList;
 import java.util.Arrays;
 
+import org.apache.helix.api.id.ParticipantId;
+import org.apache.helix.api.id.PartitionId;
+import org.apache.helix.api.id.ResourceId;
 import org.apache.helix.model.IdealState.RebalanceMode;
 
+/**
+ * IdealState builder for SEMI_AUTO mode
+ */
 public class AutoModeISBuilder extends IdealStateBuilder {
+  /**
+   * Start building a SEMI_AUTO IdealState
+   * @param resourceName the resource
+   */
   public AutoModeISBuilder(String resourceName) {
     super(resourceName);
     setRebalancerMode(RebalanceMode.SEMI_AUTO);
   }
 
-  public void add(String partitionName) {
+  /**
+   * Start building a SEMI_AUTO IdealState
+   * @param resourceId the resource
+   */
+  public AutoModeISBuilder(ResourceId resourceId) {
+    this(resourceId.stringify());
+  }
+
+  /**
+   * Add a partition; Helix will assign replicas of the partition according to preference lists
+   * @param partitionName the name of the new partition
+   * @return AutoModeISBuilder
+   */
+  public AutoModeISBuilder add(String partitionName) {
     if (_record.getListField(partitionName) == null) {
       _record.setListField(partitionName, new ArrayList<String>());
     }
+    return this;
+  }
+
+  /**
+   * Add a partition; Helix will assign replicas of the partition according to preference lists
+   * @param partitionId the id of the new partition
+   * @return AutoModeISBuilder
+   */
+  public AutoModeISBuilder add(PartitionId partitionId) {
+    if (partitionId != null) {
+      add(partitionId.stringify());
+    }
+    return this;
   }
 
+  /**
+   * Define where replicas of a partition should live
+   * @param partitionName the partition
+   * @param instanceNames ordered list of participant names
+   * @return AutoModeISBuilder
+   */
   public AutoModeISBuilder assignPreferenceList(String partitionName, String... instanceNames) {
     add(partitionName);
     _record.getListField(partitionName).addAll(Arrays.asList(instanceNames));
     return this;
   }
 
+  /**
+   * Define where replicas of a partition should live
+   * @param partitionId the partition
+   * @param participantIds ordered list of participant ids
+   * @return AutoModeISBuilder
+   */
+  public AutoModeISBuilder assignPreferenceList(PartitionId partitionId,
+      ParticipantId... participantIds) {
+    if (partitionId != null) {
+      String[] participantNames = new String[participantIds.length];
+      for (int i = 0; i < participantIds.length; i++) {
+        participantNames[i] = participantIds[i].stringify();
+      }
+      assignPreferenceList(partitionId.stringify(), participantNames);
+    }
+    return this;
+  }
+
 }

http://git-wip-us.apache.org/repos/asf/incubator-helix/blob/842035a4/helix-core/src/main/java/org/apache/helix/model/builder/AutoRebalanceModeISBuilder.java
----------------------------------------------------------------------
diff --git a/helix-core/src/main/java/org/apache/helix/model/builder/AutoRebalanceModeISBuilder.java b/helix-core/src/main/java/org/apache/helix/model/builder/AutoRebalanceModeISBuilder.java
index bfb958d..8ac3b82 100644
--- a/helix-core/src/main/java/org/apache/helix/model/builder/AutoRebalanceModeISBuilder.java
+++ b/helix-core/src/main/java/org/apache/helix/model/builder/AutoRebalanceModeISBuilder.java
@@ -21,24 +21,54 @@ package org.apache.helix.model.builder;
 
 import java.util.ArrayList;
 
+import org.apache.helix.api.id.PartitionId;
+import org.apache.helix.api.id.ResourceId;
 import org.apache.helix.model.IdealState.RebalanceMode;
 
+/**
+ * IdealState builder for FULL_AUTO mode
+ */
 public class AutoRebalanceModeISBuilder extends IdealStateBuilder {
+  /**
+   * Start building a SEMI_AUTO IdealState
+   * @param resourceName the resource
+   */
   public AutoRebalanceModeISBuilder(String resourceName) {
     super(resourceName);
     setRebalancerMode(RebalanceMode.FULL_AUTO);
   }
 
   /**
+   * Start building a SEMI_AUTO IdealState
+   * @param resourceId the resource
+   */
+  public AutoRebalanceModeISBuilder(ResourceId resourceId) {
+    this(resourceId.stringify());
+  }
+
+  /**
    * Add a partition, Helix will automatically assign the placement and state
    * for this partition at runtime.
-   * @param partitionName
+   * @param partitionName the partition to add
+   * @return AutoRebalanceModeISBuilder
    */
   public AutoRebalanceModeISBuilder add(String partitionName) {
     if (_record.getListField(partitionName) == null) {
       _record.setListField(partitionName, new ArrayList<String>());
     }
+    return this;
+  }
 
+  /**
+   * Add a partition, Helix will automatically assign the placement and state
+   * for this partition at runtime.
+   * @param partitionId the partition to add
+   * @return AutoRebalanceModeISBuilder
+   */
+  public AutoRebalanceModeISBuilder add(PartitionId partitionId) {
+    if (partitionId != null) {
+      add(partitionId.stringify());
+    }
     return this;
   }
 

http://git-wip-us.apache.org/repos/asf/incubator-helix/blob/842035a4/helix-core/src/main/java/org/apache/helix/model/builder/ClusterConstraintsBuilder.java
----------------------------------------------------------------------
diff --git a/helix-core/src/main/java/org/apache/helix/model/builder/ClusterConstraintsBuilder.java b/helix-core/src/main/java/org/apache/helix/model/builder/ClusterConstraintsBuilder.java
index f329daa..13b2a7e 100644
--- a/helix-core/src/main/java/org/apache/helix/model/builder/ClusterConstraintsBuilder.java
+++ b/helix-core/src/main/java/org/apache/helix/model/builder/ClusterConstraintsBuilder.java
@@ -25,11 +25,8 @@ import java.util.Map;
 import org.apache.helix.api.id.ConstraintId;
 import org.apache.helix.model.ClusterConstraints;
 import org.apache.helix.model.ClusterConstraints.ConstraintType;
-import org.apache.log4j.Logger;
 
 public class ClusterConstraintsBuilder {
-  private static Logger LOG = Logger.getLogger(ClusterConstraintsBuilder.class);
-
   final private ConstraintType _constraintType;
 
   /**
@@ -61,6 +58,11 @@ public class ClusterConstraintsBuilder {
     return this;
   }
 
+  public ClusterConstraintsBuilder addConstraintAttribute(String constraintId, String attribute,
+      String value) {
+    return addConstraintAttribute(ConstraintId.from(constraintId), attribute, value);
+  }
+
   public ClusterConstraints build() {
     ClusterConstraints constraints = new ClusterConstraints(_constraintType);
 

http://git-wip-us.apache.org/repos/asf/incubator-helix/blob/842035a4/helix-core/src/main/java/org/apache/helix/model/builder/CurrentStateBuilder.java
----------------------------------------------------------------------
diff --git a/helix-core/src/main/java/org/apache/helix/model/builder/CurrentStateBuilder.java b/helix-core/src/main/java/org/apache/helix/model/builder/CurrentStateBuilder.java
index 0519979..fb6235f 100644
--- a/helix-core/src/main/java/org/apache/helix/model/builder/CurrentStateBuilder.java
+++ b/helix-core/src/main/java/org/apache/helix/model/builder/CurrentStateBuilder.java
@@ -9,6 +9,7 @@ import org.apache.helix.api.id.PartitionId;
 import org.apache.helix.api.id.ResourceId;
 import org.apache.helix.api.id.SessionId;
 import org.apache.helix.api.id.StateModelDefId;
+import org.apache.helix.api.id.StateModelFactoryId;
 import org.apache.helix.model.CurrentState;
 import org.apache.helix.model.CurrentState.CurrentStateProperty;
 
@@ -39,7 +40,7 @@ public class CurrentStateBuilder {
   private final Map<PartitionId, State> _partitionStateMap;
   private SessionId _sessionId;
   private StateModelDefId _stateModelDefId;
-  private String _stateModelFactoryName;
+  private StateModelFactoryId _stateModelFactoryId;
 
   /**
    * Build a current state for a given resource
@@ -93,11 +94,11 @@ public class CurrentStateBuilder {
 
   /**
    * Set the name of the state model factory
-   * @param stateModelFactoryName state model factory identifier
+   * @param stateModelFactoryIde state model factory identifier
    * @return CurrentStateBuilder
    */
-  public CurrentStateBuilder stateModelFactory(String stateModelFactoryName) {
-    _stateModelFactoryName = stateModelFactoryName;
+  public CurrentStateBuilder stateModelFactory(StateModelFactoryId stateModelFactoryId) {
+    _stateModelFactoryId = stateModelFactoryId;
     return this;
   }
 
@@ -117,7 +118,7 @@ public class CurrentStateBuilder {
     record.setSimpleField(CurrentStateProperty.STATE_MODEL_DEF.toString(),
         _stateModelDefId.toString());
     record.setSimpleField(CurrentStateProperty.STATE_MODEL_FACTORY_NAME.toString(),
-        _stateModelFactoryName);
+        _stateModelFactoryId.toString());
     return new CurrentState(record);
   }
 }

http://git-wip-us.apache.org/repos/asf/incubator-helix/blob/842035a4/helix-core/src/main/java/org/apache/helix/model/builder/CustomModeISBuilder.java
----------------------------------------------------------------------
diff --git a/helix-core/src/main/java/org/apache/helix/model/builder/CustomModeISBuilder.java b/helix-core/src/main/java/org/apache/helix/model/builder/CustomModeISBuilder.java
index 65c40a0..566452a 100644
--- a/helix-core/src/main/java/org/apache/helix/model/builder/CustomModeISBuilder.java
+++ b/helix-core/src/main/java/org/apache/helix/model/builder/CustomModeISBuilder.java
@@ -22,31 +22,63 @@ package org.apache.helix.model.builder;
 import java.util.Map;
 import java.util.TreeMap;
 
+import org.apache.helix.api.State;
+import org.apache.helix.api.id.ParticipantId;
+import org.apache.helix.api.id.PartitionId;
+import org.apache.helix.api.id.ResourceId;
 import org.apache.helix.model.IdealState.RebalanceMode;
 
+/**
+ * IdealState builder for CUSTOMIZED mode
+ */
 public class CustomModeISBuilder extends IdealStateBuilder {
-
+  /**
+   * Start building a CUSTOMIZED IdealState
+   * @param resourceName the resource
+   */
   public CustomModeISBuilder(String resourceName) {
     super(resourceName);
     setRebalancerMode(RebalanceMode.CUSTOMIZED);
   }
 
   /**
+   * Start building a SEMI_AUTO IdealState
+   * @param resourceId the resource
+   */
+  public CustomModeISBuilder(ResourceId resourceId) {
+    this(resourceId.stringify());
+  }
+
+  /**
    * Add a sub-resource
-   * @param partitionName
+   * @param partitionName partition to add
+   * @return CustomModeISBuilder
    */
-  public void add(String partitionName) {
+  public CustomModeISBuilder add(String partitionName) {
     if (_record.getMapField(partitionName) == null) {
       _record.setMapField(partitionName, new TreeMap<String, String>());
     }
+    return this;
+  }
+
+  /**
+   * Add a sub-resource
+   * @param partitionId partition to add
+   * @return CustomModeISBuilder
+   */
+  public CustomModeISBuilder add(PartitionId partitionId) {
+    if (partitionId != null) {
+      add(partitionId.stringify());
+    }
+    return this;
   }
 
   /**
    * add an instance->state assignment
-   * @param partitionName
-   * @param instanceName
-   * @param state
-   * @return
+   * @param partitionName partition to update
+   * @param instanceName participant name
+   * @param state state the replica should be in
+   * @return CustomModeISBuilder
    */
   public CustomModeISBuilder assignInstanceAndState(String partitionName, String instanceName,
       String state) {
@@ -56,4 +88,19 @@ public class CustomModeISBuilder extends IdealStateBuilder {
     return this;
   }
 
+  /**
+   * add an instance->state assignment
+   * @param partitionId partition to update
+   * @param participantId participant to assign to
+   * @param state state the replica should be in
+   * @return CustomModeISBuilder
+   */
+  public CustomModeISBuilder assignParticipantAndState(PartitionId partitionId,
+      ParticipantId participantId, State state) {
+    if (partitionId != null && participantId != null && state != null) {
+      assignInstanceAndState(partitionId.stringify(), participantId.stringify(), state.toString());
+    }
+    return this;
+  }
+
 }

http://git-wip-us.apache.org/repos/asf/incubator-helix/blob/842035a4/helix-core/src/main/java/org/apache/helix/model/builder/ExternalViewBuilder.java
----------------------------------------------------------------------
diff --git a/helix-core/src/main/java/org/apache/helix/model/builder/ExternalViewBuilder.java b/helix-core/src/main/java/org/apache/helix/model/builder/ExternalViewBuilder.java
deleted file mode 100644
index 5fe099f..0000000
--- a/helix-core/src/main/java/org/apache/helix/model/builder/ExternalViewBuilder.java
+++ /dev/null
@@ -1,24 +0,0 @@
-package org.apache.helix.model.builder;
-
-/*
- * 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.
- */
-
-public class ExternalViewBuilder {
-
-}

http://git-wip-us.apache.org/repos/asf/incubator-helix/blob/842035a4/helix-core/src/test/java/org/apache/helix/integration/TestCustomizedIdealStateRebalancer.java
----------------------------------------------------------------------
diff --git a/helix-core/src/test/java/org/apache/helix/integration/TestCustomizedIdealStateRebalancer.java b/helix-core/src/test/java/org/apache/helix/integration/TestCustomizedIdealStateRebalancer.java
index eba322d..6b79b61 100644
--- a/helix-core/src/test/java/org/apache/helix/integration/TestCustomizedIdealStateRebalancer.java
+++ b/helix-core/src/test/java/org/apache/helix/integration/TestCustomizedIdealStateRebalancer.java
@@ -32,7 +32,7 @@ import org.apache.helix.api.Cluster;
 import org.apache.helix.api.State;
 import org.apache.helix.api.id.ParticipantId;
 import org.apache.helix.api.id.PartitionId;
-import org.apache.helix.controller.rebalancer.Rebalancer;
+import org.apache.helix.controller.rebalancer.HelixRebalancer;
 import org.apache.helix.controller.rebalancer.context.PartitionedRebalancerContext;
 import org.apache.helix.controller.rebalancer.context.RebalancerConfig;
 import org.apache.helix.controller.stages.ResourceCurrentState;
@@ -58,7 +58,7 @@ public class TestCustomizedIdealStateRebalancer extends
   static boolean testRebalancerCreated = false;
   static boolean testRebalancerInvoked = false;
 
-  public static class TestRebalancer implements Rebalancer {
+  public static class TestRebalancer implements HelixRebalancer {
 
     /**
      * Very basic mapping that evenly assigns one replica of each partition to live nodes, each of