Java Code Examples for org.apache.helix.model.IdealState#setReplicas()

@Override
public void addResource(String clusterName, String resourceName, int partitions,
    String stateModelRef, String rebalancerMode, String rebalanceStrategy, int bucketSize,
    int maxPartitionsPerInstance) {
  if (!ZKUtil.isClusterSetup(clusterName, _zkClient)) {
    throw new HelixException("cluster " + clusterName + " is not setup yet");
  }

  IdealState idealState = new IdealState(resourceName);
  idealState.setNumPartitions(partitions);
  idealState.setStateModelDefRef(stateModelRef);
  RebalanceMode mode =
      idealState.rebalanceModeFromString(rebalancerMode, RebalanceMode.SEMI_AUTO);
  idealState.setRebalanceMode(mode);
  idealState.setRebalanceStrategy(rebalanceStrategy);
  idealState.setReplicas("" + 0);
  idealState.setStateModelFactoryName(HelixConstants.DEFAULT_STATE_MODEL_FACTORY);
  if (maxPartitionsPerInstance > 0 && maxPartitionsPerInstance < Integer.MAX_VALUE) {
    idealState.setMaxPartitionsPerInstance(maxPartitionsPerInstance);
  }
  if (bucketSize > 0) {
    idealState.setBucketSize(bucketSize);
  }
  addResource(clusterName, resourceName, idealState);
}
 

@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 setupCluster(String clusterName, List<String> instanceNames, String dbName,
    int replica, int partitions, int bucketSize) {
  _gSetupTool.addCluster(clusterName, true);
  _gSetupTool.addInstancesToCluster(clusterName,
      instanceNames.toArray(new String[instanceNames.size()]));

  // add a bucketized resource
  ZKHelixDataAccessor accessor = new ZKHelixDataAccessor(clusterName, _baseAccessor);
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();
  ZNRecord idealStateRec =
      DefaultIdealStateCalculator.calculateIdealState(instanceNames, partitions, replica - 1,
          dbName,
          "MASTER", "SLAVE");
  IdealState idealState = new IdealState(idealStateRec);
  idealState.setBucketSize(bucketSize);
  idealState.setStateModelDefRef("MasterSlave");
  idealState.setRebalanceMode(IdealState.RebalanceMode.CUSTOMIZED);
  idealState.setReplicas(Integer.toString(replica));
  accessor.setProperty(keyBuilder.idealStates(dbName), idealState);

}
 

@Test(dataProvider = "IdealStateInput")
public void testIdealStateAssignment(String clusterName, List<String> instances,
    List<String> partitions, String numReplicas, String stateModeDef, String strategyName,
    Map<String, Map<String, String>> expectedMapping, List<String> disabledInstances)
    throws IllegalAccessException, InstantiationException, ClassNotFoundException {
  ClusterConfig clusterConfig = new ClusterConfig(clusterName);
  List<InstanceConfig> instanceConfigs = new ArrayList<>();
  for (String instance : instances) {
    instanceConfigs.add(new InstanceConfig(instance));
    if (disabledInstances.contains(instance)) {
      instanceConfigs.get(instanceConfigs.size() - 1).setInstanceEnabled(false);
    }
  }

  IdealState idealState = new IdealState("TestResource");
  idealState.setStateModelDefRef(stateModeDef);
  idealState.setNumPartitions(partitions.size());
  idealState.setReplicas(numReplicas);

  Map<String, Map<String, String>> idealStateMapping = HelixUtil
      .getIdealAssignmentForFullAuto(clusterConfig, instanceConfigs, instances, idealState,
          partitions, strategyName);
  Assert.assertEquals(idealStateMapping, expectedMapping);
}
 

