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

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

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

/**
 * This test was written because there is an edge case where an instance becomes disabled while a
 * partition is bootstrapping by way of pending
 * messages.
 * The newly bootstrapped partitions never get further state transitions because the instance
 * won't ever get added to instanceStateMap (this issue has been fixed). In other words, if there
 * are mapping changes while a partition is bootstrapping, the final state should go into the best
 * possible mapping for clusters to converge correctly.
 */
@Test
public void testDisabledBootstrappingPartitions() {
  String resourceName = "Test";
  String partitionName = "Test0";
  String instanceName = "localhost";
  String stateModelName = "OnlineOffline";
  StateModelDefinition stateModelDef = new OnlineOfflineSMD();

  IdealState idealState = new IdealState(resourceName);
  idealState.setStateModelDefRef(stateModelName);
  idealState.setPartitionState(partitionName, instanceName, "ONLINE");

  Resource resource = new Resource(resourceName);
  resource.addPartition(partitionName);

  CustomRebalancer customRebalancer = new CustomRebalancer();
  ResourceControllerDataProvider cache = mock(ResourceControllerDataProvider.class);
  when(cache.getStateModelDef(stateModelName)).thenReturn(stateModelDef);
  when(cache.getDisabledInstancesForPartition(resource.getResourceName(), partitionName))
      .thenReturn(ImmutableSet.of(instanceName));
  when(cache.getLiveInstances())
      .thenReturn(ImmutableMap.of(instanceName, new LiveInstance(instanceName)));

  CurrentStateOutput currOutput = new CurrentStateOutput();
  ResourceAssignment resourceAssignment =
      customRebalancer.computeBestPossiblePartitionState(cache, idealState, resource, currOutput);

  Assert.assertEquals(
      resourceAssignment.getReplicaMap(new Partition(partitionName)).get(instanceName),
      stateModelDef.getInitialState());
}
 

protected List<IdealState> setupIdealStateDeprecated(int nodes, String[] resources,
    int partitions, int replicas, IdealStateModeProperty mode) {
  List<IdealState> idealStates = new ArrayList<IdealState>();
  List<String> instances = new ArrayList<String>();
  for (int i = 0; i < nodes; i++) {
    instances.add("localhost_" + i);
  }

  for (int i = 0; i < resources.length; i++) {
    String resourceName = resources[i];
    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("localhost_" + (p + r + 1) % nodes);
      }
      record.setListField(resourceName + "_" + p, value);
    }
    IdealState idealState = new IdealState(record);
    idealState.setStateModelDefRef("MasterSlave");
    idealState.setIdealStateMode(mode.toString());
    idealState.setNumPartitions(partitions);
    idealStates.add(idealState);

    // System.out.println(idealState);

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

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

/**
 * Case where we have one resource in IdealState
 * @throws Exception
 */
@Test
public void testSimple() throws Exception {
  int nodes = 5;
  List<String> instances = new ArrayList<String>();
  for (int i = 0; i < nodes; i++) {
    instances.add("localhost_" + i);
  }
  int partitions = 10;
  int replicas = 1;
  String resourceName = "testResource";
  ZNRecord record =
      DefaultIdealStateCalculator.calculateIdealState(instances, partitions, replicas,
          resourceName, "MASTER", "SLAVE");
  IdealState idealState = new IdealState(record);
  idealState.setStateModelDefRef("MasterSlave");

  HelixDataAccessor accessor = manager.getHelixDataAccessor();
  Builder keyBuilder = accessor.keyBuilder();
  accessor.setProperty(keyBuilder.idealStates(resourceName), idealState);
  event.addAttribute(AttributeName.ControllerDataProvider.name(),
      new ResourceControllerDataProvider());
  ResourceComputationStage stage = new ResourceComputationStage();
  runStage(event, new ReadClusterDataStage());
  runStage(event, stage);

  Map<String, Resource> resource = event.getAttribute(AttributeName.RESOURCES_TO_REBALANCE.name());
  AssertJUnit.assertEquals(1, resource.size());

  AssertJUnit.assertEquals(resource.keySet().iterator().next(), resourceName);
  AssertJUnit.assertEquals(resource.values().iterator().next().getResourceName(), resourceName);
  AssertJUnit.assertEquals(resource.values().iterator().next().getStateModelDefRef(),
      idealState.getStateModelDefRef());
  AssertJUnit
      .assertEquals(resource.values().iterator().next().getPartitions().size(), partitions);
}
 

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

