Java Code Examples for org.apache.helix.model.ExternalView#getPartitionSet()

private boolean isParentTaskDone(ExternalView ev) {
  Set<String> partitionSet = ev.getPartitionSet();
  if (partitionSet.isEmpty()) {
    return false;
  }

  for (String partition : partitionSet) {
    Map<String, String> stateMap = ev.getStateMap(partition);
    for (String instance : stateMap.keySet()) {
      if (!stateMap.get(instance).equalsIgnoreCase("Online")) {
        return false;
      }
    }
  }
  return true;
}
 

private static void printState(String msg) {
  System.out.println("CLUSTER STATE: " + msg);
  ExternalView resourceExternalView = admin.getResourceExternalView(CLUSTER_NAME, RESOURCE_NAME);
  TreeSet<String> sortedSet = new TreeSet<String>(resourceExternalView.getPartitionSet());
  StringBuilder sb = new StringBuilder("\t\t");
  for (int i = 0; i < NUM_NODES; i++) {
    sb.append(INSTANCE_CONFIG_LIST.get(i).getInstanceName()).append("\t");
  }
  System.out.println(sb);
  for (String partitionName : sortedSet) {
    sb.delete(0, sb.length() - 1);
    sb.append(partitionName).append("\t");
    for (int i = 0; i < NUM_NODES; i++) {
      Map<String, String> stateMap = resourceExternalView.getStateMap(partitionName);
      if (stateMap != null && stateMap.containsKey(INSTANCE_CONFIG_LIST.get(i).getInstanceName())) {
        sb.append(stateMap.get(INSTANCE_CONFIG_LIST.get(i).getInstanceName()).charAt(0)).append(
            "\t\t");
      } else {
        sb.append("-").append("\t\t");
      }
    }
    System.out.println(sb);
  }
  System.out.println("###################################################################");
}
 

private void getPrintableResourceState(ExternalView resourceExternalView,
                                       String resourceName,
                                       StringBuilder sb) {
  TreeSet<String> sortedSet = new TreeSet<>(resourceExternalView.getPartitionSet());
  sb.append("\nCLUSTER: ");
  sb.append(getClusterName());
  sb.append("\nRESOURCE: ");
  sb.append(resourceName);
  for (String partitionName : sortedSet) {
    sb.append("\nPARTITION: ");
    sb.append(partitionName).append("\t");
    Map<String, String> states = resourceExternalView.getStateMap(partitionName);
    for (Map.Entry<String, String> stateEntry : states.entrySet()) {
      sb.append("\t");
      sb.append(stateEntry.getKey());
      sb.append("\t");
      sb.append(stateEntry.getValue());
    }
  }
}
 

@Test(dependsOnMethods = "testAddParticipant")
public void testStopAndReStartParticipant() throws InterruptedException {
  MockParticipantManager participant = _participants.get(0);
  String instance = participant.getInstanceName();

  Map<String, MasterSlaveSMD.States> affectedPartitions =
      new HashMap<String, MasterSlaveSMD.States>();

  ExternalView externalView = _accessor.getProperty(_keyBuilder.externalView(DB_NAME));

  for (String partition : externalView.getPartitionSet()) {
    Map<String, String> stateMap = externalView.getStateMap(partition);
    if (stateMap.containsKey(instance)) {
      affectedPartitions.put(partition, MasterSlaveSMD.States.valueOf(stateMap.get(instance)));
    }
  }

  stopParticipant(participant, affectedPartitions);

  // create a new participant
  participant = new MockParticipantManager(ZK_ADDR, CLUSTER_NAME, instance);
  _participants.set(0, participant);
  startParticipant(participant, affectedPartitions);
}
 

@Test
public void testCancellationWhenDisableResource() throws InterruptedException {
  // Enable cancellation
  ClusterConfig clusterConfig = _configAccessor.getClusterConfig(CLUSTER_NAME);
  clusterConfig.stateTransitionCancelEnabled(true);
  _configAccessor.setClusterConfig(CLUSTER_NAME, clusterConfig);

  // Wait for assignment done
  Thread.sleep(2000);

  // Disable the resource
  _gSetupTool.getClusterManagementTool().enableResource(CLUSTER_NAME,
      WorkflowGenerator.DEFAULT_TGT_DB, false);

  // Wait for pipeline reaching final stage
  Assert.assertTrue(_verifier.verifyByPolling());
  ExternalView externalView = _gSetupTool.getClusterManagementTool()
      .getResourceExternalView(CLUSTER_NAME, WorkflowGenerator.DEFAULT_TGT_DB);
  for (String partition : externalView.getPartitionSet()) {
    for (String currentState : externalView.getStateMap(partition).values()) {
      Assert.assertEquals(currentState, "OFFLINE");
    }
  }
}
 