@Override
public void addClusterToGrandCluster(String clusterName, String grandCluster) {
  logger.info("Add cluster {} to grand cluster {}.", clusterName, grandCluster);
  if (!ZKUtil.isClusterSetup(grandCluster, _zkClient)) {
    throw new HelixException("Grand cluster " + grandCluster + " is not setup yet");
  }

  if (!ZKUtil.isClusterSetup(clusterName, _zkClient)) {
    throw new HelixException("Cluster " + clusterName + " is not setup yet");
  }

  IdealState idealState = new IdealState(clusterName);

  idealState.setNumPartitions(1);
  idealState.setStateModelDefRef("LeaderStandby");
  idealState.setRebalanceMode(RebalanceMode.FULL_AUTO);
  idealState.setRebalancerClassName(DelayedAutoRebalancer.class.getName());
  idealState.setRebalanceStrategy(CrushEdRebalanceStrategy.class.getName());
  // TODO: Give user an option, say from RestAPI to config the number of replicas.
  idealState.setReplicas(Integer.toString(DEFAULT_SUPERCLUSTER_REPLICA));
  idealState.getRecord().setListField(clusterName, new ArrayList<String>());

  List<String> controllers = getInstancesInCluster(grandCluster);
  if (controllers.size() == 0) {
    throw new HelixException("Grand cluster " + grandCluster + " has no instances");
  }

  ZKHelixDataAccessor accessor =
      new ZKHelixDataAccessor(grandCluster, new ZkBaseDataAccessor<ZNRecord>(_zkClient));
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();

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

@Test
public void noReplicas()
    throws Exception {
  final String tableName = "myTable_REALTIME";
  List<String> allTableNames = new ArrayList<String>();
  allTableNames.add(tableName);
  IdealState idealState = new IdealState(tableName);
  idealState.setPartitionState("myTable_0", "pinot1", "OFFLINE");
  idealState.setPartitionState("myTable_0", "pinot2", "OFFLINE");
  idealState.setPartitionState("myTable_0", "pinot3", "OFFLINE");
  idealState.setReplicas("0");
  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);
  segmentStatusChecker.start();
  segmentStatusChecker.run();
  Assert.assertEquals(controllerMetrics.getValueOfTableGauge(tableName, ControllerGauge.SEGMENTS_IN_ERROR_STATE), 0);
  Assert.assertEquals(controllerMetrics.getValueOfTableGauge(tableName, ControllerGauge.NUMBER_OF_REPLICAS), 1);
  Assert.assertEquals(controllerMetrics.getValueOfTableGauge(tableName, ControllerGauge.PERCENT_OF_REPLICAS), 100);
  Assert.assertEquals(controllerMetrics.getValueOfTableGauge(tableName, ControllerGauge.PERCENT_SEGMENTS_AVAILABLE),
      100);
}
 

private void createIS(HelixDataAccessor accessor, String resourceId, String stateModelDefRef,
    RebalanceMode rebalanceMode) {
  IdealState idealState = new IdealState(resourceId);
  idealState.setRebalanceMode(rebalanceMode);
  idealState.setStateModelDefRef(stateModelDefRef);
  idealState.setNumPartitions(1);
  idealState.setReplicas("1");
  idealState.getRecord().setListField(resourceId + "_0", ImmutableList.of(PARTICIPANT));
  idealState.getRecord().setMapField(resourceId + "_0", ImmutableMap.of(PARTICIPANT, STATE));
  accessor.setProperty(accessor.keyBuilder().idealStates(resourceId), idealState);
}
 

@Test(dependsOnMethods = "test")
public void testChangeIdealState() throws InterruptedException {
  String dbName = "Test-DB-" + TestHelper.getTestMethodName();
  createResourceWithWagedRebalance(CLUSTER_NAME, dbName,
      BuiltInStateModelDefinitions.MasterSlave.name(), PARTITIONS, _replica, _replica);
  _gSetupTool.rebalanceStorageCluster(CLUSTER_NAME, dbName, _replica);
  _allDBs.add(dbName);
  Thread.sleep(300);

  validate(_replica);

  // Adjust the replica count
  IdealState is =
      _gSetupTool.getClusterManagementTool().getResourceIdealState(CLUSTER_NAME, dbName);
  int newReplicaFactor = _replica - 1;
  is.setReplicas("" + newReplicaFactor);
  _gSetupTool.getClusterManagementTool().setResourceIdealState(CLUSTER_NAME, dbName, is);
  Thread.sleep(300);

  validate(newReplicaFactor);

  // Adjust the partition list
  is = _gSetupTool.getClusterManagementTool().getResourceIdealState(CLUSTER_NAME, dbName);
  is.setNumPartitions(PARTITIONS + 1);
  _gSetupTool.getClusterManagementTool().setResourceIdealState(CLUSTER_NAME, dbName, is);
  _gSetupTool.getClusterManagementTool().rebalance(CLUSTER_NAME, dbName, newReplicaFactor);
  Thread.sleep(300);

  validate(newReplicaFactor);
  ExternalView ev =
      _gSetupTool.getClusterManagementTool().getResourceExternalView(CLUSTER_NAME, dbName);
  Assert.assertEquals(ev.getPartitionSet().size(), PARTITIONS + 1);
}
 

@Test
public void testNoRedundantDropMessage() throws Exception {
  String resourceName = "TEST_RESOURCE";
  _gSetupTool.getClusterManagementTool().addResource(CLUSTER_NAME, resourceName, 1, "MasterSlave",
      IdealState.RebalanceMode.CUSTOMIZED.name());
  String partitionName = "P_0";
  ClusterEvent event = new ClusterEvent(CLUSTER_NAME, ClusterEventType.Unknown, "ID");
  ResourceControllerDataProvider cache = new ResourceControllerDataProvider(CLUSTER_NAME);
  cache.refresh(_manager.getHelixDataAccessor());
  IdealState idealState = cache.getIdealState(resourceName);
  idealState.setReplicas("2");
  Map<String, String> stateMap = new HashMap<>();
  stateMap.put(_participants[0].getInstanceName(), "SLAVE");
  stateMap.put(_participants[1].getInstanceName(), "DROPPED");
  idealState.setInstanceStateMap(partitionName, stateMap);

  cache.setIdealStates(Arrays.asList(idealState));
  cache.setCachedIdealMapping(idealState.getResourceName(), idealState.getRecord());

  event.addAttribute(AttributeName.ControllerDataProvider.name(), cache);
  event.addAttribute(AttributeName.helixmanager.name(), _manager);

  runStage(event, new ResourceComputationStage());
  runStage(event, new CurrentStateComputationStage());
  runStage(event, new BestPossibleStateCalcStage());
  runStage(event, new IntermediateStateCalcStage());
  Assert.assertEquals(cache.getCachedIdealMapping().size(), 1);
  runStage(event, new ResourceMessageGenerationPhase());

  MessageOutput messageOutput = event.getAttribute(AttributeName.MESSAGES_ALL.name());
  Assert
      .assertEquals(messageOutput.getMessages(resourceName, new Partition(partitionName)).size(),
          1);
  Assert.assertEquals(cache.getCachedIdealMapping().size(), 0);
}
 

