Java Code Examples for org.apache.helix.HelixAdmin#addCluster()

private static void setupHelixClusterIfNeeded(String helixClusterName, String zkPath) {
  HelixAdmin admin = new ZKHelixAdmin(zkPath);
  if (admin.getClusters().contains(helixClusterName)) {
    LOGGER.info("Helix cluster: {} already exists", helixClusterName);
  } else {
    LOGGER.info("Creating a new Helix cluster: {}", helixClusterName);
    admin.addCluster(helixClusterName, false);
    // Enable Auto-Join for the cluster
    HelixConfigScope configScope =
        new HelixConfigScopeBuilder(ConfigScopeProperty.CLUSTER).forCluster(helixClusterName).build();
    Map<String, String> configMap = new HashMap<>();
    configMap.put(ZKHelixManager.ALLOW_PARTICIPANT_AUTO_JOIN, Boolean.toString(true));
    configMap.put(ENABLE_CASE_INSENSITIVE_KEY, Boolean.toString(false));
    configMap.put(DEFAULT_HYPERLOGLOG_LOG2M_KEY, Integer.toString(DEFAULT_HYPERLOGLOG_LOG2M));
    configMap.put(CommonConstants.Broker.CONFIG_OF_ENABLE_QUERY_LIMIT_OVERRIDE, Boolean.toString(false));
    admin.setConfig(configScope, configMap);
    LOGGER.info("New Helix cluster: {} created", helixClusterName);
  }
}
 

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

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

@Test
public void testAddCustomizedStateConfig() {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;

  HelixAdmin admin = new ZKHelixAdmin(ZK_ADDR);
  admin.addCluster(clusterName, true);
  CustomizedStateConfig.Builder builder =
      new CustomizedStateConfig.Builder();
  builder.addAggregationEnabledType("mockType1");
  CustomizedStateConfig customizedStateConfig = builder.build();

  admin.addCustomizedStateConfig(clusterName, customizedStateConfig);

  // Read CustomizedStateConfig from Zookeeper and check the content
  ConfigAccessor _configAccessor = new ConfigAccessor(ZK_ADDR);
  CustomizedStateConfig configFromZk =
      _configAccessor.getCustomizedStateConfig(clusterName);
  List<String> listTypesFromZk = configFromZk.getAggregationEnabledTypes();
  Assert.assertEquals(listTypesFromZk.get(0), "mockType1");
}
 

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

@Test
public void testRemoveCloudConfig() throws Exception {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;

  HelixAdmin admin = new ZKHelixAdmin(_gZkClient);
  admin.addCluster(clusterName, true);

  CloudConfig.Builder builder = new CloudConfig.Builder();
  builder.setCloudEnabled(true);
  builder.setCloudID("TestID");
  builder.addCloudInfoSource("TestURL");
  builder.setCloudProvider(CloudProvider.CUSTOMIZED);
  builder.setCloudInfoProcessorName("TestProcessor");
  CloudConfig cloudConfig = builder.build();

  admin.addCloudConfig(clusterName, cloudConfig);

  // Read CloudConfig from Zookeeper and check the content
  ConfigAccessor _configAccessor = new ConfigAccessor(_gZkClient);
  CloudConfig cloudConfigFromZk = _configAccessor.getCloudConfig(clusterName);
  Assert.assertTrue(cloudConfigFromZk.isCloudEnabled());
  Assert.assertEquals(cloudConfigFromZk.getCloudID(), "TestID");
  Assert.assertEquals(cloudConfigFromZk.getCloudProvider(), CloudProvider.CUSTOMIZED.name());
  List<String> listUrlFromZk = cloudConfigFromZk.getCloudInfoSources();
  Assert.assertEquals(listUrlFromZk.get(0), "TestURL");
  Assert.assertEquals(cloudConfigFromZk.getCloudInfoProcessorName(), "TestProcessor");

  // Remove Cloud Config and make sure it has been removed from Zookeeper
  admin.removeCloudConfig(clusterName);
  cloudConfigFromZk = _configAccessor.getCloudConfig(clusterName);
  Assert.assertNull(cloudConfigFromZk);
}
 