@Override
protected Map<String, String> getPartitionStateMap(ExternalView state) {
  Map<String, String> partitionState = new HashMap<>();

  for (String partition : state.getPartitionSet()) {
    Map<String, String> instanceStateMap = state.getStateMap(partition);
    if (instanceStateMap.containsKey(_instanceName)) {
      partitionState.put(partition, instanceStateMap.get(_instanceName));
    }
  }

  return partitionState;
}
 

public static JSONObject build(String topicName, IdealState idealStateForTopic, ExternalView externalViewForTopic) {
  JSONObject responseJson = new JSONObject();
  responseJson.put("topic", topicName);
  responseJson.put("externalView", buildExternalView(topicName, externalViewForTopic));
  responseJson.put("idealState", buildIdealState(topicName, idealStateForTopic));

  JSONObject serverToPartitionMappingJson = new JSONObject();
  JSONObject serverToNumPartitionsMappingJson = new JSONObject();

  if (externalViewForTopic != null) {
    for (String partition : externalViewForTopic.getPartitionSet()) {
      Map<String, String> stateMap = externalViewForTopic.getStateMap(partition);
      for (String server : stateMap.keySet()) {
        String state = stateMap.get(server);
        if (!serverToPartitionMappingJson.containsKey(server)) {
          serverToPartitionMappingJson.put(server, new JSONObject());
          serverToNumPartitionsMappingJson.put(server, 0);
        }
        serverToPartitionMappingJson.getJSONObject(server).put(partition, stateMap.get(server));
        if (state.equalsIgnoreCase(Constants.HELIX_ONLINE_STATE)) {
          serverToNumPartitionsMappingJson.put(server,
              serverToNumPartitionsMappingJson.getInteger(server) + 1);
        }
      }
    }
  }
  responseJson.put("serverToPartitionMapping", serverToPartitionMappingJson);
  responseJson.put("serverToNumPartitionsMapping", serverToNumPartitionsMappingJson);

  return responseJson;
}
 

/**
 * Get a set of offline instance from the external view of the resource.
 *
 * @param resourceExternalView External view of the resource
 * @return Set of string instance names of the offline instances in the external view.
 */
public static Set<String> getOfflineInstanceFromExternalView(ExternalView resourceExternalView) {
  Set<String> instanceSet = new HashSet<String>();
  for (String partition : resourceExternalView.getPartitionSet()) {
    Map<String, String> stateMap = resourceExternalView.getStateMap(partition);
    for (String instance : stateMap.keySet()) {
      if (stateMap.get(instance).equalsIgnoreCase(OFFLINE)) {
        instanceSet.add(instance);
      }
    }
  }
  return instanceSet;
}
 

/**
 * Returns the set of online instances from external view.
 *
 * @param resourceExternalView External view for the resource.
 * @return Set<String> of online instances in the external view for the resource.
 */
public static Set<String> getOnlineInstanceFromExternalView(ExternalView resourceExternalView) {
  Set<String> instanceSet = new HashSet<String>();
  if (resourceExternalView != null) {
    for (String partition : resourceExternalView.getPartitionSet()) {
      Map<String, String> stateMap = resourceExternalView.getStateMap(partition);
      for (String instance : stateMap.keySet()) {
        if (stateMap.get(instance).equalsIgnoreCase(ONLINE)) {
          instanceSet.add(instance);
        }
      }
    }
  }
  return instanceSet;
}
 

private boolean isResourceOffline(HelixAdmin helixAdmin, String resource) {
  ExternalView externalView = helixAdmin.getResourceExternalView(_helixClusterName, resource);
  // Treat deleted resource as OFFLINE
  if (externalView == null) {
    return true;
  }
  for (String partition : externalView.getPartitionSet()) {
    Map<String, String> instanceStateMap = externalView.getStateMap(partition);
    String state = instanceStateMap.get(_instanceId);
    if (StateModel.SegmentStateModel.ONLINE.equals(state) || StateModel.SegmentStateModel.CONSUMING.equals(state)) {
      return false;
    }
  }
  return true;
}
 