@Test
public void testListenerOnBucketizedResource() throws Exception {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;
  String dbName = "TestDB0";
  List<String> instanceNames =
      Arrays.asList("localhost_0", "localhost_1", "localhost_2", "localhost_3", "localhost_4");
  int n = instanceNames.size();

  ZKHelixDataAccessor accessor = new ZKHelixDataAccessor(clusterName, _baseAccessor);
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();

  setupCluster(clusterName, instanceNames, dbName, 3, 10, 2);

  // start controller
  ClusterControllerManager controller = new ClusterControllerManager(ZK_ADDR, clusterName);
  controller.syncStart();

  // start participants
  MockParticipantManager[] participants = new MockParticipantManager[n];
  for (int i = 0; i < n; i++) {
    participants[i] = new MockParticipantManager(ZK_ADDR, clusterName, instanceNames.get(i));
    participants[i].syncStart();
  }

  ZkHelixClusterVerifier _clusterVerifier =
      new BestPossibleExternalViewVerifier.Builder(clusterName).setZkAddr(ZK_ADDR).build();
  Assert.assertTrue(_clusterVerifier.verifyByPolling());

  // add an external view listener
  final TestExternalViewListener listener = new TestExternalViewListener();
  controller.addExternalViewChangeListener(listener);

  // remove "TestDB0"
  _gSetupTool.dropResourceFromCluster(clusterName, dbName);
  Assert.assertTrue(_clusterVerifier.verifyByPolling());

  // wait callback to finish
  TestHelper.verify(new TestHelper.Verifier() {
    @Override public boolean verify() throws Exception {
      return listener.cbCnt > 0;
    }
  }, 20000);
  Assert.assertTrue(listener.cbCnt > 0);

  listener.cbCnt = 0;

  // add a new db
  String newDbName = "TestDB1";
  int r = 3;
  ZNRecord idealStateRec =
      DefaultIdealStateCalculator.calculateIdealState(instanceNames, 10, r - 1, newDbName,
          "MASTER", "SLAVE");
  IdealState idealState = new IdealState(idealStateRec);
  idealState.setBucketSize(2);
  idealState.setStateModelDefRef("MasterSlave");
  idealState.setRebalanceMode(IdealState.RebalanceMode.CUSTOMIZED);
  idealState.setReplicas(Integer.toString(r));
  accessor.setProperty(keyBuilder.idealStates(newDbName), idealState);

  Assert.assertTrue(_clusterVerifier.verifyByPolling());

  TestHelper.verify(new TestHelper.Verifier() {
    @Override public boolean verify() throws Exception {
      return listener.cbCnt > 0;
    }
  }, 20000);

  Assert.assertTrue(listener.cbCnt > 0);

  // clean up
  controller.syncStop();
  for (MockParticipantManager participant : participants) {
    participant.syncStop();
  }
  deleteCluster(clusterName);
}
 

@Test
public void missingEVTest()
    throws Exception {
  final String tableName = "myTable_REALTIME";
  List<String> allTableNames = new ArrayList<String>();
  allTableNames.add(tableName);
  IdealState idealState = new IdealState(tableName);
  idealState.setPartitionState("myTable_0", "pinot1", "ONLINE");
  idealState.setPartitionState("myTable_0", "pinot2", "ONLINE");
  idealState.setPartitionState("myTable_0", "pinot3", "ONLINE");
  idealState.setPartitionState("myTable_1", "pinot1", "ONLINE");
  idealState.setPartitionState("myTable_1", "pinot2", "ONLINE");
  idealState.setPartitionState("myTable_1", "pinot3", "ONLINE");
  idealState.setPartitionState("myTable_2", "pinot3", "OFFLINE");
  idealState.setReplicas("2");
  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);
  segmentStatusChecker.start();
  segmentStatusChecker.run();
  Assert.assertEquals(controllerMetrics.getValueOfTableGauge(tableName, ControllerGauge.SEGMENTS_IN_ERROR_STATE), 0);
  Assert.assertEquals(controllerMetrics.getValueOfTableGauge(tableName, ControllerGauge.NUMBER_OF_REPLICAS), 0);
}
 

