import org.apache.helix.model.IdealState; //导入依赖的package包/类
/**
 * From IdealStates.
 *
 * @return InstanceToNumTopicPartitionMap
 */
public static Map<String, Set<TopicPartition>> getInstanceToTopicPartitionsMap(
    HelixManager helixManager) {
  Map<String, Set<TopicPartition>> instanceToNumTopicPartitionMap =
      new HashMap<String, Set<TopicPartition>>();
  HelixAdmin helixAdmin = helixManager.getClusterManagmentTool();
  String helixClusterName = helixManager.getClusterName();
  for (String topic : helixAdmin.getResourcesInCluster(helixClusterName)) {
    IdealState is = helixAdmin.getResourceIdealState(helixClusterName, topic);
    for (String partition : is.getPartitionSet()) {
      TopicPartition tpi = new TopicPartition(topic, Integer.parseInt(partition));
      for (String instance : is.getInstanceSet(partition)) {
        if (!instanceToNumTopicPartitionMap.containsKey(instance)) {
          instanceToNumTopicPartitionMap.put(instance, new HashSet<TopicPartition>());
        }
        instanceToNumTopicPartitionMap.get(instance).add(tpi);
      }
    }
  }
  return instanceToNumTopicPartitionMap;
}
 

import org.apache.helix.model.IdealState; //导入依赖的package包/类
public static IdealState buildCustomIdealStateFor(String topicName,
    int numTopicPartitions,
    PriorityQueue<InstanceTopicPartitionHolder> instanceToNumServingTopicPartitionMap) {

  final CustomModeISBuilder customModeIdealStateBuilder = new CustomModeISBuilder(topicName);

  customModeIdealStateBuilder
      .setStateModel(OnlineOfflineStateModel.name)
      .setNumPartitions(numTopicPartitions).setNumReplica(1)
      .setMaxPartitionsPerNode(numTopicPartitions);

  for (int i = 0; i < numTopicPartitions; ++i) {
    synchronized (instanceToNumServingTopicPartitionMap) {
      InstanceTopicPartitionHolder liveInstance = instanceToNumServingTopicPartitionMap.poll();
      customModeIdealStateBuilder.assignInstanceAndState(Integer.toString(i),
          liveInstance.getInstanceName(), "ONLINE");
      liveInstance.addTopicPartition(new TopicPartition(topicName, i));
      instanceToNumServingTopicPartitionMap.add(liveInstance);
    }
  }
  return customModeIdealStateBuilder.build();
}
 

import org.apache.helix.model.IdealState; //导入依赖的package包/类
public static Set<TopicPartition> getUnassignedPartitions(HelixManager helixManager) {
  Set<TopicPartition> unassignedPartitions = new HashSet<TopicPartition>();
  HelixAdmin helixAdmin = helixManager.getClusterManagmentTool();
  String helixClusterName = helixManager.getClusterName();
  for (String topic : helixAdmin.getResourcesInCluster(helixClusterName)) {
    IdealState is = helixAdmin.getResourceIdealState(helixClusterName, topic);
    int numPartitions = is.getNumPartitions();
    for (int partition = 0; partition < numPartitions; ++partition) {
      if (is.getInstanceSet(Integer.toString(partition)).isEmpty()) {
        TopicPartition tpi = new TopicPartition(topic, partition);
        unassignedPartitions.add(tpi);
      }
    }
  }
  return unassignedPartitions;
}
 

import org.apache.helix.model.IdealState; //导入依赖的package包/类
public static IdealState buildCustomIdealStateFor(String topicName,
    int numTopicPartitions,
    PriorityQueue<InstanceTopicPartitionHolder> instanceToNumServingTopicPartitionMap) {

  final CustomModeISBuilder customModeIdealStateBuilder = new CustomModeISBuilder(topicName);

  customModeIdealStateBuilder
      .setStateModel(OnlineOfflineStateModel.name)
      .setNumPartitions(numTopicPartitions).setNumReplica(1)
      .setMaxPartitionsPerNode(numTopicPartitions);

  for (int i = 0; i < numTopicPartitions; ++i) {
    InstanceTopicPartitionHolder liveInstance = instanceToNumServingTopicPartitionMap.poll();
    if (liveInstance != null) {
      customModeIdealStateBuilder.assignInstanceAndState(Integer.toString(i),
          liveInstance.getInstanceName(), "ONLINE");
      liveInstance.addTopicPartition(new TopicPartition(topicName, i));
      instanceToNumServingTopicPartitionMap.add(liveInstance);
    }
  }
  return customModeIdealStateBuilder.build();
}
 

