org.apache.helix.manager.zk.ZKHelixAdmin Java Examples

private static void addInstanceConfig(String instanceName) {
  // add node to cluster if not already added
  ZKHelixAdmin admin = new ZKHelixAdmin(ZK_ADDR);

  InstanceConfig instanceConfig = null;
  try {
    instanceConfig = admin.getInstanceConfig(_clusterName, instanceName);
  } catch (Exception ignored) {
  }
  if (instanceConfig == null) {
    InstanceConfig config = new InstanceConfig(instanceName);
    config.setHostName("localhost");
    config.setInstanceEnabled(true);
    echo("Adding InstanceConfig:" + config);
    admin.addInstance(_clusterName, config);
  }
}
 

private static void setConstraints(ZKHelixAdmin admin, int numPartitions) {

        logger.info("pause cluster");
        admin.enableCluster("PistachiosCluster", false);
        // setting partition constraints
        logger.info("setting per partition state transition constraints to 1");
        try {
            for (int constraintId = 0; constraintId < numPartitions; constraintId++) {
                java.util.HashMap<ConstraintAttribute, String>  attributes = new java.util.HashMap<ConstraintAttribute, String>();
                attributes.put(ConstraintAttribute.RESOURCE, "PistachiosResource");
                attributes.put(ConstraintAttribute.PARTITION, "PistachiosResource_"+constraintId);
                logger.info("setting per partition for {} state transition constraints to 1", "PistachiosResource_"+constraintId);
                admin.setConstraint("PistachiosCluster", ConstraintType.STATE_CONSTRAINT, "PistachiosPartitionTransitionConstraint" + constraintId,
                   new ConstraintItem(attributes,"1"));
            }

        } catch(Exception e) {
            logger.info("setting state transition constraints error, roll back and exit", e);
        }
        logger.info("resume cluster");
        admin.enableCluster("PistachiosCluster", true);
  }
 

public static void updateRouteWithValidation(String managerHelixClusterName,
    String routeForHelix,
    String instanceId, ZKHelixAdmin helixAdmin, String state, String expectedState)
    throws InterruptedException {
  if (StringUtils.isBlank(expectedState)) {
    expectedState = state;
  }
  IdealState idealState = TestUtils
      .buildManagerWorkerCustomIdealState(routeForHelix, Collections.singletonList(instanceId),
          state);
  helixAdmin.setResourceIdealState(managerHelixClusterName, routeForHelix, idealState);
  Thread.sleep(1000);
  ExternalView externalView = helixAdmin
      .getResourceExternalView(managerHelixClusterName, routeForHelix);
  Assert.assertNotNull(externalView);
  Assert.assertNotNull(externalView.getStateMap("0"));
  LOGGER.info("ExternalView: {}", externalView);
  Assert.assertEquals(externalView.getStateMap("0").get("0"), expectedState);
}
 

/**
 * Update the resources in the destination cluster with the new IdealState settings.
 * @param destZkString the destination Zookeeper server string.
 * @param destClusterName the destination cluster name.
 * @param dryRun run without actual change.
 */
static void updateResourceIdealState(String destZkString, String destClusterName, boolean dryRun) {
  HelixAdmin destAdmin = new ZKHelixAdmin(destZkString);
  Set<String> destResources = new HashSet<>(destAdmin.getResourcesInCluster(destClusterName));

  for (String resource : destResources) {
    IdealState currentIdealState = destAdmin.getResourceIdealState(destClusterName, resource);
    IdealState newIdealState = buildIdealState(resource, currentIdealState.getPartitionSet());
    if (dryRun) {
      System.out.println("Will update " + resource + " to new ideal state " + newIdealState.toString());
    } else {
      destAdmin.setResourceIdealState(destClusterName, resource, newIdealState);
      System.out.println("Updated the ideal state for resource " + resource);
      destAdmin.rebalance(destClusterName, resource, REPLICA_NUMBER, "", "");
      System.out.println("Rebalanced resource " + resource + " with REPLICA_NUM: " + REPLICA_NUMBER);
    }
  }
}
 

