Helix gateway service for Helix 2.0

.github/workflows/Helix-PR-CI.yml

@@ -1,7 +1,7 @@
 name: Helix PR CI
 on:
   pull_request:
-    branches: [ master, metaclient, ApplicationClusterManager] # TODO: remove side branch
+    branches: [ master, metaclient, ApplicationClusterManager, helix-gateway-service] # TODO: remove side branch
     paths-ignore:
       - '.github/**'
       - 'helix-front/**'

bump-snapshot.sh

@@ -45,6 +45,7 @@ mv metrics-common/metrics-common-$version-SNAPSHOT.ivy metrics-common/metrics-co
 mv zookeeper-api/zookeeper-api-$version-SNAPSHOT.ivy zookeeper-api/zookeeper-api-$new_version-SNAPSHOT.ivy
 mv helix-view-aggregator/helix-view-aggregator-$version-SNAPSHOT.ivy helix-view-aggregator/helix-view-aggregator-$new_version-SNAPSHOT.ivy
 mv meta-client/meta-client-$version-SNAPSHOT.ivy meta-client/meta-client-$new_version-SNAPSHOT.ivy
+mv helix-gateway/helix-gateway-$version-SNAPSHOT.ivy helix-gateway/helix-gateway-$new_version-SNAPSHOT.ivy
 
 
 find . -type f -name '*.ivy' -exec sed -i "s/$version/$new_version/g" {} \;

bump-up.sh

@@ -74,7 +74,7 @@ echo "bump up: $current_version -> $new_version"
 update_pom_version "pom.xml" $current_version
 
 for module in "metrics-common" "metadata-store-directory-common" "zookeeper-api" "helix-common" "helix-core" \
-              "helix-rest" "helix-lock" "helix-view-aggregator" "helix-agent" "meta-client"; do
+              "helix-rest" "helix-lock" "helix-view-aggregator" "helix-agent" "meta-client" "helix-gateway"; do
   update_ivy $module
   update_pom_version $module/pom.xml $current_version
 done

helix-core/src/main/java/org/apache/helix/controller/common/CapacityNode.java

@@ -21,23 +21,55 @@
 
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.LinkedHashMap;
 import java.util.Map;
 import java.util.Set;
 
+import org.apache.helix.controller.rebalancer.topology.Topology;
+import org.apache.helix.model.ClusterConfig;
+import org.apache.helix.model.ClusterTopologyConfig;
+import org.apache.helix.model.InstanceConfig;
+
 /**
  * A Node is an entity that can serve capacity recording purpose. It has a capacity and knowledge
  * of partitions assigned to it, so it can decide if it can receive additional partitions.
  */
