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

public static IdealState resetCustomIdealStateFor(IdealState oldIdealState,
    String topicName, String oldPartition, String newPartition, String newInstanceName) {
  final CustomModeISBuilder customModeIdealStateBuilder = new CustomModeISBuilder(topicName);

  int oldNumPartitions = oldIdealState.getNumPartitions();

  customModeIdealStateBuilder
      .setStateModel(OnlineOfflineStateModel.name)
      .setNumPartitions(oldNumPartitions).setNumReplica(1)
      .setMaxPartitionsPerNode(oldNumPartitions);

  for (String partitionName : oldIdealState.getPartitionSet()) {
    String instanceName = oldIdealState.getInstanceStateMap(partitionName).keySet().iterator().next();
    String instanceToUse = partitionName.equals(oldPartition) ? newInstanceName : instanceName;
    String partitionToUse = partitionName.equals(oldPartition) ? newPartition : oldPartition;
    customModeIdealStateBuilder.assignInstanceAndState(partitionToUse, instanceToUse, "ONLINE");
  }

  return customModeIdealStateBuilder.build();
}
 

private void setupUnbalancedDB() throws InterruptedException {
  // Start with Full-Auto mode to create the partitions, Semi-Auto won't create partitions.
  _gSetupTool.addResourceToCluster(CLUSTER_NAME, UNBALANCED_DB_NAME, 50, MASTER_SLAVE_STATE_MODEL,
      IdealState.RebalanceMode.FULL_AUTO.toString());
  _gSetupTool.rebalanceStorageCluster(CLUSTER_NAME, UNBALANCED_DB_NAME, 1);

  // Set preference list to put all partitions to one instance.
  IdealState idealState = _gSetupTool.getClusterManagementTool()
      .getResourceIdealState(CLUSTER_NAME, UNBALANCED_DB_NAME);
  Set<String> partitions = idealState.getPartitionSet();
  for (String partition : partitions) {
    idealState.setPreferenceList(partition,
        Lists.newArrayList(_blockedParticipant.getInstanceName()));
  }
  idealState.setRebalanceMode(IdealState.RebalanceMode.SEMI_AUTO);

  _gSetupTool.getClusterManagementTool().setResourceIdealState(CLUSTER_NAME, UNBALANCED_DB_NAME,
      idealState);

  Assert.assertTrue(_clusterVerifier.verifyByPolling(10000, 100));
}
 

public static IdealState resetCustomIdealStateFor(IdealState oldIdealState,
    String topicName, String partitionToReplace, String newInstanceName) {
  final CustomModeISBuilder customModeIdealStateBuilder = new CustomModeISBuilder(topicName);

  int oldNumPartitions = oldIdealState.getNumPartitions();

  customModeIdealStateBuilder
      .setStateModel(OnlineOfflineStateModel.name)
      .setNumPartitions(oldNumPartitions).setNumReplica(1)
      .setMaxPartitionsPerNode(oldNumPartitions);

  for (String partitionName : oldIdealState.getPartitionSet()) {
    String instanceName = oldIdealState.getInstanceStateMap(partitionName).keySet().iterator().next();
    String instanceToUse = partitionName.equals(partitionToReplace) ? newInstanceName : instanceName;
    customModeIdealStateBuilder.assignInstanceAndState(partitionName, instanceToUse, "ONLINE");
  }

  return customModeIdealStateBuilder.build();
}
 

private void setParticipants(String dbName) throws InterruptedException {
  _factories = new HashMap<>();
  IdealState idealState =
      _gSetupTool.getClusterManagementTool().getResourceIdealState(CLUSTER_NAME, dbName);
  for (int i = 0; i < NODE_NR; i++) {
    if (_participants[i] != null) {
      _participants[i].syncStop();
    }
    Thread.sleep(1000);
    String instanceName = PARTICIPANT_PREFIX + "_" + (START_PORT + i);
    SleepStateModelFactory factory = new SleepStateModelFactory(1000);
    _factories.put(instanceName, factory);
    for (String p : idealState.getPartitionSet()) {
      if (idealState.getPreferenceList(p).get(0).equals(instanceName)) {
        factory.addPartition(p);
      }
    }

    _participants[i] = new MockParticipantManager(ZK_ADDR, CLUSTER_NAME, instanceName);
    _participants[i].getStateMachineEngine().registerStateModelFactory("MasterSlave", factory);
    _participants[i].syncStart();
  }
}
 