import org.apache.helix.model.IdealState; //导入依赖的package包/类
/**
 * From IdealStates.
 * @param helixManager
 * @return InstanceToNumTopicPartitionMap
 */
public static Map<String, Set<TopicPartition>> getInstanceToTopicPartitionsMap(
    HelixManager helixManager) {
  Map<String, Set<TopicPartition>> instanceToNumTopicPartitionMap =
      new HashMap<String, Set<TopicPartition>>();
  HelixAdmin helixAdmin = helixManager.getClusterManagmentTool();
  String helixClusterName = helixManager.getClusterName();
  for (String topic : helixAdmin.getResourcesInCluster(helixClusterName)) {
    IdealState is = helixAdmin.getResourceIdealState(helixClusterName, topic);
    for (String partition : is.getPartitionSet()) {
      TopicPartition tpi = new TopicPartition(topic, Integer.parseInt(partition));
      for (String instance : is.getInstanceSet(partition)) {
        if (!instanceToNumTopicPartitionMap.containsKey(instance)) {
          instanceToNumTopicPartitionMap.put(instance, new HashSet<TopicPartition>());
        }
        instanceToNumTopicPartitionMap.get(instance).add(tpi);
      }
    }
  }
  return instanceToNumTopicPartitionMap;
}
 

import org.apache.helix.model.IdealState; //导入依赖的package包/类
static double calculateDeviationForResource(String resource,
                                            IdealState idealState,
                                            RoutingTableProvider routingTableProvider) {
  Set<String> partitions = idealState.getPartitionSet();
  int totalAssignments = 0, totalDeviations = 0;
  // Check if the external view has deviated from the actual view.
  for (String partition : partitions) {
    // Make a copy of the instance mapping in the ideal state.
    Set<String> idealInstances = new HashSet(idealState.getInstanceSet(partition));
    totalAssignments += idealInstances.size();
    // Now check against our real state and count the amount of deviating
    // assignments.
    List<InstanceConfig> currentInstanceConfigs = routingTableProvider.getInstances(
        resource, partition, "ONLINE");
    Set<String> currentInstances = Sets.newHashSetWithExpectedSize(
        currentInstanceConfigs.size());
    if (currentInstanceConfigs != null) {
      for (InstanceConfig instanceConfig : currentInstanceConfigs) {
        currentInstances.add(instanceConfig.getHostName());
      }
    }
    idealInstances.removeAll(currentInstances);
    totalDeviations += idealInstances.size();
  }
  return (double)totalDeviations / totalAssignments;
}
 

import org.apache.helix.model.IdealState; //导入依赖的package包/类
private void rebalanceResource(String hdfsDir, String resource, FileSetInfo fileSetInfo)
    throws Exception {
  IdealState idealState = ControllerUtil.buildIdealStateForHdfsDir(
      hdfsClient,
      hdfsDir,
      resource,
      fileSetInfo.servingInfo.partitionerType,
      configuration.getInt(Constants.NUM_SERVING_REPLICAS, 3),
      configuration.getBoolean(Constants.ENABLE_ZK_COMPRESSION,
          Constants.ENABLE_ZK_COMPRESSION_DEFAULT));

  double deviation = calculateDeviationForResource(resource, idealState, routingTableProvider);
  if (deviation > configuration.getDouble(Constants.REBALANCE_DEVIATION_THRESHOLD, 0.0)) {
    // Write the new ideal state.
    LOG.info("Writing new ideal state for " + resource);
    helixAdmin.setResourceIdealState(clusterName, resource, idealState);
  } else {
    LOG.info("Resource " + resource + " is balanced. Skipping.");
  }
}
 

