org.apache.helix.HelixDataAccessor Java Examples

@Override
public void addInstanceTag(String clusterName, String instanceName, String tag) {
  logger
      .info("Add instance tag {} for instance {} in cluster {}.", tag, instanceName, clusterName);
  if (!ZKUtil.isClusterSetup(clusterName, _zkClient)) {
    throw new HelixException("cluster " + clusterName + " is not setup yet");
  }

  if (!ZKUtil.isInstanceSetup(_zkClient, clusterName, instanceName, InstanceType.PARTICIPANT)) {
    throw new HelixException(
        "cluster " + clusterName + " instance " + instanceName + " is not setup yet");
  }
  HelixDataAccessor accessor =
      new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor<ZNRecord>(_zkClient));
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();

  InstanceConfig config = accessor.getProperty(keyBuilder.instanceConfig(instanceName));
  config.addTag(tag);
  accessor.setProperty(keyBuilder.instanceConfig(instanceName), config);
}
 

@Override
@PreFetch(enabled = false)
public void onIdealStateChange(List<IdealState> idealStates, NotificationContext changeContext) {
  logger.info(
      "START: Generic GenericClusterController.onIdealStateChange() for cluster " + _clusterName);
  notifyCaches(changeContext, ChangeType.IDEAL_STATE);
  pushToEventQueues(ClusterEventType.IdealStateChange, changeContext,
      Collections.<String, Object>emptyMap());

  if (changeContext.getType() != NotificationContext.Type.FINALIZE) {
    HelixManager manager = changeContext.getManager();
    if (manager != null) {
      HelixDataAccessor dataAccessor = changeContext.getManager().getHelixDataAccessor();
      checkRebalancingTimer(changeContext.getManager(), idealStates,
          (ClusterConfig) dataAccessor.getProperty(dataAccessor.keyBuilder().clusterConfig()));
    }
  }

  logger.info("END: GenericClusterController.onIdealStateChange() for cluster " + _clusterName);
}
 

@GET
@Path("{clusterId}/controller")
public Response getClusterController(@PathParam("clusterId") String clusterId) {
  HelixDataAccessor dataAccessor = getDataAccssor(clusterId);
  Map<String, Object> controllerInfo = new HashMap<>();
  controllerInfo.put(Properties.id.name(), clusterId);

  LiveInstance leader = dataAccessor.getProperty(dataAccessor.keyBuilder().controllerLeader());
  if (leader != null) {
    controllerInfo.put(ClusterProperties.controller.name(), leader.getInstanceName());
    controllerInfo.putAll(leader.getRecord().getSimpleFields());
  } else {
    controllerInfo.put(ClusterProperties.controller.name(), "No Lead Controller!");
  }

  return JSONRepresentation(controllerInfo);
}
 

public Map<InstanceType, List<Message>> generateMessage(final Criteria recipientCriteria,
    final Message message) {
  Map<InstanceType, List<Message>> messagesToSendMap = new HashMap<InstanceType, List<Message>>();
  InstanceType instanceType = recipientCriteria.getRecipientInstanceType();

  HelixDataAccessor targetDataAccessor = getRecipientDataAccessor(recipientCriteria);

    List<Message> messages = Collections.EMPTY_LIST;
    if (instanceType == InstanceType.CONTROLLER) {
      messages = generateMessagesForController(message);
    } else if (instanceType == InstanceType.PARTICIPANT) {
      messages =
          generateMessagesForParticipant(recipientCriteria, message, targetDataAccessor);
    }
    messagesToSendMap.put(instanceType, messages);
    return messagesToSendMap;
}
 

/**
 * Assert externalView and currentState for each participant are empty
 * @param clusterName
 * @param db
 * @param participants
 */
