org.apache.helix.HelixDefinedState Java Examples

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

/**
 * Build OnlineOfflineWithBootstrap state model definition
 *
 * @return
 */
public static OnlineOfflineWithBootstrapSMD build() {
  Builder builder = new Builder(name);
  // init state
  builder.initialState(States.OFFLINE.name());

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

  // add transitions
  builder.addTransition(States.ONLINE.name(), States.OFFLINE.name(), 0);
  builder.addTransition(States.OFFLINE.name(), States.BOOTSTRAP.name(), 1);
  builder.addTransition(States.BOOTSTRAP.name(), States.ONLINE.name(), 2);
  builder.addTransition(States.OFFLINE.name(), HelixDefinedState.DROPPED.name());

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

  return new OnlineOfflineWithBootstrapSMD(builder.build().getRecord());
}
 

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

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

  // add transitions
  builder.addTransition(States.ONLINE.name(), States.OFFLINE.name(), 0);
  builder.addTransition(States.OFFLINE.name(), States.ONLINE.name(), 1);
  builder.addTransition(States.OFFLINE.name(), HelixDefinedState.DROPPED.name());

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

  return builder.build();
}
 

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

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

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

  // add transitions
  builder.addTransition(States.MASTER.name(), States.OFFLINE.name(), 0);
  builder.addTransition(States.OFFLINE.name(), States.MASTER.name(), 1);
  builder.addTransition(States.OFFLINE.name(), HelixDefinedState.DROPPED.name());

  // bounds
  builder.dynamicUpperBound(States.MASTER.name(), "N");

  return builder.build();
}
 

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

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

  // add transitions
  builder.addTransition(States.COMPLETED.name(), States.OFFLINE.name(), 0);
  builder.addTransition(States.OFFLINE.name(), States.COMPLETED.name(), 1);
  builder.addTransition(States.OFFLINE.name(), HelixDefinedState.DROPPED.name());

  // bounds
  builder.dynamicUpperBound(States.COMPLETED.name(), "1");

  return builder.build();
}
 

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

  // add states
  builder.addState(States.MASTER.name(), 0);
  builder.addState(States.SLAVE.name(), 1);
  builder.addState(States.OFFLINE.name(), 2);
  for (HelixDefinedState state : HelixDefinedState.values()) {
    builder.addState(state.name());
  }

  // add transitions
  builder.addTransition(States.MASTER.name(), States.SLAVE.name(), 0);
  builder.addTransition(States.SLAVE.name(), States.MASTER.name(), 1);
  builder.addTransition(States.OFFLINE.name(), States.SLAVE.name(), 2);
  builder.addTransition(States.SLAVE.name(), States.OFFLINE.name(), 3);
  builder.addTransition(States.OFFLINE.name(), HelixDefinedState.DROPPED.name());

  // bounds
  builder.upperBound(States.MASTER.name(), 1);
  builder.dynamicUpperBound(States.SLAVE.name(), "R");

  return builder.build();
}
 

/**
 * This is the example used on the website, so this must work
 */
@Test
public void testBasic() {
  StateModelDefinition stateModel = new StateModelDefinition.Builder("MasterSlave")
  // OFFLINE is the state that the system starts in (initial state is REQUIRED)
      .initialState("OFFLINE")

      // Lowest number here indicates highest priority, no value indicates lowest priority
      .addState("MASTER", 1).addState("SLAVE", 2).addState("OFFLINE")

      // Note the special inclusion of the DROPPED state (REQUIRED)
      .addState(HelixDefinedState.DROPPED.toString())

      // No more than one master allowed
      .upperBound("MASTER", 1)

      // R indicates an upper bound of number of replicas for each partition
      .dynamicUpperBound("SLAVE", "R")

      // Add some high-priority transitions
      .addTransition("SLAVE", "MASTER", 1).addTransition("OFFLINE", "SLAVE", 2)

      // Using the same priority value indicates that these transitions can fire in any order
      .addTransition("MASTER", "SLAVE", 3).addTransition("SLAVE", "OFFLINE", 3)

      // Not specifying a value defaults to lowest priority
      // Notice the inclusion of the OFFLINE to DROPPED transition
      // Since every state has a path to OFFLINE, they each now have a path to DROPPED (REQUIRED)
      .addTransition("OFFLINE", HelixDefinedState.DROPPED.toString())

      // Create the StateModelDefinition instance
      .build();

  Assert.assertTrue(stateModel.isValid());
}
 