private void refresh(Collection<ExternalView> externalViewList) {
  Map<String, InstanceConfig> instanceConfigMap = new HashMap<>();
  if (externalViewList != null && !externalViewList.isEmpty()) {
    for (InstanceConfig config : _instanceConfigs) {
      instanceConfigMap.put(config.getId(), config);
    }
    for (ExternalView extView : externalViewList) {
      String resourceName = extView.getId();
      for (String partitionName : extView.getPartitionSet()) {
        Map<String, String> stateMap = extView.getStateMap(partitionName);
        for (String instanceName : stateMap.keySet()) {
          String currentState = stateMap.get(instanceName);
          if (instanceConfigMap.containsKey(instanceName)) {
            InstanceConfig instanceConfig = instanceConfigMap.get(instanceName);
            if (extView.isGroupRoutingEnabled()) {
              addEntry(resourceName, extView.getResourceGroupName(),
                  extView.getInstanceGroupTag(), partitionName, currentState, instanceConfig);
            } else {
              addEntry(resourceName, partitionName, currentState, instanceConfig);
            }
          } else {
            logger.warn(
                "Participant {} is not found with proper configuration information. It might already be removed from the cluster. "
                    + "Skip recording partition assignment entry: Partition {}, Participant {}, State {}.",
                instanceName, partitionName, instanceName, stateMap.get(instanceName));
          }
        }
      }
    }
  }
}
 

@Override
public boolean verify() {
  BaseDataAccessor<ZNRecord> baseAccessor = new ZkBaseDataAccessor<ZNRecord>(_zkClient);
  HelixDataAccessor accessor = new ZKHelixDataAccessor(_clusterName, baseAccessor);
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();
  ExternalView externalView = accessor.getProperty(keyBuilder.externalView(_resourceName));

  // verify external view empty
  if (externalView != null) {
    for (String partition : externalView.getPartitionSet()) {
      Map<String, String> stateMap = externalView.getStateMap(partition);
      if (stateMap != null && !stateMap.isEmpty()) {
        LOG.error("External view not empty for " + partition);
        return false;
      }
    }
  }

  // verify current state empty
  List<String> liveParticipants = accessor.getChildNames(keyBuilder.liveInstances());
  for (String participant : liveParticipants) {
    List<String> sessionIds = accessor.getChildNames(keyBuilder.sessions(participant));
    for (String sessionId : sessionIds) {
      CurrentState currentState =
          accessor.getProperty(keyBuilder.currentState(participant, sessionId, _resourceName));
      Map<String, String> partitionStateMap = currentState.getPartitionStateMap();
      if (partitionStateMap != null && !partitionStateMap.isEmpty()) {
        LOG.error("Current state not empty for " + participant);
        return false;
      }
    }
  }
  return true;
}
 

@Test
public void testAddParticipant() throws InterruptedException {
  String newInstance = PARTICIPANT_PREFIX + "_" + (PARTICIPANT_START_PORT + _participants.size());
  _gSetupTool.addInstanceToCluster(CLUSTER_NAME, newInstance);

  MockParticipantManager newParticipant =
      new MockParticipantManager(ZK_ADDR, CLUSTER_NAME, newInstance);
  newParticipant.syncStart();

  Thread.sleep(1000);

  List<String> instances = _accessor.getChildNames(_keyBuilder.instanceConfigs());
  Assert.assertEquals(instances.size(), _participants.size() + 1);
  Assert.assertTrue(instances.contains(newInstance));

  List<String> liveInstances = _accessor.getChildNames(_keyBuilder.liveInstances());
  Assert.assertEquals(liveInstances.size(), _participants.size() + 1);
  Assert.assertTrue(liveInstances.contains(newInstance));

  // nothing for new participant
  ExternalView externalView = _accessor.getProperty(_keyBuilder.externalView(DB_NAME));
  Assert.assertNotNull(externalView);
  for (String partition : externalView.getPartitionSet()) {
    Map<String, String> stateMap = externalView.getStateMap(partition);
    Assert.assertFalse(stateMap.containsKey(newInstance));
  }

  // clear
  newParticipant.syncStop();
  _gSetupTool.getClusterManagementTool().enableInstance(CLUSTER_NAME, newInstance, false);
  _gSetupTool.dropInstanceFromCluster(CLUSTER_NAME, newInstance);

  instances = _accessor.getChildNames(_keyBuilder.instanceConfigs());
  Assert.assertEquals(instances.size(), _participants.size());

  liveInstances = _accessor.getChildNames(_keyBuilder.liveInstances());
  Assert.assertEquals(liveInstances.size(), _participants.size());
}
 