private void assertEmptyCSandEV(String clusterName, String db,
    MockParticipantManager[] participants) throws Exception {
  HelixDataAccessor accessor =
      new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor<>(_gZkClient));
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();
  boolean isExternalViewNull = TestHelper.verify(() -> {
    ExternalView externalView = accessor.getProperty(keyBuilder.externalView(db));
    return (externalView == null);
  }, TestHelper.WAIT_DURATION);
  Assert.assertTrue(isExternalViewNull);

  for (MockParticipantManager participant : participants) {
    String instanceName = participant.getInstanceName();
    String sessionId = participant.getSessionId();
    boolean isCurrentStateNull = TestHelper.verify(() -> {
      CurrentState currentState = accessor.getProperty(keyBuilder.currentState(instanceName, sessionId, db));
      return (currentState == null);
    }, TestHelper.WAIT_DURATION);
    Assert.assertTrue(isCurrentStateNull);
  }
}
 

/**
 * Reads HISTORY ZNode from the metadata store and generates a Map object that contains the
 * pertinent history entries depending on the history type.
 * @param clusterId
 * @param historyType
 * @return
 */
private Map<String, Object> getControllerHistory(String clusterId,
    ControllerHistory.HistoryType historyType) {
  HelixDataAccessor dataAccessor = getDataAccssor(clusterId);
  Map<String, Object> history = new HashMap<>();
  history.put(Properties.id.name(), clusterId);

  ControllerHistory historyRecord =
      dataAccessor.getProperty(dataAccessor.keyBuilder().controllerLeaderHistory());

  switch (historyType) {
    case CONTROLLER_LEADERSHIP:
      history.put(Properties.history.name(),
          historyRecord != null ? historyRecord.getHistoryList() : Collections.emptyList());
      break;
    case MAINTENANCE:
      history.put(ClusterProperties.maintenanceHistory.name(),
          historyRecord != null ? historyRecord.getMaintenanceHistoryList()
              : Collections.emptyList());
      break;
  }
  return history;
}
 

@Override
protected Set<PropertyKey> PopulateParticipantKeys(HelixDataAccessor accessor,
    Map<String, LiveInstance> liveInstanceMap) {
  Set<PropertyKey> participantStateKeys = new HashSet<>();
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();
  for (String instanceName : liveInstanceMap.keySet()) {
    LiveInstance liveInstance = liveInstanceMap.get(instanceName);
    String sessionId = liveInstance.getEphemeralOwner();
    List<String> currentStateNames =
        accessor.getChildNames(keyBuilder.currentStates(instanceName, sessionId));
    for (String currentStateName : currentStateNames) {
      participantStateKeys
          .add(keyBuilder.currentState(instanceName, sessionId, currentStateName));
    }
  }
  return participantStateKeys;
}
 

@Test
public void testPartitionLevelCheck() {
  List<ExternalView> externalViews = new ArrayList<>(Arrays.asList(prepareExternalView()));
  Mock mock = new Mock();
  HelixDataAccessor accessor = mock.dataAccessor;

  when(mock.dataAccessor.keyBuilder())
      .thenReturn(new PropertyKey.Builder(TEST_CLUSTER));
  when(mock.dataAccessor
      .getProperty(new PropertyKey.Builder(TEST_CLUSTER).stateModelDef(MasterSlaveSMD.name)))
      .thenReturn(mock.stateModel);
  when(mock.stateModel.getTopState()).thenReturn("MASTER");
  List<String> failedPartitions = InstanceValidationUtil
      .perPartitionHealthCheck(externalViews, preparePartitionStateMap(), "h2", accessor);

  Assert.assertTrue(failedPartitions.size() == 1);
  Assert.assertEquals(failedPartitions.iterator().next(), "p2");
}
 