/**
 * Test success case when updating InstanceConfig in Helix after new replica is added in storage manager.
 */
@Test
public void updateInstanceConfigSuccessTest() throws Exception {
  generateConfigs(true, true);
  MockDataNodeId localNode = clusterMap.getDataNodes().get(0);
  List<ReplicaId> localReplicas = clusterMap.getReplicaIds(localNode);
  MockClusterParticipant mockHelixParticipant = new MockClusterParticipant();
  StorageManager storageManager =
      createStorageManager(localNode, metricRegistry, Collections.singletonList(mockHelixParticipant));
  storageManager.start();
  // create a new partition and get its replica on local node
  PartitionId newPartition = clusterMap.createNewPartition(Collections.singletonList(localNode));
  ReplicaId newReplica = newPartition.getReplicaIds().get(0);
  // for updating instanceConfig test, we first add an empty InstanceConfig of current node
  String instanceName =
      ClusterMapUtils.getInstanceName(clusterMapConfig.clusterMapHostName, clusterMapConfig.clusterMapPort);
  InstanceConfig instanceConfig = new InstanceConfig(instanceName);
  instanceConfig.setHostName(localNode.getHostname());
  instanceConfig.setPort(Integer.toString(localNode.getPort()));
  // for current test, we initial InstanceConfig empty, non-empty case will be tested in HelixParticipantTest
  Map<String, Map<String, String>> diskInfos = new HashMap<>();
  instanceConfig.getRecord().setMapFields(diskInfos);
  HelixAdmin helixAdmin = mockHelixParticipant.getHelixAdmin();
  helixAdmin.addCluster(CLUSTER_NAME);
  helixAdmin.addInstance(CLUSTER_NAME, instanceConfig);
  // test success case
  mockHelixParticipant.onPartitionBecomeBootstrapFromOffline(newPartition.toPathString());
  instanceConfig = helixAdmin.getInstanceConfig(CLUSTER_NAME, instanceName);
  // verify that new replica info is present in InstanceConfig
  Map<String, Map<String, String>> mountPathToDiskInfos = instanceConfig.getRecord().getMapFields();
  Map<String, String> diskInfo = mountPathToDiskInfos.get(newReplica.getMountPath());
  String replicasStr = diskInfo.get("Replicas");
  Set<String> partitionStrs = new HashSet<>();
  for (String replicaInfo : replicasStr.split(",")) {
    String[] infos = replicaInfo.split(":");
    partitionStrs.add(infos[0]);
  }
  assertTrue("New replica info is not found in InstanceConfig", partitionStrs.contains(newPartition.toPathString()));
  shutdownAndAssertStoresInaccessible(storageManager, localReplicas);
}
 

/**
 * Initialize a map of dataCenter to HelixAdmin based on the given zk Connect Strings.
 */
private void maybeAddCluster() {
  for (Map.Entry<String, HelixAdmin> entry : adminForDc.entrySet()) {
    // Add a cluster entry in every DC
    String dcName = entry.getKey();
    HelixAdmin admin = entry.getValue();
    boolean isClusterPresent = zkClientForDc.get(dcName) == null ? admin.getClusters().contains(clusterName)
        : ZKUtil.isClusterSetup(clusterName, zkClientForDc.get(dcName));
    if (!isClusterPresent) {
      info("Adding cluster {} in {}", clusterName, dcName);
      admin.addCluster(clusterName);
      admin.addStateModelDef(clusterName, stateModelDef, getStateModelDefinition(stateModelDef));
    }
  }
}
 