public static IdealState expandInstanceCustomIdealStateFor(IdealState oldIdealState,
    String topicName, String newPartition, List<String> instances, int maxNumReplica) {
  final CustomModeISBuilder customModeIdealStateBuilder = new CustomModeISBuilder(topicName);

  int oldNumPartitions = oldIdealState.getNumPartitions();

  customModeIdealStateBuilder
      .setStateModel(OnlineOfflineStateModel.name)
      .setNumPartitions(oldNumPartitions).setNumReplica(maxNumReplica)
      .setMaxPartitionsPerNode(oldNumPartitions);

  for (String partitionName : oldIdealState.getPartitionSet()) {
    for (String instanceName : oldIdealState.getInstanceStateMap(partitionName).keySet()) {
      customModeIdealStateBuilder.assignInstanceAndState(partitionName, instanceName, "ONLINE");
    }
    if (partitionName.equals(newPartition)) {
      for (String newInstanceName : instances) {
        customModeIdealStateBuilder.assignInstanceAndState(partitionName, newInstanceName, "ONLINE");
      }
    }
  }

  return customModeIdealStateBuilder.build();
}
 

@Override
public void onExternalViewChange(List<ExternalView> externalViewList,
    NotificationContext changeContext) {
  if (!trackEnabled) {
    return;
  }
  for (ExternalView ev : externalViewList) {
    IdealState is = _resourceMap.get(ev.getResourceName());
    if (is == null) {
      continue;
    }
    int replica = is.getReplicaCount(NUM_NODE);
    for (String p : is.getPartitionSet()) {
      Map<String, String> stateMap = ev.getStateMap(p);
      verifyPartitionCount(is.getResourceName(), p, stateMap, replica, "EV",
          is.getMinActiveReplicas());
    }
  }
}
 

/**
 * Validate instances for each partition is on different zone and with necessary tagged instances.
 */
private void validateZoneAndTagIsolation(IdealState is, ExternalView ev, int expectedReplica) {
  String tag = is.getInstanceGroupTag();
  for (String partition : is.getPartitionSet()) {
    Set<String> assignedZones = new HashSet<String>();

    Map<String, String> assignmentMap = ev.getRecord().getMapField(partition);
    Set<String> instancesInEV = assignmentMap.keySet();
    // TODO: preference List is not persisted in IS.
    // Assert.assertEquals(instancesInEV, instancesInIs);
    for (String instance : instancesInEV) {
      assignedZones.add(_nodeToZoneMap.get(instance));
      if (tag != null) {
        InstanceConfig config =
            _gSetupTool.getClusterManagementTool().getInstanceConfig(CLUSTER_NAME, instance);
        Assert.assertTrue(config.containsTag(tag));
      }
    }
    Assert.assertEquals(assignedZones.size(), expectedReplica);
  }
}
 

private static JSONObject buildIdealState(String topicName, IdealState idealStateForTopic) {
  JSONObject idealStatePartitionToServerMappingJson = new JSONObject();
  if (idealStateForTopic == null) {
    LOGGER.info("Ideal state for topic " + topicName + " is NULL");
  } else {
    for (String partition : idealStateForTopic.getPartitionSet()) {
      Map<String, String> stateMap = idealStateForTopic.getInstanceStateMap(partition);
      if (stateMap != null) {
        for (String server : stateMap.keySet()) {
          if (!idealStatePartitionToServerMappingJson.containsKey(partition)) {
            idealStatePartitionToServerMappingJson.put(partition, new JSONObject());
          }
          idealStatePartitionToServerMappingJson.getJSONObject(partition).put(server, stateMap.get(server));
        }
      }
    }
  }
  return idealStatePartitionToServerMappingJson;
}
 

