org.apache.helix.model.StateModelDefinition Java Examples

/**
 * Build OnlineOffline state model definition
 */
public static StateModelDefinition build() {
  Builder builder = new Builder(name);
  // init state
  builder.initialState(States.OFFLINE.name());

  // add states
  builder.addState(States.ONLINE.name(), 20);
  builder.addState(States.OFFLINE.name(), -1);
  for (HelixDefinedState state : HelixDefinedState.values()) {
    builder.addState(state.name(), -1);
  }

  // add transitions

  builder.addTransition(States.ONLINE.name(), States.OFFLINE.name(), 25);
  builder.addTransition(States.OFFLINE.name(), States.ONLINE.name(), 5);
  builder.addTransition(States.OFFLINE.name(), HelixDefinedState.DROPPED.name(), 0);

  // bounds
  builder.dynamicUpperBound(States.ONLINE.name(), "R");

  return builder.build();
}
 

public static int getStateCount(String state, StateModelDefinition stateModelDef, int liveAndEnabledSize,
    int preferenceListSize) {
  String num = stateModelDef.getNumInstancesPerState(state);
  int stateCount = -1;
  if ("N".equals(num)) {
    stateCount = liveAndEnabledSize;
  } else if ("R".equals(num)) {
    stateCount = preferenceListSize;
  } else {
    try {
      stateCount = Integer.parseInt(num);
    } catch (Exception e) {
      LOG.error("Invalid count for state:" + state + " ,count=" + num);
    }
  }

  return stateCount;
}
 

/**
 * Compute best state for partition in AUTO ideal state mode.
 * @param liveInstances
 * @param stateModelDef
 * @param preferenceList
 * @param currentStateOutput instance->state for each partition
 * @param disabledInstancesForPartition
 * @param idealState
 * @param clusterConfig
 * @param partition
 * @param monitoredResolver
 * @return
 */
protected Map<String, String> computeBestPossibleStateForPartition(Set<String> liveInstances,
    StateModelDefinition stateModelDef, List<String> preferenceList,
    CurrentStateOutput currentStateOutput, Set<String> disabledInstancesForPartition,
    IdealState idealState, ClusterConfig clusterConfig, Partition partition,
    MonitoredAbnormalResolver monitoredResolver) {
  Optional<Map<String, String>> optionalOverwrittenStates =
      computeStatesOverwriteForPartition(stateModelDef, preferenceList, currentStateOutput,
          idealState, partition, monitoredResolver);
  if (optionalOverwrittenStates.isPresent()) {
    return optionalOverwrittenStates.get();
  }

  Map<String, String> currentStateMap = new HashMap<>(
      currentStateOutput.getCurrentStateMap(idealState.getResourceName(), partition));
  return computeBestPossibleMap(preferenceList, stateModelDef, currentStateMap, liveInstances,
      disabledInstancesForPartition);
}
 

/**
 * Compute the best state for all partitions.
 * This is the default implementation, subclasses should re-implement
 * this method if its logic to generate bestpossible map for each partition is different from the default one here.
 *
 * @param cache
 * @param idealState
 * @param resource
 * @param currentStateOutput
 *          Provides the current state and pending state transitions for all partitions
 * @return
 */
@Override
public ResourceAssignment computeBestPossiblePartitionState(
    T cache, IdealState idealState, Resource resource,
    CurrentStateOutput currentStateOutput) {
  if (LOG.isDebugEnabled()) {
    LOG.debug("Processing resource:" + resource.getResourceName());
  }
  String stateModelDefName = idealState.getStateModelDefRef();
  StateModelDefinition stateModelDef = cache.getStateModelDef(stateModelDefName);
  ResourceAssignment partitionMapping = new ResourceAssignment(resource.getResourceName());
  for (Partition partition : resource.getPartitions()) {
    Set<String> disabledInstancesForPartition =
        cache.getDisabledInstancesForPartition(resource.getResourceName(), partition.toString());
    List<String> preferenceList = getPreferenceList(partition, idealState,
        Collections.unmodifiableSet(cache.getLiveInstances().keySet()));
    Map<String, String> bestStateForPartition =
        computeBestPossibleStateForPartition(cache.getLiveInstances().keySet(), stateModelDef,
            preferenceList, currentStateOutput, disabledInstancesForPartition, idealState,
            cache.getClusterConfig(), partition,
            cache.getAbnormalStateResolver(stateModelDefName));
    partitionMapping.addReplicaMap(partition, bestStateForPartition);
  }
  return partitionMapping;
}
 

