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

/**
 * Test enabledWagedRebalance by checking the rebalancer class name changed.
 */
@Test
public void testEnableWagedRebalance() {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;
  String testResourcePrefix = "TestResource";
  HelixAdmin admin = new ZKHelixAdmin(_gZkClient);
  admin.addCluster(clusterName, true);
  admin.addStateModelDef(clusterName, "MasterSlave", new MasterSlaveSMD());

  // Add an IdealState
  IdealState idealState = new IdealState(testResourcePrefix);
  idealState.setNumPartitions(3);
  idealState.setStateModelDefRef("MasterSlave");
  idealState.setRebalanceMode(IdealState.RebalanceMode.FULL_AUTO);
  admin.addResource(clusterName, testResourcePrefix, idealState);

  admin.enableWagedRebalance(clusterName, Collections.singletonList(testResourcePrefix));
  IdealState is = admin.getResourceIdealState(clusterName, testResourcePrefix);
  Assert.assertEquals(is.getRebalancerClassName(), WagedRebalancer.class.getName());
}
 

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

private void setupUnbalancedDB() throws InterruptedException {
  // Start with Full-Auto mode to create the partitions, Semi-Auto won't create partitions.
  _gSetupTool.addResourceToCluster(CLUSTER_NAME, UNBALANCED_DB_NAME, 50, MASTER_SLAVE_STATE_MODEL,
      IdealState.RebalanceMode.FULL_AUTO.toString());
  _gSetupTool.rebalanceStorageCluster(CLUSTER_NAME, UNBALANCED_DB_NAME, 1);

  // Set preference list to put all partitions to one instance.
  IdealState idealState = _gSetupTool.getClusterManagementTool()
      .getResourceIdealState(CLUSTER_NAME, UNBALANCED_DB_NAME);
  Set<String> partitions = idealState.getPartitionSet();
  for (String partition : partitions) {
    idealState.setPreferenceList(partition,
        Lists.newArrayList(_blockedParticipant.getInstanceName()));
  }
  idealState.setRebalanceMode(IdealState.RebalanceMode.SEMI_AUTO);

  _gSetupTool.getClusterManagementTool().setResourceIdealState(CLUSTER_NAME, UNBALANCED_DB_NAME,
      idealState);

  Assert.assertTrue(_clusterVerifier.verifyByPolling(10000, 100));
}
 

@BeforeMethod
public void beforeMethod() {
  _changeTypes.clear();
  _instanceConfigMap.clear();
  _idealStateMap.clear();
  _resourceConfigMap.clear();

  _changeTypes.add(HelixConstants.ChangeType.INSTANCE_CONFIG);
  _changeTypes.add(HelixConstants.ChangeType.IDEAL_STATE);
  _changeTypes.add(HelixConstants.ChangeType.RESOURCE_CONFIG);
  _changeTypes.add(HelixConstants.ChangeType.CLUSTER_CONFIG);

  _instanceConfigMap.put(INSTANCE_NAME, new InstanceConfig(INSTANCE_NAME));
  IdealState idealState = new IdealState(RESOURCE_NAME);
  idealState.setRebalanceMode(IdealState.RebalanceMode.FULL_AUTO);
  _idealStateMap.put(RESOURCE_NAME, idealState);
  _resourceConfigMap.put(RESOURCE_NAME, new ResourceConfig(RESOURCE_NAME));
  _clusterConfig = new ClusterConfig(CLUSTER_NAME);

  _dataProvider =
      getMockDataProvider(_changeTypes, _instanceConfigMap, _idealStateMap, _resourceConfigMap,
          _clusterConfig);
}
 

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

}
 

/**
 * Create an IdealState for a resource that belongs to a resource group We use
 * "resourceGroupName$resourceInstanceTag" as the IdealState znode name to differetiate different
 * resources from the same resourceGroup.
 */
public IdealState createIdealStateForResourceGroup(String resourceGroupName,
    String resourceTag, int numPartition, int replica, String rebalanceMode, String stateModelDefName) {
  String idealStateId = genIdealStateNameWithResourceTag(resourceGroupName, resourceTag);
  IdealState idealState = new IdealState(idealStateId);
  idealState.setNumPartitions(numPartition);
  idealState.setStateModelDefRef(stateModelDefName);
  IdealState.RebalanceMode mode =
      idealState.rebalanceModeFromString(rebalanceMode, IdealState.RebalanceMode.SEMI_AUTO);
  idealState.setRebalanceMode(mode);
  idealState.setReplicas("" + replica);
  idealState.setStateModelFactoryName(HelixConstants.DEFAULT_STATE_MODEL_FACTORY);
  idealState.setResourceGroupName(resourceGroupName);
  idealState.setInstanceGroupTag(resourceTag);
  idealState.enableGroupRouting(true);

  return idealState;
}
 