public static PistachioClusterInfo getClusterInfo() {
    try {
        String zookeeperConnStr = ConfigurationManager.getConfiguration().getString("Pistachio.ZooKeeper.Server");
        ZKHelixAdmin admin = new ZKHelixAdmin(zookeeperConnStr);
        IdealState idealState = admin.getResourceIdealState("PistachiosCluster", "PistachiosResource");
        PistachioClusterInfo info = new PistachioClusterInfo();
        info.numPartitions = idealState.getNumPartitions();
        info.numReplicas = Integer.parseInt(idealState.getReplicas());
        info.hostList = admin.getInstancesInCluster("PistachiosCluster");

        logger.info("num partitions: {}, num Replicas: {}, hostList: {}.", info.numPartitions,
            info.numReplicas, Joiner.on(",").join(info.hostList.toArray()));

        return info;
    } catch (Exception e) {
        logger.info("error getting cluster info", e);
        return null;
    }
}
 

/**
 * Create a helix cluster with given information.
 * @param destZkString the cluster's zk string
 * @param destClusterName the cluster's name
 */
static void createCluster(String destZkString, String destClusterName) {
  HelixZkClient destZkClient = getHelixZkClient(destZkString);
  HelixAdmin destAdmin = new ZKHelixAdmin(destZkClient);
  if (ZKUtil.isClusterSetup(destClusterName, destZkClient)) {
    errorAndExit("Failed to create cluster because " + destClusterName + " already exist.");
  }
  ClusterSetup clusterSetup = new ClusterSetup.Builder().setZkAddress(destZkString).build();
  clusterSetup.addCluster(destClusterName, true);

  // set ALLOW_PARTICIPANT_AUTO_JOIN
  HelixConfigScope configScope = new HelixConfigScopeBuilder(HelixConfigScope.ConfigScopeProperty.CLUSTER).
      forCluster(destClusterName).build();
  Map<String, String> helixClusterProperties = new HashMap<>();
  helixClusterProperties.put(ZKHelixManager.ALLOW_PARTICIPANT_AUTO_JOIN, String.valueOf(true));
  destAdmin.setConfig(configScope, helixClusterProperties);

  setClusterConfig(destZkClient, destClusterName, false);
  System.out.println("Cluster " + destClusterName + " is created successfully!");
}
 

/**
 * A method to handle failures joining Helix cluster. The method will perform the following steps before attempting
 * to re-join the cluster:
 * <li>
 *   <ul>Disconnect from Helix cluster, which would close any open clients</ul>
 *   <ul>Drop instance from Helix cluster, to remove any partial instance structure from Helix</ul>
 *   <ul>Re-construct helix manager instances, used to re-join the cluster</ul>
 * </li>
 */
private void onClusterJoinFailure() {
  logger.warn("Disconnecting Helix manager..");
  disconnectHelixManager();

  HelixAdmin admin = new ZKHelixAdmin(clusterConfig.getString(GobblinClusterConfigurationKeys.ZK_CONNECTION_STRING_KEY));
  //Drop the helix Instance
  logger.warn("Dropping instance: {} from cluster: {}", helixInstanceName, clusterName);
  HelixUtils.dropInstanceIfExists(admin, clusterName, helixInstanceName);

  if (this.taskDriverHelixManager.isPresent()) {
    String taskDriverCluster = clusterConfig.getString(GobblinClusterConfigurationKeys.TASK_DRIVER_CLUSTER_NAME_KEY);
    logger.warn("Dropping instance: {} from task driver cluster: {}", helixInstanceName, taskDriverCluster);
    HelixUtils.dropInstanceIfExists(admin, clusterName, helixInstanceName);
  }
  admin.close();

  logger.warn("Reinitializing Helix manager..");
  initHelixManager();
}
 