/**
 * The current states are not valid if there are more than one top state replicas for a single top
 * state state model.
 */
@Override
public boolean checkCurrentStates(final CurrentStateOutput currentStateOutput,
    final String resourceName, final Partition partition, StateModelDefinition stateModelDef) {
  if (!stateModelDef.isSingleTopStateModel()) {
    return true;
  }
  // TODO: Cache top state count in the ResourceControllerDataProvider and avoid repeated counting
  // TODO: here. It would be premature to do it now. But with more use case, we can improve the
  // TODO: ResourceControllerDataProvider to calculate by default.
  if (currentStateOutput.getCurrentStateMap(resourceName, partition).values().stream()
      .filter(state -> state.equals(stateModelDef.getTopState())).count() > 1) {
    return false;
  }
  return true;
}
 

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

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

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

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

/**
 * Update per-instance resource bean
 * @param stateMap partition->state
 * @tags tags instance tags
 * @param stateModelDef
 */
public synchronized void update(Map<Partition, String> stateMap, Set<String> tags,
    StateModelDefinition stateModelDef) {
  if (tags == null || tags.isEmpty()) {
    _tags = ImmutableList.of(ClusterStatusMonitor.DEFAULT_TAG);
  } else {
    _tags = Lists.newArrayList(tags);
    Collections.sort(_tags);
  }

  int cnt = 0;
  for (String state : stateMap.values()) {
    // Skip DROPPED and initial state (e.g. OFFLINE)
    if (state.equalsIgnoreCase(HelixDefinedState.DROPPED.name())
        || state.equalsIgnoreCase(stateModelDef.getInitialState())) {
      continue;
    }
    cnt++;
  }
  _partitions = cnt;
}
 

public static StateModelDefinition generatePinotStateModelDefinition() {
  StateModelDefinition.Builder builder = new StateModelDefinition.Builder(PINOT_SEGMENT_ONLINE_OFFLINE_STATE_MODEL);
  builder.initialState(OFFLINE_STATE);

  builder.addState(ONLINE_STATE);
  builder.addState(CONSUMING_STATE);
  builder.addState(OFFLINE_STATE);
  builder.addState(DROPPED_STATE);
  // Set the initial state when the node starts

  // Add transitions between the states.
  builder.addTransition(CONSUMING_STATE, ONLINE_STATE);
  builder.addTransition(OFFLINE_STATE, CONSUMING_STATE);
  builder.addTransition(OFFLINE_STATE, ONLINE_STATE);
  builder.addTransition(CONSUMING_STATE, OFFLINE_STATE);
  builder.addTransition(ONLINE_STATE, OFFLINE_STATE);
  builder.addTransition(OFFLINE_STATE, DROPPED_STATE);

  // set constraints on states.
  // static constraint
  builder.dynamicUpperBound(ONLINE_STATE, "R");
  // dynamic constraint, R means it should be derived based on the replication
  // factor.
  builder.dynamicUpperBound(CONSUMING_STATE, "R");

  StateModelDefinition statemodelDefinition = builder.build();
  return statemodelDefinition;
}
 

/**
 * Check for a partition, whether all transitions for its replicas are downward transitions. Note
 * that this function does NOT check for ERROR states.
 * @param currentStateMap
 * @param bestPossibleMap
 * @param stateModelDef
 * @return true if there are; false otherwise
 */