protected ResourceAssignment buildEmptyAssignment(String name,
    CurrentStateOutput currStateOutput) {
  ResourceAssignment assignment = new ResourceAssignment(name);
  Set<Partition> partitions = currStateOutput.getCurrentStateMappedPartitions(name);
  for (Partition partition : partitions) {
    Map<String, String> currentStateMap = currStateOutput.getCurrentStateMap(name, partition);
    Map<String, String> replicaMap = Maps.newHashMap();
    for (String instanceName : currentStateMap.keySet()) {
      replicaMap.put(instanceName, HelixDefinedState.DROPPED.toString());
    }
    assignment.addReplicaMap(partition, replicaMap);
  }
  return assignment;
}
 

private static ResourceAssignment emptyAssignment(String name, CurrentStateOutput currStateOutput) {
  ResourceAssignment assignment = new ResourceAssignment(name);
  Set<Partition> partitions = currStateOutput.getCurrentStateMappedPartitions(name);
  for (Partition partition : partitions) {
    Map<String, String> currentStateMap = currStateOutput.getCurrentStateMap(name, partition);
    Map<String, String> replicaMap = Maps.newHashMap();
    for (String instanceName : currentStateMap.keySet()) {
      replicaMap.put(instanceName, HelixDefinedState.DROPPED.toString());
    }
    assignment.addReplicaMap(partition, replicaMap);
  }
  return assignment;
}
 

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

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

  // add states
  builder.addState(States.LEADER.name(), 0);
  builder.addState(States.STANDBY.name(), 1);
  builder.addState(States.OFFLINE.name(), 2);
  for (HelixDefinedState state : HelixDefinedState.values()) {
    builder.addState(state.name());
  }

  // add transitions
  builder.addTransition(States.LEADER.name(), States.STANDBY.name(), 0);
  builder.addTransition(States.STANDBY.name(), States.LEADER.name(), 1);
  builder.addTransition(States.OFFLINE.name(), States.STANDBY.name(), 2);
  builder.addTransition(States.STANDBY.name(), States.OFFLINE.name(), 3);
  builder.addTransition(States.OFFLINE.name(), HelixDefinedState.DROPPED.name());

  // bounds
  builder.upperBound(States.LEADER.name(), 1);
  builder.dynamicUpperBound(States.STANDBY.name(), "R");

  return builder.build();
}
 

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

protected Map<String, String> computeBestPossibleMapForDroppedResource(
    final Map<String, String> currentStateMap) {
  Map<String, String> bestPossibleStateMap = new HashMap<>();
  for (String instance : currentStateMap.keySet()) {
    bestPossibleStateMap.put(instance, HelixDefinedState.DROPPED.toString());
  }
  return bestPossibleStateMap;
}
 

@Override
public void resetInstance(String clusterName, List<String> instanceNames) {
  // TODO: not mp-safe
  logger.info("Reset instances {} in cluster {}.",
      instanceNames == null ? "NULL" : HelixUtil.serializeByComma(instanceNames), clusterName);
  HelixDataAccessor accessor =
      new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor<>(_zkClient));
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();
  List<ExternalView> extViews = accessor.getChildValues(keyBuilder.externalViews(), true);

  Set<String> resetInstanceNames = new HashSet<String>(instanceNames);
  for (String instanceName : resetInstanceNames) {
    List<String> resetPartitionNames = new ArrayList<String>();
    for (ExternalView extView : extViews) {
      Map<String, Map<String, String>> stateMap = extView.getRecord().getMapFields();
      for (String partitionName : stateMap.keySet()) {
        Map<String, String> instanceStateMap = stateMap.get(partitionName);

        if (instanceStateMap.containsKey(instanceName) && instanceStateMap.get(instanceName)
            .equals(HelixDefinedState.ERROR.toString())) {
          resetPartitionNames.add(partitionName);
        }
      }
      resetPartition(clusterName, instanceName, extView.getResourceName(), resetPartitionNames);
    }
  }
}
 

/**
 * Assign the states to the instances listed in the preference list according to inputs.
 * Note that when we choose the top-state (e.g. MASTER) replica for a partition, we prefer to
 * choose it from these replicas which are already in the secondary states (e.g, SLAVE) instead
 * of in lower-state. This is because a replica in secondary state will take shorter time to
 * transition to the top-state, which could minimize the impact to the application's availability.
 * To achieve that, we sort the preferenceList based on CurrentState, by treating top-state and
 * second-states with same priority and rely on the fact that Collections.sort() is stable.
 */