@Test()
public void testRenamePartitionCustomIS() throws Exception {

  String clusterName = "CLUSTER_" + getShortClassName() + "_custom";
  System.out.println("START " + clusterName + " at " + new Date(System.currentTimeMillis()));

  TestHelper.setupCluster(clusterName, ZK_ADDR, 12918, // participant start port
      "localhost", // participant name prefix
      "TestDB", // resource name prefix
      1, // resources
      10, // partitions per resource
      5, // number of nodes
      3, // replicas
      "MasterSlave", false); // do rebalance

  // calculate idealState
  List<String> instanceNames = Arrays.asList("localhost_12918", "localhost_12919",
      "localhost_12920", "localhost_12921", "localhost_12922");
  ZNRecord destIS = DefaultIdealStateCalculator.calculateIdealState(instanceNames, 10, 3 - 1,
      "TestDB0", "MASTER", "SLAVE");
  IdealState idealState = new IdealState(destIS);
  idealState.setRebalanceMode(RebalanceMode.CUSTOMIZED);
  idealState.setReplicas("3");
  idealState.setStateModelDefRef("MasterSlave");

  ZKHelixDataAccessor accessor =
      new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor<>(_gZkClient));
  Builder keyBuilder = accessor.keyBuilder();

  accessor.setProperty(keyBuilder.idealStates("TestDB0"), idealState);

  startAndVerify(clusterName);

  Map<String, String> stateMap = idealState.getRecord().getMapFields().remove("TestDB0_0");
  idealState.getRecord().getMapFields().put("TestDB0_100", stateMap);
  accessor.setProperty(keyBuilder.idealStates("TestDB0"), idealState);

  boolean result = ClusterStateVerifier.verifyByPolling(
      new ClusterStateVerifier.BestPossAndExtViewZkVerifier(ZK_ADDR, clusterName));
  Assert.assertTrue(result);

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

@Test
public void missingEVPartitionTest()
    throws Exception {
  final String tableName = "myTable_OFFLINE";
  List<String> allTableNames = new ArrayList<String>();
  allTableNames.add(tableName);
  IdealState idealState = new IdealState(tableName);
  idealState.setPartitionState("myTable_0", "pinot1", "ONLINE");
  idealState.setPartitionState("myTable_0", "pinot2", "ONLINE");
  idealState.setPartitionState("myTable_0", "pinot3", "ONLINE");
  idealState.setPartitionState("myTable_1", "pinot1", "ONLINE");
  idealState.setPartitionState("myTable_1", "pinot2", "ONLINE");
  idealState.setPartitionState("myTable_1", "pinot3", "ONLINE");
  idealState.setPartitionState("myTable_2", "pinot3", "OFFLINE");
  idealState.setPartitionState("myTable_3", "pinot3", "ONLINE");
  idealState.setReplicas("2");
  idealState.setRebalanceMode(IdealState.RebalanceMode.CUSTOMIZED);

  ExternalView externalView = new ExternalView(tableName);
  externalView.setState("myTable_0", "pinot1", "ONLINE");
  externalView.setState("myTable_0", "pinot2", "ONLINE");
  externalView.setState("myTable_1", "pinot1", "ERROR");
  externalView.setState("myTable_1", "pinot2", "ONLINE");

  ZNRecord znrecord = new ZNRecord("myTable_0");
  znrecord.setSimpleField(CommonConstants.Segment.SEGMENT_NAME, "myTable_0");
  znrecord.setSimpleField(CommonConstants.Segment.TABLE_NAME, "myTable_OFFLINE");
  znrecord.setSimpleField(CommonConstants.Segment.INDEX_VERSION, "v1");
  znrecord.setEnumField(CommonConstants.Segment.SEGMENT_TYPE, CommonConstants.Segment.SegmentType.OFFLINE);
  znrecord.setLongField(CommonConstants.Segment.START_TIME, 1000);
  znrecord.setLongField(CommonConstants.Segment.END_TIME, 2000);
  znrecord.setSimpleField(CommonConstants.Segment.TIME_UNIT, TimeUnit.HOURS.toString());
  znrecord.setLongField(CommonConstants.Segment.TOTAL_DOCS, 10000);
  znrecord.setLongField(CommonConstants.Segment.CRC, 1234);
  znrecord.setLongField(CommonConstants.Segment.CREATION_TIME, 3000);
  znrecord.setSimpleField(CommonConstants.Segment.Offline.DOWNLOAD_URL, "http://localhost:8000/myTable_0");
  znrecord.setLongField(CommonConstants.Segment.Offline.PUSH_TIME, System.currentTimeMillis());
  znrecord.setLongField(CommonConstants.Segment.Offline.REFRESH_TIME, System.currentTimeMillis());

  ZkHelixPropertyStore<ZNRecord> propertyStore;
  {
    propertyStore = (ZkHelixPropertyStore<ZNRecord>) mock(ZkHelixPropertyStore.class);
    when(propertyStore.get("/SEGMENTS/myTable_OFFLINE/myTable_3", null, AccessOption.PERSISTENT))
        .thenReturn(znrecord);
  }

  {
    helixResourceManager = mock(PinotHelixResourceManager.class);
    when(helixResourceManager.getAllTables()).thenReturn(allTableNames);
    when(helixResourceManager.getTableIdealState(tableName)).thenReturn(idealState);
    when(helixResourceManager.getTableExternalView(tableName)).thenReturn(externalView);
    when(helixResourceManager.getOfflineSegmentZKMetadata(tableName, "myTable_3"))
        .thenReturn(new OfflineSegmentZKMetadata(znrecord));
  }
  {
    config = mock(ControllerConf.class);
    when(config.getStatusCheckerFrequencyInSeconds()).thenReturn(300);
    when(config.getStatusCheckerWaitForPushTimeInSeconds()).thenReturn(0);
  }
  {
    leadControllerManager = mock(LeadControllerManager.class);
    when(leadControllerManager.isLeaderForTable(anyString())).thenReturn(true);
  }
  metricsRegistry = new MetricsRegistry();
  controllerMetrics = new ControllerMetrics(metricsRegistry);
  segmentStatusChecker =
      new SegmentStatusChecker(helixResourceManager, leadControllerManager, config, controllerMetrics);
  segmentStatusChecker.start();
  segmentStatusChecker.run();
  Assert.assertEquals(
      controllerMetrics.getValueOfTableGauge(externalView.getId(), ControllerGauge.SEGMENTS_IN_ERROR_STATE), 1);
  Assert
      .assertEquals(controllerMetrics.getValueOfTableGauge(externalView.getId(), ControllerGauge.NUMBER_OF_REPLICAS),
          0);
  Assert.assertEquals(
      controllerMetrics.getValueOfTableGauge(externalView.getId(), ControllerGauge.PERCENT_SEGMENTS_AVAILABLE), 75);
}
 