import org.apache.helix.model.IdealState; //导入依赖的package包/类
@Override
public boolean matches(Object o) {
  if (o == null || !(o instanceof IdealState)) {
    return false;
  }
  IdealState is = (IdealState)o;
  if (is.getRebalanceMode() != IdealState.RebalanceMode.CUSTOMIZED ||
      !is.getReplicas().equals("3") ||
      is.getNumPartitions() != this.numPartitions ||
      !is.getStateModelDefRef().equals("OnlineOffline")) {
    return false;
  }
  for (Map.Entry<Integer, List<String>> entry : partitionHostMap.entrySet()) {
    Map<String, String> stateMap = is.getInstanceStateMap(
            this.resource + "$" + entry.getKey());
    if (stateMap.size() != entry.getValue().size()) {
      return false;
    }
    for (String host : entry.getValue()) {
      if (!(stateMap.containsKey(host) && stateMap.get(host).equals("ONLINE"))) {
        return false;
      }
    }
  }
  return true;
}
 

import org.apache.helix.model.IdealState; //导入依赖的package包/类
/**
 * TODO : Due to the issue of helix, if remove the segment from idealState directly, the node won't get transaction,
 * so we first disable the segment, then remove it from idealState, then enable the segment. This will trigger transactions.
 * After the helix patch, we can refine the logic here by directly drop the segment from idealState.
 *
 * @param tableName
 * @param segmentId
 * @return
 */
public synchronized PinotResourceManagerResponse deleteSegment(final String tableName, final String segmentId) {
  LOGGER.info("Trying to delete segment: {} for table: {} ", segmentId, tableName);
  final PinotResourceManagerResponse res = new PinotResourceManagerResponse();
  try {
    IdealState idealState = _helixAdmin.getResourceIdealState(_helixClusterName, tableName);
    if (idealState.getPartitionSet().contains(segmentId)) {
      LOGGER.info("Trying to delete segment: {} from IdealStates", segmentId);
      HelixHelper.removeSegmentFromIdealState(_helixZkManager, tableName, segmentId);
    } else {
      res.message = "Segment " + segmentId + " not in IDEALSTATE.";
      LOGGER.info("Segment: {} is not in IDEALSTATE", segmentId);
    }
    _segmentDeletionManager.deleteSegment(tableName, segmentId);

    res.message = "Segment successfully deleted.";
    res.status = ResponseStatus.success;
  } catch (final Exception e) {
    LOGGER.error("Caught exception while deleting segment", e);
    res.status = ResponseStatus.failure;
    res.message = e.getMessage();
  }
  return res;
}
 

import org.apache.helix.model.IdealState; //导入依赖的package包/类
public boolean isServerTenantDeletable(String tenantName) {
  Set<String> taggedInstances =
      new HashSet<String>(_helixAdmin.getInstancesInClusterWithTag(_helixClusterName,
          ControllerTenantNameBuilder.getOfflineTenantNameForTenant(tenantName)));
  taggedInstances.addAll(_helixAdmin.getInstancesInClusterWithTag(_helixClusterName,
      ControllerTenantNameBuilder.getRealtimeTenantNameForTenant(tenantName)));
  for (String tableName : getAllTableNames()) {
    if (tableName.equals(CommonConstants.Helix.BROKER_RESOURCE_INSTANCE)) {
      continue;
    }
    IdealState tableIdealState = _helixAdmin.getResourceIdealState(_helixClusterName, tableName);
    for (String partition : tableIdealState.getPartitionSet()) {
      for (String instance : tableIdealState.getInstanceSet(partition)) {
        if (taggedInstances.contains(instance)) {
          return false;
        }
      }
    }
  }
  return true;
}
 