private boolean isLoadBalanceDownwardForAllReplicas(Map<String, String> currentStateMap,
    Map<String, String> bestPossibleMap, StateModelDefinition stateModelDef) {
  Set<String> allInstances = new HashSet<>();
  allInstances.addAll(currentStateMap.keySet());
  allInstances.addAll(bestPossibleMap.keySet());
  Map<String, Integer> statePriorityMap = stateModelDef.getStatePriorityMap();

  for (String instance : allInstances) {
    String currentState = currentStateMap.get(instance);
    String bestPossibleState = bestPossibleMap.get(instance);
    if (currentState == null) {
      return false; // null -> state is upward
    }
    if (bestPossibleState != null) {
      // Compare priority values and return if an upward transition is found
      // Note that lower integer value implies higher priority
      if (!statePriorityMap.containsKey(currentState)
          || !statePriorityMap.containsKey(bestPossibleState)) {
        // If the state is not found in statePriorityMap, consider it not strictly downward by
        // default because we can't determine whether it is downward
        return false;
      }
      if (statePriorityMap.get(currentState) > statePriorityMap.get(bestPossibleState)) {
        return false;
      }
    }
  }
  return true;
}
 

/**
 * From current state output, find out the location of the top state of given resource
 * and partition
 *
 * @param currentStateOutput current state output
 * @param resourceName resource name
 * @param partition partition of the resource
 * @param stateModelDef state model def object
 * @return name of the node that contains top state, null if there is not top state recorded
 */
private String findCurrentTopStateLocation(CurrentStateOutput currentStateOutput,
    String resourceName, Partition partition, StateModelDefinition stateModelDef) {
  Map<String, String> stateMap = currentStateOutput.getCurrentStateMap(resourceName, partition);
  for (String instance : stateMap.keySet()) {
    if (stateMap.get(instance).equals(stateModelDef.getTopState())) {
      return instance;
    }
  }
  return null;
}
 

private Message generateCancellationMessageForPendingMessage(final String desiredState, final String currentState,
    final String nextState, final Message pendingMessage, final HelixManager manager,
    final Resource resource, final Partition partition, final Map<String, String> sessionIdMap,
    final String instanceName, final StateModelDefinition stateModelDef,
    final Message cancellationMessage, final boolean isCancellationEnabled) {

  Message message = null;

  if (pendingMessage != null) {
    String pendingState = pendingMessage.getToState();
    if (nextState.equalsIgnoreCase(pendingState)) {
      LogUtil.logInfo(logger, _eventId,
          "Message already exists for " + instanceName + " to transit " + resource
              .getResourceName() + "." + partition.getPartitionName() + " from "
              + currentState + " to " + nextState + ", isRelay: " + pendingMessage.isRelayMessage());
    } else if (currentState.equalsIgnoreCase(pendingState)) {
      LogUtil.logInfo(logger, _eventId,
          "Message hasn't been removed for " + instanceName + " to transit " + resource
              .getResourceName() + "." + partition.getPartitionName() + " to "
              + pendingState + ", desiredState: " + desiredState + ", isRelay: " + pendingMessage.isRelayMessage());
    } else {
      LogUtil.logInfo(logger, _eventId,
          "IdealState changed before state transition completes for " + resource
              .getResourceName() + "." + partition.getPartitionName() + " on "
              + instanceName + ", pendingState: " + pendingState + ", currentState: "
              + currentState + ", nextState: " + nextState + ", isRelay: " + pendingMessage.isRelayMessage());

      message = createStateTransitionCancellationMessage(manager, resource,
          partition.getPartitionName(), instanceName, sessionIdMap.get(instanceName),
          stateModelDef.getId(), pendingMessage.getFromState(), pendingState, nextState,
          cancellationMessage, isCancellationEnabled, currentState);
    }
  }
  return message;
}
 