@Test
public void testDisableCancellationWhenDisableResource() throws InterruptedException {
  // Disable cancellation
  ClusterConfig clusterConfig = _configAccessor.getClusterConfig(CLUSTER_NAME);
  clusterConfig.stateTransitionCancelEnabled(false);
  _configAccessor.setClusterConfig(CLUSTER_NAME, clusterConfig);

  // Reenable resource
  stateCleanUp();
  _gSetupTool.getClusterManagementTool().enableResource(CLUSTER_NAME,
      WorkflowGenerator.DEFAULT_TGT_DB, true);

  // Wait for assignment done
  Thread.sleep(2000);

  // Disable the resource
  _gSetupTool.getClusterManagementTool().enableResource(CLUSTER_NAME,
      WorkflowGenerator.DEFAULT_TGT_DB, false);

  // Wait for pipeline reaching final stage
  Thread.sleep(2000L);
  ExternalView externalView = _gSetupTool.getClusterManagementTool()
      .getResourceExternalView(CLUSTER_NAME, WorkflowGenerator.DEFAULT_TGT_DB);
  for (String partition : externalView.getPartitionSet()) {
    Assert.assertTrue(externalView.getStateMap(partition).values().contains("SLAVE"));
  }
}
 

/**
 * Check if sibling nodes of the instance meet min active replicas constraint
 * Two instances are sibling of each other if they host the same partition
 * WARNING: The check uses ExternalView to reduce network traffic but suffer from accuracy
 * due to external view propagation latency
 *
 * TODO: Use in memory cache and query instance's currentStates
 *
 * @param dataAccessor
 * @param instanceName
 * @return
 */
public static boolean siblingNodesActiveReplicaCheck(HelixDataAccessor dataAccessor, String instanceName) {
  PropertyKey.Builder propertyKeyBuilder = dataAccessor.keyBuilder();
  List<String> resources = dataAccessor.getChildNames(propertyKeyBuilder.idealStates());

  for (String resourceName : resources) {
    IdealState idealState = dataAccessor.getProperty(propertyKeyBuilder.idealStates(resourceName));
    if (idealState == null || !idealState.isEnabled() || !idealState.isValid()
        || TaskConstants.STATE_MODEL_NAME.equals(idealState.getStateModelDefRef())) {
      continue;
    }
    ExternalView externalView =
        dataAccessor.getProperty(propertyKeyBuilder.externalView(resourceName));
    if (externalView == null) {
      throw new HelixException(
          String.format("Resource %s does not have external view!", resourceName));
    }
    // Get the minActiveReplicas constraint for the resource
    int minActiveReplicas = externalView.getMinActiveReplicas();
    if (minActiveReplicas == -1) {
      _logger.warn("Resource " + resourceName
          + " is missing minActiveReplica field. Skip the sibling check");
      continue;
    }
    String stateModeDef = externalView.getStateModelDefRef();
    StateModelDefinition stateModelDefinition =
        dataAccessor.getProperty(propertyKeyBuilder.stateModelDef(stateModeDef));
    Set<String> unhealthyStates = new HashSet<>(UNHEALTHY_STATES);
    if (stateModelDefinition != null) {
      unhealthyStates.add(stateModelDefinition.getInitialState());
    }
    for (String partition : externalView.getPartitionSet()) {
      Map<String, String> stateByInstanceMap = externalView.getStateMap(partition);
      // found the resource hosted on the instance
      if (stateByInstanceMap.containsKey(instanceName)) {
        int numHealthySiblings = 0;
        for (Map.Entry<String, String> entry : stateByInstanceMap.entrySet()) {
          if (!entry.getKey().equals(instanceName)
              && !unhealthyStates.contains(entry.getValue())) {
            numHealthySiblings++;
          }
        }
        if (numHealthySiblings < minActiveReplicas) {
          _logger.info(
              "Partition {} doesn't have enough active replicas in sibling nodes. NumHealthySiblings: {}, minActiveReplicas: {}",
              partition, numHealthySiblings, minActiveReplicas);
          return false;
        }
      }
    }
  }

  return true;
}
 

