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

@Override
public boolean matches(Object o) {
  if (o == null || !(o instanceof IdealState)) {
    return false;
  }
  IdealState is = (IdealState)o;
  if (is.getRebalanceMode() != IdealState.RebalanceMode.CUSTOMIZED ||
      !is.getReplicas().equals("3") ||
      is.getNumPartitions() != this.numPartitions ||
      !is.getStateModelDefRef().equals("OnlineOffline")) {
    return false;
  }
  for (Map.Entry<Integer, List<String>> entry : partitionHostMap.entrySet()) {
    Map<String, String> stateMap = is.getInstanceStateMap(
            this.resource + "$" + entry.getKey());
    if (stateMap.size() != entry.getValue().size()) {
      return false;
    }
    for (String host : entry.getValue()) {
      if (!(stateMap.containsKey(host) && stateMap.get(host).equals("ONLINE"))) {
        return false;
      }
    }
  }
  return true;
}
 

private boolean checkBestPossibleStateCalculation(IdealState idealState) {
  // If replicas is 0, indicate the resource is not fully initialized or ready to be rebalanced
  if (idealState.getRebalanceMode() == IdealState.RebalanceMode.FULL_AUTO && !idealState
      .getReplicas().equals("0")) {
    Map<String, List<String>> preferenceLists = idealState.getPreferenceLists();
    if (preferenceLists == null || preferenceLists.isEmpty()) {
      return false;
    }
    int emptyListCount = 0;
    for (List<String> preferenceList : preferenceLists.values()) {
      if (preferenceList.isEmpty()) {
        emptyListCount++;
      }
    }
    // If all lists are empty, rebalance fails completely
    return emptyListCount != preferenceLists.values().size();
  } else {
    // For non FULL_AUTO RebalanceMode, rebalancing is not controlled by Helix
    return true;
  }
}
 

@Override
protected Map<FieldType, Set<String>> getNonTrimmableFields(IdealState idealState) {
  Map<FieldType, Set<String>> nonTrimmableFields =
      new HashMap<>(STATIC_TOPOLOGY_RELATED_FIELD_MAP);
  // Also don't trim the application defined assignment in the list and map fields.
  // They are fixed and considered as part of the cluster topology.
  switch (idealState.getRebalanceMode()) {
  case CUSTOMIZED:
    // For CUSTOMZIED resources, both list and map fields are user configured partition state
    // assignment. So they are not trimmable.
    nonTrimmableFields.put(FieldType.MAP_FIELD, idealState.getRecord().getMapFields().keySet());
    nonTrimmableFields.put(FieldType.LIST_FIELD, idealState.getRecord().getListFields().keySet());
    break;
  case SEMI_AUTO:
    // For SEMI_AUTO resources, list fields are user configured partition placement. So it is not
    // trimmable.
    nonTrimmableFields.put(FieldType.LIST_FIELD, idealState.getRecord().getListFields().keySet());
    break;
  case FULL_AUTO:
    // For FULL_AUTO resources, both map fields and list fields are trimmable since they are
    // re-computed and updated in every controller rebalance pipelines.
  default:
    break;
  }
  return nonTrimmableFields;
}
 

private Rebalancer<ResourceControllerDataProvider> getRebalancer(IdealState idealState,
    String resourceName, boolean isMaintenanceModeEnabled) {
  Rebalancer<ResourceControllerDataProvider> rebalancer = null;
  switch (idealState.getRebalanceMode()) {
  case FULL_AUTO:
    if (isMaintenanceModeEnabled) {
      rebalancer = new MaintenanceRebalancer();
    } else {
      Rebalancer<ResourceControllerDataProvider> customizedRebalancer =
          getCustomizedRebalancer(idealState.getRebalancerClassName(), resourceName);
      if (customizedRebalancer != null) {
        rebalancer = customizedRebalancer;
      } else {
        rebalancer = new AutoRebalancer();
      }
    }
    break;
  case SEMI_AUTO:
    rebalancer = new SemiAutoRebalancer<>();
    break;
  case CUSTOMIZED:
    rebalancer = new CustomRebalancer();
    break;
  case USER_DEFINED:
  case TASK:
    rebalancer = getCustomizedRebalancer(idealState.getRebalancerClassName(), resourceName);
    break;
  default:
    LogUtil.logError(logger, _eventId,
        "Fail to find the rebalancer, invalid rebalance mode " + idealState.getRebalanceMode());
    break;
  }
  return rebalancer;
}
 