private void assignStatesToInstances(final List<String> preferenceList,
    final StateModelDefinition stateModelDef, final Map<String, String> currentStateMap,
    final Set<String> liveInstances, final Set<String> disabledInstancesForPartition,
    Map<String, String> bestPossibleStateMap) {
  // Record the assigned instances to avoid double calculating or conflict assignment.
  Set<String> assignedInstances = new HashSet<>();

  Set<String> liveAndEnabled =
      liveInstances.stream().filter(instance -> !disabledInstancesForPartition.contains(instance))
          .collect(Collectors.toSet());

  Queue<String> preferredActiveInstanceQueue = new LinkedList<>(preferenceList);
  preferredActiveInstanceQueue.retainAll(liveAndEnabled);
  int totalCandidateCount = preferredActiveInstanceQueue.size();

  // Sort the preferred instances based on replicas' state priority in the current state.
  // Note that if one instance exists in the current states but not in the preference list, then
  // it won't show in the prioritized list.
  List<String> currentStatePrioritizedList = new ArrayList<>(preferredActiveInstanceQueue);
  currentStatePrioritizedList.sort(new StatePriorityComparator(currentStateMap, stateModelDef));
  Iterator<String> currentStatePrioritizedInstanceIter = currentStatePrioritizedList.iterator();

  // Assign the states to the instances that appear in the preference list.
  for (String state : stateModelDef.getStatesPriorityList()) {
    int stateCount =
        getStateCount(state, stateModelDef, liveAndEnabled.size(), preferenceList.size());
    while (!preferredActiveInstanceQueue.isEmpty()) {
      if (stateCount <= 0) {
        break; // continue assigning for the next state
      }
      String peekInstance = preferredActiveInstanceQueue.peek();
      if (assignedInstances.contains(peekInstance)) {
        preferredActiveInstanceQueue.poll();
        continue; // continue checking for the next available instance
      }
      String proposedInstance = adjustInstanceIfNecessary(state, peekInstance,
          currentStateMap.getOrDefault(peekInstance, stateModelDef.getInitialState()),
          stateModelDef, assignedInstances, totalCandidateCount - assignedInstances.size(),
          stateCount, currentStatePrioritizedInstanceIter);

      if (proposedInstance.equals(peekInstance)) {
        // If the peeked instance is the final decision, then poll it from the queue.
        preferredActiveInstanceQueue.poll();
      }
      // else, if we found a different instance for the partition placement, then we need to
      // check the same instance again or it will not be assigned with any partitions.

      // Assign the desired state to the proposed instance if not on ERROR state.
      if (HelixDefinedState.ERROR.toString().equals(currentStateMap.get(proposedInstance))) {
        bestPossibleStateMap.put(proposedInstance, HelixDefinedState.ERROR.toString());
      } else {
        bestPossibleStateMap.put(proposedInstance, state);
        stateCount--;
      }
      // Note that in either case, the proposed instance is considered to be assigned with a state
      // by now.
      if (!assignedInstances.add(proposedInstance)) {
        throw new AssertionError(String
            .format("The proposed instance %s has been already assigned before.",
                proposedInstance));
      }
    }
  }
}
 

/**
 * Check if the StateModelDefinition passes all validation checks
 * @return true if state model definition is valid, false otherwise
 */
public boolean isStateModelDefinitionValid() {
  // has a name
  if (_stateModelDef.getId() == null || _stateModelDef.getId().isEmpty()) {
    _logger.error("State model does not have a name");
    return false;
  }

  // has an initial state
  if (_stateModelDef.getInitialState() == null || _stateModelDef.getInitialState().isEmpty()) {
    _logger
        .error("State model does not contain init state, statemodel:" + _stateModelDef.getId());
    return false;
  }

  // has states
  if (_statePriorityList == null || _statePriorityList.isEmpty()) {
    _logger.error("CurrentState does not contain StatesPriorityList, state model : "
        + _stateModelDef.getId());
    return false;
  }

  // initial state is a state
  if (!_stateSet.contains(_stateModelDef.getInitialState())) {
    _logger.error("Defined states does not include the initial state, state model: "
        + _stateModelDef.getId());
    return false;
  }

  // has a dropped state
  if (!_stateSet.contains(HelixDefinedState.DROPPED.toString())) {
    _logger.error("Defined states does not include the DROPPED state, state model: "
        + _stateModelDef.getId());
    return false;
  }

  // make sure individual checks all pass
  if (!areStateCountsValid() || !areNextStatesValid() || !isTransitionPriorityListValid()
      || !arePathsValid()) {
    return false;
  }

  return true;
}
 