@Test
public void realtimeBasicTest()
    throws Exception {
  final String tableName = "myTable_REALTIME";
  final String rawTableName = TableNameBuilder.extractRawTableName(tableName);
  List<String> allTableNames = new ArrayList<String>();
  allTableNames.add(tableName);
  final LLCSegmentName seg1 = new LLCSegmentName(rawTableName, 1, 0, System.currentTimeMillis());
  final LLCSegmentName seg2 = new LLCSegmentName(rawTableName, 1, 1, System.currentTimeMillis());
  final LLCSegmentName seg3 = new LLCSegmentName(rawTableName, 2, 1, System.currentTimeMillis());
  IdealState idealState = new IdealState(tableName);
  idealState.setPartitionState(seg1.getSegmentName(), "pinot1", "ONLINE");
  idealState.setPartitionState(seg1.getSegmentName(), "pinot2", "ONLINE");
  idealState.setPartitionState(seg1.getSegmentName(), "pinot3", "ONLINE");
  idealState.setPartitionState(seg2.getSegmentName(), "pinot1", "ONLINE");
  idealState.setPartitionState(seg2.getSegmentName(), "pinot2", "ONLINE");
  idealState.setPartitionState(seg2.getSegmentName(), "pinot3", "ONLINE");
  idealState.setPartitionState(seg3.getSegmentName(), "pinot1", "CONSUMING");
  idealState.setPartitionState(seg3.getSegmentName(), "pinot2", "CONSUMING");
  idealState.setPartitionState(seg3.getSegmentName(), "pinot3", "OFFLINE");
  idealState.setReplicas("3");
  idealState.setRebalanceMode(IdealState.RebalanceMode.CUSTOMIZED);

  ExternalView externalView = new ExternalView(tableName);
  externalView.setState(seg1.getSegmentName(), "pinot1", "ONLINE");
  externalView.setState(seg1.getSegmentName(), "pinot2", "ONLINE");
  externalView.setState(seg1.getSegmentName(), "pinot3", "ONLINE");
  externalView.setState(seg2.getSegmentName(), "pinot1", "CONSUMING");
  externalView.setState(seg2.getSegmentName(), "pinot2", "ONLINE");
  externalView.setState(seg2.getSegmentName(), "pinot3", "CONSUMING");
  externalView.setState(seg3.getSegmentName(), "pinot1", "CONSUMING");
  externalView.setState(seg3.getSegmentName(), "pinot2", "CONSUMING");
  externalView.setState(seg3.getSegmentName(), "pinot3", "OFFLINE");

  {
    helixResourceManager = mock(PinotHelixResourceManager.class);
    when(helixResourceManager.getAllTables()).thenReturn(allTableNames);
    when(helixResourceManager.getTableIdealState(tableName)).thenReturn(idealState);
    when(helixResourceManager.getTableExternalView(tableName)).thenReturn(externalView);
  }
  {
    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);
  segmentStatusChecker.start();
  segmentStatusChecker.run();
  Assert.assertEquals(
      controllerMetrics.getValueOfTableGauge(externalView.getId(), ControllerGauge.SEGMENTS_IN_ERROR_STATE), 0);
  Assert
      .assertEquals(controllerMetrics.getValueOfTableGauge(externalView.getId(), ControllerGauge.NUMBER_OF_REPLICAS),
          3);
  Assert
      .assertEquals(controllerMetrics.getValueOfTableGauge(externalView.getId(), ControllerGauge.PERCENT_OF_REPLICAS),
          100);
  Assert.assertEquals(
      controllerMetrics.getValueOfTableGauge(externalView.getId(), ControllerGauge.PERCENT_SEGMENTS_AVAILABLE), 100);
}
 