public static void setup() {
  admin = new ZKHelixAdmin(ZK_ADDRESS);
  // create cluster
  echo("Creating cluster: " + CLUSTER_NAME);
  admin.addCluster(CLUSTER_NAME, true);

  // Add nodes to the cluster
  echo("Adding " + NUM_NODES + " participants to the cluster");
  for (int i = 0; i < NUM_NODES; i++) {
    admin.addInstance(CLUSTER_NAME, INSTANCE_CONFIG_LIST.get(i));
    echo("\t Added participant: " + INSTANCE_CONFIG_LIST.get(i).getInstanceName());
  }

  // Add a state model
  StateModelDefinition myStateModel = defineStateModel();
  echo("Configuring StateModel: " + "MyStateModel  with 1 Master and 1 Slave");
  admin.addStateModelDef(CLUSTER_NAME, STATE_MODEL_NAME, myStateModel);

  // Add a resource with 6 partitions and 2 replicas
  echo("Adding a resource MyResource: " + "with 6 partitions and 2 replicas");
  admin.addResource(CLUSTER_NAME, RESOURCE_NAME, NUM_PARTITIONS, STATE_MODEL_NAME, "AUTO");
  // this will set up the ideal state, it calculates the preference list for
  // each partition similar to consistent hashing
  admin.rebalance(CLUSTER_NAME, RESOURCE_NAME, NUM_REPLICAS);
}
 

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

public static ZKHelixAdmin initHelixClustersForWorkerTest(Properties properties, String route1,
    String route2) throws InterruptedException {
  String zkRoot = properties.getProperty("zkServer");
  Thread.sleep(500);
  ZkClient zkClient = ZkUtils.createZkClient(ZkStarter.DEFAULT_ZK_STR, 1000, 1000);
  zkClient.createPersistent("/ureplicator");
  zkClient.close();
  ZKHelixAdmin helixAdmin = new ZKHelixAdmin(zkRoot);
  String deployment = properties.getProperty("federated.deployment.name");
  String managerHelixClusterName = WorkerUtils.getManagerWorkerHelixClusterName(deployment);
  String controllerHelixClusterName = WorkerUtils.getControllerWorkerHelixClusterName(route1);
  if (StringUtils.isNotBlank(route2)) {
    String controllerHelixClusterName2 = WorkerUtils.getControllerWorkerHelixClusterName(route2);
    HelixSetupUtils.setup(controllerHelixClusterName2, zkRoot, "0");
  }

  HelixSetupUtils.setup(managerHelixClusterName, zkRoot, "0");
  HelixSetupUtils.setup(controllerHelixClusterName, zkRoot, "0");

  return helixAdmin;
}
 

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

  ZKHelixAdmin admin = new ZKHelixAdmin(ZK_ADDR);
  admin.addCluster(clusterName, true);
  ConfigAccessor configAccessor = new ConfigAccessor(ZK_ADDR);
  HelixConfigScope scope =
      new HelixConfigScopeBuilder(ConfigScopeProperty.REST).forCluster(clusterName).build();
  Assert.assertNull(configAccessor.getRESTConfig(clusterName));

  RESTConfig restConfig = new RESTConfig(clusterName);
  restConfig.set(RESTConfig.SimpleFields.CUSTOMIZED_HEALTH_URL, "TEST_URL");
  configAccessor.setRESTConfig(clusterName, restConfig);
  Assert.assertEquals(restConfig, configAccessor.getRESTConfig(clusterName));
}
 

private void verifyHelixAdminMsdsEndpoint(
    RealmAwareZkClient.RealmAwareZkConnectionConfig connectionConfig) {
  System.out.println("Start " + TestHelper.getTestMethodName());

  HelixAdmin firstHelixAdmin = new ZKHelixAdmin.Builder().build();
  HelixAdmin secondHelixAdmin =
      new ZKHelixAdmin.Builder().setRealmAwareZkConnectionConfig(connectionConfig).build();

  try {
    verifyMsdsZkRealm(CLUSTER_ONE, true,
        () -> firstHelixAdmin.enableCluster(CLUSTER_ONE, true));
    verifyMsdsZkRealm(CLUSTER_FOUR, false,
        () -> firstHelixAdmin.enableCluster(CLUSTER_FOUR, true));

    verifyMsdsZkRealm(CLUSTER_FOUR, true,
        () -> secondHelixAdmin.enableCluster(CLUSTER_FOUR, true));
    verifyMsdsZkRealm(CLUSTER_ONE, false,
        () -> secondHelixAdmin.enableCluster(CLUSTER_ONE, true));
  } finally {
    firstHelixAdmin.close();
    secondHelixAdmin.close();
  }
}
 