@Override
public void process(ClusterEvent event) throws Exception {
  _eventId = event.getEventId();
  CurrentStateOutput currentStateOutput = event.getAttribute(AttributeName.CURRENT_STATE.name());
  final Map<String, Resource> resourceMap =
      event.getAttribute(AttributeName.RESOURCES_TO_REBALANCE.name());
  final ClusterStatusMonitor clusterStatusMonitor =
      event.getAttribute(AttributeName.clusterStatusMonitor.name());
  ResourceControllerDataProvider cache =
      event.getAttribute(AttributeName.ControllerDataProvider.name());

  if (currentStateOutput == null || resourceMap == null || cache == null) {
    throw new StageException(
        "Missing attributes in event:" + event + ". Requires CURRENT_STATE|RESOURCES|DataCache");
  }

  final BestPossibleStateOutput bestPossibleStateOutput =
      compute(event, resourceMap, currentStateOutput);
  event.addAttribute(AttributeName.BEST_POSSIBLE_STATE.name(), bestPossibleStateOutput);

  final Map<String, InstanceConfig> instanceConfigMap = cache.getInstanceConfigMap();
  final Map<String, StateModelDefinition> stateModelDefMap = cache.getStateModelDefMap();
  asyncExecute(cache.getAsyncTasksThreadPool(), new Callable<Object>() {
    @Override
    public Object call() {
      try {
        if (clusterStatusMonitor != null) {
          clusterStatusMonitor
              .setPerInstanceResourceStatus(bestPossibleStateOutput, instanceConfigMap,
                  resourceMap, stateModelDefMap);
        }
      } catch (Exception e) {
        LogUtil.logError(logger, _eventId, "Could not update cluster status metrics!", e);
      }
      return null;
    }
  });
}
 

@Override
public Representation post(Representation entity) {
  try {
    String clusterName = (String) getRequest().getAttributes().get("clusterName");
    String modelName = (String) getRequest().getAttributes().get("modelName");
    ZkClient zkClient =
        (ZkClient) getContext().getAttributes().get(RestAdminApplication.ZKCLIENT);

    JsonParameters jsonParameters = new JsonParameters(entity);
    String command = jsonParameters.getCommand();

    if (command.equalsIgnoreCase(ClusterSetup.addStateModelDef)) {
      ZNRecord newStateModel =
          jsonParameters.getExtraParameter(JsonParameters.NEW_STATE_MODEL_DEF);
      HelixDataAccessor accessor =
          ClusterRepresentationUtil.getClusterDataAccessor(zkClient, clusterName);

      accessor.setProperty(accessor.keyBuilder().stateModelDef(newStateModel.getId()),
          new StateModelDefinition(newStateModel));
    } else {
      throw new HelixException("Unsupported command: " + command + ". Should be one of ["
          + ClusterSetup.addStateModelDef + "]");
    }
    getResponse().setEntity(getStateModelRepresentation(clusterName, modelName));
    getResponse().setStatus(Status.SUCCESS_OK);
  } catch (Exception e) {
    getResponse().setEntity(ClusterRepresentationUtil.getErrorAsJsonStringFromException(e),
        MediaType.APPLICATION_JSON);
    getResponse().setStatus(Status.SUCCESS_OK);
    LOG.error("Error in posting " + entity, e);
  }
  return null;
}
 

/**
 * Get a StateModelDefinition without transitions. The auto rebalancer doesn't take transitions
 * into account when computing mappings, so this is acceptable.
 * @param modelName name to give the model
 * @param initialState initial state for all nodes
 * @param states ordered map of state to count
 * @return incomplete StateModelDefinition for rebalancing
 */
private StateModelDefinition getIncompleteStateModelDef(String modelName, String initialState,
    LinkedHashMap<String, Integer> states) {
  StateModelDefinition.Builder builder = new StateModelDefinition.Builder(modelName);
  builder.initialState(initialState);
  int i = states.size();
  for (String state : states.keySet()) {
    builder.addState(state, i);
    builder.upperBound(state, states.get(state));
    i--;
  }
  return builder.build();
}
 

/**
 * Compute the best state for all partitions.
 * This is the default implementation, subclasses should re-implement
 * this method if its logic to generate bestpossible map for each partition is different from the default one here.
 *
 * @param cache
 * @param idealState
 * @param resource
 * @param currentStateOutput Provides the current state and pending state transitions for all partitions
 * @return
 */