@Test
public void testGetResourcesWithTag() {
  String TEST_TAG = "TestTAG";

  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;

  HelixAdmin tool = new ZKHelixAdmin(_gZkClient);
  tool.addCluster(clusterName, true);
  Assert.assertTrue(ZKUtil.isClusterSetup(clusterName, _gZkClient));

  tool.addStateModelDef(clusterName, "OnlineOffline",
      new StateModelDefinition(StateModelConfigGenerator.generateConfigForOnlineOffline()));

  for (int i = 0; i < 4; i++) {
    String instanceName = "host" + i + "_9999";
    InstanceConfig config = new InstanceConfig(instanceName);
    config.setHostName("host" + i);
    config.setPort("9999");
    // set tag to two instances
    if (i < 2) {
      config.addTag(TEST_TAG);
    }
    tool.addInstance(clusterName, config);
    tool.enableInstance(clusterName, instanceName, true);
    String path = PropertyPathBuilder.instance(clusterName, instanceName);
    AssertJUnit.assertTrue(_gZkClient.exists(path));
  }

  for (int i = 0; i < 4; i++) {
    String resourceName = "database_" + i;
    IdealState is = new IdealState(resourceName);
    is.setStateModelDefRef("OnlineOffline");
    is.setNumPartitions(2);
    is.setRebalanceMode(IdealState.RebalanceMode.FULL_AUTO);
    is.setReplicas("1");
    is.enable(true);
    if (i < 2) {
      is.setInstanceGroupTag(TEST_TAG);
    }
    tool.addResource(clusterName, resourceName, is);
  }

  List<String> allResources = tool.getResourcesInCluster(clusterName);
  List<String> resourcesWithTag = tool.getResourcesInClusterWithTag(clusterName, TEST_TAG);
  AssertJUnit.assertEquals(allResources.size(), 4);
  AssertJUnit.assertEquals(resourcesWithTag.size(), 2);

  tool.dropCluster(clusterName);
}
 

/**
 * Creates a setup where the health API can be tested.
 * @param clusterName
 * @param resourceName
 * @param idealStateParams
 * @param partitionReplicaStates maps partitionName to its replicas' states
 * @throws Exception
 */
private void createDummyMapping(String clusterName, String resourceName,
    Map<String, String> idealStateParams, Map<String, List<String>> partitionReplicaStates)
    throws Exception {
  IdealState idealState = new IdealState(resourceName);
  idealState.setMinActiveReplicas(Integer.parseInt(idealStateParams.get("MinActiveReplicas"))); // 2
  idealState.setStateModelDefRef(idealStateParams.get("StateModelDefRef")); // MasterSlave
  idealState.setMaxPartitionsPerInstance(
      Integer.parseInt(idealStateParams.get("MaxPartitionsPerInstance"))); // 3
  idealState.setReplicas(idealStateParams.get("Replicas")); // 3
  idealState.setNumPartitions(Integer.parseInt(idealStateParams.get("NumPartitions"))); // 3
  idealState.enable(false);

  Map<String, List<String>> partitionNames = new LinkedHashMap<>();
  List<String> dummyPrefList = new ArrayList<>();

  for (int i = 0; i < Integer.parseInt(idealStateParams.get("MaxPartitionsPerInstance")); i++) {
    dummyPrefList.add(ANY_INSTANCE);
    partitionNames.put("p" + i, dummyPrefList);
  }
  idealState.getRecord().getListFields().putAll(partitionNames);

  if (!_gSetupTool.getClusterManagementTool().getClusters().contains(clusterName)) {
    _gSetupTool.getClusterManagementTool().addCluster(clusterName);
  }
  _gSetupTool.getClusterManagementTool().setResourceIdealState(clusterName, resourceName,
      idealState);

  // Set ExternalView's replica states for a given parameter map
  ExternalView externalView = new ExternalView(resourceName);

  Map<String, Map<String, String>> mappingCurrent = new LinkedHashMap<>();

  List<String> partitionReplicaStatesList = new ArrayList<>(partitionReplicaStates.keySet());
  for (int k = 0; k < partitionReplicaStatesList.size(); k++) {
    Map<String, String> replicaStatesForPartition = new LinkedHashMap<>();
    List<String> replicaStateList = partitionReplicaStates.get(partitionReplicaStatesList.get(k));
    for (int i = 0; i < replicaStateList.size(); i++) {
      replicaStatesForPartition.put("r" + i, replicaStateList.get(i));
    }
    mappingCurrent.put("p" + k, replicaStatesForPartition);
  }

  externalView.getRecord().getMapFields().putAll(mappingCurrent);

  HelixManager helixManager = HelixManagerFactory.getZKHelixManager(clusterName, "p1",
      InstanceType.ADMINISTRATOR, ZK_ADDR);
  helixManager.connect();
  HelixDataAccessor helixDataAccessor = helixManager.getHelixDataAccessor();
  helixDataAccessor.setProperty(helixDataAccessor.keyBuilder().externalView(resourceName),
      externalView);
  System.out.println("End test :" + TestHelper.getTestMethodName());
}
 