public static void waitForExternalViewUpdate(String zkAddress, final String clusterName, long timeoutInMilliseconds) {
  final ZKHelixAdmin helixAdmin = new ZKHelixAdmin(zkAddress);

  List<String> allResourcesInCluster = helixAdmin.getResourcesInCluster(clusterName);
  Set<String> tableAndBrokerResources = new HashSet<>();
  for (String resourceName : allResourcesInCluster) {
    // Only check table resources and broker resource
    if (TableNameBuilder.isTableResource(resourceName) || resourceName
        .equals(CommonConstants.Helix.BROKER_RESOURCE_INSTANCE)) {
      tableAndBrokerResources.add(resourceName);
    }
  }

  StrictMatchExternalViewVerifier verifier =
      new StrictMatchExternalViewVerifier.Builder(clusterName).setZkAddr(zkAddress)
          .setResources(tableAndBrokerResources).build();

  boolean success = verifier.verify(timeoutInMilliseconds);
  if (success) {
    LOGGER.info("Cluster is ready to serve queries");
  }
}
 

@Test()
public void testDisableNode() throws Exception {
  String command =
      "-zkSvr " + ZK_ADDR + " -enableInstance " + CLUSTER_NAME + " " + PARTICIPANT_PREFIX
          + "_12918" + " TestDB TestDB_0 false";
  ClusterSetup.processCommandLineArgs(command.split(" "));
  boolean result =
      ClusterStateVerifier.verifyByPolling(new ClusterStateVerifier.BestPossAndExtViewZkVerifier(
          ZK_ADDR, CLUSTER_NAME));
  Assert.assertTrue(result);

  ZKHelixAdmin tool = new ZKHelixAdmin(_gZkClient);
  tool.enableInstance(CLUSTER_NAME, PARTICIPANT_PREFIX + "_12918", true);

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

}
 

@Deprecated
public ClusterSetup(String zkServerAddress) {
  // If the multi ZK config is enabled, use FederatedZkClient on multi-realm mode
  if (Boolean.parseBoolean(System.getProperty(SystemPropertyKeys.MULTI_ZK_ENABLED))) {
    try {
      _zkClient = new FederatedZkClient(
          new RealmAwareZkClient.RealmAwareZkConnectionConfig.Builder().build(),
          new RealmAwareZkClient.RealmAwareZkClientConfig()
              .setZkSerializer(new ZNRecordSerializer()));
    } catch (IOException | InvalidRoutingDataException | IllegalStateException e) {
      throw new HelixException("Failed to create ConfigAccessor!", e);
    }
  } else {
    _zkClient = SharedZkClientFactory.getInstance()
        .buildZkClient(new HelixZkClient.ZkConnectionConfig(zkServerAddress));
    _zkClient.setZkSerializer(new ZNRecordSerializer());
  }

  _admin = new ZKHelixAdmin(_zkClient);
  _usesExternalZkClient = false;
}
 

public TaskCluster(String zkAddr, String clusterName) throws Exception {
  _clusterName = clusterName;
  _zkclient =
      new ZkClient(zkAddr, ZkClient.DEFAULT_SESSION_TIMEOUT, ZkClient.DEFAULT_CONNECTION_TIMEOUT,
          new ZNRecordSerializer());
  _admin = new ZKHelixAdmin(_zkclient);
}
 

public PinotZKChanger(String zkAddress, String clusterName) {
  this.clusterName = clusterName;
  helixAdmin = new ZKHelixAdmin(zkAddress);
  helixManager = HelixManagerFactory
      .getZKHelixManager(clusterName, "PinotNumReplicaChanger", InstanceType.ADMINISTRATOR, zkAddress);
  try {
    helixManager.connect();
  } catch (Exception e) {
    throw new RuntimeException(e);
  }
  ZNRecordSerializer serializer = new ZNRecordSerializer();
  String path = PropertyPathConfig.getPath(PropertyType.PROPERTYSTORE, clusterName);
  propertyStore = new ZkHelixPropertyStore<>(zkAddress, serializer, path);
}
 

/**
 * A method to verify resources and partitions in src cluster and dest cluster are same.
 */
