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

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

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

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

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

protected List<IdealState> setupIdealState(String clusterName, int[] nodes, String[] resources,
    int partitions, int replicas) {
  ZKHelixDataAccessor accessor =
      new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor<>(_gZkClient));
  Builder keyBuilder = accessor.keyBuilder();

  List<IdealState> idealStates = new ArrayList<>();
  List<String> instances = new ArrayList<>();
  for (int i : nodes) {
    instances.add("localhost_" + i);
  }

  for (String resourceName : resources) {
    IdealState idealState = new IdealState(resourceName);
    for (int p = 0; p < partitions; p++) {
      List<String> value = new ArrayList<>();
      for (int r = 0; r < replicas; r++) {
        int n = nodes[(p + r) % nodes.length];
        value.add("localhost_" + n);
      }
      idealState.getRecord().setListField(resourceName + "_" + p, value);
    }

    idealState.setReplicas(Integer.toString(replicas));
    idealState.setStateModelDefRef("MasterSlave");
    idealState.setRebalanceMode(IdealState.RebalanceMode.SEMI_AUTO);
    idealState.setNumPartitions(partitions);
    idealStates.add(idealState);

    // System.out.println(idealState);
    accessor.setProperty(keyBuilder.idealStates(resourceName), idealState);
  }
  return idealStates;
}
 

public static Map<String, IdealState> getIdealStatesFromAssignment(
    Set<InstanceTopicPartitionHolder> newAssignment,
    Set<TopicPartition> blacklistedTopicPartitions) {
  Map<String, CustomModeISBuilder> idealStatesBuilderMap = new HashMap<>();
  for (InstanceTopicPartitionHolder instance : newAssignment) {
    for (TopicPartition tpi : instance.getServingTopicPartitionSet()) {
      String topicName = tpi.getTopic();
      String partition = Integer.toString(tpi.getPartition());
      if (!idealStatesBuilderMap.containsKey(topicName)) {
        final CustomModeISBuilder customModeIdealStateBuilder =
            new CustomModeISBuilder(topicName);
        customModeIdealStateBuilder
            .setStateModel(OnlineOfflineStateModel.name)
            .setNumReplica(1);

        idealStatesBuilderMap.put(topicName, customModeIdealStateBuilder);
      }
      String state = Constants.HELIX_ONLINE_STATE;
      if (blacklistedTopicPartitions.contains(tpi)) {
        state = Constants.HELIX_OFFLINE_STATE;
      }
      idealStatesBuilderMap.get(topicName).assignInstanceAndState(partition,
          instance.getInstanceName(),
          state);
    }
  }
  Map<String, IdealState> idealStatesMap = new HashMap<>();
  for (String topic : idealStatesBuilderMap.keySet()) {
    IdealState idealState = idealStatesBuilderMap.get(topic).build();
    idealState.setMaxPartitionsPerInstance(idealState.getPartitionSet().size());
    idealState.setNumPartitions(idealState.getPartitionSet().size());
    idealStatesMap.put(topic, idealState);
  }
  return idealStatesMap;
}
 

private void updateIdealState() {
  HelixDataAccessor accessor = _manager.getHelixDataAccessor();
  Builder keyBuilder = accessor.keyBuilder();
  List<String> idealStates = accessor.getChildNames(keyBuilder.idealStates());
  IdealState idealState = accessor.getProperty(keyBuilder.idealStates(idealStates.get(0)));
  idealState.setNumPartitions((int)(System.currentTimeMillis()%50L));
  accessor.setProperty(keyBuilder.idealStates(idealState.getId()), idealState);
}
 

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

@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 = "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 addResourceWithWeight() and validateResourcesForWagedRebalance() by trying to add a resource with incomplete ResourceConfig.
 */
