org.apache.helix.model.IdealState Java Examples

public static IdealState buildManagerWorkerCustomIdealState(String routeName,
    List<String> instanceIds,
    String state) {
  final CustomModeISBuilder customModeIdealStateBuilder = new CustomModeISBuilder(routeName);

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

  for (String instanceId : instanceIds) {
    customModeIdealStateBuilder.assignInstanceAndState("0", instanceId, state);
  }

  return customModeIdealStateBuilder.build();
}
 

/**
 * 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 idealstate", resourceTag);
  HelixHelper
      .updateIdealState(helixManager, CommonConstants.Helix.BROKER_RESOURCE_INSTANCE, updater, DEFAULT_RETRY_POLICY);
}
 

public static IdealState expandCustomIdealStateFor(IdealState oldIdealState,
    String topicName, String newPartition, List<String> instances, int maxNumReplica) {
  final CustomModeISBuilder customModeIdealStateBuilder = new CustomModeISBuilder(topicName);

  int oldNumPartitions = oldIdealState.getNumPartitions();

  customModeIdealStateBuilder
      .setStateModel(OnlineOfflineStateModel.name)
      .setNumPartitions(oldNumPartitions + 1).setNumReplica(maxNumReplica)
      .setMaxPartitionsPerNode(oldNumPartitions + 1);

  for (String partitionName : oldIdealState.getPartitionSet()) {
    for (String instanceName : oldIdealState.getInstanceStateMap(partitionName).keySet()) {
      customModeIdealStateBuilder.assignInstanceAndState(partitionName, instanceName, "ONLINE");
    }
  }

  for (String instance : instances) {
    customModeIdealStateBuilder.assignInstanceAndState(newPartition, instance, "ONLINE");
  }

  return customModeIdealStateBuilder.build();
}
 

private void preSetup(StateTransitionThrottleConfig.RebalanceType rebalanceType,
    Set<String> resourceSet, int numOfLiveInstances, int numOfReplicas, String stateModelName,
    int minActiveReplica) {
  setupIdealState(numOfLiveInstances, resourceSet.toArray(new String[resourceSet.size()]),
      numOfLiveInstances, numOfReplicas, IdealState.RebalanceMode.FULL_AUTO, stateModelName,
      DelayedAutoRebalancer.class.getName(), CrushEdRebalanceStrategy.class.getName(),
      minActiveReplica);
  setupStateModel();
  setupLiveInstances(numOfLiveInstances);

  // Set up cluster configs
  _clusterConfig = accessor.getProperty(accessor.keyBuilder().clusterConfig());
  StateTransitionThrottleConfig throttleConfig = new StateTransitionThrottleConfig(rebalanceType,
      StateTransitionThrottleConfig.ThrottleScope.CLUSTER, Integer.MAX_VALUE);
  _clusterConfig.setStateTransitionThrottleConfigs(Collections.singletonList(throttleConfig));
  setClusterConfig(_clusterConfig);
}
 

@Override
public void onExternalViewChange(List<ExternalView> externalViewList,
    NotificationContext changeContext) {
  if (!trackEnabled) {
    return;
  }
  for (ExternalView ev : externalViewList) {
    IdealState is = _resourceMap.get(ev.getResourceName());
    if (is == null) {
      continue;
    }
    int replica = is.getReplicaCount(NUM_NODE);
    for (String p : is.getPartitionSet()) {
      Map<String, String> stateMap = ev.getStateMap(p);
      verifyPartitionCount(is.getResourceName(), p, stateMap, replica, "EV",
          is.getMinActiveReplicas());
    }
  }
}
 

@Test
public void testMetricValuePropagation()
    throws JMException, HelixRebalanceException, IOException {
  _metadataStore.reset();
  _metricCollector = new WagedRebalancerMetricCollector(TEST_STRING);
  WagedRebalancer rebalancer =
      new WagedRebalancer(_metadataStore, _algorithm, Optional.of(_metricCollector));

  // Generate the input for the rebalancer.
  ResourceControllerDataProvider clusterData = setupClusterDataCache();
  Map<String, Resource> resourceMap = clusterData.getIdealStates().entrySet().stream()
      .collect(Collectors.toMap(entry -> entry.getKey(), entry -> {
        Resource resource = new Resource(entry.getKey());
        entry.getValue().getPartitionSet().stream()
            .forEach(partition -> resource.addPartition(partition));
        return resource;
      }));
  Map<String, IdealState> newIdealStates =
      rebalancer.computeNewIdealStates(clusterData, resourceMap, new CurrentStateOutput());

  // Check that there exists a non-zero value in the metrics
  Assert.assertTrue(_metricCollector.getMetricMap().values().stream()
      .anyMatch(metric -> (long) metric.getLastEmittedMetricValue() > 0L));
}
 

public static IdealState resetCustomIdealStateFor(IdealState oldIdealState,
    String topicName, String partitionToReplace, String newInstanceName) {
  final CustomModeISBuilder customModeIdealStateBuilder = new CustomModeISBuilder(topicName);

  int oldNumPartitions = oldIdealState.getNumPartitions();

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

  for (String partitionName : oldIdealState.getPartitionSet()) {
    String instanceName = oldIdealState.getInstanceStateMap(partitionName).keySet().iterator().next();
    String instanceToUse = partitionName.equals(partitionToReplace) ? newInstanceName : instanceName;
    customModeIdealStateBuilder.assignInstanceAndState(partitionName, instanceToUse, "ONLINE");
  }

  return customModeIdealStateBuilder.build();
}
 

@Override
public IdealState computeNewIdealState(String resourceName, IdealState currentIdealState,
    CurrentStateOutput currentStateOutput, ResourceControllerDataProvider clusterData) {
  testRebalancerInvoked = true;
  List<String> liveNodes = Lists.newArrayList(clusterData.getLiveInstances().keySet());
  int i = 0;
  for (String partition : currentIdealState.getPartitionSet()) {
    int index = i++ % liveNodes.size();
    String instance = liveNodes.get(index);
    currentIdealState.getPreferenceList(partition).clear();
    currentIdealState.getPreferenceList(partition).add(instance);

    currentIdealState.getInstanceStateMap(partition).clear();
    currentIdealState.getInstanceStateMap(partition).put(instance, "MASTER");
  }
  currentIdealState.setReplicas("1");
  return currentIdealState;
}
 

private void verifyUserDefinedPreferenceLists(String db,
    Map<String, List<String>> userDefinedPreferenceLists, List<String> userDefinedPartitions)
    throws InterruptedException {
  IdealState is = _gSetupTool.getClusterManagementTool().getResourceIdealState(CLUSTER_NAME, db);
  for (String p : userDefinedPreferenceLists.keySet()) {
    List<String> userDefined = userDefinedPreferenceLists.get(p);
    List<String> preferenceListInIs = is.getPreferenceList(p);
    if (userDefinedPartitions.contains(p)) {
      Assert.assertTrue(userDefined.equals(preferenceListInIs));
    } else {
      if (userDefined.equals(preferenceListInIs)) {
        Assert.fail("Something is not good!");
      }
      Assert.assertFalse(userDefined.equals(preferenceListInIs), String
          .format("Partition %s, List in Is: %s, List as defined in config: %s", p, preferenceListInIs,
              userDefined));
    }
  }
}
 

public static IdealState resetCustomIdealStateFor(IdealState oldIdealState,
    String topicName, List<String> instanceToReplace, String availableInstance, int maxNumReplica) {
  final CustomModeISBuilder customModeIdealStateBuilder = new CustomModeISBuilder(topicName);

  int oldNumPartitions = oldIdealState.getNumPartitions();

  customModeIdealStateBuilder
      .setStateModel(OnlineOfflineStateModel.name)
      .setNumPartitions(oldNumPartitions).setNumReplica(maxNumReplica)
      .setMaxPartitionsPerNode(oldNumPartitions);
  for (String partitionName : oldIdealState.getPartitionSet()) {
    for (String instanceName : oldIdealState.getInstanceStateMap(partitionName).keySet()) {
      String instanceToUse = instanceToReplace.contains(instanceName) ? availableInstance : instanceName;
      customModeIdealStateBuilder.assignInstanceAndState(partitionName, instanceToUse, "ONLINE");
      if (instanceToReplace.contains(instanceName)) {
        LOGGER.info("replaceing: route: {}@{}, old {}, new {}",
            topicName, partitionName, instanceName, instanceToUse);
      }
    }
  }

  return customModeIdealStateBuilder.build();
}
 

public static void buildLowLevelRealtimeIdealStateFor(PinotLLCRealtimeSegmentManager pinotLLCRealtimeSegmentManager,
    String realtimeTableName, TableConfig realtimeTableConfig, IdealState idealState,
    boolean enableBatchMessageMode) {

  // Validate replicasPerPartition here.
  final String replicasPerPartitionStr = realtimeTableConfig.getValidationConfig().getReplicasPerPartition();
  if (replicasPerPartitionStr == null || replicasPerPartitionStr.isEmpty()) {
    throw new RuntimeException("Null or empty value for replicasPerPartition, expected a number");
  }
  final int nReplicas;
  try {
    nReplicas = Integer.valueOf(replicasPerPartitionStr);
  } catch (NumberFormatException e) {
    throw new PinotHelixResourceManager.InvalidTableConfigException(
        "Invalid value for replicasPerPartition, expected a number: " + replicasPerPartitionStr, e);
  }
  if (idealState == null) {
    idealState = buildEmptyRealtimeIdealStateFor(realtimeTableName, nReplicas, enableBatchMessageMode);
  }
  pinotLLCRealtimeSegmentManager.setUpNewTable(realtimeTableConfig, idealState);
}
 

private boolean hasInstanceMapChanged(Map<Partition, Map<String, String>> newAssiments,
    IdealState idealState) {
  Set<Partition> partitions = new HashSet<Partition>(newAssiments.keySet());
  for (String p : idealState.getPartitionSet()) {
    partitions.add(new Partition(p));
  }

  for (Partition partition : partitions) {
    Map<String, String> instanceMap = newAssiments.get(partition);
    Map<String, String> existInstanceMap =
        idealState.getInstanceStateMap(partition.getPartitionName());
    if (instanceMap == null && existInstanceMap == null) {
      continue;
    }
    if (instanceMap == null || existInstanceMap == null || !instanceMap
        .equals(existInstanceMap)) {
      return true;
    }
  }

  return false;
}
 

/**
 * Validate instances for each partition is on different zone and with necessary tagged instances.
 */