@Test
public void testRebalancingCauseCancellation() throws InterruptedException {
  // Enable cancellation
  ClusterConfig clusterConfig = _configAccessor.getClusterConfig(CLUSTER_NAME);
  clusterConfig.stateTransitionCancelEnabled(true);
  clusterConfig.setPersistBestPossibleAssignment(true);
  _configAccessor.setClusterConfig(CLUSTER_NAME, clusterConfig);

  // Reenable resource
  stateCleanUp();
  _gSetupTool.getClusterManagementTool().enableResource(CLUSTER_NAME,
      WorkflowGenerator.DEFAULT_TGT_DB, true);

  // Wait for assignment done
  Thread.sleep(2000);
  int numNodesToStart = 10;
  for (int i = 0; i < numNodesToStart; i++) {
    String storageNodeName = PARTICIPANT_PREFIX + "_" + (_startPort + _numNodes + i);
    _gSetupTool.addInstanceToCluster(CLUSTER_NAME, storageNodeName);
  }
  MockParticipantManager[] newParticipants = new MockParticipantManager[numNodesToStart];
  registerParticipants(newParticipants, numNodesToStart, _startPort + _numNodes, 1000, -3000000L);

  // Wait for pipeline reaching final stage
  Thread.sleep(2000L);
  int numOfMasters = 0;
  ExternalView externalView = _gSetupTool.getClusterManagementTool()
      .getResourceExternalView(CLUSTER_NAME, WorkflowGenerator.DEFAULT_TGT_DB);
  for (String partition : externalView.getPartitionSet()) {
    if (externalView.getStateMap(partition).values().contains("MASTER")) {
      numOfMasters++;
    }
  }

  for (MockParticipantManager participant : newParticipants) {
    participant.syncStop();
  }

  // Either partial of state transitions have been cancelled or all the Slave -> Master
  // reassigned to other cluster
  Assert.assertTrue((numOfMasters > 0 && numOfMasters <= _numPartitions));
}
 

@Test
public void testSubMessageFailed() throws Exception {
  TestOnlineOfflineStateModel._numOfSuccessBeforeFailure.set(6);

  // Let one instance handle all the batch messages.
  _participants[0].getStateMachineEngine().registerStateModelFactory("OnlineOffline",
      new TestOnlineOfflineStateModelFactory(), "TestFactory");
  for (int i = 1; i < _participants.length; i++) {
    _participants[i].syncStop();
  }

  HelixDataAccessor dataAccessor = new ZKHelixDataAccessor(CLUSTER_NAME, _baseAccessor);
  // Check that the Participants really stopped
  boolean result = TestHelper.verify(() -> {
    List<String> liveInstances =
        dataAccessor.getChildNames(dataAccessor.keyBuilder().liveInstances());
    for (int i = 1; i < _participants.length; i++) {
      if (_participants[i].isConnected()
          || liveInstances.contains(_participants[i].getInstanceName())) {
        return false;
      }
    }
    return true;
  }, TestHelper.WAIT_DURATION);
  Assert.assertTrue(result);

  // Add 1 db with batch message enabled. Each db has 10 partitions.
  // So it will have 1 batch message and 10 sub messages.
  String dbName = "TestDBSubMessageFail";
  IdealState idealState = new FullAutoModeISBuilder(dbName).setStateModel("OnlineOffline")
      .setStateModelFactoryName("TestFactory").setNumPartitions(10).setNumReplica(1).build();
  idealState.setBatchMessageMode(true);
  _gSetupTool.getClusterManagementTool().addResource(CLUSTER_NAME, dbName, idealState);

  // Check that IdealState has really been added
  result = TestHelper.verify(
      () -> dataAccessor.getPropertyStat(dataAccessor.keyBuilder().idealStates(dbName)) != null,
      TestHelper.WAIT_DURATION);
  Assert.assertTrue(result);

  for (int i = 0; i < 5; i++) {
    IdealState is =
        _gSetupTool.getClusterManagementTool().getResourceIdealState(CLUSTER_NAME, dbName);
    if (!idealState.equals(is)) {
      Thread.sleep(1000L);
    }
  }
  _gSetupTool.rebalanceStorageCluster(CLUSTER_NAME, dbName, 1);
  Assert.assertTrue(_clusterVerifier.verifyByPolling());
  Thread.sleep(2000L);

  int numOfOnlines = 0;
  int numOfErrors = 0;
  ExternalView externalView =
      _gSetupTool.getClusterManagementTool().getResourceExternalView(CLUSTER_NAME, dbName);
  for (String partition : externalView.getPartitionSet()) {
    if (externalView.getStateMap(partition).values().contains("ONLINE")) {
      numOfOnlines++;
    }
    if (externalView.getStateMap(partition).values().contains("ERROR")) {
      numOfErrors++;
    }
  }
  Assert.assertEquals(numOfErrors, 4);
  Assert.assertEquals(numOfOnlines, 6);
}
 