@Test
public void testUpdateCustomizedStateConfig() throws Exception {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;

  HelixAdmin admin = new ZKHelixAdmin(ZK_ADDR);
  admin.addCluster(clusterName, true);
  CustomizedStateConfig.Builder builder =
      new CustomizedStateConfig.Builder();
  builder.addAggregationEnabledType("mockType1");
  CustomizedStateConfig customizedStateConfig = builder.build();

  admin.addCustomizedStateConfig(clusterName, customizedStateConfig);

  // Read CustomizedStateConfig from Zookeeper and check the content
  ConfigAccessor _configAccessor = new ConfigAccessor(ZK_ADDR);
  CustomizedStateConfig configFromZk =
      _configAccessor.getCustomizedStateConfig(clusterName);
  List<String> listTypesFromZk = configFromZk.getAggregationEnabledTypes();
  Assert.assertEquals(listTypesFromZk.get(0), "mockType1");

  admin.addTypeToCustomizedStateConfig(clusterName, "mockType2");
  admin.addTypeToCustomizedStateConfig(clusterName, "mockType3");
  configFromZk =
      _configAccessor.getCustomizedStateConfig(clusterName);
  listTypesFromZk = configFromZk.getAggregationEnabledTypes();
  Assert.assertEquals(listTypesFromZk.get(0), "mockType1");
  Assert.assertEquals(listTypesFromZk.get(1), "mockType2");
  Assert.assertEquals(listTypesFromZk.get(2), "mockType3");

  admin.removeTypeFromCustomizedStateConfig(clusterName, "mockType1");
  configFromZk =
      _configAccessor.getCustomizedStateConfig(clusterName);
  listTypesFromZk = configFromZk.getAggregationEnabledTypes();
  Assert.assertEquals(listTypesFromZk.get(0), "mockType2");
  Assert.assertEquals(listTypesFromZk.get(1), "mockType3");
}
 

@Test
public void testRemoveCustomizedStateConfig() throws Exception {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;

  HelixAdmin admin = new ZKHelixAdmin(ZK_ADDR);
  admin.addCluster(clusterName, true);
  CustomizedStateConfig.Builder builder =
      new CustomizedStateConfig.Builder();
  builder.addAggregationEnabledType("mockType1");
  CustomizedStateConfig customizedStateConfig = builder.build();

  admin.addCustomizedStateConfig(clusterName, customizedStateConfig);

  // Read CustomizedStateConfig from Zookeeper and check the content
  ConfigAccessor _configAccessor = new ConfigAccessor(ZK_ADDR);
  CustomizedStateConfig configFromZk =
      _configAccessor.getCustomizedStateConfig(clusterName);
  List<String> listTypesFromZk = configFromZk.getAggregationEnabledTypes();
  Assert.assertEquals(listTypesFromZk.get(0), "mockType1");

  // Remove CustomizedStateConfig Config and make sure it has been removed from
  // Zookeeper
  admin.removeCustomizedStateConfig(clusterName);
  configFromZk = _configAccessor.getCustomizedStateConfig(clusterName);
  Assert.assertNull(configFromZk);
}
 

public static void createHelixClusterIfNeeded(String helixClusterName, String zkPath) {
  final HelixAdmin admin = new ZKHelixAdmin(zkPath);

  if (admin.getClusters().contains(helixClusterName)) {
    LOGGER.info("cluster already exist, skipping it.. ********************************************* ");
    return;
  }

  LOGGER.info("Creating a new cluster, as the helix cluster : " + helixClusterName
      + " was not found ********************************************* ");
  admin.addCluster(helixClusterName, false);

  LOGGER.info("Enable mirror maker machines auto join.");
  final HelixConfigScope scope = new HelixConfigScopeBuilder(ConfigScopeProperty.CLUSTER)
      .forCluster(helixClusterName).build();

  final Map<String, String> props = new HashMap<String, String>();
  props.put(ZKHelixManager.ALLOW_PARTICIPANT_AUTO_JOIN, String.valueOf(true));
  props.put(MessageType.STATE_TRANSITION + "." + HelixTaskExecutor.MAX_THREADS,
      String.valueOf(100));

  admin.setConfig(scope, props);

  LOGGER.info("Adding state model definition named : OnlineOffline generated using : "
      + OnlineOfflineStateModel.class.toString()
      + " ********************************************** ");

  // add state model definition
  admin.addStateModelDef(helixClusterName, "OnlineOffline", OnlineOfflineStateModel.build());
  LOGGER.info("New Cluster setup completed... ********************************************** ");
}
 