private void validateZoneAndTagIsolation(IdealState is, ExternalView ev, int expectedReplica) {
  String tag = is.getInstanceGroupTag();
  for (String partition : is.getPartitionSet()) {
    Set<String> assignedZones = new HashSet<String>();

    Map<String, String> assignmentMap = ev.getRecord().getMapField(partition);
    Set<String> instancesInEV = assignmentMap.keySet();
    // TODO: preference List is not persisted in IS.
    // Assert.assertEquals(instancesInEV, instancesInIs);
    for (String instance : instancesInEV) {
      assignedZones.add(_nodeToZoneMap.get(instance));
      if (tag != null) {
        InstanceConfig config =
            _gSetupTool.getClusterManagementTool().getInstanceConfig(CLUSTER_NAME, instance);
        Assert.assertTrue(config.containsTag(tag));
      }
    }
    Assert.assertEquals(assignedZones.size(), expectedReplica);
  }
}
 

/**
 * Test when two nodes go offline,  the minimal active replica should be maintained.
 * @throws Exception
 */
@Test(dependsOnMethods = {"testDelayedPartitionMovement"})
public void testMinimalActiveReplicaMaintain() throws Exception {
  setDelayTimeInCluster(_gZkClient, CLUSTER_NAME, 1000000);
  Map<String, ExternalView> externalViewsBefore = createTestDBs(-1);
  validateDelayedMovements(externalViewsBefore);

  // bring down another node, the minimal active replica for each partition should be maintained.
  _participants.get(3).syncStop();
  Thread.sleep(DEFAULT_REBALANCE_PROCESSING_WAIT_TIME);
  Assert.assertTrue(_clusterVerifier.verifyByPolling());
  for (String db : _testDBs) {
    ExternalView ev =
        _gSetupTool.getClusterManagementTool().getResourceExternalView(CLUSTER_NAME, db);
    IdealState is = _gSetupTool.getClusterManagementTool().getResourceIdealState(CLUSTER_NAME, db);
    validateMinActiveAndTopStateReplica(is, ev, _minActiveReplica, NUM_NODE);
  }
  setDelayTimeInCluster(_gZkClient, CLUSTER_NAME, -1);
}
 

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();
}
 

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.");
  }
}
 