public static Map<String, Integer> assertIdealServerPartitionCount(
    HelixMirrorMakerManager helixMirrorMakerManager,
    int numTotalPartitions) {
  Map<String, Integer> serverToPartitionMapping = new HashMap<>();
  int assginedPartitions = 0;
  for (String topicName : helixMirrorMakerManager.getTopicLists()) {
    IdealState idealStateForTopic =
        helixMirrorMakerManager.getIdealStateForTopic(topicName);
    LOGGER.info("IdealState: " + idealStateForTopic.toString());
    for (String partition : idealStateForTopic.getPartitionSet()) {
      String instanceName =
          idealStateForTopic.getInstanceStateMap(partition).keySet().iterator().next();
      if (!serverToPartitionMapping.containsKey(instanceName)) {
        serverToPartitionMapping.put(instanceName, 0);
      }
      serverToPartitionMapping.put(instanceName, serverToPartitionMapping.get(instanceName) + 1);
      assginedPartitions++;
    }
  }
  Assert.assertEquals(assginedPartitions, numTotalPartitions, "assignedPartitions not match with numTotalPartitions");
  return serverToPartitionMapping;
}
 

private boolean verify(String dbName) {
  IdealState idealState =
      _gSetupTool.getClusterManagementTool().getResourceIdealState(CLUSTER_NAME, dbName);
  HelixDataAccessor accessor = _manager.getHelixDataAccessor();
  ExternalView ev = accessor.getProperty(accessor.keyBuilder().externalView(dbName));
  for (String p : idealState.getPartitionSet()) {
    String idealMaster = idealState.getPreferenceList(p).get(0);
    if(!ev.getStateMap(p).get(idealMaster).equals("ERROR")) {
      return false;
    }

    TimeOutStateModel model = _factories.get(idealMaster).getStateModel(dbName, p);
    if (model._errorCallcount != 1 || model._error.getCode() != ErrorCode.TIMEOUT) {
      return false;
    }
  }

  return true;
}
 

static double calculateDeviationForResource(String resource,
                                            IdealState idealState,
                                            RoutingTableProvider routingTableProvider) {
  Set<String> partitions = idealState.getPartitionSet();
  int totalAssignments = 0, totalDeviations = 0;
  // Check if the external view has deviated from the actual view.
  for (String partition : partitions) {
    // Make a copy of the instance mapping in the ideal state.
    Set<String> idealInstances = new HashSet(idealState.getInstanceSet(partition));
    totalAssignments += idealInstances.size();
    // Now check against our real state and count the amount of deviating
    // assignments.
    List<InstanceConfig> currentInstanceConfigs = routingTableProvider.getInstances(
        resource, partition, "ONLINE");
    Set<String> currentInstances = Sets.newHashSetWithExpectedSize(
        currentInstanceConfigs.size());
    if (currentInstanceConfigs != null) {
      for (InstanceConfig instanceConfig : currentInstanceConfigs) {
        currentInstances.add(instanceConfig.getHostName());
      }
    }
    idealInstances.removeAll(currentInstances);
    totalDeviations += idealInstances.size();
  }
  return (double)totalDeviations / totalAssignments;
}
 

private boolean hasInstanceMapChanged(Map<Partition, Map<String, String>> newAssiments,
    IdealState idealState) {
  Set<Partition> partitions = new HashSet<Partition>(newAssiments.keySet());
  for (String p : idealState.getPartitionSet()) {
    partitions.add(new Partition(p));
  }

  for (Partition partition : partitions) {
    Map<String, String> instanceMap = newAssiments.get(partition);
    Map<String, String> existInstanceMap =
        idealState.getInstanceStateMap(partition.getPartitionName());
    if (instanceMap == null && existInstanceMap == null) {
      continue;
    }
    if (instanceMap == null || existInstanceMap == null || !instanceMap
        .equals(existInstanceMap)) {
      return true;
    }
  }

  return false;
}
 