@Override
public List<String> getResourcesInClusterWithTag(String clusterName, String tag) {
  List<String> resourcesWithTag = new ArrayList<String>();

  HelixDataAccessor accessor =
      new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor<ZNRecord>(_zkClient));
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();

  for (String resourceName : getResourcesInCluster(clusterName)) {
    IdealState is = accessor.getProperty(keyBuilder.idealStates(resourceName));
    if (is != null && is.getInstanceGroupTag() != null && is.getInstanceGroupTag().equals(tag)) {
      resourcesWithTag.add(resourceName);
    }
  }

  return resourcesWithTag;
}
 

@Override
public void dropCluster(String clusterName) {
  logger.info("Deleting cluster {}.", clusterName);
  HelixDataAccessor accessor =
      new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor<ZNRecord>(_zkClient));
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();

  String root = "/" + clusterName;
  if (accessor.getChildNames(keyBuilder.liveInstances()).size() > 0) {
    throw new HelixException(
        "There are still live instances in the cluster, shut them down first.");
  }

  if (accessor.getProperty(keyBuilder.controllerLeader()) != null) {
    throw new HelixException("There are still LEADER in the cluster, shut them down first.");
  }

  _zkClient.deleteRecursively(root);
}
 

@GET @Path("resources/{resourceName}")
public Response getResourceOnInstance(@PathParam("clusterId") String clusterId,
    @PathParam("instanceName") String instanceName,
    @PathParam("resourceName") String resourceName) throws IOException {
  HelixDataAccessor accessor = getDataAccssor(clusterId);
  List<String> sessionIds = accessor.getChildNames(accessor.keyBuilder().sessions(instanceName));
  if (sessionIds == null || sessionIds.size() == 0) {
    return notFound();
  }

  // Only get resource list from current session id
  String currentSessionId = sessionIds.get(0);
  CurrentState resourceCurrentState = accessor.getProperty(
      accessor.keyBuilder().currentState(instanceName, currentSessionId, resourceName));
  if (resourceCurrentState != null) {
    return JSONRepresentation(resourceCurrentState.getRecord());
  }

  return notFound();
}
 

/**
 * This function is deprecated. Please use RebalanceUtil.scheduleInstantPipeline method instead.
 * Trigger the controller to perform rebalance for a given resource.
 * @param accessor Helix data accessor
 * @param resource the name of the resource changed to triggering the execution
 */
@Deprecated
public static void invokeRebalanceForResourceConfig(HelixDataAccessor accessor, String resource) {
  LOG.info("invoke rebalance for " + resource);
  PropertyKey key = accessor.keyBuilder().resourceConfig(resource);
  ResourceConfig cfg = accessor.getProperty(key);
  if (cfg != null) {
    // Here it uses the updateProperty function with no-op DataUpdater. Otherwise, it will use default
    // ZNRecordUpdater which will duplicate elements for listFields.
    if (!accessor.updateProperty(key, znRecord -> znRecord, cfg)) {
      LOG.warn("Failed to invoke rebalance on resource config {}", resource);
    }
  } else {
    LOG.warn("Can't find resource config for {}", resource);
  }
}
 

private void sendNopMessage() {
  if (_manager.isConnected()) {
    try {
      Message nopMsg = new Message(MessageType.NO_OP, UUID.randomUUID().toString());
      nopMsg.setSrcName(_manager.getInstanceName());

      HelixDataAccessor accessor = _manager.getHelixDataAccessor();
      Builder keyBuilder = accessor.keyBuilder();

      if (_manager.getInstanceType() == InstanceType.CONTROLLER
          || _manager.getInstanceType() == InstanceType.CONTROLLER_PARTICIPANT) {
        nopMsg.setTgtName(InstanceType.CONTROLLER.name());
        accessor.setProperty(keyBuilder.controllerMessage(nopMsg.getId()), nopMsg);
      }

      if (_manager.getInstanceType() == InstanceType.PARTICIPANT
          || _manager.getInstanceType() == InstanceType.CONTROLLER_PARTICIPANT) {
        nopMsg.setTgtName(_manager.getInstanceName());
        accessor.setProperty(keyBuilder.message(nopMsg.getTgtName(), nopMsg.getId()), nopMsg);
      }
      logger.info("Send NO_OP message to " + nopMsg.getTgtName() + ", msgId: " + nopMsg.getId());
    } catch (Exception e) {
      logger.error(e.toString());
    }
  }
}
 