@Override
public ResourceAssignment computeBestPossiblePartitionState(ResourceControllerDataProvider cache,
    IdealState idealState, Resource resource, CurrentStateOutput currentStateOutput) {
  if (LOG.isDebugEnabled()) {
    LOG.debug("Processing resource:" + resource.getResourceName());
  }

  Set<String> allNodes = cache.getEnabledInstances();
  Set<String> liveNodes = cache.getLiveInstances().keySet();

  ClusterConfig clusterConfig = cache.getClusterConfig();
  long delayTime = DelayedRebalanceUtil.getRebalanceDelay(idealState, clusterConfig);
  Set<String> activeNodes = DelayedRebalanceUtil
      .getActiveNodes(allNodes, idealState, liveNodes, cache.getInstanceOfflineTimeMap(),
          cache.getLiveInstances().keySet(), cache.getInstanceConfigMap(), delayTime,
          clusterConfig);

  String stateModelDefName = idealState.getStateModelDefRef();
  StateModelDefinition stateModelDef = cache.getStateModelDef(stateModelDefName);
  ResourceAssignment partitionMapping = new ResourceAssignment(resource.getResourceName());
  for (Partition partition : resource.getPartitions()) {
    Set<String> disabledInstancesForPartition =
        cache.getDisabledInstancesForPartition(resource.getResourceName(), partition.toString());
    List<String> preferenceList = getPreferenceList(partition, idealState, activeNodes);
    Map<String, String> bestStateForPartition =
        computeBestPossibleStateForPartition(liveNodes, stateModelDef, preferenceList,
            currentStateOutput, disabledInstancesForPartition, idealState, clusterConfig,
            partition, cache.getAbnormalStateResolver(stateModelDefName));

    partitionMapping.addReplicaMap(partition, bestStateForPartition);
  }

  if (LOG.isDebugEnabled()) {
    LOG.debug("Best possible mapping for resource  " + resource.getResourceName() + ": "
        + partitionMapping);
  }

  return partitionMapping;
}
 

protected Map<String, String> computeBestPossibleMapForDisabledResource(
    final Map<String, String> currentStateMap, StateModelDefinition stateModelDef) {
  Map<String, String> bestPossibleStateMap = new HashMap<>();
  for (String instance : currentStateMap.keySet()) {
    if (!HelixDefinedState.ERROR.name().equals(currentStateMap.get(instance))) {
      bestPossibleStateMap.put(instance, stateModelDef.getInitialState());
    }
  }
  return bestPossibleStateMap;
}
 

@Override
public void addStateModelDef(String clusterName, String stateModelDefName,
    String stateModelDefFile)
    throws IOException {
  ZNRecord record =
      (ZNRecord) (new ZNRecordSerializer().deserialize(readFile(stateModelDefFile)));
  if (record == null || record.getId() == null || !record.getId().equals(stateModelDefName)) {
    throw new IllegalArgumentException(
        "state model definition must have same id as state model def name");
  }
  addStateModelDef(clusterName, stateModelDefName, new StateModelDefinition(record), false);
}
 

/**
 * If the proposed instance may cause unnecessary state transition (according to the current
 * state), check and return a alternative instance to avoid.
 *
 * @param requestedState                      The requested state.
 * @param proposedInstance                    The current proposed instance to host the replica
 *                                            with the specified state.
 * @param currentState                        The current state of the proposed Instance, or init
 *                                            state if the proposed instance does not have an
 *                                            assignment.
 * @param stateModelDef
 * @param assignedInstances
 * @param remainCandidateCount                The count of the remaining unassigned instances
 * @param remainRequestCount                  The count of the remaining replicas that need to be
 *                                            assigned with the given state.
 * @param currentStatePrioritizedInstanceIter The iterator of the prioritized instance list which
 *                                            can be used to find a better alternative instance.
 * @return The alternative instance, or the original proposed instance if adjustment is not
 * necessary.
 */
