public static Set<TopicPartition> getUnassignedPartitions(HelixManager helixManager) {
  Set<TopicPartition> unassignedPartitions = new HashSet<TopicPartition>();
  HelixAdmin helixAdmin = helixManager.getClusterManagmentTool();
  String helixClusterName = helixManager.getClusterName();
  for (String topic : helixAdmin.getResourcesInCluster(helixClusterName)) {
    IdealState is = helixAdmin.getResourceIdealState(helixClusterName, topic);
    int numPartitions = is.getNumPartitions();
    for (int partition = 0; partition < numPartitions; ++partition) {
      if (is.getInstanceSet(Integer.toString(partition)).isEmpty()) {
        TopicPartition tpi = new TopicPartition(topic, partition);
        unassignedPartitions.add(tpi);
      }
    }
  }
  return unassignedPartitions;
}
 

private static HelixManager getMockHelixManager(Config config) {
  HelixManager helixManager = Mockito.mock(HelixManager.class);
  HelixAdmin helixAdmin = Mockito.mock(HelixAdmin.class);
  HelixDataAccessor helixDataAccessor = Mockito.mock(HelixDataAccessor.class);
  PropertyKey propertyKey = Mockito.mock(PropertyKey.class);
  PropertyKey.Builder propertyKeyBuilder = Mockito.mock(PropertyKey.Builder.class);

  Mockito.when(helixManager.getInstanceName()).thenReturn("helixInstance1");
  Mockito.when(helixManager.getClusterName()).thenReturn(config.getString(GobblinClusterConfigurationKeys.HELIX_CLUSTER_NAME_KEY));
  Mockito.doNothing().when(helixAdmin).enableInstance(Mockito.anyString(), Mockito.anyString(), Mockito.anyBoolean());
  Mockito.when(helixManager.getHelixDataAccessor()).thenReturn(helixDataAccessor);
  Mockito.when(helixDataAccessor.keyBuilder()).thenReturn(propertyKeyBuilder);
  Mockito.when(propertyKeyBuilder.liveInstance(Mockito.anyString())).thenReturn(propertyKey);
  Mockito.when(helixDataAccessor.getProperty(propertyKey)).thenReturn(null);

  return helixManager;
}
 

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

private Set<String> getCurrentPartitionsOnInstance(String cluster, String dbName,
    String instance) {
  HelixAdmin admin = _gSetupTool.getClusterManagementTool();
  Set<String> partitionSet = new HashSet<>();

  ExternalView ev = admin.getResourceExternalView(cluster, dbName);
  for (String partition : ev.getRecord().getMapFields().keySet()) {
    Map<String, String> assignments = ev.getRecord().getMapField(partition);
    for (String ins : assignments.keySet()) {
      if (ins.equals(instance)) {
        partitionSet.add(partition);
      }
    }
  }
  return partitionSet;
}
 

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

/**
 * Reset the partition on specific node.
 */
private void resetPartition() {
  if (adminForDc.size() != 1) {
    throw new IllegalStateException("The dc count is not 1 for resetting partition operation");
  }
  HelixAdmin helixAdmin = adminForDc.values().iterator().next();
  String instanceName;
  if (portNum == null) {
    Optional<DataNodeId> optionalDataNode =
        staticClusterMap.getDataNodeIds().stream().filter(node -> node.getHostname().equals(hostName)).findFirst();
    if (!optionalDataNode.isPresent()) {
      throw new IllegalStateException("Host " + hostName + " is not found in static clustermap");
    }
    DataNodeId dataNodeId = optionalDataNode.get();
    instanceName = getInstanceName(dataNodeId);
  } else {
    instanceName = ClusterMapUtils.getInstanceName(hostName, portNum);
  }
  String resourceName = getResourceNameOfPartition(helixAdmin, clusterName, partitionName);
  info("Resetting partition {} under resource {} on node {}", partitionName, resourceName, hostName);
  helixAdmin.resetPartition(clusterName, instanceName, resourceName, Collections.singletonList(partitionName));
  partitionsReset.getAndIncrement();
}
 

/**
 * Add new state model def to ambry cluster in enabled datacenter(s).
 */