import org.apache.helix.model.IdealState; //导入依赖的package包/类
private void handleBrokerResource(AbstractTableConfig tableConfig) {
  try {
    String brokerTenant =
        ControllerTenantNameBuilder.getBrokerTenantNameForTenant(tableConfig.getTenantConfig().getBroker());
    if (_helixAdmin.getInstancesInClusterWithTag(_helixClusterName, brokerTenant).isEmpty()) {
      throw new RuntimeException("broker tenant : " + tableConfig.getTenantConfig().getBroker() + " is not existed!");
    }
    LOGGER.info("Trying to update BrokerDataResource IdealState!");
    final IdealState idealState =
        _helixAdmin.getResourceIdealState(_helixClusterName, CommonConstants.Helix.BROKER_RESOURCE_INSTANCE);
    String tableName = tableConfig.getTableName();
    for (String instanceName : _helixAdmin.getInstancesInClusterWithTag(_helixClusterName, brokerTenant)) {
      idealState.setPartitionState(tableName, instanceName, BrokerOnlineOfflineStateModel.ONLINE);
    }
    if (idealState != null) {
      _helixAdmin
          .setResourceIdealState(_helixClusterName, CommonConstants.Helix.BROKER_RESOURCE_INSTANCE, idealState);
    }
  } catch (final Exception e) {
    LOGGER.warn("Caught exception while creating broker", e);
  }
}
 

import org.apache.helix.model.IdealState; //导入依赖的package包/类
public Map<String, List<String>> getInstanceToSegmentsInATableMap(String tableName) {
  Map<String, List<String>> instancesToSegmentsMap = new HashMap<String, List<String>>();
  IdealState is = _helixAdmin.getResourceIdealState(_helixClusterName, tableName);
  Set<String> segments = is.getPartitionSet();

  for (String segment : segments) {
    Set<String> instances = is.getInstanceSet(segment);
    for (String instance : instances) {
      if (instancesToSegmentsMap.containsKey(instance)) {
        instancesToSegmentsMap.get(instance).add(segment);
      } else {
        List<String> a = new ArrayList<String>();
        a.add(segment);
        instancesToSegmentsMap.put(instance, a);
      }
    }
  }

  return instancesToSegmentsMap;
}
 

import org.apache.helix.model.IdealState; //导入依赖的package包/类
public static IdealState buildInitialRealtimeIdealStateFor(String realtimeTableName,
    AbstractTableConfig realtimeTableConfig, HelixAdmin helixAdmin, String helixClusterName,
    ZkHelixPropertyStore<ZNRecord> zkHelixPropertyStore) {
  KafkaStreamMetadata kafkaStreamMetadata =
      new KafkaStreamMetadata(realtimeTableConfig.getIndexingConfig().getStreamConfigs());
  String realtimeServerTenant =
      ControllerTenantNameBuilder.getRealtimeTenantNameForTenant(realtimeTableConfig.getTenantConfig().getServer());
  switch (kafkaStreamMetadata.getConsumerType()) {
    case highLevel:
      IdealState idealState =
          buildInitialKafkaHighLevelConsumerRealtimeIdealStateFor(realtimeTableName, helixAdmin, helixClusterName,
              zkHelixPropertyStore);
      List<String> realtimeInstances = helixAdmin.getInstancesInClusterWithTag(helixClusterName, realtimeServerTenant);
      if (realtimeInstances.size() % Integer.parseInt(realtimeTableConfig.getValidationConfig().getReplication()) != 0) {
        throw new RuntimeException("Number of instance in current tenant should be an integer multiples of the number of replications");
      }
      setupInstanceConfigForKafkaHighLevelConsumer(realtimeTableName, realtimeInstances.size(),
          Integer.parseInt(realtimeTableConfig.getValidationConfig().getReplication()), realtimeTableConfig
              .getIndexingConfig().getStreamConfigs(), zkHelixPropertyStore, realtimeInstances);
      return idealState;
    case simple:
    default:
      throw new UnsupportedOperationException("Not support kafka consumer type: "
          + kafkaStreamMetadata.getConsumerType());
  }
}
 