private String adjustInstanceIfNecessary(String requestedState, String proposedInstance,
    String currentState, StateModelDefinition stateModelDef, Set<String> assignedInstances,
    int remainCandidateCount, int remainRequestCount,
    Iterator<String> currentStatePrioritizedInstanceIter) {
  String adjustedInstance = proposedInstance;
  // Check and alternate the assignment for reducing top state handoff.
  // 1. If the requested state is not the top state, then it does not worth it to adjust.
  // 2. If all remaining candidates need to be assigned with the the state, then there is no need
  // to adjust.
  // 3. If the proposed instance already has the top state or a secondary state, then adjustment
  // is not necessary.
  if (remainRequestCount < remainCandidateCount && requestedState
      .equals(stateModelDef.getTopState()) && !requestedState.equals(currentState)
      && !stateModelDef.getSecondTopStates().contains(currentState)) {
    // If the desired state is the top state, but the instance cannot be transited to the
    // top state in one hop, try to keep the top state on current host or a host with a closer
    // state.
    while (currentStatePrioritizedInstanceIter.hasNext()) {
      // Note that it is safe to check the prioritized instance items only once here.
      // Since the only possible condition when we don't use an instance in this list is that
      // it has been assigned with a state. And this is not revertable.
      String currentStatePrioritizedInstance = currentStatePrioritizedInstanceIter.next();
      if (!assignedInstances.contains(currentStatePrioritizedInstance)) {
        adjustedInstance = currentStatePrioritizedInstance;
        break;
      }
    }
    // Note that if all the current top state instances are not assignable, then we fallback
    // to the default logic that assigning the state according to preference list order.
  }
  return adjustedInstance;
}
 

@Override
public ResourceAssignment computeBestPossiblePartitionState(ResourceControllerDataProvider cache,
    IdealState idealState, Resource resource, CurrentStateOutput currentStateOutput) {
  // Looking for cached BestPossible mapping for this resource, if it is already there, do not recompute it again.
  // The cached mapping will be cleared in ResourceControllerDataProvider if there is anything changed in cluster state that can
  // cause the potential changes in BestPossible state.
  ResourceAssignment partitionMapping =
      cache.getCachedResourceAssignment(resource.getResourceName());
  if (partitionMapping != null) {
    return partitionMapping;
  }

  LOG.info("Computing BestPossibleMapping for " + resource.getResourceName());

  String stateModelDefName = idealState.getStateModelDefRef();
  StateModelDefinition stateModelDef = cache.getStateModelDef(stateModelDefName);
  partitionMapping = new ResourceAssignment(resource.getResourceName());
  for (Partition partition : resource.getPartitions()) {
    Map<String, String> currentStateMap =
        currentStateOutput.getCurrentStateMap(resource.getResourceName(), partition);
    Set<String> disabledInstancesForPartition =
        cache.getDisabledInstancesForPartition(resource.getResourceName(), partition.toString());
    Map<String, String> idealStateMap =
        idealState.getInstanceStateMap(partition.getPartitionName());
    Map<String, String> bestStateForPartition =
        computeCustomizedBestStateForPartition(cache, stateModelDef, idealStateMap,
            currentStateMap, disabledInstancesForPartition, idealState.isEnabled());
    partitionMapping.addReplicaMap(partition, bestStateForPartition);
  }

  cache.setCachedResourceAssignment(resource.getResourceName(), partitionMapping);
  if (LOG.isDebugEnabled()) {
    LOG.debug(String.format("Processing resource: %s", resource.getResourceName()));
    LOG.debug(String.format("Final Mapping of resource : %s", partitionMapping.toString()));
  }
  return partitionMapping;
}
 

public HelixStateMachineEngine(HelixManager manager) {
  _stateModelParser = new StateModelParser();
  _manager = manager;

  _stateModelFactoryMap =
      new ConcurrentHashMap<String, Map<String, StateModelFactory<? extends StateModel>>>();
  _stateModelDefs = new ConcurrentHashMap<String, StateModelDefinition>();
}
 

/**
 * Instantiate a validator instance
 * @param stateModelDef the state model definition to validate
 */
private StateModelDefinitionValidator(StateModelDefinition stateModelDef) {
  _stateModelDef = stateModelDef;
  _statePriorityList = stateModelDef.getStatesPriorityList();
  _transitionPriorityList = stateModelDef.getStateTransitionPriorityList();
  _stateSet = Sets.newHashSet(_statePriorityList);
}
 