/**
 * Returns the set of all partition ids for a job.
 * If a set of partition ids was explicitly specified in the config, that is used. Otherwise, we
 * use the list of all partition ids from the target resource.
 * return empty set if target resource does not exist.
 */
private static Set<Integer> getAllTaskPartitions(IdealState tgtResourceIs, JobConfig jobCfg,
    JobContext taskCtx) {
  Map<String, List<Integer>> currentTargets = taskCtx.getPartitionsByTarget();
  SortedSet<String> targetPartitions = Sets.newTreeSet();
  if (jobCfg.getTargetPartitions() != null) {
    targetPartitions.addAll(jobCfg.getTargetPartitions());
  } else {
    if (tgtResourceIs != null) {
      targetPartitions.addAll(tgtResourceIs.getPartitionSet());
    } else {
      LOG.warn("Missing target resource for the scheduled job!");
    }
  }

  Set<Integer> taskPartitions = Sets.newTreeSet();
  for (String targetPartition : targetPartitions) {
    taskPartitions
        .addAll(getPartitionsForTargetPartition(targetPartition, currentTargets, taskCtx));
  }
  return taskPartitions;
}
 

@Test
public void testDisableResource() {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;
  System.out.println("START " + clusterName + " at " + new Date(System.currentTimeMillis()));
  HelixAdmin admin = new ZKHelixAdmin(_gZkClient);
  admin.addCluster(clusterName, true);
  Assert.assertTrue(ZKUtil.isClusterSetup(clusterName, _gZkClient), "Cluster should be setup");
  String resourceName = "TestDB";
  admin.addStateModelDef(clusterName, "MasterSlave",
      new StateModelDefinition(StateModelConfigGenerator.generateConfigForMasterSlave()));
  admin.addResource(clusterName, resourceName, 4, "MasterSlave");
  admin.enableResource(clusterName, resourceName, false);
  BaseDataAccessor<ZNRecord> baseAccessor = new ZkBaseDataAccessor<>(_gZkClient);
  HelixDataAccessor accessor = new ZKHelixDataAccessor(clusterName, baseAccessor);
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();
  IdealState idealState = accessor.getProperty(keyBuilder.idealStates(resourceName));
  Assert.assertFalse(idealState.isEnabled());
  admin.enableResource(clusterName, resourceName, true);
  idealState = accessor.getProperty(keyBuilder.idealStates(resourceName));
  Assert.assertTrue(idealState.isEnabled());

  admin.dropCluster(clusterName);
  System.out.println("END " + clusterName + " at " + new Date(System.currentTimeMillis()));
}
 