StringRepresentation getHostedEntitiesRepresentation(String clusterName)
    throws JsonGenerationException, JsonMappingException, IOException {
  // Get all resources
  ZkClient zkClient =
      ResourceUtil.getAttributeFromCtx(getContext(), ResourceUtil.ContextKey.ZKCLIENT);
  HelixDataAccessor accessor =
      ClusterRepresentationUtil.getClusterDataAccessor(zkClient, clusterName);
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();
  Map<String, HelixProperty> resourceConfigMap =
      accessor.getChildValuesMap(keyBuilder.resourceConfigs());

  // Create the result
  ZNRecord hostedEntitiesRecord = new ZNRecord("JobQueues");

  // Filter out non-workflow resources
  Iterator<Map.Entry<String, HelixProperty>> it = resourceConfigMap.entrySet().iterator();
  while (it.hasNext()) {
    Map.Entry<String, HelixProperty> e = it.next();
    HelixProperty resource = e.getValue();
    Map<String, String> simpleFields = resource.getRecord().getSimpleFields();
    boolean isTerminable = resource.getRecord()
        .getBooleanField(WorkflowConfig.WorkflowConfigProperty.Terminable.name(), true);
    if (!simpleFields.containsKey(WorkflowConfig.WorkflowConfigProperty.TargetState.name())
        || !simpleFields.containsKey(WorkflowConfig.WorkflowConfigProperty.Dag.name())
        || isTerminable) {
      it.remove();
    }
  }

  // Populate the result
  List<String> allResources = Lists.newArrayList(resourceConfigMap.keySet());
  hostedEntitiesRecord.setListField("JobQueues", allResources);

  StringRepresentation representation =
      new StringRepresentation(ClusterRepresentationUtil.ZNRecordToJson(hostedEntitiesRecord),
          MediaType.APPLICATION_JSON);

  return representation;
}
 

private boolean waitForMessageProcessed(HelixDataAccessor accessor, String messageId)
    throws InterruptedException {
  String path =
      accessor.keyBuilder().message(_participant.getInstanceName(), messageId).getPath();
  long startTime = System.currentTimeMillis();
  while (accessor.getBaseDataAccessor().exists(path, 0)) {
    if (System.currentTimeMillis() - startTime > DEFAULT_TIMEOUT) {
      return false;
    }
    Thread.sleep(200);
  }
  return true;
}
 

@Override
public boolean verify() {
  try {
    HelixDataAccessor accessor =
        new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor<ZNRecord>(zkClient));

    return verifyBestPossAndExtView(accessor, errStates, clusterName, resources);
  } catch (Exception e) {
    LOG.error("exception in verification", e);
  }
  return false;
}
 

private void doRefreshWithSelectiveUpdate(final HelixDataAccessor accessor) {
  SelectivePropertyRefreshInputs<T> input =
      genSelectiveUpdateInput(accessor, _objCache, _keyFuncs);
  Map<PropertyKey, T> updatedData = refreshProperties(accessor, new HashSet<>(input.getReloadKeys()),
      input.getCachedKeys(), input.getCachedPropertyMap(), new HashSet<>());
  _objCache = propertyKeyMapToStringMap(updatedData, _keyFuncs);

  // need to separate keys so we can potentially update cache map asynchronously while
  // keeping snapshot unchanged
  _objMap = new HashMap<>(_objCache);
}
 