public void noSegmentsInternal(final int nReplicas)
    throws Exception {
  final String tableName = "myTable_REALTIME";
  String nReplicasStr = Integer.toString(nReplicas);
  int nReplicasExpectedValue = nReplicas;
  if (nReplicas < 0) {
    nReplicasStr = "abc";
    nReplicasExpectedValue = 1;
  }
  List<String> allTableNames = new ArrayList<String>();
  allTableNames.add(tableName);
  IdealState idealState = new IdealState(tableName);
  idealState.setReplicas(nReplicasStr);
  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);
  segmentStatusChecker.start();
  segmentStatusChecker.run();
  Assert.assertEquals(controllerMetrics.getValueOfTableGauge(tableName, ControllerGauge.SEGMENTS_IN_ERROR_STATE),
      Long.MIN_VALUE);
  Assert.assertEquals(controllerMetrics.getValueOfTableGauge(tableName, ControllerGauge.NUMBER_OF_REPLICAS),
      nReplicasExpectedValue);
  Assert.assertEquals(controllerMetrics.getValueOfTableGauge(tableName, ControllerGauge.PERCENT_OF_REPLICAS), 100);
  Assert.assertEquals(controllerMetrics.getValueOfTableGauge(tableName, ControllerGauge.PERCENT_SEGMENTS_AVAILABLE),
      100);
}
 

@Override
protected ResourceControllerDataProvider setupClusterDataCache() throws IOException {
  ResourceControllerDataProvider testCache = super.setupClusterDataCache();

  // Set up mock idealstate
  Map<String, IdealState> isMap = new HashMap<>();
  for (String resource : _resourceNames) {
    IdealState is = new IdealState(resource);
    is.setNumPartitions(_partitionNames.size());
    is.setRebalanceMode(IdealState.RebalanceMode.FULL_AUTO);
    is.setStateModelDefRef("MasterSlave");
    is.setReplicas("100");
    is.setRebalancerClassName(WagedRebalancer.class.getName());
    _partitionNames.stream()
        .forEach(partition -> is.setPreferenceList(partition, Collections.emptyList()));
    isMap.put(resource, is);
  }
  when(testCache.getIdealState(anyString())).thenAnswer(
      (Answer<IdealState>) invocationOnMock -> isMap.get(invocationOnMock.getArguments()[0]));
  when(testCache.getIdealStates()).thenReturn(isMap);
  when(testCache.getAsyncTasksThreadPool()).thenReturn(Executors.newSingleThreadExecutor());

  // Set up 2 more instances
  for (int i = 1; i < 3; i++) {
    String instanceName = _testInstanceId + i;
    _instances.add(instanceName);
    // 1. Set up the default instance information with capacity configuration.
    InstanceConfig testInstanceConfig = createMockInstanceConfig(instanceName);
    Map<String, InstanceConfig> instanceConfigMap = testCache.getInstanceConfigMap();
    instanceConfigMap.put(instanceName, testInstanceConfig);
    when(testCache.getInstanceConfigMap()).thenReturn(instanceConfigMap);
    // 2. Mock the live instance node for the default instance.
    LiveInstance testLiveInstance = createMockLiveInstance(instanceName);
    Map<String, LiveInstance> liveInstanceMap = testCache.getLiveInstances();
    liveInstanceMap.put(instanceName, testLiveInstance);
    when(testCache.getLiveInstances()).thenReturn(liveInstanceMap);
    when(testCache.getEnabledInstances()).thenReturn(liveInstanceMap.keySet());
    when(testCache.getEnabledLiveInstances()).thenReturn(liveInstanceMap.keySet());
    when(testCache.getAllInstances()).thenReturn(_instances);
  }

  return testCache;
}
 