/**
 * This function is deprecated. Please use RebalanceUtil.scheduleInstantPipeline method instead.
 * Trigger the controller to perform rebalance for a given resource.
 * @param accessor Helix data accessor
 * @param resource the name of the resource changed to triggering the execution
 */
@Deprecated
public static void invokeRebalance(HelixDataAccessor accessor, String resource) {
  LOG.info("invoke rebalance for " + resource);
  PropertyKey key = accessor.keyBuilder().idealStates(resource);
  IdealState is = accessor.getProperty(key);
  if (is != null) {
    // Here it uses the updateProperty function with no-op DataUpdater. Otherwise, it will use default
    // ZNRecordUpdater which will duplicate elements for listFields.
    if (!accessor.updateProperty(key, znRecord -> znRecord, is)) {
      LOG.warn("Failed to invoke rebalance on resource {}", resource);
    }
  } else {
    LOG.warn("Can't find ideal state for {}", resource);
  }
}
 

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();
}
 

/**
 * Validate each partition is different instances and with necessary tagged instances.
 */
private void validateIsolation(IdealState is, ExternalView ev, int expectedReplica) {
  String tag = is.getInstanceGroupTag();
  for (String partition : is.getPartitionSet()) {
    Map<String, String> assignmentMap = ev.getRecord().getMapField(partition);
    Set<String> instancesInEV = assignmentMap.keySet();
    Assert.assertEquals(instancesInEV.size(), expectedReplica);
    for (String instance : instancesInEV) {
      if (tag != null) {
        InstanceConfig config =
            _gSetupTool.getClusterManagementTool().getInstanceConfig(CLUSTER_NAME, instance);
        Assert.assertTrue(config.containsTag(tag));
      }
    }
  }
}
 