public void postHandleMessage() {
  if (_message.getBatchMessageMode() == true && _batchMsgWrapper != null) {
    _batchMsgWrapper.end(_message, _notificationContext);
  }

  // update currentState
  HelixManager manager = _notificationContext.getManager();
  HelixDataAccessor accessor = manager.getHelixDataAccessor();
  ConcurrentHashMap<String, CurrentStateUpdate> csUpdateMap =
      (ConcurrentHashMap<String, CurrentStateUpdate>) _notificationContext
          .get(MapKey.CURRENT_STATE_UPDATE.toString());

  if (csUpdateMap != null) {
    Map<PropertyKey, CurrentState> csUpdate = mergeCurStateUpdate(csUpdateMap);

    // TODO: change to use asyncSet
    for (PropertyKey key : csUpdate.keySet()) {
      // logger.info("updateCS: " + key);
      // System.out.println("\tupdateCS: " + key.getPath() + ", " +
      // curStateMap.get(key));
      if(!accessor.updateProperty(key, csUpdate.get(key))) {
        LOG.error(
            "Fails to persist current state to ZK for key " + key);
      }
    }
  }
}
 

@GET
@Path("history")
public Response getHistoryOnInstance(@PathParam("clusterId") String clusterId,
    @PathParam("instanceName") String instanceName) throws IOException {
  HelixDataAccessor accessor = getDataAccssor(clusterId);
  ParticipantHistory history =
      accessor.getProperty(accessor.keyBuilder().participantHistory(instanceName));
  if (history != null) {
    return JSONRepresentation(history.getRecord());
  }
  return notFound();
}
 

private void callParticipantCode(NotificationContext context) {
  // since ZkClient.unsubscribe() does not immediately remove listeners
  // from zk, it is possible that two listeners exist when leadership transfers
  // therefore, double check to make sure only one participant invokes the code
  if (context.getType() == Type.CALLBACK) {
    HelixManager manager = context.getManager();
    // DataAccessor accessor = manager.getDataAccessor();
    HelixDataAccessor accessor = manager.getHelixDataAccessor();
    Builder keyBuilder = accessor.keyBuilder();

    String instance = manager.getInstanceName();
    String sessionId = manager.getSessionId();

    // get resource name from partition key: "PARTICIPANT_LEADER_XXX_0"
    String resourceName = _partitionKey.substring(0, _partitionKey.lastIndexOf('_'));

    CurrentState curState =
        accessor.getProperty(keyBuilder.currentState(instance, sessionId, resourceName));
    if (curState == null) {
      return;
    }

    String state = curState.getState(_partitionKey);
    if (state == null || !state.equalsIgnoreCase("LEADER")) {
      return;
    }
  }

  try {
    _callback.onCallback(context);
  } catch (Exception e) {
    LOG.error("Error invoking callback:" + _callback, e);
  }
}
 

private static void addSegmentStateModelDefinitionIfNeeded(String helixClusterName, HelixAdmin helixAdmin,
    HelixDataAccessor helixDataAccessor, boolean isUpdateStateModel) {
  String segmentStateModelName =
      PinotHelixSegmentOnlineOfflineStateModelGenerator.PINOT_SEGMENT_ONLINE_OFFLINE_STATE_MODEL;
  StateModelDefinition stateModelDefinition = helixAdmin.getStateModelDef(helixClusterName, segmentStateModelName);
  if (stateModelDefinition == null || isUpdateStateModel) {
    if (stateModelDefinition == null) {
      LOGGER.info("Adding state model: {} with CONSUMING state", segmentStateModelName);
    } else {
      LOGGER.info("Updating state model: {} to contain CONSUMING state", segmentStateModelName);
    }
    helixDataAccessor
        .createStateModelDef(PinotHelixSegmentOnlineOfflineStateModelGenerator.generatePinotStateModelDefinition());
  }
}
 