-public class CapacityNode {
+public class CapacityNode implements Comparable<CapacityNode> {
   private int _currentlyAssigned;
   private int _capacity;
-  private final String _id;
+  private final String _instanceName;
+  private final String _logicaId;
+  private final String _faultZone;
   private final Map<String, Set<String>> _partitionMap;
 
-  public CapacityNode(String id) {
-    _partitionMap = new HashMap<>();
-    _currentlyAssigned = 0;
-    this._id = id;
+  /**
+   * Constructor used for non-topology-aware use case
+   * @param instanceName  The instance name of this node
+   * @param capacity  The capacity of this node
+   */
+  public CapacityNode(String instanceName, int capacity) {
+    this(instanceName, null, null, null);
+    this._capacity = capacity;
+  }
+
+  /**
+   * Constructor used for topology-aware use case
+   * @param instanceName  The instance name of this node
+   * @param clusterConfig  The cluster config for current helix cluster
+   * @param clusterTopologyConfig  The cluster topology config for current helix cluster
+   * @param instanceConfig  The instance config for current instance
+   */
+  public CapacityNode(String instanceName, ClusterConfig clusterConfig,
+      ClusterTopologyConfig clusterTopologyConfig, InstanceConfig instanceConfig) {
+    this._instanceName = instanceName;
+    this._logicaId = clusterTopologyConfig != null ? instanceConfig.getLogicalId(
+        clusterTopologyConfig.getEndNodeType()) : instanceName;
+    this._faultZone =
+        clusterConfig != null ? computeFaultZone(clusterConfig, instanceConfig) : null;
+    this._partitionMap = new HashMap<>();
+    this._capacity =
+        clusterConfig != null ? clusterConfig.getGlobalMaxPartitionAllowedPerInstance() : 0;
+    this._currentlyAssigned = 0;
   }
 
   /**
@@ -52,11 +84,25 @@ public boolean canAdd(String resource, String partition) {
         && _partitionMap.get(resource).contains(partition))) {
       return false;
     }
+
+    // Add the partition to the resource's set of partitions in this node
     _partitionMap.computeIfAbsent(resource, k -> new HashSet<>()).add(partition);
     _currentlyAssigned++;
     return true;
   }
 
+  /**
+   * Checks if a specific resource + partition is assigned to this node.
+   *
+   * @param resource  the name of the resource
+   * @param partition the partition
+   * @return {@code true} if the resource + partition is assigned to this node, {@code false} otherwise
+   */
+  public boolean hasPartition(String resource, String partition) {
+    Set<String> partitions = _partitionMap.get(resource);
+    return partitions != null && partitions.contains(partition);
+  }
+
   /**
    * Set the capacity of this node
    * @param capacity  The capacity to set
@@ -66,11 +112,27 @@ public void setCapacity(int capacity) {
   }
 
   /**
-   * Get the ID of this node
-   * @return The ID of this node
+   * Get the instance name of this node
+   * @return The instance name of this node
    */
-  public String getId() {
-    return _id;
+  public String getInstanceName() {
+    return _instanceName;
+  }
+
+  /**
+   * Get the logical id of this node
+   * @return The logical id of this node
+   */
+  public String getLogicalId() {
+    return _logicaId;
+  }
+
+  /**
+   * Get the fault zone of this node
+   * @return The fault zone of this node
+   */
+  public String getFaultZone() {
+    return _faultZone;
   }
 
   /**
@@ -84,8 +146,40 @@ public int getCurrentlyAssigned() {
   @Override
   public String toString() {
     StringBuilder sb = new StringBuilder();
-    sb.append("##########\nname=").append(_id).append("\nassigned:").append(_currentlyAssigned)
-        .append("\ncapacity:").append(_capacity);
+    sb.append("##########\nname=").append(_instanceName).append("\nassigned:")
+        .append(_currentlyAssigned).append("\ncapacity:").append(_capacity).append("\nlogicalId:")
+        .append(_logicaId).append("\nfaultZone:").append(_faultZone);
     return sb.toString();
   }
+
+  @Override
+  public int compareTo(CapacityNode o) {
+    if (_logicaId != null) {
+      return _logicaId.compareTo(o.getLogicalId());
+    }
+    return _instanceName.compareTo(o.getInstanceName());
+  }
+
+  /**
+   * Computes the fault zone id based on the domain and fault zone type when topology is enabled.
+   * For example, when
+   * the domain is "zone=2, instance=testInstance" and the fault zone type is "zone", this function
+   * returns "2".
+   * If cannot find the fault zone type, this function leaves the fault zone id as the instance name.
+   * TODO: change the return value to logical id when no fault zone type found. Also do the same for
+   *  waged rebalancer in helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/model/AssignableNode.java
+   */
+  private String computeFaultZone(ClusterConfig clusterConfig, InstanceConfig instanceConfig) {
+    LinkedHashMap<String, String> instanceTopologyMap =
+        Topology.computeInstanceTopologyMap(clusterConfig, instanceConfig.getInstanceName(),
+            instanceConfig, true /*earlyQuitTillFaultZone*/);
+
+    StringBuilder faultZoneStringBuilder = new StringBuilder();
+    for (Map.Entry<String, String> entry : instanceTopologyMap.entrySet()) {
+      faultZoneStringBuilder.append(entry.getValue());
+      faultZoneStringBuilder.append('/');
+    }
+    faultZoneStringBuilder.setLength(faultZoneStringBuilder.length() - 1);
+    return faultZoneStringBuilder.toString();
+  }
 }

helix-core/src/main/java/org/apache/helix/controller/dataproviders/ResourceControllerDataProvider.java

@@ -41,16 +41,20 @@
 import org.apache.helix.controller.LogUtil;
 import org.apache.helix.controller.common.CapacityNode;
 import org.apache.helix.controller.pipeline.Pipeline;
-import org.apache.helix.controller.rebalancer.strategy.GreedyRebalanceStrategy;
+import org.apache.helix.controller.rebalancer.strategy.StickyRebalanceStrategy;
 import org.apache.helix.controller.rebalancer.waged.WagedInstanceCapacity;
 import org.apache.helix.controller.rebalancer.waged.WagedResourceWeightsProvider;
+import org.apache.helix.controller.stages.CurrentStateOutput;
 import org.apache.helix.controller.stages.MissingTopStateRecord;
+import org.apache.helix.model.ClusterConfig;
+import org.apache.helix.model.ClusterTopologyConfig;
 import org.apache.helix.model.CustomizedState;
 import org.apache.helix.model.CustomizedStateConfig;
 import org.apache.helix.model.CustomizedView;
 import org.apache.helix.model.ExternalView;
 import org.apache.helix.model.IdealState;
 import org.apache.helix.model.InstanceConfig;
+import org.apache.helix.model.Partition;
 import org.apache.helix.model.ResourceAssignment;
 import org.apache.helix.zookeeper.datamodel.ZNRecord;
 import org.slf4j.Logger;
@@ -73,6 +77,9 @@ public class ResourceControllerDataProvider extends BaseControllerDataProvider {
   // a map from customized state type to customized view cache
   private final Map<String, CustomizedViewCache> _customizedViewCacheMap;
 
+  // maintain a cache of ideal state (preference list + best possible assignment) which will be managed ondemand in rebalancer
+  private final Map<String, ZNRecord> _ondemandIdealStateCache;
+
   // maintain a cache of bestPossible assignment across pipeline runs
   // TODO: this is only for customRebalancer, remove it and merge it with _idealMappingCache.
   private Map<String, ResourceAssignment> _resourceAssignmentCache;
@@ -149,6 +156,7 @@ public String getObjName(ExternalView obj) {
     _refreshedChangeTypes = ConcurrentHashMap.newKeySet();
     _customizedStateCache = new CustomizedStateCache(this, _aggregationEnabledTypes);
     _customizedViewCacheMap = new HashMap<>();
+    _ondemandIdealStateCache = new HashMap<>();
   }
 
   public synchronized void refresh(HelixDataAccessor accessor) {
@@ -182,16 +190,17 @@ public synchronized void refresh(HelixDataAccessor accessor) {
     refreshStablePartitionList(getIdealStates());
     refreshDisabledInstancesForAllPartitionsSet();
 
-    if (getClusterConfig().getGlobalMaxPartitionAllowedPerInstance() != -1) {
-      buildSimpleCapacityMap(getClusterConfig().getGlobalMaxPartitionAllowedPerInstance());
+    if (getClusterConfig() != null
+        && getClusterConfig().getGlobalMaxPartitionAllowedPerInstance() != -1) {
+      buildSimpleCapacityMap();
       // Remove all cached IdealState because it is a global computation cannot partially be
       // performed for some resources. The computation is simple as well not taking too much resource
       // to recompute the assignments.
-      Set<String> cachedGreedyIdealStates = _idealMappingCache.values().stream().filter(
+      Set<String> cachedStickyIdealStates = _idealMappingCache.values().stream().filter(
               record -> record.getSimpleField(IdealState.IdealStateProperty.REBALANCE_STRATEGY.name())
-                  .equals(GreedyRebalanceStrategy.class.getName())).map(ZNRecord::getId)
+                  .equals(StickyRebalanceStrategy.class.getName())).map(ZNRecord::getId)
           .collect(Collectors.toSet());
-      _idealMappingCache.keySet().removeAll(cachedGreedyIdealStates);
+      _idealMappingCache.keySet().removeAll(cachedStickyIdealStates);
     }
 
     LogUtil.logInfo(logger, getClusterEventId(), String.format(
@@ -388,6 +397,28 @@ public Map<String, Map<String, String>> getLastTopStateLocationMap() {
     return _lastTopStateLocationMap;
   }
 
+  /**
+   * Get cached ideal state (preference list + best possible assignment) for a resource
+   * @param resource
+   * @return
+   */
+  public ZNRecord getCachedOndemandIdealState(String resource) {
+    return _ondemandIdealStateCache.get(resource);
+  }
+
+  /**
+   * Cache ideal state (preference list + best possible assignment) for a resource
+   * @param resource
+   * @return
+   */
+  public void setCachedOndemandIdealState(String resource, ZNRecord idealState) {
+    _ondemandIdealStateCache.put(resource, idealState);
+  }
+
+  public void clearCachedOndemandIdealStates() {
+    _ondemandIdealStateCache.clear();
+  }
+
   /**
    * Get cached resourceAssignment (bestPossible mapping) for a resource
    * @param resource
@@ -552,11 +583,16 @@ public WagedInstanceCapacity getWagedInstanceCapacity() {
     return _wagedInstanceCapacity;
   }
 
-  private void buildSimpleCapacityMap(int globalMaxPartitionAllowedPerInstance) {
+  private void buildSimpleCapacityMap() {
+    ClusterConfig clusterConfig = getClusterConfig();
+    ClusterTopologyConfig clusterTopologyConfig =
+        ClusterTopologyConfig.createFromClusterConfig(clusterConfig);
+    Map<String, InstanceConfig> instanceConfigMap = getAssignableInstanceConfigMap();
     _simpleCapacitySet = new HashSet<>();
-    for (String instance : getEnabledLiveInstances()) {
-      CapacityNode capacityNode = new CapacityNode(instance);
-      capacityNode.setCapacity(globalMaxPartitionAllowedPerInstance);
+    for (String instanceName : getAssignableInstances()) {
+      CapacityNode capacityNode =
+          new CapacityNode(instanceName, clusterConfig, clusterTopologyConfig,
+              instanceConfigMap.getOrDefault(instanceName, new InstanceConfig(instanceName)));
       _simpleCapacitySet.add(capacityNode);
     }
   }
@@ -565,6 +601,35 @@ public Set<CapacityNode> getSimpleCapacitySet() {
     return _simpleCapacitySet;
   }
 
+  public void populateSimpleCapacitySetUsage(final Set<String> resourceNameSet,
+      final CurrentStateOutput currentStateOutput) {
+    // Convert the assignableNodes to map for quick lookup
+    Map<String, CapacityNode> simpleCapacityMap = new HashMap<>();
+    for (CapacityNode node : _simpleCapacitySet) {
+      simpleCapacityMap.put(node.getInstanceName(), node);
+    }
+    for (String resourceName : resourceNameSet) {
+      // Process current state mapping
+      populateCapacityNodeUsageFromStateMap(resourceName, simpleCapacityMap,
+          currentStateOutput.getCurrentStateMap(resourceName));
+      // Process pending state mapping
+      populateCapacityNodeUsageFromStateMap(resourceName, simpleCapacityMap,
+          currentStateOutput.getPendingMessageMap(resourceName));
+    }
+  }
+
+  private <T> void populateCapacityNodeUsageFromStateMap(String resourceName,
+      Map<String, CapacityNode> simpleCapacityMap, Map<Partition, Map<String, T>> stateMap) {
+    for (Map.Entry<Partition, Map<String, T>> entry : stateMap.entrySet()) {
+      for (String instanceName : entry.getValue().keySet()) {
+        CapacityNode node = simpleCapacityMap.get(instanceName);
+        if (node != null) {
+          node.canAdd(resourceName, entry.getKey().getPartitionName());
+        }
+      }
+    }
+  }
+
   private void refreshDisabledInstancesForAllPartitionsSet() {
     _disabledInstancesForAllPartitionsSet.clear();
     Collection<InstanceConfig> allConfigs = getInstanceConfigMap().values();