private boolean shouldDeleteInProgressLLCSegment(String segmentName, IdealState idealState,
    RealtimeSegmentZKMetadata realtimeSegmentZKMetadata) {
  if (idealState == null) {
    return false;
  }
  // delete a segment only if it is old enough (5 days) or else,
  // 1. latest segment could get deleted in the middle of repair by RealtimeSegmentValidationManager
  // 2. for a brand new segment, if this code kicks in after new metadata is created but ideal state entry is not yet created (between step 2 and 3),
  // the latest segment metadata could get marked for deletion
  if (System.currentTimeMillis() - realtimeSegmentZKMetadata.getCreationTime()
      <= OLD_LLC_SEGMENTS_RETENTION_IN_MILLIS) {
    return false;
  }
  Map<String, String> stateMap = idealState.getInstanceStateMap(segmentName);
  if (stateMap == null) {
    // Segment is in property store but not in ideal state, delete it
    return true;
  } else {
    // Delete segment if all of its replicas are OFFLINE
    Set<String> states = new HashSet<>(stateMap.values());
    return states.size() == 1 && states.contains(CommonConstants.Helix.StateModel.SegmentStateModel.OFFLINE);
  }
}
 

private void applyUserDefinedPreferenceList(ResourceConfig resourceConfig,
    IdealState idealState) {
  if (resourceConfig != null) {
    Map<String, List<String>> userDefinedPreferenceList = resourceConfig.getPreferenceLists();
    if (!userDefinedPreferenceList.isEmpty()) {
      LOG.info("Using user defined preference list for partitions.");
      for (String partition : userDefinedPreferenceList.keySet()) {
        idealState.setPreferenceList(partition, userDefinedPreferenceList.get(partition));
      }
    }
  }
}
 

public boolean isBrokerTenantDeletable(String tenantName) {
  String brokerTag = TagNameUtils.getBrokerTagForTenant(tenantName);
  Set<String> taggedInstances = new HashSet<>(HelixHelper.getInstancesWithTag(_helixZkManager, brokerTag));
  String brokerName = Helix.BROKER_RESOURCE_INSTANCE;
  IdealState brokerIdealState = _helixAdmin.getResourceIdealState(_helixClusterName, brokerName);
  for (String partition : brokerIdealState.getPartitionSet()) {
    for (String instance : brokerIdealState.getInstanceSet(partition)) {
      if (taggedInstances.contains(instance)) {
        return false;
      }
    }
  }
  return true;
}
 

@Test
public void disabledEmptyTableTest()
    throws Exception {

  final String tableName = "myTable_OFFLINE";
  List<String> allTableNames = Lists.newArrayList(tableName);
  IdealState idealState = new IdealState(tableName);
  // disable table in idealstate
  idealState.enable(false);
  idealState.setReplicas("1");
  idealState.setRebalanceMode(IdealState.RebalanceMode.CUSTOMIZED);

  {
    helixResourceManager = mock(PinotHelixResourceManager.class);
    when(helixResourceManager.getAllTables()).thenReturn(allTableNames);
    when(helixResourceManager.getTableIdealState(tableName)).thenReturn(idealState);
    when(helixResourceManager.getTableExternalView(tableName)).thenReturn(null);
  }
  {
    config = mock(ControllerConf.class);
    when(config.getStatusCheckerFrequencyInSeconds()).thenReturn(300);
    when(config.getStatusCheckerWaitForPushTimeInSeconds()).thenReturn(300);
  }
  {
    leadControllerManager = mock(LeadControllerManager.class);
    when(leadControllerManager.isLeaderForTable(anyString())).thenReturn(true);
  }
  metricsRegistry = new MetricsRegistry();
  controllerMetrics = new ControllerMetrics(metricsRegistry);
  segmentStatusChecker =
      new SegmentStatusChecker(helixResourceManager, leadControllerManager, config, controllerMetrics);
  // verify state before test
  Assert.assertEquals(controllerMetrics.getValueOfGlobalGauge(ControllerGauge.DISABLED_TABLE_COUNT), 0);
  // update metrics
  segmentStatusChecker.start();
  segmentStatusChecker.run();
  Assert.assertEquals(controllerMetrics.getValueOfGlobalGauge(ControllerGauge.DISABLED_TABLE_COUNT), 1);
}
 