public void expandResource(String clusterName, String resourceName) {
  IdealState idealState = _admin.getResourceIdealState(clusterName, resourceName);
  if (idealState.getRebalanceMode() == RebalanceMode.FULL_AUTO
      || idealState.getRebalanceMode() == RebalanceMode.CUSTOMIZED) {
    _logger.info("Skipping idealState " + idealState.getResourceName() + " "
        + idealState.getRebalanceMode());
    return;
  }
  boolean anyLiveInstance = false;
  for (List<String> list : idealState.getRecord().getListFields().values()) {
    if (list.contains(IdealState.IdealStateConstants.ANY_LIVEINSTANCE.toString())) {
      _logger.info("Skipping idealState " + idealState.getResourceName()
          + " with ANY_LIVEINSTANCE");
      anyLiveInstance = true;
      continue;
    }
  }
  if (anyLiveInstance) {
    return;
  }
  try {
    int replica = Integer.parseInt(idealState.getReplicas());
  } catch (Exception e) {
    _logger.error("", e);
    return;
  }
  if (idealState.getRecord().getListFields().size() == 0) {
    _logger.warn("Resource " + resourceName + " not balanced, skip");
    return;
  }
  balanceIdealState(clusterName, idealState);
}
 

private boolean verifyExternalView(ResourceControllerDataProvider dataCache, ExternalView externalView,
    IdealState idealState) {
  Map<String, Map<String, String>> mappingInExtview = externalView.getRecord().getMapFields();
  Map<String, Map<String, String>> idealPartitionState;

  switch (idealState.getRebalanceMode()) {
  case FULL_AUTO:
    ClusterConfig clusterConfig = new ConfigAccessor(_zkClient).getClusterConfig(dataCache.getClusterName());
    if (!clusterConfig.isPersistBestPossibleAssignment() && !clusterConfig.isPersistIntermediateAssignment()) {
      throw new HelixException(String.format("Full-Auto IdealState verifier requires "
          + "ClusterConfig.PERSIST_BEST_POSSIBLE_ASSIGNMENT or ClusterConfig.PERSIST_INTERMEDIATE_ASSIGNMENT "
          + "is enabled."));
    }
    for (String partition : idealState.getPartitionSet()) {
      if (idealState.getPreferenceList(partition) == null || idealState.getPreferenceList(partition).isEmpty()) {
        return false;
      }
    }
    idealPartitionState = computeIdealPartitionState(dataCache, idealState);
    break;
  case SEMI_AUTO:
  case USER_DEFINED:
    idealPartitionState = computeIdealPartitionState(dataCache, idealState);
    break;
  case CUSTOMIZED:
    idealPartitionState = idealState.getRecord().getMapFields();
    break;
  case TASK:
    // ignore jobs
  default:
    return true;
  }

  return mappingInExtview.equals(idealPartitionState);
}
 

private void verifyAssignmentInIdealStateWithPersistDisabled(IdealState idealState,
    Set<String> excludedInstances) {
  boolean mapFieldEmpty = true;
  boolean assignmentNotChanged = false;
  for (String partition : idealState.getPartitionSet()) {
    Map<String, String> instanceStateMap = idealState.getInstanceStateMap(partition);
    if (instanceStateMap == null || instanceStateMap.isEmpty()) {
      continue;
    }
    mapFieldEmpty = false;
    Set<String> instancesInMap = instanceStateMap.keySet();
    for (String ins : excludedInstances) {
      if (instancesInMap.contains(ins)) {
        // if at least one excluded instance is included, it means assignment was not updated.
        assignmentNotChanged = true;
      }
      if (idealState.getRebalanceMode() == RebalanceMode.FULL_AUTO) {
        List<String> instanceList = idealState.getPreferenceList(partition);
        if (instanceList.contains(ins)) {
          assignmentNotChanged = true;
        }
      }
    }
  }

  Assert.assertTrue((mapFieldEmpty || assignmentNotChanged),
      "BestPossible assignment was updated.");
}
 

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