@Test
public void testGetResources() throws IOException {
  final String clusterName = "TestTagAwareness_testGetResources";
  final String TAG = "tag";
  final String URL_BASE =
      "http://localhost:" + ADMIN_PORT + "/clusters/" + clusterName + "/resourceGroups";

  _gSetupTool.addCluster(clusterName, true);
  HelixAdmin admin = _gSetupTool.getClusterManagementTool();

  // Add a tagged resource
  IdealState taggedResource = new IdealState("taggedResource");
  taggedResource.setInstanceGroupTag(TAG);
  taggedResource.setStateModelDefRef("OnlineOffline");
  admin.addResource(clusterName, taggedResource.getId(), taggedResource);

  // Add an untagged resource
  IdealState untaggedResource = new IdealState("untaggedResource");
  untaggedResource.setStateModelDefRef("OnlineOffline");
  admin.addResource(clusterName, untaggedResource.getId(), untaggedResource);

  // Now make a REST call for all resources
  Reference resourceRef = new Reference(URL_BASE);
  Request request = new Request(Method.GET, resourceRef);
  Response response = _gClient.handle(request);
  ZNRecord responseRecord =
      ClusterRepresentationUtil.JsonToObject(ZNRecord.class, response.getEntityAsText());

  // Ensure that the tagged resource has information and the untagged one doesn't
  Assert.assertNotNull(responseRecord.getMapField("ResourceTags"));
  Assert
      .assertEquals(TAG, responseRecord.getMapField("ResourceTags").get(taggedResource.getId()));
  Assert.assertFalse(responseRecord.getMapField("ResourceTags").containsKey(
      untaggedResource.getId()));
}
 

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

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

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

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