@Test
public void testAddResourceWithWeightAndValidation()
    throws IOException {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;
  String mockInstance = "MockInstance";
  String testResourcePrefix = "TestResource";
  HelixAdmin admin = new ZKHelixAdmin(_gZkClient);
  admin.addCluster(clusterName, true);
  admin.addStateModelDef(clusterName, "MasterSlave", new MasterSlaveSMD());

  // Create a dummy instance
  InstanceConfig instanceConfig = new InstanceConfig(mockInstance);
  Map<String, Integer> mockInstanceCapacity =
      ImmutableMap.of("WCU", 100, "RCU", 100, "STORAGE", 100);
  instanceConfig.setInstanceCapacityMap(mockInstanceCapacity);
  admin.addInstance(clusterName, instanceConfig);
  MockParticipantManager mockParticipantManager =
      new MockParticipantManager(ZK_ADDR, clusterName, mockInstance);
  mockParticipantManager.syncStart();

  IdealState idealState = new IdealState(testResourcePrefix);
  idealState.setNumPartitions(3);
  idealState.setStateModelDefRef("MasterSlave");
  idealState.setRebalanceMode(IdealState.RebalanceMode.FULL_AUTO);

  ResourceConfig resourceConfig = new ResourceConfig(testResourcePrefix);
  // validate
  Map<String, Boolean> validationResult = admin.validateResourcesForWagedRebalance(clusterName,
      Collections.singletonList(testResourcePrefix));
  Assert.assertEquals(validationResult.size(), 1);
  Assert.assertFalse(validationResult.get(testResourcePrefix));
  try {
    admin.addResourceWithWeight(clusterName, idealState, resourceConfig);
    Assert.fail();
  } catch (HelixException e) {
    // OK since resourceConfig is empty
  }

  // Set PARTITION_CAPACITY_MAP
  Map<String, String> capacityDataMap =
      ImmutableMap.of("WCU", "1", "RCU", "2", "STORAGE", "3");
  resourceConfig.getRecord()
      .setMapField(ResourceConfig.ResourceConfigProperty.PARTITION_CAPACITY_MAP.name(),
          Collections.singletonMap(ResourceConfig.DEFAULT_PARTITION_KEY,
              OBJECT_MAPPER.writeValueAsString(capacityDataMap)));

  // validate
  validationResult = admin.validateResourcesForWagedRebalance(clusterName,
      Collections.singletonList(testResourcePrefix));
  Assert.assertEquals(validationResult.size(), 1);
  Assert.assertFalse(validationResult.get(testResourcePrefix));

  // Add the capacity key to ClusterConfig
  HelixDataAccessor dataAccessor = new ZKHelixDataAccessor(clusterName, _baseAccessor);
  PropertyKey.Builder keyBuilder = dataAccessor.keyBuilder();
  ClusterConfig clusterConfig = dataAccessor.getProperty(keyBuilder.clusterConfig());
  clusterConfig.setInstanceCapacityKeys(Arrays.asList("WCU", "RCU", "STORAGE"));
  dataAccessor.setProperty(keyBuilder.clusterConfig(), clusterConfig);

  // Should succeed now
  Assert.assertTrue(admin.addResourceWithWeight(clusterName, idealState, resourceConfig));
  // validate
  validationResult = admin.validateResourcesForWagedRebalance(clusterName,
      Collections.singletonList(testResourcePrefix));
  Assert.assertEquals(validationResult.size(), 1);
  Assert.assertTrue(validationResult.get(testResourcePrefix));
}
 

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

/**
 * This method will be invoked when there is a change in any subscribed
 * notificationTypes
 * @throws Exception
 */
public void start() throws Exception {
  if (_callback == null || _notificationTypes == null || _notificationTypes.size() == 0
      || _resourceName == null) {
    throw new IllegalArgumentException("Require callback | notificationTypes | resourceName");
  }

  LOG.info("Register participantLeader on " + _notificationTypes + " using " + _resourceName);

  _stateModelFty = new GenericLeaderStandbyStateModelFactory(_callback, _notificationTypes);

  StateMachineEngine stateMach = _manager.getStateMachineEngine();
  stateMach.registerStateModelFactory(LEADER_STANDBY, _stateModelFty, _resourceName);
  HelixZkClient zkClient = null;
  try {
    // manually add ideal state for participant leader using LeaderStandby
    // model
    HelixZkClient.ZkClientConfig clientConfig = new HelixZkClient.ZkClientConfig();
    clientConfig.setZkSerializer(new ZNRecordSerializer());
    zkClient = SharedZkClientFactory.getInstance()
        .buildZkClient(new HelixZkClient.ZkConnectionConfig(_zkAddr), clientConfig);

    HelixDataAccessor accessor =
        new ZKHelixDataAccessor(_manager.getClusterName(), new ZkBaseDataAccessor<>(zkClient));
    Builder keyBuilder = accessor.keyBuilder();

    IdealState idealState = new IdealState(_resourceName);
    idealState.setRebalanceMode(RebalanceMode.SEMI_AUTO);
    idealState.setReplicas(IdealState.IdealStateConstants.ANY_LIVEINSTANCE.toString());
    idealState.setNumPartitions(1);
    idealState.setStateModelDefRef(LEADER_STANDBY);
    idealState.setStateModelFactoryName(_resourceName);
    List<String> prefList = new ArrayList<String>(
        Arrays.asList(IdealState.IdealStateConstants.ANY_LIVEINSTANCE.toString()));
    idealState.getRecord().setListField(_resourceName + "_0", prefList);

    List<String> idealStates = accessor.getChildNames(keyBuilder.idealStates());
    while (idealStates == null || !idealStates.contains(_resourceName)) {
      accessor.setProperty(keyBuilder.idealStates(_resourceName), idealState);
      idealStates = accessor.getChildNames(keyBuilder.idealStates());
    }

    LOG.info(
        "Set idealState for participantLeader:" + _resourceName + ", idealState:" + idealState);
  } finally {
    if (zkClient != null && !zkClient.isClosed()) {
      zkClient.close();
    }
  }
}
 

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

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

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

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