@Override
public boolean checkCurrentStates(final CurrentStateOutput currentStateOutput,
    final String resourceName, final Partition partition,
    final StateModelDefinition stateModelDef) {
  // By default, all current states are valid.
  return true;
}
 

private boolean verifyExternalView(ExternalView externalView,
    PartitionStateMap bestPossibleState, StateModelDefinition stateModelDef) {
  Set<String> ignoreStates = new HashSet<>(
      Arrays.asList(stateModelDef.getInitialState(), HelixDefinedState.DROPPED.toString()));

  Map<String, Map<String, String>> bestPossibleStateMap =
      convertBestPossibleState(bestPossibleState);

  removeEntryWithIgnoredStates(bestPossibleStateMap.entrySet().iterator(), ignoreStates);

  Map<String, Map<String, String>> externalViewMap = externalView.getRecord().getMapFields();
  removeEntryWithIgnoredStates(externalViewMap.entrySet().iterator(), ignoreStates);

  return externalViewMap.equals(bestPossibleStateMap);
}
 

@GET
@Path("{clusterId}/statemodeldefs/{statemodel}")
public Response getClusterStateModelDefinition(@PathParam("clusterId") String clusterId,
    @PathParam("statemodel") String statemodel) {
  HelixDataAccessor dataAccessor = getDataAccssor(clusterId);
  StateModelDefinition stateModelDef =
      dataAccessor.getProperty(dataAccessor.keyBuilder().stateModelDef(statemodel));

  if (stateModelDef == null) {
    return badRequest("Statemodel not found!");
  }
  return JSONRepresentation(stateModelDef.getRecord());
}
 

/**
 * Validate there should be always minimal active replica and top state replica for each
 * partition.
 * Also make sure there is always some partitions with only active replica count.
 */
protected void validateMinActiveAndTopStateReplica(IdealState is, ExternalView ev,
    int minActiveReplica, int numNodes) {
  StateModelDefinition stateModelDef =
      BuiltInStateModelDefinitions.valueOf(is.getStateModelDefRef()).getStateModelDefinition();
  String topState = stateModelDef.getStatesPriorityList().get(0);
  int replica = Integer.valueOf(is.getReplicas());

  Map<String, Integer> stateCount = stateModelDef.getStateCountMap(numNodes, replica);
  Set<String> activeStates = stateCount.keySet();

  for (String partition : is.getPartitionSet()) {
    Map<String, String> assignmentMap = ev.getRecord().getMapField(partition);
    Assert.assertNotNull(assignmentMap,
        is.getResourceName() + "'s best possible assignment is null for partition " + partition);
    Assert.assertTrue(!assignmentMap.isEmpty(),
        is.getResourceName() + "'s partition " + partition + " has no best possible map in IS.");

    boolean hasTopState = false;
    int activeReplica = 0;
    for (String state : assignmentMap.values()) {
      if (topState.equalsIgnoreCase(state)) {
        hasTopState = true;
      }
      if (activeStates.contains(state)) {
        activeReplica++;
      }
    }

    if (activeReplica < minActiveReplica) {
      int a = 0;
    }

    Assert.assertTrue(hasTopState, String.format("%s missing %s replica", partition, topState));
    Assert.assertTrue(activeReplica >= minActiveReplica,
        String.format("%s has less active replica %d then required %d", partition, activeReplica,
            minActiveReplica));
  }
}
 