import org.apache.helix.model.IdealState; //导入依赖的package包/类
/**
 * Remove a resource (offline/realtime table) from the Broker's ideal state.
 *
 * @param helixManager The HelixManager object for accessing helix cluster.
 * @param resourceTag Name of the resource that needs to be removed from Broker ideal state.
 */
public static void removeResourceFromBrokerIdealState(HelixManager helixManager, final String resourceTag) {
  Function<IdealState, IdealState> updater = new Function<IdealState, IdealState>() {
    @Override
    public IdealState apply(IdealState idealState) {
      if (idealState.getPartitionSet().contains(resourceTag)) {
        idealState.getPartitionSet().remove(resourceTag);
        return idealState;
      } else {
        return null;
      }
    }
  };

  // Removing partitions from ideal state
  LOGGER.info("Trying to remove resource from idealstats");
  HelixHelper.updateIdealState(helixManager, CommonConstants.Helix.BROKER_RESOURCE_INSTANCE, updater,
      DEFAULT_RETRY_POLICY);
}
 

import org.apache.helix.model.IdealState; //导入依赖的package包/类
/**
 * Remove the segment from the cluster.
 *
 * @param helixManager The HelixManager object to access the helix cluster.
 * @param tableName Name of the table to which the new segment is to be added.
 * @param segmentName Name of the new segment to be added
 */
public static void removeSegmentFromIdealState(HelixManager helixManager, String tableName, final String segmentName) {
  Function<IdealState, IdealState> updater = new Function<IdealState, IdealState>() {
    @Override
    public IdealState apply(IdealState idealState) {
      final Set<String> currentInstanceSet = idealState.getInstanceSet(segmentName);

      if (!currentInstanceSet.isEmpty() && idealState.getPartitionSet().contains(segmentName)) {
        idealState.getPartitionSet().remove(segmentName);
        return idealState;
      } else {
        return null;
      }
    }
  };

  updateIdealState(helixManager, tableName, updater, DEFAULT_RETRY_POLICY);
}
 

import org.apache.helix.model.IdealState; //导入依赖的package包/类
/**
 * Add the new specified segment to the idealState of the specified table in the specified cluster.
 *
 * @param helixManager The HelixManager object to access the helix cluster.
 * @param tableName Name of the table to which the new segment is to be added.
 * @param segmentName Name of the new segment to be added
 * @param getInstancesForSegment Callable returning list of instances where the segment should be uploaded.
 */
public static void addSegmentToIdealState(HelixManager helixManager, String tableName, final String segmentName,
    final Callable<List<String>> getInstancesForSegment) {

  Function<IdealState, IdealState> updater = new Function<IdealState, IdealState>() {
    @Override
    public IdealState apply(IdealState idealState) {
      List<String> targetInstances = null;
      try {
        targetInstances = getInstancesForSegment.call();
      } catch (Exception e) {
        LOGGER.error("Unable to get new instances for segment uploading.");
        return null;
      }

      for (final String instance : targetInstances) {
        idealState.setPartitionState(segmentName, instance, ONLINE);
      }

      idealState.setNumPartitions(idealState.getNumPartitions() + 1);
      return idealState;
    }
  };

  updateIdealState(helixManager, tableName, updater, DEFAULT_RETRY_POLICY);
}
 

import org.apache.helix.model.IdealState; //导入依赖的package包/类
public static PistachioClusterInfo getClusterInfo() {
    try {
        String zookeeperConnStr = ConfigurationManager.getConfiguration().getString("Pistachio.ZooKeeper.Server");
        ZKHelixAdmin admin = new ZKHelixAdmin(zookeeperConnStr);
        IdealState idealState = admin.getResourceIdealState("PistachiosCluster", "PistachiosResource");
        PistachioClusterInfo info = new PistachioClusterInfo();
        info.numPartitions = idealState.getNumPartitions();
        info.numReplicas = Integer.parseInt(idealState.getReplicas());
        info.hostList = admin.getInstancesInCluster("PistachiosCluster");

        logger.info("num partitions: {}, num Replicas: {}, hostList: {}.", info.numPartitions,
            info.numReplicas, Joiner.on(",").join(info.hostList.toArray()));

        return info;
    } catch (Exception e) {
        logger.info("error getting cluster info", e);
        return null;
    }
}
 