public static IdealState addNewRealtimeSegmentToIdealState(String segmentId, IdealState state, String instanceName) {
  state.setPartitionState(segmentId, instanceName, ONLINE);
  state.setNumPartitions(state.getNumPartitions() + 1);
  return state;
}
 

public static IdealState calculateRelayIdealState(List<String> partitions,
    List<String> instances, String resultRecordName, int replica, String firstValue,
    String restValue, String stateModelName) {
  Collections.sort(partitions);
  Collections.sort(instances);
  if (instances.size() % replica != 0) {
    throw new HelixException("Instances must be divided by replica");
  }

  IdealState result = new IdealState(resultRecordName);
  result.setNumPartitions(partitions.size());
  result.setReplicas("" + replica);
  result.setStateModelDefRef(stateModelName);

  int groups = instances.size() / replica;
  int remainder = instances.size() % replica;

  int remainder2 = partitions.size() % groups;
  int storageNodeGroupSize = partitions.size() / groups;

  for (int i = 0; i < groups; i++) {
    int relayStart = 0, relayEnd = 0, storageNodeStart = 0, storageNodeEnd = 0;
    if (i < remainder) {
      relayStart = (replica + 1) * i;
      relayEnd = (replica + 1) * (i + 1);
    } else {
      relayStart = (replica + 1) * remainder + replica * (i - remainder);
      relayEnd = relayStart + replica;
    }
    // System.out.println("relay start :" + relayStart + " relayEnd:" + relayEnd);
    if (i < remainder2) {
      storageNodeStart = (storageNodeGroupSize + 1) * i;
      storageNodeEnd = (storageNodeGroupSize + 1) * (i + 1);
    } else {
      storageNodeStart =
          (storageNodeGroupSize + 1) * remainder2 + storageNodeGroupSize * (i - remainder2);
      storageNodeEnd = storageNodeStart + storageNodeGroupSize;
    }

    // System.out.println("storageNodeStart :" + storageNodeStart + " storageNodeEnd:" +
    // storageNodeEnd);
    List<String> snBatch = partitions.subList(storageNodeStart, storageNodeEnd);
    List<String> relayBatch = instances.subList(relayStart, relayEnd);

    Map<String, List<String>> sublistFields =
        calculateSubIdealState(snBatch, relayBatch, replica);

    result.getRecord().getListFields().putAll(sublistFields);
  }

  for (String snName : result.getRecord().getListFields().keySet()) {
    Map<String, String> mapField = new TreeMap<String, String>();
    List<String> relayCandidates = result.getRecord().getListField(snName);
    mapField.put(relayCandidates.get(0), firstValue);
    for (int i = 1; i < relayCandidates.size(); i++) {
      mapField.put(relayCandidates.get(i), restValue);
    }
    result.getRecord().getMapFields().put(snName, mapField);
  }
  System.out.println();
  return result;
}
 

/**
 * @return
 */
public IdealState build() {
  IdealState idealstate = new IdealState(_record);
  idealstate.setNumPartitions(numPartitions);
  idealstate.setStateModelDefRef(stateModel);
  idealstate.setStateModelFactoryName(stateModelFactoryName);
  idealstate.setRebalanceMode(rebalancerMode);
  idealstate.setReplicas("" + numReplica);

  if (minActiveReplica >= 0) {
    idealstate.setMinActiveReplicas(minActiveReplica);
  }

  if (rebalancerClassName != null) {
    idealstate.setRebalancerClassName(rebalancerClassName);
  }

  if (rebalanceStrategy != null) {
    idealstate.setRebalanceStrategy(rebalanceStrategy);
  }

  if (maxPartitionsPerNode > 0) {
    idealstate.setMaxPartitionsPerInstance(maxPartitionsPerNode);
  }

  if (disableExternalView != null) {
    idealstate.setDisableExternalView(disableExternalView);
  }

  if (enableGroupRouting != null) {
    idealstate.enableGroupRouting(enableGroupRouting);
  }

  if (resourceGroupName != null) {
    idealstate.setResourceGroupName(resourceGroupName);
  }

  if (resourceType != null) {
    idealstate.setResourceType(resourceType);
  }

  if (rebalanceDelayInMs >= 0) {
    idealstate.setRebalanceDelay(rebalanceDelayInMs);
  }

  if (delayRebalanceEnabled != null) {
    idealstate.setDelayRebalanceEnabled(delayRebalanceEnabled);
  }

  if (!idealstate.isValid()) {
    throw new HelixException("invalid ideal-state: " + idealstate);
  }
  return idealstate;
}