/**
 * Check if the "next" states in the state model definition are valid. These check the next values
 * at a single level. To check full paths, use {@link #arePathsValid()}.
 * @return true if next states are properly defined, false otherwise
 */
private boolean areNextStatesValid() {
  for (String state : _statePriorityList) {
    // all states can reach DROPPED
    if (!state.equals(HelixDefinedState.DROPPED.toString())
        && _stateModelDef.getNextStateForTransition(state, HelixDefinedState.DROPPED.toString()) == null) {
      _logger.error("State " + state + " cannot reach the DROPPED state, state model: "
          + _stateModelDef.getId());
      return false;
    }

    // initial state should reach all states (other than error)
    if (!state.equals(_stateModelDef.getInitialState())
        && !state.equals(HelixDefinedState.ERROR.toString())
        && _stateModelDef.getNextStateForTransition(_stateModelDef.getInitialState(), state) == null) {
      _logger.error("Initial state " + _stateModelDef.getInitialState()
          + " should be able to reach all states, state model: " + _stateModelDef.getId());
      return false;
    }

    // validate "next" states
    for (String destState : _statePriorityList) {
      if (state.equals(destState)) {
        continue;
      }
      // the next state should exist
      String intermediate = _stateModelDef.getNextStateForTransition(state, destState);
      if (intermediate != null && !_stateSet.contains(intermediate)) {
        _logger.error("Intermediate state " + intermediate + " for transition " + state + "-"
            + destState + " is not a valid state, state model: " + _stateModelDef.getId());
        return false;
      }

      // the next state should not allow a self loop
      if (intermediate != null && intermediate.equals(state)) {
        _logger.error("Intermediate state " + intermediate + " for transition " + state + "-"
            + destState + " should never be the from state, state model: "
            + _stateModelDef.getId());
        return false;
      }
    }
  }
  return true;
}
 

/**
 * Check that the state model does not have loops or unreachable states and that next states
 * actually help make progress
 * @return true if the transitions are valid, false otherwise
 */
private boolean arePathsValid() {
  // create a map for memoized path checking
  Map<String, Set<String>> alreadyChecked = Maps.newHashMap();
  for (String state : _statePriorityList) {
    alreadyChecked.put(state, new HashSet<String>());
  }

  // check all pairs for paths
  for (String from : _statePriorityList) {
    for (String to : _statePriorityList) {
      // ignore self transitions
      if (from.equals(to)) {
        continue;
      }

      // see if a path is claimed to exist
      Set<String> used = Sets.newHashSet(from);
      String next = _stateModelDef.getNextStateForTransition(from, to);
      if (next == null) {
        if (from.equals(_stateModelDef.getInitialState())
            && !to.equals(HelixDefinedState.ERROR.toString())) {
          _logger.error("Initial state " + from + " cannot reach " + to + ", state model: "
              + _stateModelDef.getId());
          return false;
        }
        continue;
      }
      // if a path exists, follow it all the way
      while (!to.equals(next)) {
        // no need to proceed if this path has already been traversed
        if (alreadyChecked.get(next).contains(to)) {
          break;
        }
        if (used.contains(next)) {
          _logger.error("Path from " + from + " to " + to
              + " contains an infinite loop, state model: " + _stateModelDef.getId());
          return false;
        }
        alreadyChecked.get(next).add(to);
        used.add(next);
        next = _stateModelDef.getNextStateForTransition(next, to);
        if (next == null) {
          _logger.error("Path from " + from + " to " + to + " is incomplete, state model: "
              + _stateModelDef.getId());
          return false;
        }
      }
      alreadyChecked.get(from).add(to);
    }
  }
  return true;
}
 

/**
 * Check that the intermediateStateOutput assignment does not violate maxPartitionPerInstance
 * threshold.
 * @param cache
 * @param intermediateStateOutput
 * @return true if violation is found, false otherwise.
 */