@Test
public void testStateMismatch() throws InterruptedException {
  // String controllerName = CONTROLLER_PREFIX + "_0";

  HelixManager manager = _controller; // _startCMResultMap.get(controllerName)._manager;
  HelixDataAccessor accessor = manager.getHelixDataAccessor();
  Builder kb = accessor.keyBuilder();
  ExternalView ev = accessor.getProperty(kb.externalView(TEST_DB));
  Map<String, LiveInstance> liveinstanceMap =
      accessor.getChildValuesMap(accessor.keyBuilder().liveInstances(), true);

  for (String instanceName : liveinstanceMap.keySet()) {
    String sessionid = liveinstanceMap.get(instanceName).getEphemeralOwner();
    for (String partition : ev.getPartitionSet()) {
      if (ev.getStateMap(partition).containsKey(instanceName)) {
        String uuid = UUID.randomUUID().toString();
        Message message = new Message(MessageType.STATE_TRANSITION, uuid);
        boolean rand = new Random().nextInt(10) > 5;
        if (ev.getStateMap(partition).get(instanceName).equals("MASTER")) {
          message.setSrcName(manager.getInstanceName());
          message.setTgtName(instanceName);
          message.setMsgState(MessageState.NEW);
          message.setPartitionName(partition);
          message.setResourceName(TEST_DB);
          message.setFromState(rand ? "SLAVE" : "OFFLINE");
          message.setToState(rand ? "MASTER" : "SLAVE");
          message.setTgtSessionId(sessionid);
          message.setSrcSessionId(manager.getSessionId());
          message.setStateModelDef("MasterSlave");
          message.setStateModelFactoryName("DEFAULT");
        } else if (ev.getStateMap(partition).get(instanceName).equals("SLAVE")) {
          message.setSrcName(manager.getInstanceName());
          message.setTgtName(instanceName);
          message.setMsgState(MessageState.NEW);
          message.setPartitionName(partition);
          message.setResourceName(TEST_DB);
          message.setFromState(rand ? "MASTER" : "OFFLINE");
          message.setToState(rand ? "SLAVE" : "SLAVE");
          message.setTgtSessionId(sessionid);
          message.setSrcSessionId(manager.getSessionId());
          message.setStateModelDef("MasterSlave");
          message.setStateModelFactoryName("DEFAULT");
        }
        accessor.setProperty(accessor.keyBuilder().message(instanceName, message.getMsgId()),
            message);
      }
    }
  }
  Thread.sleep(3000);
  ExternalView ev2 = accessor.getProperty(kb.externalView(TEST_DB));
  Assert.assertTrue(ev.equals(ev2));
}
 

/**
 * Get the problematic partitions on the to-be-stop instance
 * Requirement:
 *  If the instance gets stopped and the partitions on the instance are OFFLINE,
 *  the cluster still have enough "healthy" replicas on other sibling instances
 *
 *  - sibling instances mean those who share the same partition (replicas) of the to-be-stop instance
 *
 * @param globalPartitionHealthStatus (instance => (partition name, health status))
 * @param instanceToBeStop The instance to be stopped
 * @param dataAccessor The data accessor
 * @return A list of problematic partitions if the instance is stopped
 */