private void addStateModelDef() {
  for (Map.Entry<String, HelixAdmin> entry : adminForDc.entrySet()) {
    // Add a cluster entry in every enabled 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) {
      throw new IllegalStateException("Cluster " + clusterName + " in " + dcName + " doesn't exist!");
    }
    if (!admin.getStateModelDefs(clusterName).contains(stateModelDef)) {
      info("Adding state model def {} in {} for cluster {}", stateModelDef, dcName, clusterName);
      admin.addStateModelDef(clusterName, stateModelDef, getStateModelDefinition(stateModelDef));
    } else {
      info("{} in {} already has state model def {}, skip adding operation", clusterName, dcName, stateModelDef);
    }
  }
}
 

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

@DELETE
@Path("{clusterId}/customized-state-config")
public Response removeCustomizedStateConfig(@PathParam("clusterId") String clusterId) {
  if (!doesClusterExist(clusterId)) {
    return notFound(String.format("Cluster %s does not exist", clusterId));
  }

  HelixAdmin admin = getHelixAdmin();
  try {
    admin.removeCustomizedStateConfig(clusterId);
  } catch (Exception ex) {
    LOG.error(
        "Cannot remove CustomizedStateConfig from cluster: {}, Exception: {}",
        clusterId, ex);
    return serverError(ex);
  }

  return OK();
}
 

@GET
@Path("/cluster/configs")
@Produces(MediaType.APPLICATION_JSON)
@ApiOperation(value = "List cluster configurations", notes = "List cluster level configurations")
@ApiResponses(value = {@ApiResponse(code = 200, message = "Success"), @ApiResponse(code = 500, message = "Internal server error")})
public String listClusterConfigs() {
  HelixAdmin helixAdmin = pinotHelixResourceManager.getHelixAdmin();
  HelixConfigScope configScope = new HelixConfigScopeBuilder(HelixConfigScope.ConfigScopeProperty.CLUSTER)
      .forCluster(pinotHelixResourceManager.getHelixClusterName()).build();
  List<String> configKeys = helixAdmin.getConfigKeys(configScope);
  ObjectNode ret = JsonUtils.newObjectNode();
  Map<String, String> configs = helixAdmin.getConfig(configScope, configKeys);
  for (String key : configs.keySet()) {
    ret.put(key, configs.get(key));
  }
  return ret.toString();
}
 

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

protected void addFakeBrokerInstanceToAutoJoinHelixCluster(String instanceId, boolean isSingleTenant)
    throws Exception {
  HelixManager helixManager = HelixManagerFactory
      .getZKHelixManager(getHelixClusterName(), instanceId, InstanceType.PARTICIPANT, ZkStarter.DEFAULT_ZK_STR);
  helixManager.getStateMachineEngine()
      .registerStateModelFactory(FakeBrokerResourceOnlineOfflineStateModelFactory.STATE_MODEL_DEF,
          FakeBrokerResourceOnlineOfflineStateModelFactory.FACTORY_INSTANCE);
  helixManager.connect();
  HelixAdmin helixAdmin = helixManager.getClusterManagmentTool();
  if (isSingleTenant) {
    helixAdmin.addInstanceTag(getHelixClusterName(), instanceId, TagNameUtils.getBrokerTagForTenant(null));
  } else {
    helixAdmin.addInstanceTag(getHelixClusterName(), instanceId, UNTAGGED_BROKER_INSTANCE);
  }
  _fakeInstanceHelixManagers.add(helixManager);
}
 

protected void addFakeServerInstanceToAutoJoinHelixCluster(String instanceId, boolean isSingleTenant, int adminPort)
    throws Exception {
  HelixManager helixManager = HelixManagerFactory
      .getZKHelixManager(getHelixClusterName(), instanceId, InstanceType.PARTICIPANT, ZkStarter.DEFAULT_ZK_STR);
  helixManager.getStateMachineEngine()
      .registerStateModelFactory(FakeSegmentOnlineOfflineStateModelFactory.STATE_MODEL_DEF,
          FakeSegmentOnlineOfflineStateModelFactory.FACTORY_INSTANCE);
  helixManager.connect();
  HelixAdmin helixAdmin = helixManager.getClusterManagmentTool();
  if (isSingleTenant) {
    helixAdmin.addInstanceTag(getHelixClusterName(), instanceId, TagNameUtils.getOfflineTagForTenant(null));
    helixAdmin.addInstanceTag(getHelixClusterName(), instanceId, TagNameUtils.getRealtimeTagForTenant(null));
  } else {
    helixAdmin.addInstanceTag(getHelixClusterName(), instanceId, UNTAGGED_SERVER_INSTANCE);
  }
  HelixConfigScope configScope =
      new HelixConfigScopeBuilder(HelixConfigScope.ConfigScopeProperty.PARTICIPANT, getHelixClusterName())
          .forParticipant(instanceId).build();
  helixAdmin.setConfig(configScope, Collections.singletonMap(ADMIN_PORT_KEY, Integer.toString(adminPort)));
  _fakeInstanceHelixManagers.add(helixManager);
}
 