protected List<IdealState> setupIdealState(int nodes, String[] resources, int partitions,
    int replicas, RebalanceMode rebalanceMode, String stateModelName, String rebalanceClassName,
    String rebalanceStrategyName, int minActiveReplica) {
  List<IdealState> idealStates = new ArrayList<IdealState>();
  for (String resourceName : resources) {
    ZNRecord record = new ZNRecord(resourceName);
    for (int p = 0; p < partitions; p++) {
      List<String> value = new ArrayList<String>();
      for (int r = 0; r < replicas; r++) {
        value.add(HOSTNAME_PREFIX + (p + r + 1) % nodes);
      }
      record.setListField(resourceName + "_" + p, value);
    }
    IdealState idealState = new IdealState(record);
    idealState.setStateModelDefRef(stateModelName);
    if (rebalanceClassName != null) {
      idealState.setRebalancerClassName(rebalanceClassName);
    }
    if (rebalanceStrategyName != null) {
      idealState.setRebalanceStrategy(rebalanceStrategyName);
    }
    idealState.setRebalanceMode(rebalanceMode);
    idealState.setNumPartitions(partitions);
    idealStates.add(idealState);
    idealState.setReplicas(String.valueOf(replicas));

    if (minActiveReplica > 0) {
      idealState.setMinActiveReplicas(minActiveReplica);
    }

    Builder keyBuilder = accessor.keyBuilder();

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

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

@BeforeClass
public void setUp() {
  startZk();
  startController();

  IdealState idealState = new IdealState(RESOURCE_NAME);
  idealState.setStateModelDefRef("OnlineOffline");
  idealState.setRebalanceMode(IdealState.RebalanceMode.CUSTOMIZED);
  idealState.setReplicas("0");
  _helixAdmin.addResource(getHelixClusterName(), RESOURCE_NAME, 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]));

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

  return testCache;
}
 

private IdealState generateNewIdealState(String resourceName, IdealState currentIdealState,
    ZNRecord newMapping) {
  IdealState newIdealState = new IdealState(resourceName);
  newIdealState.getRecord().setSimpleFields(currentIdealState.getRecord().getSimpleFields());
  newIdealState.setRebalanceMode(currentIdealState.getRebalanceMode());
  newIdealState.getRecord().setListFields(newMapping.getListFields());

  return newIdealState;
}
 

@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(dependsOnMethods = "deleteFromResourceIdealState")
public void testAddResourceWithWeight() throws IOException {
  // Test case 1: Add a valid resource with valid weights
  // Create a resource with IdealState and ResourceConfig
  String wagedResourceName = "newWagedResource";

  // Create an IdealState on full-auto with 1 partition
  IdealState idealState = new IdealState(wagedResourceName);
  idealState.getRecord().getSimpleFields().putAll(_gSetupTool.getClusterManagementTool()
      .getResourceIdealState(CLUSTER_NAME, RESOURCE_NAME).getRecord().getSimpleFields());
  idealState.setRebalanceMode(IdealState.RebalanceMode.FULL_AUTO);
  idealState.setRebalancerClassName(WagedRebalancer.class.getName());
  idealState.setNumPartitions(1); // 1 partition for convenience of testing

  // Create a ResourceConfig with FOO and BAR at 100 respectively
  ResourceConfig resourceConfig = new ResourceConfig(wagedResourceName);
  Map<String, Map<String, Integer>> partitionCapacityMap = new HashMap<>();
  Map<String, Integer> partitionCapacity = ImmutableMap.of("FOO", 100, "BAR", 100);
  partitionCapacityMap.put(wagedResourceName + "_0", partitionCapacity);
  // Also add a default key
  partitionCapacityMap.put(ResourceConfig.DEFAULT_PARTITION_KEY, partitionCapacity);
  resourceConfig.setPartitionCapacityMap(partitionCapacityMap);

  // Put both IdealState and ResourceConfig into a map as required
  Map<String, ZNRecord> inputMap = ImmutableMap.of(
      ResourceAccessor.ResourceProperties.idealState.name(), idealState.getRecord(),
      ResourceAccessor.ResourceProperties.resourceConfig.name(), resourceConfig.getRecord());

  // Create an entity using the inputMap
  Entity entity =
      Entity.entity(OBJECT_MAPPER.writeValueAsString(inputMap), MediaType.APPLICATION_JSON_TYPE);

  // Make a HTTP call to the REST endpoint
  put("clusters/" + CLUSTER_NAME + "/resources/" + wagedResourceName,
      ImmutableMap.of("command", "addWagedResource"), entity, Response.Status.OK.getStatusCode());

  // Test case 2: Add a resource with invalid weights
  String invalidResourceName = "invalidWagedResource";
  ResourceConfig invalidWeightResourceConfig = new ResourceConfig(invalidResourceName);
  IdealState invalidWeightIdealState = new IdealState(invalidResourceName);

  Map<String, ZNRecord> invalidInputMap = ImmutableMap.of(
      ResourceAccessor.ResourceProperties.idealState.name(), invalidWeightIdealState.getRecord(),
      ResourceAccessor.ResourceProperties.resourceConfig.name(),
      invalidWeightResourceConfig.getRecord());

  // Create an entity using invalidInputMap
  entity = Entity.entity(OBJECT_MAPPER.writeValueAsString(invalidInputMap),
      MediaType.APPLICATION_JSON_TYPE);

  // Make a HTTP call to the REST endpoint
  put("clusters/" + CLUSTER_NAME + "/resources/" + invalidResourceName,
      ImmutableMap.of("command", "addWagedResource"), entity,
      Response.Status.BAD_REQUEST.getStatusCode());
}
 