private boolean isSrcDestSync(String srcZkString, String srcClusterName, String destZkString,
    String destClusterName) {

  HelixAdmin srcAdmin = new ZKHelixAdmin(srcZkString);
  Set<String> srcResources = new HashSet<>(srcAdmin.getResourcesInCluster(srcClusterName));
  HelixAdmin destAdmin = new ZKHelixAdmin(destZkString);
  Set<String> destResources = new HashSet<>(destAdmin.getResourcesInCluster(destClusterName));

  for (String resource : srcResources) {
    if (HelixVcrPopulateTool.ignoreResourceKeyWords.stream().anyMatch(resource::contains)) {
      System.out.println("Resource " + resource + " from src cluster is ignored");
      continue;
    }
    if (destResources.contains(resource)) {
      // check if every partition exist.
      Set<String> srcPartitions = srcAdmin.getResourceIdealState(srcClusterName, resource).getPartitionSet();
      Set<String> destPartitions = destAdmin.getResourceIdealState(destClusterName, resource).getPartitionSet();
      for (String partition : srcPartitions) {
        if (!destPartitions.contains(partition)) {
          return false;
        }
      }
    } else {
      return false;
    }
  }
  return true;
}
 

public AutoAddInvertedIndex(@Nonnull String zkAddress, @Nonnull String clusterName, @Nonnull String controllerAddress,
    @Nonnull String brokerAddress, @Nonnull Strategy strategy, @Nonnull Mode mode) {
  _clusterName = clusterName;
  _controllerAddress = controllerAddress;
  _brokerAddress = brokerAddress;
  _helixAdmin = new ZKHelixAdmin(zkAddress);
  _propertyStore = new ZkHelixPropertyStore<>(zkAddress, new ZNRecordSerializer(),
      PropertyPathConfig.getPath(PropertyType.PROPERTYSTORE, clusterName));
  _strategy = strategy;
  _mode = mode;
}
 

/**
 * Initializes the cache loader using the given data source config.
 *
 * @param pinotThirdEyeDataSourceConfig the data source config that provides controller's information.
 *
 * @throws Exception when an error occurs connecting to the Pinot controller.
 */
private void init(PinotThirdEyeDataSourceConfig pinotThirdEyeDataSourceConfig) throws Exception {
  if (pinotThirdEyeDataSourceConfig.getBrokerUrl() != null
      && pinotThirdEyeDataSourceConfig.getBrokerUrl().trim().length() > 0) {
    ZkClient zkClient = new ZkClient(pinotThirdEyeDataSourceConfig.getZookeeperUrl());
    zkClient.setZkSerializer(new ZNRecordSerializer());
    zkClient.waitUntilConnected(CONNECTION_TIMEOUT, TimeUnit.MILLISECONDS);
    ZKHelixAdmin helixAdmin = new ZKHelixAdmin(zkClient);
    List<String> thirdeyeBrokerList = helixAdmin.getInstancesInClusterWithTag(
        pinotThirdEyeDataSourceConfig.getClusterName(), pinotThirdEyeDataSourceConfig.getTag());

    String[] thirdeyeBrokers = new String[thirdeyeBrokerList.size()];
    for (int i = 0; i < thirdeyeBrokerList.size(); i++) {
      String instanceName = thirdeyeBrokerList.get(i);
      InstanceConfig instanceConfig =
          helixAdmin.getInstanceConfig(pinotThirdEyeDataSourceConfig.getClusterName(), instanceName);
      thirdeyeBrokers[i] = instanceConfig.getHostName().replaceAll(BROKER_PREFIX, "") + ":"
          + instanceConfig.getPort();
    }
    this.connections = fromHostList(thirdeyeBrokers);
    LOG.info("Created PinotControllerResponseCacheLoader with brokers {}", thirdeyeBrokers);
  } else {
    this.connections = fromZookeeper(pinotThirdEyeDataSourceConfig);
    LOG.info("Created PinotControllerResponseCacheLoader with controller {}:{}",
        pinotThirdEyeDataSourceConfig.getControllerHost(), pinotThirdEyeDataSourceConfig.getControllerPort());
  }
}
 

/**
 * Populate info on ZooKeeper server and start {@link HelixControllerManager}.
 * @param zkConnectString zk connect string to zk server.
 * @param vcrClusterName the vcr cluster name.
 * @param clusterMap the {@link ClusterMap} to use.
 * @return the created {@link HelixControllerManager}.
 */