/**
 * 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 testInvocation() throws Exception {

  HelixManager testHelixManager =
      HelixManagerFactory.getZKHelixManager(clusterName, "localhost_8900",
          InstanceType.PARTICIPANT, ZK_ADDR);
  testHelixManager.connect();

  TestZKCallback test = new TestZKCallback();

  TestZKCallback.TestCallbackListener testListener = test.new TestCallbackListener();

  testHelixManager.addMessageListener(testListener, "localhost_8900");
  testHelixManager.addCurrentStateChangeListener(testListener, "localhost_8900",
      testHelixManager.getSessionId());
  testHelixManager.addCustomizedStateRootChangeListener(testListener, "localhost_8900");
  testHelixManager.addConfigChangeListener(testListener);
  testHelixManager.addIdealStateChangeListener(testListener);
  testHelixManager.addExternalViewChangeListener(testListener);
  testHelixManager.addLiveInstanceChangeListener(testListener);
  // Initial add listener should trigger the first execution of the
  // listener callbacks
  AssertJUnit.assertTrue(testListener.configChangeReceived
      & testListener.currentStateChangeReceived & testListener.externalViewChangeReceived
      & testListener.idealStateChangeReceived & testListener.liveInstanceChangeReceived
      & testListener.messageChangeReceived);

  testListener.Reset();
  HelixDataAccessor accessor = testHelixManager.getHelixDataAccessor();
  Builder keyBuilder = accessor.keyBuilder();

  ExternalView extView = new ExternalView("db-12345");
  accessor.setProperty(keyBuilder.externalView("db-12345"), extView);
  Thread.sleep(100);
  AssertJUnit.assertTrue(testListener.externalViewChangeReceived);
  testListener.Reset();

  CurrentState curState = new CurrentState("db-12345");
  curState.setSessionId("sessionId");
  curState.setStateModelDefRef("StateModelDef");
  accessor.setProperty(keyBuilder.currentState("localhost_8900", testHelixManager.getSessionId(),
      curState.getId()), curState);
  Thread.sleep(100);
  AssertJUnit.assertTrue(testListener.currentStateChangeReceived);
  testListener.Reset();

  IdealState idealState = new IdealState("db-1234");
  idealState.setNumPartitions(400);
  idealState.setReplicas(Integer.toString(2));
  idealState.setStateModelDefRef("StateModeldef");
  accessor.setProperty(keyBuilder.idealStates("db-1234"), idealState);
  Thread.sleep(100);
  AssertJUnit.assertTrue(testListener.idealStateChangeReceived);
  testListener.Reset();

  // dummyRecord = new ZNRecord("db-12345");
  // dataAccessor.setProperty(PropertyType.IDEALSTATES, idealState, "db-12345"
  // );
  // Thread.sleep(100);
  // AssertJUnit.assertTrue(testListener.idealStateChangeReceived);
  // testListener.Reset();

  // dummyRecord = new ZNRecord("localhost:8900");
  // List<ZNRecord> recList = new ArrayList<ZNRecord>();
  // recList.add(dummyRecord);

  testListener.Reset();
  Message message = new Message(MessageType.STATE_TRANSITION, UUID.randomUUID().toString());
  message.setTgtSessionId("*");
  message.setResourceName("testResource");
  message.setPartitionName("testPartitionKey");
  message.setStateModelDef("MasterSlave");
  message.setToState("toState");
  message.setFromState("fromState");
  message.setTgtName("testTarget");
  message.setStateModelFactoryName(HelixConstants.DEFAULT_STATE_MODEL_FACTORY);

  accessor.setProperty(keyBuilder.message("localhost_8900", message.getId()), message);
  Thread.sleep(500);
  AssertJUnit.assertTrue(testListener.messageChangeReceived);

  // dummyRecord = new ZNRecord("localhost_9801");
  LiveInstance liveInstance = new LiveInstance("localhost_9801");
  liveInstance.setSessionId(UUID.randomUUID().toString());
  liveInstance.setHelixVersion(UUID.randomUUID().toString());
  accessor.setProperty(keyBuilder.liveInstance("localhost_9801"), liveInstance);
  Thread.sleep(500);
  AssertJUnit.assertTrue(testListener.liveInstanceChangeReceived);
  testListener.Reset();

  // dataAccessor.setNodeConfigs(recList); Thread.sleep(100);
  // AssertJUnit.assertTrue(testListener.configChangeReceived);
  // testListener.Reset();

  accessor.removeProperty(keyBuilder.liveInstance("localhost_8900"));
  accessor.removeProperty(keyBuilder.liveInstance("localhost_9801"));
}
 

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

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

void handleMessageUsingScheduledTaskQueue(Criteria recipientCriteria, Message messageTemplate,
    String controllerMsgId) {
  HelixDataAccessor accessor = _manager.getHelixDataAccessor();
  Builder keyBuilder = accessor.keyBuilder();

  String clusterName = recipientCriteria.getClusterName();
  if (clusterName != null && !clusterName.equals(_manager.getClusterName())) {
    throw new HelixException(String.format(
        "ScheduledTaskQueue cannot send message to another cluster. Local cluster name %s, remote cluster name %s.",
        _manager.getClusterName(), clusterName));
  }

  Map<String, String> sendSummary = new HashMap<String, String>();
  sendSummary.put("MessageCount", "0");
  Map<InstanceType, List<Message>> messages =
      _manager.getMessagingService().generateMessage(recipientCriteria, messageTemplate);

  // Calculate tasks, and put them into the idealState of the SCHEDULER_TASK_QUEUE resource.
  // List field are the destination node, while the Message parameters are stored in the
  // mapFields
  // task throttling can be done on SCHEDULER_TASK_QUEUE resource
  if (messages.size() > 0) {
    String taskQueueName = _message.getRecord().getSimpleField(SCHEDULER_TASK_QUEUE);
    if (taskQueueName == null) {
      throw new HelixException("SchedulerTaskMessage need to have " + SCHEDULER_TASK_QUEUE
          + " specified.");
    }
    IdealState newAddedScheduledTasks = new IdealState(taskQueueName);
    newAddedScheduledTasks.setBucketSize(TASKQUEUE_BUCKET_NUM);
    newAddedScheduledTasks.setStateModelDefRef(SCHEDULER_TASK_QUEUE);

    synchronized (_manager) {
      int existingTopPartitionId = 0;
      IdealState currentTaskQueue =
          _manager.getHelixDataAccessor().getProperty(
              accessor.keyBuilder().idealStates(newAddedScheduledTasks.getId()));
      if (currentTaskQueue != null) {
        existingTopPartitionId = findTopPartitionId(currentTaskQueue) + 1;
      }

      List<Message> taskMessages = (List<Message>) (messages.values().toArray()[0]);
      for (Message task : taskMessages) {
        String partitionId = taskQueueName + "_" + existingTopPartitionId;
        existingTopPartitionId++;
        String instanceName = task.getTgtName();
        newAddedScheduledTasks.setPartitionState(partitionId, instanceName, "COMPLETED");
        task.getRecord().setSimpleField(instanceName, "COMPLETED");
        task.getRecord().setSimpleField(CONTROLLER_MSG_ID, controllerMsgId);

        List<String> priorityList = new LinkedList<String>();
        priorityList.add(instanceName);
        newAddedScheduledTasks.getRecord().setListField(partitionId, priorityList);
        newAddedScheduledTasks.getRecord().setMapField(partitionId,
            task.getRecord().getSimpleFields());
        _logger.info("Scheduling for controllerMsg " + controllerMsgId + " , sending task "
            + partitionId + " " + task.getMsgId() + " to " + instanceName);

        if (_logger.isDebugEnabled()) {
          _logger.debug(task.getRecord().getSimpleFields().toString());
        }
      }
      _manager.getHelixDataAccessor().updateProperty(
          accessor.keyBuilder().idealStates(newAddedScheduledTasks.getId()),
          newAddedScheduledTasks);
      sendSummary.put("MessageCount", "" + taskMessages.size());
    }
  }
  // Record the number of messages sent into scheduler message status updates

  ZNRecord statusUpdate =
      accessor.getProperty(
          keyBuilder.controllerTaskStatus(MessageType.SCHEDULER_MSG.name(),
              _message.getMsgId())).getRecord();

  statusUpdate.getMapFields().put("SentMessageCount", sendSummary);
  accessor.updateProperty(keyBuilder.controllerTaskStatus(MessageType.SCHEDULER_MSG.name(),
      _message.getMsgId()), new StatusUpdate(statusUpdate));
}