public static List<String> perPartitionHealthCheck(List<ExternalView> externalViews,
    Map<String, Map<String, Boolean>> globalPartitionHealthStatus, String instanceToBeStop,
    HelixDataAccessor dataAccessor) {
  List<String> unhealthyPartitions = new ArrayList<>();

  for (ExternalView externalView : externalViews) {
    // Skip ANY_LIVEINSTANCES resources check, since ANY_LIVEINSTANCES resources have single partition
    // with 1 replica. There is no need to check sibiling replicas.
    if (ResourceConfig.ResourceConfigConstants.ANY_LIVEINSTANCE.name()
        .equals(externalView.getReplicas())) {
      continue;
    }

    StateModelDefinition stateModelDefinition = dataAccessor
        .getProperty(dataAccessor.keyBuilder().stateModelDef(externalView.getStateModelDefRef()));
    for (String partition : externalView.getPartitionSet()) {
      Map<String, String> stateMap = externalView.getStateMap(partition);
      // Only check if instance holds top state
      if (stateMap.containsKey(instanceToBeStop)
          && stateMap.get(instanceToBeStop).equals(stateModelDefinition.getTopState())) {
        for (String siblingInstance : stateMap.keySet()) {
          // Skip this self check
          if (siblingInstance.equals(instanceToBeStop)) {
            continue;
          }

          // We are checking sibling partition healthy status. So if partition health does not
          // exist or it is not healthy. We should mark this partition is unhealthy.
          if (!globalPartitionHealthStatus.containsKey(siblingInstance)
              || !globalPartitionHealthStatus.get(siblingInstance).containsKey(partition)
              || !globalPartitionHealthStatus.get(siblingInstance).get(partition)) {
            unhealthyPartitions.add(partition);
            break;
          }
        }
      }
    }
  }

  return unhealthyPartitions;
}
 

@Test(timeout = 180000)
public void testHAControllerDistributedAggregation() throws Exception {
  MetricCollectorHAController haController = new MetricCollectorHAController(configuration);
  haController.initializeHAController();
  // Wait for task assignment
  Thread.sleep(10000);

  Assert.assertTrue(haController.isInitialized());
  Assert.assertEquals(1, haController.getLiveInstanceHostNames().size());
  Assert.assertTrue(haController.getAggregationTaskRunner().performsClusterAggregation());
  Assert.assertTrue(haController.getAggregationTaskRunner().performsHostAggregation());

  // Add new instance
  InstanceConfig instanceConfig2 = new InstanceConfig("h2_12001");
  haController.admin.addInstance(CLUSTER_NAME, instanceConfig2);
  HelixManager manager2 = HelixManagerFactory.getZKHelixManager(CLUSTER_NAME,
    instanceConfig2.getInstanceName(),
    InstanceType.PARTICIPANT, haController.zkConnectUrl);
  manager2.getStateMachineEngine().registerStateModelFactory(DEFAULT_STATE_MODEL,
    new OnlineOfflineStateModelFactory(instanceConfig2.getInstanceName(),
      new AggregationTaskRunner(instanceConfig2.getInstanceName(), "", CLUSTER_NAME)));
  manager2.connect();
  haController.admin.rebalance(CLUSTER_NAME, METRIC_AGGREGATORS, 1);

  // Wait on re-assignment of partitions
  Thread.sleep(10000);
  Assert.assertEquals(2, haController.getLiveInstanceHostNames().size());

  ExternalView view = haController.admin.getResourceExternalView(CLUSTER_NAME, METRIC_AGGREGATORS);

  Map<String, String> partitionInstanceMap = new HashMap<>();

  for (String partition : view.getPartitionSet()) {
    Map<String, String> states = view.getStateMap(partition);
    // (instance, state) pairs
    for (Map.Entry<String, String> stateEntry : states.entrySet()) {
      partitionInstanceMap.put(partition, stateEntry.getKey());
      Assert.assertEquals("ONLINE", stateEntry.getValue());
    }
  }
  // Re-assigned partitions
  Assert.assertEquals(2, partitionInstanceMap.size());

  haController.getAggregationTaskRunner().stop();
  haController.manager.disconnect();
}