@GET
@Path("{clusterId}")
public Response getClusterInfo(@PathParam("clusterId") String clusterId) {
  if (!doesClusterExist(clusterId)) {
    return notFound();
  }

  HelixDataAccessor dataAccessor = getDataAccssor(clusterId);
  PropertyKey.Builder keyBuilder = dataAccessor.keyBuilder();

  Map<String, Object> clusterInfo = new HashMap<>();
  clusterInfo.put(Properties.id.name(), clusterId);

  LiveInstance controller = dataAccessor.getProperty(keyBuilder.controllerLeader());
  if (controller != null) {
    clusterInfo.put(ClusterProperties.controller.name(), controller.getInstanceName());
  } else {
    clusterInfo.put(ClusterProperties.controller.name(), "No Lead Controller!");
  }

  boolean paused = dataAccessor.getBaseDataAccessor()
      .exists(keyBuilder.pause().getPath(), AccessOption.PERSISTENT);
  clusterInfo.put(ClusterProperties.paused.name(), paused);
  boolean maintenance = getHelixAdmin().isInMaintenanceMode(clusterId);
  clusterInfo.put(ClusterProperties.maintenance.name(), maintenance);

  List<String> idealStates = dataAccessor.getChildNames(keyBuilder.idealStates());
  clusterInfo.put(ClusterProperties.resources.name(), idealStates);
  List<String> instances = dataAccessor.getChildNames(keyBuilder.instanceConfigs());
  clusterInfo.put(ClusterProperties.instances.name(), instances);
  List<String> liveInstances = dataAccessor.getChildNames(keyBuilder.liveInstances());
  clusterInfo.put(ClusterProperties.liveInstances.name(), liveInstances);

  return JSONRepresentation(clusterInfo);
}
 

@Override
public void run() {
  try {
    if (_shouldRefreshCacheOption.orElse(
        _clusterEventType.equals(ClusterEventType.PeriodicalRebalance) || _clusterEventType
            .equals(ClusterEventType.OnDemandRebalance))) {
      requestDataProvidersFullRefresh();

      HelixDataAccessor accessor = _manager.getHelixDataAccessor();
      PropertyKey.Builder keyBuilder = accessor.keyBuilder();
      List<LiveInstance> liveInstances =
          accessor.getChildValues(keyBuilder.liveInstances(), true);

      if (liveInstances != null && !liveInstances.isEmpty()) {
        NotificationContext changeContext = new NotificationContext(_manager);
        changeContext.setType(NotificationContext.Type.CALLBACK);
        synchronized (_manager) {
          checkLiveInstancesObservation(liveInstances, changeContext);
        }
      }
    }
    forceRebalance(_manager, _clusterEventType);
  } catch (Throwable ex) {
    logger.error("Time task failed. Rebalance task type: " + _clusterEventType + ", cluster: "
        + _clusterName, ex);
  }
}
 

private HelixDataAccessor getSrcClusterDataAccessor(final Message message) {
  HelixDataAccessor helixDataAccessor = _manager.getHelixDataAccessor();
  String clusterName = message.getSrcClusterName();
  if (clusterName != null && !clusterName.equals(_manager.getClusterName())) {
    // for cross cluster message, create different HelixDataAccessor for replying message.
    /*
      TODO On frequent cross clsuter messaging request, keeping construct data accessor may cause
      performance issue. We should consider adding cache in this class or HelixManager. --JJ
     */
    helixDataAccessor = new ZKHelixDataAccessor(clusterName, helixDataAccessor.getBaseDataAccessor());
  }
  return helixDataAccessor;
}
 

/**
 * Update an ideal state so that partitions will have a new instance at the tails of their
 * preference lists
 * @param accessor
 * @param instanceName
 * @param resourceName
 * @param partitions
 */