public static HelixControllerManager populateZkInfoAndStartController(String zkConnectString, String vcrClusterName,
    ClusterMap clusterMap) {
  HelixZkClient zkClient = DedicatedZkClientFactory.getInstance()
      .buildZkClient(new HelixZkClient.ZkConnectionConfig(zkConnectString), new HelixZkClient.ZkClientConfig());
  try {
    zkClient.setZkSerializer(new ZNRecordSerializer());
    ClusterSetup clusterSetup = new ClusterSetup(zkClient);
    clusterSetup.addCluster(vcrClusterName, true);
    HelixAdmin admin = new ZKHelixAdmin(zkClient);
    // set ALLOW_PARTICIPANT_AUTO_JOIN
    HelixConfigScope configScope = new HelixConfigScopeBuilder(HelixConfigScope.ConfigScopeProperty.CLUSTER).
        forCluster(vcrClusterName).build();
    Map<String, String> helixClusterProperties = new HashMap<>();
    helixClusterProperties.put(ZKHelixManager.ALLOW_PARTICIPANT_AUTO_JOIN, String.valueOf(true));
    admin.setConfig(configScope, helixClusterProperties);
    // set PersistBestPossibleAssignment
    ConfigAccessor configAccessor = new ConfigAccessor(zkClient);
    ClusterConfig clusterConfig = configAccessor.getClusterConfig(vcrClusterName);
    clusterConfig.setPersistBestPossibleAssignment(true);
    configAccessor.setClusterConfig(vcrClusterName, clusterConfig);

    FullAutoModeISBuilder builder = new FullAutoModeISBuilder(helixResource);
    builder.setStateModel(LeaderStandbySMD.name);
    for (PartitionId partitionId : clusterMap.getAllPartitionIds(null)) {
      builder.add(partitionId.toPathString());
    }
    builder.setRebalanceStrategy(CrushEdRebalanceStrategy.class.getName());
    IdealState idealState = builder.build();
    admin.addResource(vcrClusterName, helixResource, idealState);
    admin.rebalance(vcrClusterName, helixResource, 3, "", "");
    HelixControllerManager helixControllerManager = new HelixControllerManager(zkConnectString, vcrClusterName);
    helixControllerManager.syncStart();
    return helixControllerManager;
  } finally {
    zkClient.close();
  }
}
 

public static void setupPinotCluster(String helixClusterName, String zkPath, boolean isUpdateStateModel,
    boolean enableBatchMessageMode, String leadControllerResourceRebalanceStrategy) {
  HelixZkClient zkClient = null;
  try {
    zkClient = SharedZkClientFactory.getInstance().buildZkClient(new HelixZkClient.ZkConnectionConfig(zkPath),
        new HelixZkClient.ZkClientConfig().setZkSerializer(new ZNRecordSerializer())
            .setConnectInitTimeout(TimeUnit.SECONDS.toMillis(ZkClient.DEFAULT_CONNECT_TIMEOUT_SEC)));
    zkClient.waitUntilConnected(ZkClient.DEFAULT_CONNECT_TIMEOUT_SEC, TimeUnit.SECONDS);
    HelixAdmin helixAdmin = new ZKHelixAdmin(zkClient);
    HelixDataAccessor helixDataAccessor =
        new ZKHelixDataAccessor(helixClusterName, new ZkBaseDataAccessor<>(zkClient));
    ConfigAccessor configAccessor = new ConfigAccessor(zkClient);

    Preconditions.checkState(helixAdmin.getClusters().contains(helixClusterName),
        String.format("Helix cluster: %s hasn't been set up", helixClusterName));

    // Add segment state model definition if needed
    addSegmentStateModelDefinitionIfNeeded(helixClusterName, helixAdmin, helixDataAccessor, isUpdateStateModel);

    // Add broker resource if needed
    createBrokerResourceIfNeeded(helixClusterName, helixAdmin, enableBatchMessageMode);

    // Add lead controller resource if needed
    createLeadControllerResourceIfNeeded(helixClusterName, helixAdmin, configAccessor, enableBatchMessageMode,
        leadControllerResourceRebalanceStrategy);
  } finally {
    if (zkClient != null) {
      zkClient.close();
    }
  }
}
 

/**
 * Configure the instance, the configuration of each node is available to
 * other nodes.
 * @param instanceName
 */