private void checkInstanceState(HelixAdmin helixAdmin) {
  String expectedInstanceState = "MASTER";
  TestUtils.waitForCondition(aVoid -> {
    ExternalView leadControllerResourceExternalView =
        helixAdmin.getResourceExternalView(getHelixClusterName(), Helix.LEAD_CONTROLLER_RESOURCE_NAME);
    Set<String> partitionSet = leadControllerResourceExternalView.getPartitionSet();
    if (partitionSet.size() != Helix.NUMBER_OF_PARTITIONS_IN_LEAD_CONTROLLER_RESOURCE) {
      return false;
    }
    for (String partition : partitionSet) {
      Map<String, String> stateMap = leadControllerResourceExternalView.getStateMap(partition);
      if (stateMap.size() != 1 || !stateMap.values().contains(expectedInstanceState)) {
        return false;
      }
    }
    return true;
  }, TIMEOUT_IN_MS, "Failed to pick only one instance as: " + expectedInstanceState);
}
 

private void assignIdealStates(HelixManager helixManager,
    Map<String, IdealState> idealStatesFromAssignment) {
  HelixAdmin helixAdmin = helixManager.getClusterManagmentTool();
  String helixClusterName = helixManager.getClusterName();
  for (String topic : idealStatesFromAssignment.keySet()) {
    IdealState idealState = idealStatesFromAssignment.get(topic);
    helixAdmin.setResourceIdealState(helixClusterName, topic, idealState);
  }
}
 

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

/**
 * From IdealStates.
 *
 * @return InstanceToNumTopicPartitionMap
 */
public static Map<String, Set<TopicPartition>> getInstanceToTopicPartitionsMap(
    HelixManager helixManager,
    Map<String, KafkaBrokerTopicObserver> clusterToObserverMap) {
  Map<String, Set<TopicPartition>> instanceToNumTopicPartitionMap = new HashMap<>();
  HelixAdmin helixAdmin = helixManager.getClusterManagmentTool();
  String helixClusterName = helixManager.getClusterName();
  for (String topic : helixAdmin.getResourcesInCluster(helixClusterName)) {
    IdealState is = helixAdmin.getResourceIdealState(helixClusterName, topic);
    for (String partition : is.getPartitionSet()) {
      TopicPartition tpi;
      if (partition.startsWith("@")) {
        if (clusterToObserverMap != null) {
          TopicPartition topicParition = clusterToObserverMap.get(getSrcFromRoute(partition))
              .getTopicPartitionWithRefresh(topic);
          int trueNumPartition = topicParition != null ? topicParition.getPartition() : -1;
          tpi = new TopicPartition(topic, trueNumPartition, partition);
        } else {
          tpi = new TopicPartition(topic, -1, partition);
        }
      } else {
        // route
        tpi = new TopicPartition(topic, Integer.parseInt(partition));
      }
      for (String instance : is.getInstanceSet(partition)) {
        instanceToNumTopicPartitionMap.putIfAbsent(instance, new HashSet<>());
        instanceToNumTopicPartitionMap.get(instance).add(tpi);
      }
    }
  }
  return instanceToNumTopicPartitionMap;
}
 

/**
 * Get sealed partitions from given datacenter.
 * @param dc the datacenter where sealed partitions come from.
 * @param dcToSealedPartitions a map to track sealed partitions in each dc. This entry associated with given dc will
 *                             be populated in this method.
 * @param nodeToNonExistentReplicas a map to track if any replica is in sealed list but not actually on local node.
 */