public static String[] parseStates(String clusterName, StateModelDefinition stateModDef) {
  String[] result = new String[2];
  String masterStateValue = null, slaveStateValue = null;

  // StateModelDefinition def = new StateModelDefinition(stateModDef);

  List<String> statePriorityList = stateModDef.getStatesPriorityList();

  for (String state : statePriorityList) {
    String count = stateModDef.getNumInstancesPerState(state);
    if (count.equals("1")) {
      if (masterStateValue != null) {
        throw new HelixException("Invalid or unsupported state model definition");
      }
      masterStateValue = state;
    } else if (count.equalsIgnoreCase("R")) {
      if (slaveStateValue != null) {
        throw new HelixException("Invalid or unsupported state model definition");
      }
      slaveStateValue = state;
    } else if (count.equalsIgnoreCase("N")) {
      if (!(masterStateValue == null && slaveStateValue == null)) {
        throw new HelixException("Invalid or unsupported state model definition");
      }
      masterStateValue = slaveStateValue = state;
    }
  }
  if (masterStateValue == null && slaveStateValue == null) {
    throw new HelixException("Invalid or unsupported state model definition");
  }

  if (masterStateValue == null) {
    masterStateValue = slaveStateValue;
  }
  result[0] = masterStateValue;
  result[1] = slaveStateValue;
  return result;
}
 

/**
 * Get state model definition based on given name.
 * @param stateModelDefName the name of state model definition that would be employed by Ambry cluster.
 * @return {@link StateModelDefinition}
 */
private StateModelDefinition getStateModelDefinition(String stateModelDefName) {
  StateModelDefinition stateModelDefinition;
  switch (stateModelDefName) {
    case ClusterMapConfig.OLD_STATE_MODEL_DEF:
      stateModelDefinition = LeaderStandbySMD.build();
      break;
    case ClusterMapConfig.AMBRY_STATE_MODEL_DEF:
      stateModelDefinition = AmbryStateModelDefinition.getDefinition();
      break;
    default:
      throw new IllegalArgumentException("Unsupported state model def: " + stateModelDefName);
  }
  return stateModelDefinition;
}
 

private static StateModelDefinition defineStateModel() {
  StateModelDefinition.Builder builder = new StateModelDefinition.Builder("Bootstrap");
  // Add states and their rank to indicate priority. Lower the rank higher the priority
  builder.addState("MASTER", 1);
  builder.addState("SLAVE", 2);
  builder.addState("BOOTSTRAP", 3);
  builder.addState("OFFLINE");
  builder.addState("DROPPED");
  // Set the initial state when the node starts
  builder.initialState("OFFLINE");

  // Add transitions between the states.
  builder.addTransition("OFFLINE", "BOOTSTRAP", 3);
  builder.addTransition("BOOTSTRAP", "SLAVE", 2);
  builder.addTransition("SLAVE", "MASTER", 1);
  builder.addTransition("MASTER", "SLAVE", 4);
  builder.addTransition("SLAVE", "OFFLINE", 5);
  builder.addTransition("OFFLINE", "DROPPED", 6);

  // set constraints on states.
  // static constraint
  builder.upperBound("MASTER", 1);
  // dynamic constraint, R means it should be derived based on the replication factor.
  builder.dynamicUpperBound("SLAVE", "R");

  StateModelDefinition statemodelDefinition = builder.build();

  assert (statemodelDefinition.isValid());

  return statemodelDefinition;
}
 

/**
 * 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 void validateMap(String resource, String partition, int replica,
    Map<String, String> instanceStateMap, StateModelDefinition stateModelDef) {
  if (instanceStateMap == null || instanceStateMap.isEmpty()) {
    _testSuccess = false;
    Assert.fail(
        String.format("Resource %s partition %s has no active replica!", resource, partition));
  }
  if (instanceStateMap.size() < replica) {
    _testSuccess = false;
    Assert.fail(
        String.format("Resource %s partition %s has %d active replica, less than required %d!",
            resource, partition, instanceStateMap.size(), replica));
  }

  Map<String, Integer> stateCountMap = stateModelDef.getStateCountMap(NUM_NODE, replica);
  String topState = stateModelDef.getStatesPriorityList().get(0);
  if (stateCountMap.get(topState) == 1) {
    int topCount = 0;
    for (String val : instanceStateMap.values()) {
      if (topState.equals(val)) {
        topCount++;
      }
    }
    if (topCount > 1) {
      _testSuccess = false;
      Assert.fail(String.format("Resource %s partition %s has %d replica in %s, more than 1!",
          resource, partition, topCount, topState));
    }
  }
}