private void configureInstance(String instanceName) {
  ZKHelixAdmin helixAdmin = new ZKHelixAdmin(zkAddress);

  List<String> instancesInCluster = helixAdmin.getInstancesInCluster(clusterName);
  if (instancesInCluster == null || !instancesInCluster.contains(instanceName)) {
    InstanceConfig config = new InstanceConfig(instanceName);
    config.setHostName("localhost");
    config.setPort("12000");
    helixAdmin.addInstance(clusterName, config);
  }
}
 

@Test()
public void testDisablePartition() throws Exception {
  LOG.info("START testDisablePartition() at " + new Date(System.currentTimeMillis()));

  // localhost_12919 is MASTER for TestDB_0
  String command = "--zkSvr " + ZK_ADDR + " --enablePartition false " + CLUSTER_NAME
      + " localhost_12919 TestDB TestDB_0 TestDB_9";
  ClusterSetup.processCommandLineArgs(command.split("\\s+"));
  Map<String, Set<String>> map = new HashMap<>();
  map.put("TestDB_0", TestHelper.setOf("localhost_12919"));
  map.put("TestDB_9", TestHelper.setOf("localhost_12919"));

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

  TestHelper.verifyState(CLUSTER_NAME, ZK_ADDR, map, "OFFLINE");

  ZKHelixAdmin tool = new ZKHelixAdmin(_gZkClient);
  tool.enablePartition(true, CLUSTER_NAME, "localhost_12919", "TestDB",
      Collections.singletonList("TestDB_9"));

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

  map.clear();
  map.put("TestDB_0", TestHelper.setOf("localhost_12919"));
  TestHelper.verifyState(CLUSTER_NAME, ZK_ADDR, map, "OFFLINE");

  map.clear();
  map.put("TestDB_9", TestHelper.setOf("localhost_12919"));
  TestHelper.verifyState(CLUSTER_NAME, ZK_ADDR, map, "MASTER");

  LOG.info("STOP testDisablePartition() at " + new Date(System.currentTimeMillis()));

}
 

@Override
public void run() {
  ZkClient zkclient = null;
  try {
    // add node to cluster if not already added
    zkclient =
        new ZkClient(_zkAddr, ZkClient.DEFAULT_SESSION_TIMEOUT,
            ZkClient.DEFAULT_CONNECTION_TIMEOUT, new ZNRecordSerializer());
    ZKHelixAdmin admin = new ZKHelixAdmin(zkclient);

    List<String> nodes = admin.getInstancesInCluster(_clusterName);
    if (!nodes.contains(_instanceName)) {
      InstanceConfig config = new InstanceConfig(_instanceName);
      config.setHostName("localhost");
      config.setInstanceEnabled(true);
      admin.addInstance(_clusterName, config);
    }

    Runtime.getRuntime().addShutdownHook(new Thread() {
      @Override
      public void run() {
        System.out.println("Shutting down " + _instanceName);
        disconnect();
      }
    });

    connect();
  } finally {
    if (zkclient != null) {
      zkclient.close();
    }
  }
}
 

public static void main(String[] args) {
  if (args.length < 1) {
    System.err.println("USAGE: java SetupConsumerCluster zookeeperAddress (e.g. localhost:2181)");
    System.exit(1);
  }

  final String zkAddr = args[0];
  final String clusterName = DEFAULT_CLUSTER_NAME;

  ZkClient zkclient = null;
  try {
    zkclient =
        new ZkClient(zkAddr, ZkClient.DEFAULT_SESSION_TIMEOUT,
            ZkClient.DEFAULT_CONNECTION_TIMEOUT, new ZNRecordSerializer());
    ZKHelixAdmin admin = new ZKHelixAdmin(zkclient);

    // add cluster
    admin.addCluster(clusterName, true);

    // add state model definition
    admin.addStateModelDef(clusterName, DEFAULT_STATE_MODEL,
        new StateModelDefinition(StateModelConfigGenerator.generateConfigForOnlineOffline()));

    // add resource "topic" which has 6 partitions
    String resourceName = DEFAULT_RESOURCE_NAME;
    admin.addResource(clusterName, resourceName, DEFAULT_PARTITION_NUMBER, DEFAULT_STATE_MODEL,
        RebalanceMode.FULL_AUTO.toString());

    admin.rebalance(clusterName, resourceName, 1);

  } finally {
    if (zkclient != null) {
      zkclient.close();
    }
  }
}
 