@Test
public void testStateTransitionTimeOut() throws Exception {
  Map<String, SleepStateModelFactory> factories = new HashMap<String, SleepStateModelFactory>();
  IdealState idealState =
      _gSetupTool.getClusterManagementTool().getResourceIdealState(CLUSTER_NAME, TEST_DB);
  for (int i = 0; i < NODE_NR; i++) {
    String instanceName = PARTICIPANT_PREFIX + "_" + (START_PORT + i);
    SleepStateModelFactory factory = new SleepStateModelFactory(1000);
    factories.put(instanceName, factory);
    for (String p : idealState.getPartitionSet()) {
      if (idealState.getPreferenceList(p).get(0).equals(instanceName)) {
        factory.addPartition(p);
      }
    }

    _participants[i] = new MockParticipantManager(ZK_ADDR, CLUSTER_NAME, instanceName);
    _participants[i].getStateMachineEngine().registerStateModelFactory("MasterSlave", factory);
    _participants[i].syncStart();
  }
  String controllerName = CONTROLLER_PREFIX + "_0";
  _controller =
      new ClusterControllerManager(ZK_ADDR, CLUSTER_NAME, controllerName);
  _controller.syncStart();

  boolean result =
      ClusterStateVerifier
          .verifyByPolling(new MasterNbInExtViewVerifier(ZK_ADDR, CLUSTER_NAME));
  Assert.assertTrue(result);
  HelixDataAccessor accessor = _participants[0].getHelixDataAccessor();

  TestHelper.verify(() -> verify(accessor, idealState, factories), 5000);
  Assert.assertTrue(verify(accessor, idealState, factories));
}
 

private void addInstanceToBrokerIdealState(String brokerTenantTag, String instanceName) {
  IdealState tableIdealState = _helixAdmin.getResourceIdealState(_helixClusterName, Helix.BROKER_RESOURCE_INSTANCE);
  for (String tableNameWithType : tableIdealState.getPartitionSet()) {
    TableConfig tableConfig = ZKMetadataProvider.getTableConfig(_propertyStore, tableNameWithType);
    Preconditions.checkNotNull(tableConfig);
    String brokerTag = TagNameUtils.extractBrokerTag(tableConfig.getTenantConfig());
    if (brokerTag.equals(brokerTenantTag)) {
      tableIdealState.setPartitionState(tableNameWithType, instanceName, BrokerResourceStateModel.ONLINE);
    }
  }
  _helixAdmin.setResourceIdealState(_helixClusterName, Helix.BROKER_RESOURCE_INSTANCE, tableIdealState);
}
 

/**
 * Returns all instances for the given resource.
 *
 * @param idealState IdealState of the resource for which to return the instances of.
 * @return Returns a Set of strings containing the instance names for the given cluster.
 */
public static Set<String> getAllInstancesForResource(IdealState idealState) {
  final Set<String> instances = new HashSet<String>();

  for (final String partition : idealState.getPartitionSet()) {
    for (final String instance : idealState.getInstanceSet(partition)) {
      instances.add(instance);
    }
  }
  return instances;
}
 

private int findTopPartitionId(IdealState currentTaskQueue) {
  int topId = 0;
  for (String partitionName : currentTaskQueue.getPartitionSet()) {
    try {
      String partitionNumStr = partitionName.substring(partitionName.lastIndexOf('_') + 1);
      int num = Integer.parseInt(partitionNumStr);
      if (topId < num) {
        topId = num;
      }
    } catch (Exception e) {
      _logger.error("", e);
    }
  }
  return topId;
}
 

public void verifyReplication(HelixZkClient zkClient, String clusterName, String resource,
    int repl) {
  ZKHelixDataAccessor accessor =
      new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor<ZNRecord>(zkClient));
  Builder keyBuilder = accessor.keyBuilder();

  IdealState idealState = accessor.getProperty(keyBuilder.idealStates(resource));
  for (String partitionName : idealState.getPartitionSet()) {
    if (idealState.getRebalanceMode() == IdealState.RebalanceMode.SEMI_AUTO) {
      AssertJUnit.assertEquals(repl, idealState.getPreferenceList(partitionName).size());
    } else if (idealState.getRebalanceMode() == IdealState.RebalanceMode.CUSTOMIZED) {
      AssertJUnit.assertEquals(repl, idealState.getInstanceStateMap(partitionName).size());
    }
  }
}
 

protected void createDBInSemiAuto(ClusterSetup clusterSetup, String clusterName, String dbName,
    List<String> preferenceList, String stateModelDef, int numPartition, int replica) {
  clusterSetup.addResourceToCluster(clusterName, dbName, numPartition, stateModelDef,
      IdealState.RebalanceMode.SEMI_AUTO.toString());
  clusterSetup.rebalanceStorageCluster(clusterName, dbName, replica);

  IdealState is =
      _gSetupTool.getClusterManagementTool().getResourceIdealState(clusterName, dbName);
  for (String p : is.getPartitionSet()) {
    is.setPreferenceList(p, preferenceList);
  }
  clusterSetup.getClusterManagementTool().setResourceIdealState(clusterName, dbName, is);
}
 