import org.apache.helix.model.IdealState; //导入依赖的package包/类
public IdealState doInitialAssignment(String clusterName, List<String> instanceNames,
    int replicationFactor) {
  IdealState idealState = new IdealState(MySQLConstants.MASTER_SLAVE_RESOURCE_NAME);
  idealState.setStateModelDefRef("MasterSlave");
  idealState.setRebalanceMode(RebalanceMode.CUSTOMIZED);
  if (instanceNames.size() % replicationFactor != 0) {
    LOG.error(String.format(
        "Number of instances (%s) in the cluster must be a multiple of replication factor (%s)",
        instanceNames.size(), replicationFactor));
    return null;
  }
  int numSlices = instanceNames.size() / replicationFactor;
  idealState.setNumPartitions(numSlices);
  idealState.setReplicas(String.valueOf(replicationFactor));
  Collections.sort(instanceNames);
  for (int i = 0; i < numSlices; i++) {
    for (int j = 0; j < replicationFactor; j++) {
      idealState.setPartitionState(MySQLConstants.MASTER_SLAVE_RESOURCE_NAME + "_" + i,
          instanceNames.get(i * replicationFactor + j), (j == 0) ? "MASTER" : "SLAVE");
    }
  }
  LOG.info("Creating initial assignment \n" + idealState);
  _helixAdmin.setResourceIdealState(clusterName, MySQLConstants.MASTER_SLAVE_RESOURCE_NAME,
      idealState);
  return idealState;
}
 

import org.apache.helix.model.IdealState; //导入依赖的package包/类
@Override
public void addResource(String clusterName, String resourceName,
    IdealState idealstate) {
  String stateModelRef = idealstate.getStateModelDefRef();
  String stateModelDefPath = PropertyPathBuilder.stateModelDef(clusterName, stateModelRef);
  if (!_zkClient.exists(stateModelDefPath)) {
    throw new HelixException(
        "State model " + stateModelRef + " not found in the cluster STATEMODELDEFS path");
  }

  String idealStatePath = PropertyPathBuilder.idealState(clusterName);
  String resourceIdealStatePath = idealStatePath + "/" + resourceName;
  if (_zkClient.exists(resourceIdealStatePath)) {
    throw new HelixException("Skip the operation. Resource ideal state directory already exists:"
        + resourceIdealStatePath);
  }

  ZKUtil.createChildren(_zkClient, idealStatePath, idealstate.getRecord());
}
 

import org.apache.helix.model.IdealState; //导入依赖的package包/类
@Test (dependsOnMethods = {"testDisableDelayRebalanceInResource"})
public void testDisableDelayRebalanceInCluster() throws Exception {
  enableDelayRebalanceInCluster(_gZkClient, CLUSTER_NAME, true);

  Map<String, ExternalView> externalViewsBefore = createTestDBs(1000000);
  validateDelayedMovements(externalViewsBefore);

  // disable delay rebalance for the entire cluster.
  enableDelayRebalanceInCluster(_gZkClient, CLUSTER_NAME, false);
  // TODO: remove this once controller is listening on cluster config change.
  RebalanceScheduler.invokeRebalance(_controller.getHelixDataAccessor(), _testDBs.get(0));
  Thread.sleep(500);
  Assert.assertTrue(_clusterVerifier.verify());
  for (String db : _testDBs) {
    ExternalView ev =
        _setupTool.getClusterManagementTool().getResourceExternalView(CLUSTER_NAME, db);
    IdealState is = _setupTool.getClusterManagementTool().getResourceIdealState(CLUSTER_NAME, db);
    validateMinActiveAndTopStateReplica(is, ev, _replica, NUM_NODE);
  }

  enableDelayRebalanceInCluster(_gZkClient, CLUSTER_NAME, true);
}