void rebalance(String clusterName, String resourceName, int replica, String keyPrefix,
    List<String> instanceNames, String groupId) {
  logger.info("Rebalance resource {} with replica {} in cluster {}.", resourceName, replica,
      clusterName);
  // ensure we get the same idealState with the same set of instances
  Collections.sort(instanceNames);

  IdealState idealState = getResourceIdealState(clusterName, resourceName);
  if (idealState == null) {
    throw new HelixException("Resource: " + resourceName + " has NOT been added yet");
  }

  if (groupId != null && groupId.length() > 0) {
    idealState.setInstanceGroupTag(groupId);
  }
  idealState.setReplicas(Integer.toString(replica));
  int partitions = idealState.getNumPartitions();
  String stateModelName = idealState.getStateModelDefRef();
  StateModelDefinition stateModDef = getStateModelDef(clusterName, stateModelName);

  if (stateModDef == null) {
    throw new HelixException("cannot find state model: " + stateModelName);
  }
  // StateModelDefinition def = new StateModelDefinition(stateModDef);

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

  String masterStateValue = null;
  String slaveStateValue = null;
  replica--;

  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");
      }
      replica = instanceNames.size() - 1;
      masterStateValue = slaveStateValue = state;
    }
  }
  if (masterStateValue == null && slaveStateValue == null) {
    throw new HelixException("Invalid or unsupported state model definition");
  }

  if (masterStateValue == null) {
    masterStateValue = slaveStateValue;
  }
  if (idealState.getRebalanceMode() != RebalanceMode.FULL_AUTO
      && idealState.getRebalanceMode() != RebalanceMode.USER_DEFINED) {
    ZNRecord newIdealState = DefaultIdealStateCalculator
        .calculateIdealState(instanceNames, partitions, replica, keyPrefix, masterStateValue,
            slaveStateValue);

    // for now keep mapField in SEMI_AUTO mode and remove listField in CUSTOMIZED mode
    if (idealState.getRebalanceMode() == RebalanceMode.SEMI_AUTO) {
      idealState.getRecord().setListFields(newIdealState.getListFields());
      // TODO: need consider to remove this.
      idealState.getRecord().setMapFields(newIdealState.getMapFields());
    }
    if (idealState.getRebalanceMode() == RebalanceMode.CUSTOMIZED) {
      idealState.getRecord().setMapFields(newIdealState.getMapFields());
    }
  } else {
    for (int i = 0; i < partitions; i++) {
      String partitionName = keyPrefix + "_" + i;
      idealState.getRecord().setMapField(partitionName, new HashMap<String, String>());
      idealState.getRecord().setListField(partitionName, new ArrayList<String>());
    }
  }
  setResourceIdealState(clusterName, resourceName, idealState);
}
 

private void persistAssignment(final Resource resource, final ResourceControllerDataProvider cache,
    final ClusterEvent event, final BestPossibleStateOutput bestPossibleAssignment,
    final ClusterConfig clusterConfig, final HelixDataAccessor accessor,
    final PropertyKey.Builder keyBuilder) {
  String resourceId = resource.getResourceName();
  if (resource != null) {
    final IdealState idealState = cache.getIdealState(resourceId);
    if (idealState == null) {
      LogUtil.logWarn(LOG, event.getEventId(), "IdealState not found for resource " + resourceId);
      return;
    }
    IdealState.RebalanceMode mode = idealState.getRebalanceMode();
    if (!mode.equals(IdealState.RebalanceMode.SEMI_AUTO) && !mode
        .equals(IdealState.RebalanceMode.FULL_AUTO)) {
      // do not persist assignment for resource in neither semi or full auto.
      return;
    }

    boolean needPersist = false;
    if (mode.equals(IdealState.RebalanceMode.FULL_AUTO)) {
      // persist preference list in ful-auto mode.
      Map<String, List<String>> newLists = bestPossibleAssignment.getPreferenceLists(resourceId);
      if (newLists != null && hasPreferenceListChanged(newLists, idealState)) {
        idealState.setPreferenceLists(newLists);
        needPersist = true;
      }
    }

    PartitionStateMap partitionStateMap = bestPossibleAssignment.getPartitionStateMap(resourceId);
    if (clusterConfig.isPersistIntermediateAssignment()) {
      IntermediateStateOutput intermediateAssignment =
          event.getAttribute(AttributeName.INTERMEDIATE_STATE.name());
      partitionStateMap = intermediateAssignment.getPartitionStateMap(resourceId);
    }

    //TODO: temporary solution for Espresso/Dbus backcompatible, should remove this.
    Map<Partition, Map<String, String>> assignmentToPersist =
        convertAssignmentPersisted(resource, idealState, partitionStateMap.getStateMap());

    if (assignmentToPersist != null && hasInstanceMapChanged(assignmentToPersist, idealState)) {
      for (Partition partition : assignmentToPersist.keySet()) {
        Map<String, String> instanceMap = assignmentToPersist.get(partition);
        idealState.setInstanceStateMap(partition.getPartitionName(), instanceMap);
      }
      needPersist = true;
    }

    if (needPersist) {
      // Update instead of set to ensure any intermediate changes that the controller does not update are kept.
      accessor.updateProperty(keyBuilder.idealStates(resourceId), new DataUpdater<ZNRecord>() {
        @Override
        public ZNRecord update(ZNRecord current) {
          if (current != null) {
            // Overwrite MapFields and ListFields items with the same key.
            // Note that default merge will keep old values in the maps or lists unchanged, which is not desired.
            current.getMapFields().clear();
            current.getMapFields().putAll(idealState.getRecord().getMapFields());
            current.getListFields().putAll(idealState.getRecord().getListFields());
          }
          return current;
        }
      }, idealState);
    }
  }
}