@Override
public void setResourceIdealState(String clusterName, String resourceName,
    IdealState idealState) {
  logger
      .info("Set IdealState for resource {} in cluster {} with new IdealState {}.", resourceName,
          clusterName, idealState == null ? "NULL" : idealState.toString());
  HelixDataAccessor accessor =
      new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor<ZNRecord>(_zkClient));
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();

  accessor.setProperty(keyBuilder.idealStates(resourceName), idealState);
}
 

@Override
public IdealState getResourceIdealState(String clusterName, String resourceName) {
  HelixDataAccessor accessor =
      new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor<ZNRecord>(_zkClient));
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();

  return accessor.getProperty(keyBuilder.idealStates(resourceName));
}
 

private int getNumPartitionsFromIdealState(IdealState idealState) {
  int numPartitions = 0;
  for (String segmentName : idealState.getRecord().getMapFields().keySet()) {
    if (LLCSegmentName.isLowLevelConsumerSegmentName(segmentName)) {
      numPartitions = Math.max(numPartitions, new LLCSegmentName(segmentName).getPartitionId() + 1);
    }
  }
  return numPartitions;
}
 

private void updateBestPossibleStateOutput(BestPossibleStateOutput output, Resource resource,
    IdealState computedIdealState) {
  output.setPreferenceLists(resource.getResourceName(), computedIdealState.getPreferenceLists());
  for (Partition partition : resource.getPartitions()) {
    Map<String, String> newStateMap =
        computedIdealState.getInstanceStateMap(partition.getPartitionName());
    output.setState(resource.getResourceName(), partition, newStateMap);
  }
}
 

@Test
public void testLeadControllerResource() {
  IdealState leadControllerResourceIdealState = _helixResourceManager.getHelixAdmin()
      .getResourceIdealState(getHelixClusterName(), CommonConstants.Helix.LEAD_CONTROLLER_RESOURCE_NAME);
  Assert.assertTrue(leadControllerResourceIdealState.isValid());
  Assert.assertTrue(leadControllerResourceIdealState.isEnabled());
  Assert.assertEquals(leadControllerResourceIdealState.getInstanceGroupTag(),
      CommonConstants.Helix.CONTROLLER_INSTANCE);
  Assert.assertEquals(leadControllerResourceIdealState.getNumPartitions(),
      CommonConstants.Helix.NUMBER_OF_PARTITIONS_IN_LEAD_CONTROLLER_RESOURCE);
  Assert.assertEquals(leadControllerResourceIdealState.getReplicas(),
      Integer.toString(LEAD_CONTROLLER_RESOURCE_REPLICA_COUNT));
  Assert.assertEquals(leadControllerResourceIdealState.getRebalanceMode(), IdealState.RebalanceMode.FULL_AUTO);
  Assert.assertTrue(leadControllerResourceIdealState
      .getInstanceSet(leadControllerResourceIdealState.getPartitionSet().iterator().next()).isEmpty());

  TestUtils.waitForCondition(aVoid -> {
    ExternalView leadControllerResourceExternalView = _helixResourceManager.getHelixAdmin()
        .getResourceExternalView(getHelixClusterName(), CommonConstants.Helix.LEAD_CONTROLLER_RESOURCE_NAME);
    for (String partition : leadControllerResourceExternalView.getPartitionSet()) {
      Map<String, String> stateMap = leadControllerResourceExternalView.getStateMap(partition);
      Map.Entry<String, String> entry = stateMap.entrySet().iterator().next();
      boolean result =
          (LeadControllerUtils.generateParticipantInstanceId(LOCAL_HOST, _controllerPort)).equals(entry.getKey());
      result &= MasterSlaveSMD.States.MASTER.name().equals(entry.getValue());
      if (!result) {
        return false;
      }
    }
    return true;
  }, TIMEOUT_IN_MS, "Failed to assign controller hosts to lead controller resource in " + TIMEOUT_IN_MS + " ms.");
}