protected void setupInstances(String clusterName, int[] instances) {
  HelixAdmin admin = new ZKHelixAdmin(_gZkClient);
  for (int i = 0; i < instances.length; i++) {
    String instance = "localhost_" + instances[i];
    InstanceConfig instanceConfig = new InstanceConfig(instance);
    instanceConfig.setHostName("localhost");
    instanceConfig.setPort("" + instances[i]);
    instanceConfig.setInstanceEnabled(true);
    admin.addInstance(clusterName, instanceConfig);
  }
}
 

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

/***
 * A helper method to verify IdealState for given replica's partition. It checks number of replicas (represented by
 * instances) under certain partition and verifies new added replica exists or removed old replica is no longer present.
 * @param replica the replica to check if it exists (could be either added replica or removed replica)
 * @param shouldExist if {@code true}, it means the given replica is newly added and should exist in IdealState.
 *                    {@code false} otherwise.
 * @param expectedReplicaCountForPartition expected number of replicas under certain partition
 * @param expectedResourceCount expected total number resource count in this cluster.
 */
private void verifyIdealStateForPartition(ReplicaId replica, boolean shouldExist,
    int expectedReplicaCountForPartition, long expectedResourceCount) {
  String dcName = replica.getDataNodeId().getDatacenterName();
  ZkInfo zkInfo = dcsToZkInfo.get(dcName);
  String clusterName = CLUSTER_NAME_PREFIX + CLUSTER_NAME_IN_STATIC_CLUSTER_MAP;
  ZKHelixAdmin admin = new ZKHelixAdmin("localhost:" + zkInfo.getPort());
  if (!activeDcSet.contains(dcName)) {
    Assert.assertFalse("Cluster should not be present, as dc " + dcName + " is not enabled",
        admin.getClusters().contains(clusterName));
  } else {
    List<String> resources = admin.getResourcesInCluster(clusterName);
    assertEquals("Resource count mismatch", expectedResourceCount, resources.size());
    String partitionName = replica.getPartitionId().toPathString();
    boolean resourceFound = false;
    for (String resource : resources) {
      IdealState idealState = admin.getResourceIdealState(clusterName, resource);
      if (idealState.getPartitionSet().contains(partitionName)) {
        resourceFound = true;
        Set<String> instances = idealState.getInstanceSet(partitionName);
        assertEquals("Replica number is not expected", expectedReplicaCountForPartition, instances.size());
        String instanceNameOfNewReplica = getInstanceName(replica.getDataNodeId());
        if (shouldExist) {
          assertTrue("Instance set doesn't include new instance that holds added replica",
              instances.contains(instanceNameOfNewReplica));
        } else {
          assertFalse("Instance that holds deleted replica should not exist in the set",
              instances.contains(instanceNameOfNewReplica));
        }
        break;
      }
    }
    assertTrue("No corresponding resource found for partition " + partitionName, resourceFound);
  }
}
 

/**
 * Verify that the number of resources in Helix is as expected.
 * @param hardwareLayout the {@link HardwareLayout} of the static clustermap.
 * @param expectedResourceCount the expected number of resources in Helix.
 */
private void verifyResourceCount(HardwareLayout hardwareLayout, long expectedResourceCount) {
  for (Datacenter dc : hardwareLayout.getDatacenters()) {
    ZkInfo zkInfo = dcsToZkInfo.get(dc.getName());
    ZKHelixAdmin admin = new ZKHelixAdmin("localhost:" + zkInfo.getPort());
    if (!activeDcSet.contains(dc.getName())) {
      Assert.assertFalse("Cluster should not be present, as dc " + dc.getName() + " is not enabled",
          admin.getClusters().contains(CLUSTER_NAME_PREFIX + CLUSTER_NAME_IN_STATIC_CLUSTER_MAP));
    } else {
      assertEquals("Resource count mismatch", expectedResourceCount,
          admin.getResourcesInCluster(CLUSTER_NAME_PREFIX + CLUSTER_NAME_IN_STATIC_CLUSTER_MAP).size());
    }
  }
}