/**
 * Validate instances for each partition is on different zone and with necessary tagged
 * instances.
 */
private void validateNoZeroReplica(IdealState is, ExternalView ev) {
  int replica = is.getReplicaCount(NUM_NODE);
  StateModelDefinition stateModelDef =
      BuiltInStateModelDefinitions.valueOf(is.getStateModelDefRef()).getStateModelDefinition();

  for (String partition : is.getPartitionSet()) {
    Map<String, String> evStateMap = ev.getRecord().getMapField(partition);
    Map<String, String> isStateMap = is.getInstanceStateMap(partition);
    validateMap(is.getResourceName(), partition, replica, evStateMap, stateModelDef);
    validateMap(is.getResourceName(), partition, replica, isStateMap, stateModelDef);
  }
}
 

private Map<String, String> computePartitionHealth(String clusterId, String resourceName) {
  HelixAdmin admin = getHelixAdmin();
  IdealState idealState = admin.getResourceIdealState(clusterId, resourceName);
  ExternalView externalView = admin.getResourceExternalView(clusterId, resourceName);
  StateModelDefinition stateModelDef =
      admin.getStateModelDef(clusterId, idealState.getStateModelDefRef());
  String initialState = stateModelDef.getInitialState();
  List<String> statesPriorityList = stateModelDef.getStatesPriorityList();
  statesPriorityList = statesPriorityList.subList(0, statesPriorityList.indexOf(initialState)); // Trim
                                                                                                // stateList
                                                                                                // to
                                                                                                // initialState
                                                                                                // and
                                                                                                // above
  int minActiveReplicas = idealState.getMinActiveReplicas();

  // Start the logic that determines the health status of each partition
  Map<String, String> partitionHealthResult = new HashMap<>();
  Set<String> allPartitionNames = idealState.getPartitionSet();
  if (!allPartitionNames.isEmpty()) {
    for (String partitionName : allPartitionNames) {
      int replicaCount =
          idealState.getReplicaCount(idealState.getPreferenceList(partitionName).size());
      // Simplify expectedStateCountMap by assuming that all instances are available to reduce
      // computation load on this REST endpoint
      LinkedHashMap<String, Integer> expectedStateCountMap =
          stateModelDef.getStateCountMap(replicaCount, replicaCount);
      // Extract all states into Collections from ExternalView
      Map<String, String> stateMapInExternalView = externalView.getStateMap(partitionName);
      Collection<String> allReplicaStatesInExternalView =
          (stateMapInExternalView != null && !stateMapInExternalView.isEmpty())
              ? stateMapInExternalView.values()
              : Collections.<String> emptyList();
      int numActiveReplicasInExternalView = 0;
      HealthStatus status = HealthStatus.HEALTHY;

      // Go through all states that are "active" states (higher priority than InitialState)
      for (int statePriorityIndex = 0; statePriorityIndex < statesPriorityList
          .size(); statePriorityIndex++) {
        String currentState = statesPriorityList.get(statePriorityIndex);
        int currentStateCountInIdealState = expectedStateCountMap.get(currentState);
        int currentStateCountInExternalView =
            Collections.frequency(allReplicaStatesInExternalView, currentState);
        numActiveReplicasInExternalView += currentStateCountInExternalView;
        // Top state counts must match, if not, unhealthy
        if (statePriorityIndex == 0
            && currentStateCountInExternalView != currentStateCountInIdealState) {
          status = HealthStatus.UNHEALTHY;
          break;
        } else if (currentStateCountInExternalView < currentStateCountInIdealState) {
          // For non-top states, if count in ExternalView is less than count in IdealState,
          // partially healthy
          status = HealthStatus.PARTIAL_HEALTHY;
        }
      }
      if (numActiveReplicasInExternalView < minActiveReplicas) {
        // If this partition does not satisfy the number of minimum active replicas, unhealthy
        status = HealthStatus.UNHEALTHY;
      }
      partitionHealthResult.put(partitionName, status.name());
    }
  }
  return partitionHealthResult;
}