private boolean violatesMaxPartitionsPerInstance(ResourceControllerDataProvider cache,
    IntermediateStateOutput intermediateStateOutput) {
  int maxPartitionPerInstance = cache.getClusterConfig().getMaxPartitionsPerInstance();
  if (maxPartitionPerInstance <= 0) {
    // Config is not set; return
    return false;
  }
  Map<String, PartitionStateMap> resourceStatesMap =
      intermediateStateOutput.getResourceStatesMap();
  Map<String, Integer> instancePartitionCounts = new HashMap<>();

  for (String resource : resourceStatesMap.keySet()) {
    IdealState idealState = cache.getIdealState(resource);
    if (idealState != null
        && idealState.getStateModelDefRef().equals(BuiltInStateModelDefinitions.Task.name())) {
      // Ignore task here. Task has its own throttling logic
      continue;
    }

    PartitionStateMap partitionStateMap = resourceStatesMap.get(resource);
    Map<Partition, Map<String, String>> stateMaps = partitionStateMap.getStateMap();
    for (Partition p : stateMaps.keySet()) {
      Map<String, String> stateMap = stateMaps.get(p);
      for (String instance : stateMap.keySet()) {
        // If this replica is in DROPPED state, do not count it in the partition count since it is
        // to be dropped
        String state = stateMap.get(instance);
        if (state.equals(HelixDefinedState.DROPPED.name())) {
          continue;
        }
        if (!instancePartitionCounts.containsKey(instance)) {
          instancePartitionCounts.put(instance, 0);
        }
        // Number of replicas (from different partitions) held in this instance
        int partitionCount = instancePartitionCounts.get(instance);
        partitionCount++;
        if (partitionCount > maxPartitionPerInstance) {
          // There exists an instance whose intermediate state assignment violates the maximum
          // partitions per instance threshold, return!
          return true;
        }
        instancePartitionCounts.put(instance, partitionCount);
      }
    }
  }
  // No violation found
  return false;
}
 

private void checkRelayHost(Message relayMessage, Map<String, LiveInstance> liveInstanceMap,
    Map<String, Map<String, Map<String, CurrentState>>> currentStateMap, Message hostedMessage) {

  long currentTime = System.currentTimeMillis();

  String sessionId = hostedMessage.getTgtSessionId();
  String relayInstance = hostedMessage.getTgtName();
  String resourceName = hostedMessage.getResourceName();
  String partitionName = hostedMessage.getPartitionName();

  if (!liveInstanceMap.containsKey(relayInstance)) {
    // If the p2p forwarding host is no longer live, we should not remove the relay message immediately
    // since we do not know whether the relay message was forwarded before the instance went offline.
    setMessageRelayTime(relayMessage, currentTime);
    return;
  }
  String instanceSessionId = liveInstanceMap.get(relayInstance).getEphemeralOwner();
  if (!instanceSessionId.equals(sessionId)) {
    LOG.info("Relay instance sessionId {} does not match sessionId {} in hosted message {}, "
            + "set relay message {} to be expired.", instanceSessionId, sessionId,
        relayMessage.getId(), hostedMessage.getMsgId());
    setMessageRelayTime(relayMessage, currentTime);
    return;
  }

  Map<String, Map<String, CurrentState>> instanceCurrentStateMap =
      currentStateMap.get(relayInstance);
  if (instanceCurrentStateMap == null || !instanceCurrentStateMap.containsKey(sessionId)) {
    LOG.warn(
        "CurrentStateMap null for {}, session {}, set relay messages {} to be expired. Hosted message {}.",
        relayInstance, sessionId, relayMessage.getId(), hostedMessage.getId());
    setMessageRelayTime(relayMessage, currentTime);
    return;
  }

  Map<String, CurrentState> sessionCurrentStateMap = instanceCurrentStateMap.get(sessionId);
  CurrentState currentState = sessionCurrentStateMap.get(resourceName);

  if (currentState == null) {
    LOG.info("No currentState found for {} on {}, set relay message {} to be expired.",
        resourceName, relayInstance, relayMessage.getId());
    setMessageRelayTime(relayMessage, currentTime);
    return;
  }

  String partitionState = currentState.getState(partitionName);
  String targetState = hostedMessage.getToState();
  String fromState = hostedMessage.getFromState();

  // The relay host partition state has been changed after relay message was created.
  if (!fromState.equals(partitionState)) {
    // If the partition on the relay host turned to ERROR while transited from top state,
    // we can remove the cached relay message right away since participant won't forward the relay message anyway.
    if (HelixDefinedState.ERROR.name().equals(partitionState) && fromState
        .equals(currentState.getPreviousState(partitionName))) {
      LOG.info("Partition {} got to ERROR from the top state, "
              + "expiring relay message {} immediately. Hosted message {}.", partitionName,
          relayMessage.getId(), hostedMessage.getId());
      relayMessage.setExpired(true);
      return;
    }

    // If the partition completed the transition, set the relay time to be the actual time when state transition completed.
    if (targetState.equals(partitionState) && fromState
        .equals(currentState.getPreviousState(partitionName))) {
      // The relay host already completed the state transition.
      long completeTime = currentState.getEndTime(partitionName);
      if (completeTime > relayMessage.getCreateTimeStamp()) {
        setMessageRelayTime(relayMessage, completeTime);
        LOG.error("Target state for partition {} matches the hosted message's target state, "
            + "set relay message {} to be expired.", partitionName, relayMessage.getId());
        return;
      }
    }

    // For all other situations, set relay time to be current time.
    setMessageRelayTime(relayMessage, currentTime);
    // the state has been changed after it completed the required state transition (maybe another state-transition happened).
    LOG.info("Current state {} for partition {} does not match hosted message's from state, "
            + "set relay message {} to be expired.", partitionState, partitionName,
        relayMessage.getId());
  }
}
 