private void getSealedPartitionsInDc(Datacenter dc, Map<String, Set<String>> dcToSealedPartitions,
    Map<String, Set<String>> nodeToNonExistentReplicas) {
  String dcName = dc.getName();
  dcToSealedPartitions.put(dcName, new HashSet<>());
  HelixAdmin admin = adminForDc.get(dcName);
  Set<String> allInstancesInHelix = new HashSet<>(admin.getInstancesInCluster(clusterName));
  for (DataNodeId dataNodeId : dc.getDataNodes()) {
    DataNode dataNode = (DataNode) dataNodeId;
    Set<String> replicasOnNode = staticClusterMap.getReplicas(dataNode)
        .stream()
        .map(replicaId -> replicaId.getPartitionId().toPathString())
        .collect(Collectors.toSet());
    String instanceName = getInstanceName(dataNode);
    ensureOrThrow(allInstancesInHelix.contains(instanceName), "Instance not present in Helix " + instanceName);
    InstanceConfig instanceConfig = admin.getInstanceConfig(clusterName, instanceName);
    List<String> sealedReplicas = instanceConfig.getRecord().getListField(ClusterMapUtils.SEALED_STR);
    if (sealedReplicas != null) {
      for (String sealedReplica : sealedReplicas) {
        info("Replica {} is sealed on {}", sealedReplica, instanceName);
        dcToSealedPartitions.get(dcName).add(sealedReplica);
        if (!replicasOnNode.contains(sealedReplica)) {
          logger.warn("Replica {} is in sealed list but not on node {}", sealedReplica, instanceName);
          nodeToNonExistentReplicas.computeIfAbsent(instanceName, key -> new HashSet<>()).add(sealedReplica);
        }
      }
    }
  }
}
 

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

/**
 * A method to disable pre-existing live instances in a Helix cluster. This can happen when a previous Yarn application
 * leaves behind orphaned Yarn worker processes. Since Helix does not provide an API to drop a live instance, we use
 * the disable instance API to fence off these orphaned instances and prevent them from becoming participants in the
 * new cluster.
 *
 * NOTE: this is a workaround for an existing YARN bug. Once YARN has a fix to guarantee container kills on application
 * completion, this method should be removed.
 */
void disableLiveHelixInstances() {
  String clusterName = this.helixManager.getClusterName();
  HelixAdmin helixAdmin = this.helixManager.getClusterManagmentTool();
  List<String> liveInstances = HelixUtils.getLiveInstances(this.helixManager);
  LOGGER.warn("Found {} live instances in the cluster.", liveInstances.size());
  for (String instanceName: liveInstances) {
    LOGGER.warn("Disabling instance: {}", instanceName);
    helixAdmin.enableInstance(clusterName, instanceName, false);
  }
}
 

public static void dropInstanceIfExists(HelixAdmin admin, String clusterName, String helixInstanceName) {
  try {
    admin.dropInstance(clusterName, new InstanceConfig(helixInstanceName));
  } catch (HelixException e) {
    log.error("Could not drop instance: {} due to: {}", helixInstanceName, e);
  }
}
 

@Override
public synchronized HelixAdmin getClusterManagmentTool() {
  checkConnected(_waitForConnectedTimeout);
  if (_zkclient != null && !_zkclient.isClosed()) {
    return new ZKHelixAdmin(_zkclient);
  }

  LOG.error("Couldn't get ZKClusterManagementTool because zkclient is null");
  return null;
}
 

void disablePartition() {
  String instanceName = _manager.getInstanceName();
  String resourceName = _message.getResourceName();
  String partitionName = _message.getPartitionName();
  String clusterName = _manager.getClusterName();
  HelixAdmin admin = _manager.getClusterManagmentTool();
  admin.enablePartition(false, clusterName, instanceName, resourceName,
      Arrays.asList(partitionName));
  logger.info("error in transit from ERROR to " + _message.getToState() + " for partition: "
      + partitionName + ". disable it on " + instanceName);

}
 

public TaskDriver(HelixAdmin admin, HelixDataAccessor accessor,
    HelixPropertyStore<ZNRecord> propertyStore, String clusterName) {
  _admin = admin;
  _accessor = accessor;
  _propertyStore = propertyStore;
  _clusterName = clusterName;
}
 

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

/**
 * Verify that the information in Helix and the information in the static clustermap are equivalent.
 * @param hardwareLayout the {@link HardwareLayout} of the static clustermap.
 * @param partitionLayout the {@link PartitionLayout} of the static clustermap.
 */