@Test
public void testAddCloudConfig() {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;

  HelixAdmin admin = new ZKHelixAdmin(_gZkClient);
  admin.addCluster(clusterName, true);

  CloudConfig.Builder builder = new CloudConfig.Builder();
  builder.setCloudEnabled(true);
  builder.setCloudID("TestID");
  builder.addCloudInfoSource("TestURL");
  builder.setCloudProvider(CloudProvider.CUSTOMIZED);
  builder.setCloudInfoProcessorName("TestProcessor");
  CloudConfig cloudConfig = builder.build();

  admin.addCloudConfig(clusterName, cloudConfig);

  // Read CloudConfig from Zookeeper and check the content
  ConfigAccessor _configAccessor = new ConfigAccessor(_gZkClient);
  CloudConfig cloudConfigFromZk = _configAccessor.getCloudConfig(clusterName);
  Assert.assertTrue(cloudConfigFromZk.isCloudEnabled());
  Assert.assertEquals(cloudConfigFromZk.getCloudID(), "TestID");
  Assert.assertEquals(cloudConfigFromZk.getCloudProvider(), CloudProvider.CUSTOMIZED.name());
  List<String> listUrlFromZk = cloudConfigFromZk.getCloudInfoSources();
  Assert.assertEquals(listUrlFromZk.get(0), "TestURL");
  Assert.assertEquals(cloudConfigFromZk.getCloudInfoProcessorName(), "TestProcessor");
}
 

@Test
public void testEnableDisablePartitions() {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;
  String instanceName = "TestInstance";
  String testResourcePrefix = "TestResource";
  System.out.println("START " + clusterName + " at " + new Date(System.currentTimeMillis()));
  HelixAdmin admin = new ZKHelixAdmin(_gZkClient);
  admin.addCluster(clusterName, true);
  admin.addInstance(clusterName, new InstanceConfig(instanceName));

  // Test disable instances with resources
  admin.enablePartition(false, clusterName, instanceName, testResourcePrefix + "0",
      Arrays.asList("1", "2"));
  admin.enablePartition(false, clusterName, instanceName, testResourcePrefix + "1",
      Arrays.asList("2", "3", "4"));
  InstanceConfig instanceConfig = admin.getInstanceConfig(clusterName, instanceName);
  Assert.assertEquals(instanceConfig.getDisabledPartitions(testResourcePrefix + "0").size(), 2);
  Assert.assertEquals(instanceConfig.getDisabledPartitions(testResourcePrefix + "1").size(), 3);

  // Test disable partition across resources
  // TODO : Remove this part once setInstanceEnabledForPartition(partition, enabled) is removed
  instanceConfig.setInstanceEnabledForPartition("10", false);
  Assert.assertEquals(instanceConfig.getDisabledPartitions(testResourcePrefix + "0").size(), 3);
  Assert.assertEquals(instanceConfig.getDisabledPartitions(testResourcePrefix + "1").size(), 4);
  admin.dropCluster(clusterName);
}
 