/**
 * For a partition, given its preferenceList, bestPossibleState, and currentState, determine which
 * type of rebalance is needed to model IdealState's states defined by the state model definition.
 * @return RebalanceType needed to bring the replicas to idea states
 *         RECOVERY_BALANCE - not all required states (replicas) are available through all
 *         replicas, or the partition is disabled
 *         NONE - current state matches the ideal state
 *         LOAD_BALANCE - although all replicas required exist, Helix needs to optimize the
 *         allocation
 */
private RebalanceType getRebalanceType(ResourceControllerDataProvider cache,
    Map<String, String> bestPossibleMap, List<String> preferenceList,
    StateModelDefinition stateModelDef, Map<String, String> currentStateMap,
    IdealState idealState, String partitionName) {
  if (preferenceList == null) {
    preferenceList = Collections.emptyList();
  }

  // If there is a minimum active replica number specified in IS, we should respect it.
  // TODO: We should implement the per replica level throttling with generated message
  // Issue: https://github.com/apache/helix/issues/343
  int replica = idealState.getMinActiveReplicas() == -1
      ? idealState.getReplicaCount(preferenceList.size())
      : idealState.getMinActiveReplicas();
  Set<String> activeList = new HashSet<>(preferenceList);
  activeList.retainAll(cache.getEnabledLiveInstances());

  // For each state, check that this partition currently has the required number of that state as
  // required by StateModelDefinition.
  LinkedHashMap<String, Integer> expectedStateCountMap =
      stateModelDef.getStateCountMap(activeList.size(), replica); // StateModelDefinition's counts
  // Current counts without disabled partitions or disabled instances
  Map<String, String> currentStateMapWithoutDisabled = new HashMap<>(currentStateMap);
  currentStateMapWithoutDisabled.keySet().removeAll(
      cache.getDisabledInstancesForPartition(idealState.getResourceName(), partitionName));
  Map<String, Integer> currentStateCounts =
      StateModelDefinition.getStateCounts(currentStateMapWithoutDisabled);

  // Go through each state and compare counts
  for (String state : expectedStateCountMap.keySet()) {
    Integer expectedCount = expectedStateCountMap.get(state);
    Integer currentCount = currentStateCounts.get(state);
    expectedCount = expectedCount == null ? 0 : expectedCount;
    currentCount = currentCount == null ? 0 : currentCount;

    // If counts do not match up, this partition requires recovery
    if (currentCount < expectedCount) {
      // Recovery is not needed in cases where this partition just started, was dropped, or is in
      // error
      if (!state.equals(HelixDefinedState.DROPPED.name())
          && !state.equals(HelixDefinedState.ERROR.name())
          && !state.equals(stateModelDef.getInitialState())) {
        return RebalanceType.RECOVERY_BALANCE;
      }
    }
  }
  // No recovery needed, all expected replicas exist
  // Check if this partition is actually in the BestPossibleState
  if (currentStateMap.equals(bestPossibleMap)) {
    return RebalanceType.NONE; // No further action required
  } else {
    return RebalanceType.LOAD_BALANCE; // Required state counts are satisfied, but in order to
    // achieve BestPossibleState, load balance may be required
    // to shift replicas around
  }
}
 

/**
 * Go through every instance in the assignment and check that each instance does NOT have more
 * replicas for partitions assigned to it than maxPartitionsPerInstance. If the assignment
 * violates this, put the cluster on maintenance mode.
 * This logic could be integrated with compute() for IntermediateState calculation but appended
 * separately for visibility and testing. Additionally, performing validation after compute()
 * ensures that we have a full intermediate state mapping complete prior to validation.
 * @param event
 * @param cache
 * @param intermediateStateOutput
 * @param maxPartitionPerInstance
 */