void rebalance(String clusterName, String resourceName, int replica, String keyPrefix,
    List<String> instanceNames, String groupId) {
  logger.info("Rebalance resource {} with replica {} in cluster {}.", resourceName, replica,
      clusterName);
  // ensure we get the same idealState with the same set of instances
  Collections.sort(instanceNames);

  IdealState idealState = getResourceIdealState(clusterName, resourceName);
  if (idealState == null) {
    throw new HelixException("Resource: " + resourceName + " has NOT been added yet");
  }

  if (groupId != null && groupId.length() > 0) {
    idealState.setInstanceGroupTag(groupId);
  }
  idealState.setReplicas(Integer.toString(replica));
  int partitions = idealState.getNumPartitions();
  String stateModelName = idealState.getStateModelDefRef();
  StateModelDefinition stateModDef = getStateModelDef(clusterName, stateModelName);

  if (stateModDef == null) {
    throw new HelixException("cannot find state model: " + stateModelName);
  }
  // StateModelDefinition def = new StateModelDefinition(stateModDef);

  List<String> statePriorityList = stateModDef.getStatesPriorityList();

  String masterStateValue = null;
  String slaveStateValue = null;
  replica--;

  for (String state : statePriorityList) {
    String count = stateModDef.getNumInstancesPerState(state);
    if (count.equals("1")) {
      if (masterStateValue != null) {
        throw new HelixException("Invalid or unsupported state model definition");
      }
      masterStateValue = state;
    } else if (count.equalsIgnoreCase("R")) {
      if (slaveStateValue != null) {
        throw new HelixException("Invalid or unsupported state model definition");
      }
      slaveStateValue = state;
    } else if (count.equalsIgnoreCase("N")) {
      if (!(masterStateValue == null && slaveStateValue == null)) {
        throw new HelixException("Invalid or unsupported state model definition");
      }
      replica = instanceNames.size() - 1;
      masterStateValue = slaveStateValue = state;
    }
  }
  if (masterStateValue == null && slaveStateValue == null) {
    throw new HelixException("Invalid or unsupported state model definition");
  }

  if (masterStateValue == null) {
    masterStateValue = slaveStateValue;
  }
  if (idealState.getRebalanceMode() != RebalanceMode.FULL_AUTO
      && idealState.getRebalanceMode() != RebalanceMode.USER_DEFINED) {
    ZNRecord newIdealState = DefaultIdealStateCalculator
        .calculateIdealState(instanceNames, partitions, replica, keyPrefix, masterStateValue,
            slaveStateValue);

    // for now keep mapField in SEMI_AUTO mode and remove listField in CUSTOMIZED mode
    if (idealState.getRebalanceMode() == RebalanceMode.SEMI_AUTO) {
      idealState.getRecord().setListFields(newIdealState.getListFields());
      // TODO: need consider to remove this.
      idealState.getRecord().setMapFields(newIdealState.getMapFields());
    }
    if (idealState.getRebalanceMode() == RebalanceMode.CUSTOMIZED) {
      idealState.getRecord().setMapFields(newIdealState.getMapFields());
    }
  } else {
    for (int i = 0; i < partitions; i++) {
      String partitionName = keyPrefix + "_" + i;
      idealState.getRecord().setMapField(partitionName, new HashMap<String, String>());
      idealState.getRecord().setListField(partitionName, new ArrayList<String>());
    }
  }
  setResourceIdealState(clusterName, resourceName, idealState);
}
 

@Override
protected ResourceControllerDataProvider setupClusterDataCache() throws IOException {
  ResourceControllerDataProvider testCache = super.setupClusterDataCache();

  // Set up mock idealstate
  Map<String, IdealState> isMap = new HashMap<>();
  for (String resource : _resourceNames) {
    IdealState is = new IdealState(resource);
    is.setNumPartitions(_partitionNames.size());
    is.setRebalanceMode(IdealState.RebalanceMode.FULL_AUTO);
    is.setStateModelDefRef("MasterSlave");
    is.setReplicas("3");
    is.setRebalancerClassName(WagedRebalancer.class.getName());
    _partitionNames.stream()
        .forEach(partition -> is.setPreferenceList(partition, Collections.emptyList()));
    isMap.put(resource, is);
  }
  when(testCache.getIdealState(anyString())).thenAnswer(
      (Answer<IdealState>) invocationOnMock -> isMap.get(invocationOnMock.getArguments()[0]));
  when(testCache.getIdealStates()).thenReturn(isMap);

  // Set up 2 more instances
  for (int i = 1; i < 3; i++) {
    String instanceName = _testInstanceId + i;
    _instances.add(instanceName);
    // 1. Set up the default instance information with capacity configuration.
    InstanceConfig testInstanceConfig = createMockInstanceConfig(instanceName);
    Map<String, InstanceConfig> instanceConfigMap = testCache.getInstanceConfigMap();
    instanceConfigMap.put(instanceName, testInstanceConfig);
    when(testCache.getInstanceConfigMap()).thenReturn(instanceConfigMap);
    // 2. Mock the live instance node for the default instance.
    LiveInstance testLiveInstance = createMockLiveInstance(instanceName);
    Map<String, LiveInstance> liveInstanceMap = testCache.getLiveInstances();
    liveInstanceMap.put(instanceName, testLiveInstance);
    when(testCache.getLiveInstances()).thenReturn(liveInstanceMap);
    when(testCache.getEnabledInstances()).thenReturn(liveInstanceMap.keySet());
    when(testCache.getEnabledLiveInstances()).thenReturn(liveInstanceMap.keySet());
    when(testCache.getAllInstances()).thenReturn(_instances);
  }

  return testCache;
}