@Test
public void testDropResource() {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;

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

  HelixAdmin tool = new ZKHelixAdmin(_gZkClient);
  tool.addCluster(clusterName, true);
  Assert.assertTrue(ZKUtil.isClusterSetup(clusterName, _gZkClient), "Cluster should be setup");

  tool.addStateModelDef(clusterName, "MasterSlave",
      new StateModelDefinition(StateModelConfigGenerator.generateConfigForMasterSlave()));
  tool.addResource(clusterName, "test-db", 4, "MasterSlave");
  Map<String, String> resourceConfig = new HashMap<>();
  resourceConfig.put("key1", "value1");
  tool.setConfig(new HelixConfigScopeBuilder(ConfigScopeProperty.RESOURCE).forCluster(clusterName)
      .forResource("test-db").build(), resourceConfig);

  PropertyKey.Builder keyBuilder = new PropertyKey.Builder(clusterName);
  Assert.assertTrue(_gZkClient.exists(keyBuilder.idealStates("test-db").getPath()),
      "test-db ideal-state should exist");
  Assert.assertTrue(_gZkClient.exists(keyBuilder.resourceConfig("test-db").getPath()),
      "test-db resource config should exist");

  tool.dropResource(clusterName, "test-db");
  Assert.assertFalse(_gZkClient.exists(keyBuilder.idealStates("test-db").getPath()),
      "test-db ideal-state should be dropped");
  Assert.assertFalse(_gZkClient.exists(keyBuilder.resourceConfig("test-db").getPath()),
      "test-db resource config should be dropped");

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

private void startAdmin() {
  HelixAdmin admin = new ZKHelixAdmin(ZK_ADDR);

  // create cluster
  System.out.println("Creating cluster: " + _clusterName);
  admin.addCluster(_clusterName, true);

  // add MasterSlave state mode definition
  admin.addStateModelDef(_clusterName, "MasterSlave",
      new StateModelDefinition(generateConfigForMasterSlave()));

  // ideal-state znrecord
  ZNRecord record = new ZNRecord(_resourceName);
  record.setSimpleField("IDEAL_STATE_MODE", "AUTO");
  record.setSimpleField("NUM_PARTITIONS", "1");
  record.setSimpleField("REPLICAS", "2");
  record.setSimpleField("STATE_MODEL_DEF_REF", "MasterSlave");
  record.setListField(_resourceName, Arrays.asList("node1", "node2"));

  admin.setResourceIdealState(_clusterName, _resourceName, new IdealState(record));

  ConstraintItemBuilder builder = new ConstraintItemBuilder();

  // limit one transition message at a time across the entire cluster
  builder.addConstraintAttribute("MESSAGE_TYPE", "STATE_TRANSITION")
      // .addConstraintAttribute("INSTANCE", ".*") // un-comment this line if using instance-level
      // constraint
      .addConstraintAttribute("CONSTRAINT_VALUE", "1");
  admin.setConstraint(_clusterName, ClusterConstraints.ConstraintType.MESSAGE_CONSTRAINT,
      "constraint1", builder.build());
}
 

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

@Test
public void test() throws Exception {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;

  System.out.println("START " + clusterName + " at " + new Date(System.currentTimeMillis()));
  HelixAdmin admin = new ZKHelixAdmin(_gZkClient);
  admin.addCluster(clusterName);
  admin.addStateModelDef(clusterName, BuiltInStateModelDefinitions.MasterSlave.getStateModelDefinition().getId(),
                         BuiltInStateModelDefinitions.MasterSlave.getStateModelDefinition());
  ClusterControllerManager controller = new ClusterControllerManager(ZK_ADDR, clusterName);
  controller.syncStart();

  // controller shall create all built-in state model definitions
  final BaseDataAccessor<ZNRecord> baseAccessor = new ZkBaseDataAccessor<ZNRecord>(_gZkClient);
  final PropertyKey.Builder keyBuilder = new PropertyKey.Builder(clusterName);
  boolean ret = TestHelper.verify(new TestHelper.Verifier() {

    @Override
    public boolean verify() throws Exception {
      for (BuiltInStateModelDefinitions def : BuiltInStateModelDefinitions.values()) {
        String path = keyBuilder.stateModelDef(def.getStateModelDefinition().getId()).getPath();
        boolean exist = baseAccessor.exists(path, 0);
        if (!exist) {
          return false;
        }

        // make sure MasterSlave is not over-written
        if (def == BuiltInStateModelDefinitions.MasterSlave) {
          Stat stat = new Stat();
          baseAccessor.get(path, stat, 0);
          if (stat.getVersion() != 0) {
            return false;
          }
        }
      }
      return true;
    }
  }, 10 * 1000);
  Assert.assertTrue(ret);
  controller.syncStop();
  admin.dropCluster(clusterName);
  System.out.println("END " + clusterName + " at " + new Date(System.currentTimeMillis()));
}
 

/**
 * LockManagerDemo clusterName, numInstances, lockGroupName, numLocks
 * @param args
 * @throws Exception
 */
public static void main(String[] args) throws Exception {
  final String zkAddress = "localhost:2199";
  final String clusterName = "lock-manager-demo";
  final String lockGroupName = "lock-group";
  final int numInstances = 3;
  final int numPartitions = 12;
  final boolean startController = false;
  HelixManager controllerManager = null;
  Thread[] processArray;
  processArray = new Thread[numInstances];
  try {
    startLocalZookeeper(2199);
    HelixAdmin admin = new ZKHelixAdmin(zkAddress);
    admin.addCluster(clusterName, true);
    StateModelConfigGenerator generator = new StateModelConfigGenerator();
    admin.addStateModelDef(clusterName, "OnlineOffline",
        new StateModelDefinition(generator.generateConfigForOnlineOffline()));
    admin.addResource(clusterName, lockGroupName, numPartitions, "OnlineOffline",
        RebalanceMode.FULL_AUTO.toString());
    admin.rebalance(clusterName, lockGroupName, 1);
    for (int i = 0; i < numInstances; i++) {
      final String instanceName = "localhost_" + (12000 + i);
      processArray[i] = new Thread(new Runnable() {

        @Override
        public void run() {
          LockProcess lockProcess = null;

          try {
            lockProcess = new LockProcess(clusterName, zkAddress, instanceName, startController);
            lockProcess.start();
            Thread.currentThread().join();
          } catch (InterruptedException e) {
            System.out.println(instanceName + "Interrupted");
            if (lockProcess != null) {
              lockProcess.stop();
            }
          } catch (Exception e) {
            e.printStackTrace();
          }
        }

      });
      processArray[i].start();
    }
    Thread.sleep(3000);
    controllerManager =
        HelixControllerMain.startHelixController(zkAddress, clusterName, "controller",
            HelixControllerMain.STANDALONE);
    Thread.sleep(5000);
    printStatus(admin, clusterName, lockGroupName);
    System.out.println("Stopping localhost_12000");
    processArray[0].interrupt();
    Thread.sleep(3000);
    printStatus(admin, clusterName, lockGroupName);
    Thread.currentThread().join();
  } catch (Exception e) {
    e.printStackTrace();
  } finally {
    if (controllerManager != null) {
      controllerManager.disconnect();
    }
    for (Thread process : processArray) {
      if (process != null) {
        process.interrupt();
      }
    }
  }
}
 

@Test
public void testAddRemoveMsgConstraint() {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;

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

  HelixAdmin tool = new ZKHelixAdmin(_gZkClient);
  tool.addCluster(clusterName, true);
  Assert.assertTrue(ZKUtil.isClusterSetup(clusterName, _gZkClient), "Cluster should be setup");

  // test admin.getMessageConstraints()
  ClusterConstraints constraints =
      tool.getConstraints(clusterName, ConstraintType.MESSAGE_CONSTRAINT);
  Assert.assertNull(constraints, "message-constraint should NOT exist for cluster: " + className);

  // remove non-exist constraint
  try {
    tool.removeConstraint(clusterName, ConstraintType.MESSAGE_CONSTRAINT, "constraint1");
    // will leave a null message-constraint znode on zk
  } catch (Exception e) {
    Assert.fail("Should not throw exception when remove a non-exist constraint.");
  }

  // add a message constraint
  ConstraintItemBuilder builder = new ConstraintItemBuilder();
  builder.addConstraintAttribute(ConstraintAttribute.RESOURCE.toString(), "MyDB")
      .addConstraintAttribute(ConstraintAttribute.CONSTRAINT_VALUE.toString(), "1");
  tool.setConstraint(clusterName, ConstraintType.MESSAGE_CONSTRAINT, "constraint1",
      builder.build());

  HelixDataAccessor accessor =
      new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor<>(_gZkClient));
  PropertyKey.Builder keyBuilder = new PropertyKey.Builder(clusterName);
  constraints =
      accessor.getProperty(keyBuilder.constraint(ConstraintType.MESSAGE_CONSTRAINT.toString()));
  Assert.assertNotNull(constraints, "message-constraint should exist");
  ConstraintItem item = constraints.getConstraintItem("constraint1");
  Assert.assertNotNull(item, "message-constraint for constraint1 should exist");
  Assert.assertEquals(item.getConstraintValue(), "1");
  Assert.assertEquals(item.getAttributeValue(ConstraintAttribute.RESOURCE), "MyDB");

  // test admin.getMessageConstraints()
  constraints = tool.getConstraints(clusterName, ConstraintType.MESSAGE_CONSTRAINT);
  Assert.assertNotNull(constraints, "message-constraint should exist");
  item = constraints.getConstraintItem("constraint1");
  Assert.assertNotNull(item, "message-constraint for constraint1 should exist");
  Assert.assertEquals(item.getConstraintValue(), "1");
  Assert.assertEquals(item.getAttributeValue(ConstraintAttribute.RESOURCE), "MyDB");

  // remove a exist message-constraint
  tool.removeConstraint(clusterName, ConstraintType.MESSAGE_CONSTRAINT, "constraint1");
  constraints =
      accessor.getProperty(keyBuilder.constraint(ConstraintType.MESSAGE_CONSTRAINT.toString()));
  Assert.assertNotNull(constraints, "message-constraint should exist");
  item = constraints.getConstraintItem("constraint1");
  Assert.assertNull(item, "message-constraint for constraint1 should NOT exist");

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

/**
 * Test both replica info addition and removal cases when updating node info in Helix cluster.
 * @throws Exception
 */
@Test
public void testUpdateNodeInfoInCluster() throws Exception {
  // test setup: 3 disks on local node, each disk has 3 replicas
  MockClusterMap clusterMap = new MockClusterMap(false, true, 1, 3, 3, false, false);
  MockDataNodeId localNode = clusterMap.getDataNodes().get(0);
  List<ReplicaId> localReplicas = clusterMap.getReplicaIds(localNode);
  ReplicaId existingReplica = localReplicas.get(0);
  // override some props for current test
  props.setProperty("clustermap.update.datanode.info", Boolean.toString(true));
  props.setProperty("clustermap.port", String.valueOf(localNode.getPort()));
  ClusterMapConfig clusterMapConfig = new ClusterMapConfig(new VerifiableProperties(props));
  HelixParticipant participant = new HelixParticipant(clusterMapConfig, helixManagerFactory, new MetricRegistry(),
      getDefaultZkConnectStr(clusterMapConfig), true);
  // create InstanceConfig for local node. Also, put existing replica into sealed list
  List<String> sealedList = new ArrayList<>();
  sealedList.add(existingReplica.getPartitionId().toPathString());
  InstanceConfig instanceConfig = generateInstanceConfig(clusterMap, localNode, sealedList);
  HelixAdmin helixAdmin = participant.getHelixAdmin();
  helixAdmin.addCluster(clusterMapConfig.clusterMapClusterName);
  helixAdmin.addInstance(clusterMapConfig.clusterMapClusterName, instanceConfig);
  String instanceName = ClusterMapUtils.getInstanceName(localNode.getHostname(), localNode.getPort());
  // generate exactly same config for comparison
  InstanceConfig initialInstanceConfig = generateInstanceConfig(clusterMap, localNode, sealedList);
  // 1. add existing replica's info to Helix should be no-op
  assertTrue("Adding existing replica's info should succeed",
      participant.updateDataNodeInfoInCluster(existingReplica, true));
  assertEquals("InstanceConfig should stay unchanged", initialInstanceConfig,
      helixAdmin.getInstanceConfig(clusterMapConfig.clusterMapClusterName, instanceName));
  // create two new partitions on the same disk of local node
  PartitionId newPartition1 = clusterMap.createNewPartition(Collections.singletonList(localNode), 0);
  PartitionId newPartition2 = clusterMap.createNewPartition(Collections.singletonList(localNode), 0);
  // 2. add new partition2 (id = 10, replicaFromPartition2) into InstanceConfig
  ReplicaId replicaFromPartition2 = newPartition2.getReplicaIds().get(0);
  assertTrue("Adding new replica info into InstanceConfig should succeed.",
      participant.updateDataNodeInfoInCluster(replicaFromPartition2, true));
  // verify new added replica (id = 10, replicaFromPartition2) info is present in InstanceConfig
  instanceConfig = helixAdmin.getInstanceConfig(clusterMapConfig.clusterMapClusterName, instanceName);
  verifyReplicaInfoInInstanceConfig(instanceConfig, replicaFromPartition2, true);
  // 3. add new partition1 (id = 9, replicaFromPartition1) into InstanceConfig
  ReplicaId replicaFromPartition1 = newPartition1.getReplicaIds().get(0);
  assertTrue("Adding new replica info into InstanceConfig should succeed.",
      participant.updateDataNodeInfoInCluster(replicaFromPartition1, true));
  // verify new added replica (id = 9, replicaFromPartition1) info is present in InstanceConfig
  instanceConfig = helixAdmin.getInstanceConfig(clusterMapConfig.clusterMapClusterName, instanceName);
  verifyReplicaInfoInInstanceConfig(instanceConfig, replicaFromPartition1, true);
  // ensure previous added replica (id = 10, replicaFromPartition2) still exists
  verifyReplicaInfoInInstanceConfig(instanceConfig, replicaFromPartition2, true);
  // 4. remove recently added new replica (id = 9, replicaFromPartition1)
  assertTrue("Removing replica info from InstanceConfig should succeed.",
      participant.updateDataNodeInfoInCluster(replicaFromPartition1, false));
  instanceConfig = helixAdmin.getInstanceConfig(clusterMapConfig.clusterMapClusterName, instanceName);
  verifyReplicaInfoInInstanceConfig(instanceConfig, replicaFromPartition1, false);
  verifyReplicaInfoInInstanceConfig(instanceConfig, replicaFromPartition2, true);
  // 5. remove same replica again (id = 9, replicaFromPartition1) should be no-op
  assertTrue("Removing non-found replica info from InstanceConfig should succeed.",
      participant.updateDataNodeInfoInCluster(replicaFromPartition1, false));
  // 6. remove an existing replica should succeed
  assertTrue("Removing replica info from InstanceConfig should succeed.",
      participant.updateDataNodeInfoInCluster(existingReplica, false));
  instanceConfig = helixAdmin.getInstanceConfig(clusterMapConfig.clusterMapClusterName, instanceName);
  verifyReplicaInfoInInstanceConfig(instanceConfig, existingReplica, false);
  verifyReplicaInfoInInstanceConfig(instanceConfig, replicaFromPartition2, true);
  // reset props
  props.setProperty("clustermap.update.datanode.info", Boolean.toString(false));
  props.setProperty("clustermap.port", "2200");
}