private void validateMaxPartitionsPerInstance(ClusterEvent event,
    ResourceControllerDataProvider cache, IntermediateStateOutput intermediateStateOutput,
    int maxPartitionPerInstance) {
  Map<String, PartitionStateMap> resourceStatesMap =
      intermediateStateOutput.getResourceStatesMap();
  Map<String, Integer> instancePartitionCounts = new HashMap<>();

  for (String resource : resourceStatesMap.keySet()) {
    IdealState idealState = cache.getIdealState(resource);
    if (idealState != null
        && idealState.getStateModelDefRef().equals(BuiltInStateModelDefinitions.Task.name())) {
      // Ignore task here. Task has its own throttling logic
      continue;
    }

    PartitionStateMap partitionStateMap = resourceStatesMap.get(resource);
    Map<Partition, Map<String, String>> stateMaps = partitionStateMap.getStateMap();
    for (Partition p : stateMaps.keySet()) {
      Map<String, String> stateMap = stateMaps.get(p);
      for (String instance : stateMap.keySet()) {
        // If this replica is in DROPPED state, do not count it in the partition count since it is
        // to be dropped
        String state = stateMap.get(instance);
        if (state.equals(HelixDefinedState.DROPPED.name())) {
          continue;
        }

        if (!instancePartitionCounts.containsKey(instance)) {
          instancePartitionCounts.put(instance, 0);
        }
        int partitionCount = instancePartitionCounts.get(instance); // Number of replicas (from
        // different partitions) held
        // in this instance
        partitionCount++;
        if (partitionCount > maxPartitionPerInstance) {
          HelixManager manager = event.getAttribute(AttributeName.helixmanager.name());
          String errMsg = String.format(
              "Problem: according to this assignment, instance %s contains more "
                  + "replicas/partitions than the maximum number allowed (%d). Pipeline will "
                  + "stop the rebalance and put the cluster %s into maintenance mode",
              instance, maxPartitionPerInstance, cache.getClusterName());
          if (manager != null) {
            if (manager.getHelixDataAccessor()
                .getProperty(manager.getHelixDataAccessor().keyBuilder().maintenance()) == null) {
              manager.getClusterManagmentTool().autoEnableMaintenanceMode(
                  manager.getClusterName(), true, errMsg,
                  MaintenanceSignal.AutoTriggerReason.MAX_PARTITION_PER_INSTANCE_EXCEEDED);
            }
            LogUtil.logWarn(logger, _eventId, errMsg);
          } else {
            LogUtil.logError(logger, _eventId,
                "HelixManager is not set/null! Failed to pause this cluster/enable maintenance"
                    + " mode due to an instance being assigned more replicas/partitions than "
                    + "the limit.");
          }

          ClusterStatusMonitor clusterStatusMonitor =
              event.getAttribute(AttributeName.clusterStatusMonitor.name());
          if (clusterStatusMonitor != null) {
            clusterStatusMonitor.setResourceRebalanceStates(Collections.singletonList(resource),
                ResourceMonitor.RebalanceStatus.INTERMEDIATE_STATE_CAL_FAILED);
          }
          // Throw an exception here so that messages won't be sent out based on this mapping
          throw new HelixException(errMsg);
        }
        instancePartitionCounts.put(instance, partitionCount);
      }
    }
  }
}
 