private void verifyEquivalencyWithStaticClusterMap(HardwareLayout hardwareLayout, PartitionLayout partitionLayout)
    throws Exception {
  String clusterNameInStaticClusterMap = hardwareLayout.getClusterName();
  info("Verifying equivalency of static cluster: " + clusterNameInStaticClusterMap + " with the "
      + "corresponding cluster in Helix: " + clusterName);
  CountDownLatch verificationLatch = new CountDownLatch(adminForDc.size());
  AtomicInteger errorCount = new AtomicInteger();
  for (Datacenter dc : hardwareLayout.getDatacenters()) {
    HelixAdmin admin = adminForDc.get(dc.getName());
    if (admin == null) {
      info("Skipping {}", dc.getName());
      continue;
    }
    ensureOrThrow(zkClientForDc.get(dc.getName()) == null ? admin.getClusters().contains(clusterName)
            : ZKUtil.isClusterSetup(clusterName, zkClientForDc.get(dc.getName())),
        "Cluster not found in ZK " + dataCenterToZkAddress.get(dc.getName()));
    Utils.newThread(() -> {
      try {
        verifyResourcesAndPartitionEquivalencyInDc(dc, clusterName, partitionLayout);
        verifyDataNodeAndDiskEquivalencyInDc(dc, clusterName, partitionLayout);
      } catch (Throwable t) {
        logger.error("[{}] error message: {}", dc.getName().toUpperCase(), t.getMessage());
        errorCount.getAndIncrement();
      } finally {
        verificationLatch.countDown();
      }
    }, false).start();
  }
  verificationLatch.await(10, TimeUnit.MINUTES);
  ensureOrThrow(errorCount.get() == 0, "Error occurred when verifying equivalency with static cluster map");
  info("Successfully verified equivalency of static cluster: " + clusterNameInStaticClusterMap
      + " with the corresponding cluster in Helix: " + clusterName);
}
 

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
public void allPassed() {
  HelixAdmin helixAdmin = makeHelixAdmin();
  ZkHelixPropertyStore<ZNRecord> propertyStore = makePropertyStore();
  FakeDeletionManager deletionManager = new FakeDeletionManager(helixAdmin, propertyStore);
  Set<String> segments = new HashSet<>();
  segments.addAll(segmentsThatShouldBeDeleted());
  deletionManager.deleteSegmentsFromPropertyStoreAndLocal(tableName, segments);

  Assert.assertEquals(deletionManager.segmentsToRetry.size(), 0);
  Assert.assertEquals(deletionManager.segmentsRemovedFromStore.size(), segments.size());
  Assert.assertTrue(deletionManager.segmentsRemovedFromStore.containsAll(segments));
}
 

@Test
public void testControllerLeaderExists() {
  HelixManager helixManager = mock(HelixManager.class);
  HelixDataAccessor helixDataAccessor = mock(HelixDataAccessor.class);
  HelixAdmin helixAdmin = mock(HelixAdmin.class);
  final String leaderHost = "host";
  final int leaderPort = 12345;

  // Lead controller resource disabled.
  ConfigAccessor configAccessor = mock(ConfigAccessor.class);
  ResourceConfig resourceConfig = mock(ResourceConfig.class);
  when(helixManager.getConfigAccessor()).thenReturn(configAccessor);
  when(configAccessor.getResourceConfig(any(), anyString())).thenReturn(resourceConfig);
  when(resourceConfig.getSimpleConfig(anyString())).thenReturn("false");

  // Mocks the helix leader
  when(helixManager.getHelixDataAccessor()).thenReturn(helixDataAccessor);
  PropertyKey.Builder keyBuilder = mock(PropertyKey.Builder.class);
  when(helixDataAccessor.keyBuilder()).thenReturn(keyBuilder);
  PropertyKey controllerLeader = mock(PropertyKey.class);
  when(keyBuilder.controllerLeader()).thenReturn(controllerLeader);
  LiveInstance liveInstance = mock(LiveInstance.class);
  when(helixDataAccessor.getProperty(controllerLeader)).thenReturn(liveInstance);
  when(liveInstance.getInstanceName()).thenReturn(leaderHost + "_" + leaderPort);

  when(helixManager.getClusterName()).thenReturn("myCluster");
  when(helixManager.getClusterManagmentTool()).thenReturn(helixAdmin);
  when(helixAdmin.getResourceExternalView(anyString(), anyString())).thenReturn(null);

  // Create Controller Leader Locator
  FakeControllerLeaderLocator.create(helixManager);
  ControllerLeaderLocator controllerLeaderLocator = FakeControllerLeaderLocator.getInstance();

  Pair<String, Integer> expectedLeaderLocation = new Pair<>(leaderHost, leaderPort);
  Assert.assertEquals(controllerLeaderLocator.getControllerLeader(testTable).getFirst(),
      expectedLeaderLocation.getFirst());
  Assert.assertEquals(controllerLeaderLocator.getControllerLeader(testTable).getSecond(),
      expectedLeaderLocation.getSecond());
}