private void addInstanceToPreferences(HelixDataAccessor accessor, final String instanceName,
    final String resourceName, final List<String> partitions) {
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();
  String idealStatePath = keyBuilder.idealStates(resourceName).getPath();
  synchronized (_prefListHistory) {
    // Updater for ideal state
    final List<String> prefList = Lists.newLinkedList();
    DataUpdater<ZNRecord> idealStateUpdater = currentData -> {
      for (String partitionName : partitions) {
        List<String> preferenceList = currentData.getListField(partitionName);
        int numReplicas =
            Integer.valueOf(currentData.getSimpleField(IdealStateProperty.REPLICAS.toString()));
        List<String> newPrefList =
            addInstanceToPreferenceList(preferenceList, instanceName, numReplicas);
        currentData.setListField(partitionName, newPrefList);
        prefList.clear();
        prefList.addAll(newPrefList);
      }
      return currentData;
    };

    // Send update requests together
    List<DataUpdater<ZNRecord>> updaters = Lists.newArrayList();
    updaters.add(idealStateUpdater);
    accessor.updateChildren(Collections.singletonList(idealStatePath), updaters,
        AccessOption.PERSISTENT);
    _prefListHistory.add(prefList);
  }
}
 

private long getTopStateUpdateTime(ExternalView ev, String partition, String state) {
  String topStateHost = ev.getStateMap(partition).entrySet().stream()
      .filter(entry -> entry.getValue().equals(state)).findFirst().get().getKey();
  MockParticipantManager participant = Arrays.stream(_participants)
      .filter(instance -> instance.getInstanceName().equals(topStateHost)).findFirst().get();

  HelixDataAccessor accessor = _controller.getHelixDataAccessor();
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();
  CurrentState currentState = accessor.getProperty(keyBuilder
      .currentState(participant.getInstanceName(), participant.getSessionId(),
          ev.getResourceName()));
  return currentState.getEndTime(partition);
}
 

@Override
public Map<String, Boolean> validateResourcesForWagedRebalance(String clusterName,
    List<String> resourceNames) {
  // Null checks
  if (clusterName == null || clusterName.isEmpty()) {
    throw new HelixException("Cluster name is invalid!");
  }
  if (resourceNames == null || resourceNames.isEmpty()) {
    throw new HelixException("Resource name list is invalid!");
  }

  // Ensure that all instances are valid
  HelixDataAccessor accessor =
      new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor<>(_zkClient));
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();
  List<String> instances = accessor.getChildNames(keyBuilder.instanceConfigs());
  if (validateInstancesForWagedRebalance(clusterName, instances).containsValue(false)) {
    throw new HelixException(String
        .format("Instance capacities haven't been configured properly for cluster %s",
            clusterName));
  }

  Map<String, Boolean> result = new HashMap<>();
  ClusterConfig clusterConfig = _configAccessor.getClusterConfig(clusterName);
  for (String resourceName : resourceNames) {
    IdealState idealState = getResourceIdealState(clusterName, resourceName);
    if (idealState == null || !idealState.isValid()) {
      result.put(resourceName, false);
      continue;
    }
    ResourceConfig resourceConfig = _configAccessor.getResourceConfig(clusterName, resourceName);
    result.put(resourceName,
        validateWeightForResourceConfig(clusterConfig, resourceConfig, idealState));
  }
  return result;
}
 

/**
 * Remove workflow or job config.
 * @param accessor
 * @param workflowJobResource the workflow or job name
 */
private static boolean removeWorkflowJobConfig(HelixDataAccessor accessor,
    String workflowJobResource) {
  PropertyKey cfgKey = accessor.keyBuilder().resourceConfig(workflowJobResource);
  if (accessor.getPropertyStat(cfgKey) != null) {
    if (!accessor.removeProperty(cfgKey)) {
      LOG.warn(String.format(
          "Error occurred while trying to remove config for %s. Failed to remove node %s.",
          workflowJobResource, cfgKey));
      return false;
    }
  }

  return true;
}
 

@Override
public StateModelDefinition getStateModelDef(String clusterName, String stateModelName) {
  HelixDataAccessor accessor =
      new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor<ZNRecord>(_zkClient));
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();

  return accessor.getProperty(keyBuilder.stateModelDef(stateModelName));
}
 

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