@Test
public void TestOFFLINEToDROPPED() throws Exception {

  ClusterEvent event = new ClusterEvent(TEST_CLUSTER, ClusterEventType.Unknown);


  // Set current state to event
  CurrentStateOutput currentStateOutput = mock(CurrentStateOutput.class);
  Partition partition = mock(Partition.class);
  when(partition.getPartitionName()).thenReturn(TEST_PARTITION);
  when(currentStateOutput.getCurrentState(TEST_RESOURCE, partition, TEST_INSTANCE)).thenReturn(null);
  Message message = mock(Message.class);
  when(message.getFromState()).thenReturn("OFFLINE");
  when(message.getToState()).thenReturn("SLAVE");
  when(currentStateOutput.getPendingMessage(TEST_RESOURCE, partition, TEST_INSTANCE)).thenReturn(message);
  event.addAttribute(AttributeName.CURRENT_STATE.name(), currentStateOutput);

  // Set helix manager to event
  event.addAttribute(AttributeName.helixmanager.name(), mock(HelixManager.class));

  // Set controller data provider to event
  BaseControllerDataProvider cache = mock(BaseControllerDataProvider.class);
  StateModelDefinition stateModelDefinition = new StateModelDefinition(StateModelConfigGenerator.generateConfigForMasterSlave());
  when(cache.getStateModelDef(TaskConstants.STATE_MODEL_NAME)).thenReturn(stateModelDefinition);
  Map<String, LiveInstance> liveInstances= mock(Map.class);
  LiveInstance mockLiveInstance = mock(LiveInstance.class);
  when(mockLiveInstance.getInstanceName()).thenReturn(TEST_INSTANCE);
  when(mockLiveInstance.getEphemeralOwner()).thenReturn("TEST");
  when(liveInstances.values()).thenReturn(Arrays.asList(mockLiveInstance));
  when(cache.getLiveInstances()).thenReturn(liveInstances);
  ClusterConfig clusterConfig = mock(ClusterConfig.class);
  when(cache.getClusterConfig()).thenReturn(clusterConfig);
  when(clusterConfig.isStateTransitionCancelEnabled()).thenReturn(true);
  event.addAttribute(AttributeName.ControllerDataProvider.name(), cache);


  // Set resources to rebalance to event
  Map<String, Resource> resourceMap = new HashMap<>();
  Resource resource = mock(Resource.class);
  when(resource.getResourceName()).thenReturn(TEST_RESOURCE);
  List<Partition> partitions = Arrays.asList(partition);
  when(resource.getPartitions()).thenReturn(partitions);
  when(resource.getStateModelDefRef()).thenReturn(TaskConstants.STATE_MODEL_NAME);
  resourceMap.put(TEST_RESOURCE, resource);
  event.addAttribute(AttributeName.RESOURCES_TO_REBALANCE.name(), resourceMap);

  // set up resource state map
  ResourcesStateMap resourcesStateMap = new ResourcesStateMap();
  PartitionStateMap partitionStateMap = new PartitionStateMap(TEST_RESOURCE);
  Map<Partition, Map<String, String>> stateMap = partitionStateMap.getStateMap();
  Map<String, String> instanceStateMap = new HashMap<>();
  instanceStateMap.put(TEST_INSTANCE, HelixDefinedState.DROPPED.name());
  stateMap.put(partition, instanceStateMap);
  resourcesStateMap.setState(TEST_RESOURCE, partition, instanceStateMap);

  processEvent(event, resourcesStateMap);
  MessageOutput output = event.getAttribute(AttributeName.MESSAGES_ALL.name());
  Assert.assertEquals(output.getMessages(TEST_RESOURCE, partition).size(), 1);
}
 

static StateModelDefinition getDefinition() {

    StateModelDefinition.Builder builder = new StateModelDefinition.Builder(AMBRY_LEADER_STANDBY_MODEL);
    // Init state
    builder.initialState(ReplicaState.OFFLINE.name());

    /*
     * States and their priority which are managed by Helix controller.
     *
     * The implication is that transition to a state which has a higher number is
     * considered a "downward" transition by Helix. Downward transitions are not
     * throttled.
     */
    builder.addState(ReplicaState.LEADER.name(), 0);
    builder.addState(ReplicaState.STANDBY.name(), 1);
    builder.addState(ReplicaState.BOOTSTRAP.name(), 2);
    builder.addState(ReplicaState.INACTIVE.name(), 2);
    builder.addState(ReplicaState.OFFLINE.name(), 3);
    // HelixDefinedState adds two additional states: ERROR and DROPPED. The priority is Integer.MAX_VALUE
    for (HelixDefinedState state : HelixDefinedState.values()) {
      builder.addState(state.name());
    }

    // Add valid transitions between the states.
    builder.addTransition(ReplicaState.OFFLINE.name(), ReplicaState.BOOTSTRAP.name(), 3);
    builder.addTransition(ReplicaState.BOOTSTRAP.name(), ReplicaState.STANDBY.name(), 2);
    builder.addTransition(ReplicaState.STANDBY.name(), ReplicaState.LEADER.name(), 1);
    builder.addTransition(ReplicaState.LEADER.name(), ReplicaState.STANDBY.name(), 0);
    builder.addTransition(ReplicaState.STANDBY.name(), ReplicaState.INACTIVE.name(), 2);
    builder.addTransition(ReplicaState.INACTIVE.name(), ReplicaState.OFFLINE.name(), 3);
    builder.addTransition(ReplicaState.OFFLINE.name(), HelixDefinedState.DROPPED.name());

    // States constraints
    /*
     * Static constraint: number of leader replica for certain partition shouldn't exceed 1 at any time.
     */
    builder.upperBound(ReplicaState.LEADER.name(), 1);
    /*
     * Dynamic constraint: R means it should be derived based on the replication factor for the cluster
     * this allows a different replication factor for each resource without having to define a new state model.
     */
    builder.dynamicUpperBound(ReplicaState.STANDBY.name(), UPPER_BOUND_REPLICATION_FACTOR);

    return builder.build();
  }