@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 testDisableNodeCustomIS() throws Exception {
  // Logger.getRootLogger().setLevel(Level.INFO);
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;
  final int n = 5;
  String disableNode = "localhost_12918";

  System.out.println("START " + clusterName + " at " + new Date(System.currentTimeMillis()));

  MockParticipantManager[] participants = new MockParticipantManager[n];

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

  // set ideal state to customized mode
  ZkBaseDataAccessor<ZNRecord> baseAccessor = new ZkBaseDataAccessor<>(_gZkClient);
  ZKHelixDataAccessor accessor = new ZKHelixDataAccessor(clusterName, baseAccessor);
  Builder keyBuilder = accessor.keyBuilder();
  IdealState idealState = accessor.getProperty(keyBuilder.idealStates("TestDB0"));
  idealState.setRebalanceMode(RebalanceMode.CUSTOMIZED);
  accessor.setProperty(keyBuilder.idealStates("TestDB0"), idealState);

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

  // start participants
  for (int i = 0; i < n; i++) {
    String instanceName = "localhost_" + (12918 + i);

    participants[i] = new MockParticipantManager(ZK_ADDR, clusterName, instanceName);
    participants[i].syncStart();
  }

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

  // disable localhost_12918
  String command =
      "--zkSvr " + ZK_ADDR + " --enableInstance " + clusterName + " " + disableNode + " false";
  ClusterSetup.processCommandLineArgs(command.split("\\s+"));
  Assert.assertTrue(_clusterVerifier.verifyByPolling());

  // make sure localhost_12918 is in OFFLINE state
  Map<String, Map<String, String>> expectStateMap = new HashMap<>();
  Map<String, String> expectInstanceStateMap = new HashMap<>();
  expectInstanceStateMap.put(disableNode, "OFFLINE");
  expectStateMap.put(".*", expectInstanceStateMap);
  boolean result =
      ZkTestHelper.verifyState(_gZkClient, clusterName, "TestDB0", expectStateMap, "==");
  Assert.assertTrue(result, disableNode + " should be in OFFLINE");

  // re-enable localhost_12918
  command =
      "--zkSvr " + ZK_ADDR + " --enableInstance " + clusterName + " " + disableNode + " true";
  ClusterSetup.processCommandLineArgs(command.split("\\s+"));
  Assert.assertTrue(_clusterVerifier.verifyByPolling());

  // make sure localhost_12918 is NOT in OFFLINE state
  result = ZkTestHelper.verifyState(_gZkClient, clusterName, "TestDB0", expectStateMap, "!=");
  Assert.assertTrue(result, disableNode + " should NOT be in OFFLINE");

  // clean up
  // wait for all zk callbacks done
  controller.syncStop();
  for (int i = 0; i < 5; i++) {
    participants[i].syncStop();
  }

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

@Test
public void testIgnoreTrimmableFieldChanges() {
  // Fill mock data to initialize the detector
  ResourceChangeDetector detector = new ResourceChangeDetector(true);
  detector.updateSnapshots(_dataProvider);

  // Verify that all the trimmable field changes will not be detected
  // 1. Cluster Config
  changeTrimmableValuesAndVerifyDetector(FieldType.values(), _clusterConfig, detector,
      _dataProvider);
  // 2. Ideal States
  for (IdealState idealState : _idealStateMap.values()) {
    changeTrimmableValuesAndVerifyDetector(FieldType.values(), idealState, detector,
        _dataProvider);
    // Additional test to ensure Ideal State map/list fields are detected correctly according to
    // the rebalance mode.

    // SEMI_AUTO: List fields are non-trimmable
    idealState.setRebalanceMode(IdealState.RebalanceMode.SEMI_AUTO);
    // refresh the detector cache after modification to avoid unexpected change detected.
    detector.updateSnapshots(_dataProvider);
    changeTrimmableValuesAndVerifyDetector(
        new FieldType[] { FieldType.SIMPLE_FIELD, FieldType.MAP_FIELD }, idealState, detector,
        _dataProvider);

    // CUSTOMZIED: List and Map fields are non-trimmable
    idealState.setRebalanceMode(IdealState.RebalanceMode.CUSTOMIZED);
    // refresh the detector cache after modification to avoid unexpected change detected.
    detector.updateSnapshots(_dataProvider);
    changeTrimmableValuesAndVerifyDetector(new FieldType[] { FieldType.SIMPLE_FIELD }, idealState,
        detector, _dataProvider);
  }
  // 3. Resource Config
  for (ResourceConfig resourceConfig : _resourceConfigMap.values()) {
    // Preference lists in the list fields are non-trimmable
    changeTrimmableValuesAndVerifyDetector(
        new FieldType[] { FieldType.SIMPLE_FIELD, FieldType.MAP_FIELD }, resourceConfig, detector,
        _dataProvider);
  }
  // 4. Instance Config
  for (InstanceConfig instanceConfig : _instanceConfigMap.values()) {
    changeTrimmableValuesAndVerifyDetector(FieldType.values(), instanceConfig, detector,
        _dataProvider);
  }
}
 

private void testMultipleResourcesAndPercent(double percentReady) {
  final long now = System.currentTimeMillis();
  random = new Random(now);
  final String clusterName = "noSuchCluster";
  final List<String> tables = new ArrayList<>();
  final int tableCount = 2500 + random.nextInt(100);
  int readyTables = 0;
  Map<String, IdealState> idealStates = new HashMap<>();
  Map<String, ExternalView> externalViews = new HashMap<>();

  for (int i = 1; i <= tableCount; i++) {
    final String tableName = generateRandomString(10) + String.valueOf(i);
    tables.add(tableName);
    final String segmentName = "segment1";
    String evState;
    if (random.nextDouble() * 100 < percentReady) {
      evState = "ONLINE";
      readyTables++;
    } else {
      evState = "OFFLINE";
    }
    IdealState idealState = new IdealState(tableName);
    idealState.setRebalanceMode(IdealState.RebalanceMode.CUSTOMIZED);
    idealState.setPartitionState(segmentName, INSTANCE_NAME, "ONLINE");
    ExternalView externalView = new ExternalView(tableName);
    externalView.setState(segmentName, INSTANCE_NAME, evState);
    idealStates.put(tableName, idealState);
    externalViews.put(tableName, externalView);
  }

  final double actualReadyPercent = (double) readyTables * 100 / tableCount;
  double lowestReadyPercent = (int) Math.round(actualReadyPercent);
  lowestReadyPercent = lowestReadyPercent > 2 ? lowestReadyPercent - 2 : 1;  // Should be 2 below  percentReady

  // Create ServiceCallback objects with minReadyPercent set to values between lowestReadyPercent and 100.
  // Call getServiceStatus() enough number of times so that we are only left with the tables
  // that are not ready yet. We need to call getServiceStatus() at most tableCount times.
  for (double minReadyPercent = lowestReadyPercent; minReadyPercent <= 100; minReadyPercent += 0.1) {
    TestMultiResourceISAndEVMatchCB callback =
        new TestMultiResourceISAndEVMatchCB(clusterName, INSTANCE_NAME, tables, minReadyPercent);
    callback.setIdealStates(idealStates);
    callback.setExternalViews(externalViews);

    ServiceStatus.Status status = callback.getServiceStatus();
    // we need to call getServiceStatus() at most the number of bad tables plus 1,
    // to get a STARTED condition if we can. After that, the return value should
    // never change.
    final int nBadTables = tableCount - readyTables;
    for (int i = 0; i <= nBadTables; i++) {
      status = callback.getServiceStatus();
    }

    ServiceStatus.Status expected =
        minReadyPercent > actualReadyPercent ? ServiceStatus.Status.STARTING : ServiceStatus.Status.GOOD;
    String errorMsg = "Mismatch at " + minReadyPercent + "%, tableCount=" + tableCount + ", percentTablesReady="
        + actualReadyPercent + ":" + callback.getStatusDescription();
    Assert.assertEquals(status, expected, errorMsg);

    // The status should never change going forward from here.
    for (int i = nBadTables + 1; i < tableCount; i++) {
      ServiceStatus.Status laterStatus = callback.getServiceStatus();
      String msg = "Mismatch at " + minReadyPercent + "%, tableCount=" + tableCount + ", percentTablesReady="
          + actualReadyPercent + ", i=" + i + ":" + callback.getStatusDescription();
      Assert.assertEquals(laterStatus, status, msg);
    }
  }
}
 

@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
@PrepareForTest({ControllerUtil.class})
public void testLoadFileSet() throws Exception {
  PowerMockito.mockStatic(ControllerUtil.class);
  TerrapinLoadRequest request =
      new TerrapinLoadRequest("fileset", "/terrapin/data/fileset/12345564534", 200);
  Options requestOptions = new Options();
  requestOptions.setNumVersionsToKeep(2);
  request.setOptions(requestOptions);
  FileSetInfo fileSetInfo = new FileSetInfo(
      request.getFileSet(), "/terrapin/data/fileset/12345563434",
      request.getExpectedNumPartitions(),
      Lists.newArrayList(mock(FileSetInfo.ServingInfo.class)), new Options()
  );
  String resourceName = TerrapinUtil.hdfsDirToHelixResource(request.getHdfsDirectory());
  CustomModeISBuilder idealStateBuilder = new CustomModeISBuilder(resourceName);
  idealStateBuilder.assignInstanceAndState(resourceName + "$0", "host0", "ONLINE");
  idealStateBuilder.assignInstanceAndState(resourceName + "$1", "host1", "ONLINE");
  idealStateBuilder.setStateModel("OnlineOffline");
  idealStateBuilder.setNumReplica(2);
  idealStateBuilder.setNumPartitions(request.getExpectedNumPartitions());
  IdealState is = idealStateBuilder.build();
  is.setBucketSize(2);
  is.setRebalanceMode(IdealState.RebalanceMode.CUSTOMIZED);
  ViewInfo viewInfo1 = mock(ViewInfo.class);
  ViewInfo viewInfo2 = mock(ViewInfo.class);

  when(configuration.getInt(eq(Constants.NUM_SERVING_REPLICAS), eq(3))).thenReturn(3);
  when(zkManager.getFileSetInfo(eq(request.getFileSet()))).thenReturn(fileSetInfo);
  when(zkManager.getViewInfo(eq(resourceName))).thenReturn(viewInfo1).thenReturn(viewInfo2);
  doNothing().when(zkManager).setFileSetInfo(eq(request.getFileSet()), any(FileSetInfo.class));
  when(viewInfo1.getNumOnlinePartitions()).thenReturn(request.getExpectedNumPartitions() / 2);
  when(viewInfo2.getNumOnlinePartitions()).thenReturn(request.getExpectedNumPartitions());
  when(helixAdmin.getResourcesInCluster(CLUSTER)).thenReturn(new ArrayList<String>());
  when(ControllerUtil.buildIdealStateForHdfsDir(any(DFSClient.class),
      anyString(), anyString(), any(PartitionerType.class), anyInt(), anyBoolean())).
          thenReturn(is);
  doNothing().when(helixAdmin).addResource(eq(CLUSTER), eq(resourceName), eq(is));
  doNothing().when(helixAdmin).addResource(eq(CLUSTER), eq(resourceName),
      eq(is.getNumPartitions()), eq("OnlineOffline"), eq("CUSTOMIZED"), eq(is.getBucketSize()));
  doNothing().when(helixAdmin).setResourceIdealState(eq(CLUSTER), eq(resourceName), eq(is));

  serviceImpl.loadFileSet(request).apply();

  ArgumentCaptor<FileSetInfo> fileSetInfoCaptor = ArgumentCaptor.forClass(FileSetInfo.class);
  verify(zkManager).setFileSetInfo(eq(request.getFileSet()), fileSetInfoCaptor.capture());
  FileSetInfo capturedInfo = fileSetInfoCaptor.getValue();
  assertEquals(request.getFileSet(), capturedInfo.fileSetName);
  assertEquals(request.getExpectedNumPartitions(), capturedInfo.servingInfo.numPartitions);
  assertEquals(Lists.newArrayList(fileSetInfo.servingInfo),
      (ArrayList)capturedInfo.oldServingInfoList);
}
 

@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;
}