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

/**
 * has the preference list changed from the one persisted in current IdealState
 */
private boolean hasPreferenceListChanged(Map<String, List<String>> newLists,
    IdealState idealState) {
  Map<String, List<String>> existLists = idealState.getPreferenceLists();

  Set<String> partitions = new HashSet<String>(newLists.keySet());
  partitions.addAll(existLists.keySet());

  for (String partition : partitions) {
    List<String> assignedInstances = newLists.get(partition);
    List<String> existingInstances = existLists.get(partition);
    if (assignedInstances == null && existingInstances == null) {
      continue;
    }
    if (assignedInstances == null || existingInstances == null || !assignedInstances
        .equals(existingInstances)) {
      return true;
    }
  }

  return false;
}
 

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

@Test (dependsOnMethods = {"testDisableDelayRebalanceInCluster"})
public void testDisableDelayRebalanceInInstance() throws Exception {
  setDelayTimeInCluster(_gZkClient, CLUSTER_NAME, 1000000);
  Map<String, ExternalView> externalViewsBefore = createTestDBs(-1);
  validateDelayedMovements(externalViewsBefore);

  String disabledInstanceName = _participants.get(0).getInstanceName();
  enableDelayRebalanceInInstance(_gZkClient, CLUSTER_NAME, disabledInstanceName, false);
  Thread.sleep(DEFAULT_REBALANCE_PROCESSING_WAIT_TIME);
  Assert.assertTrue(_clusterVerifier.verifyByPolling());
  for (String db : _testDBs) {
    IdealState is = _gSetupTool.getClusterManagementTool().getResourceIdealState(CLUSTER_NAME, db);
    Map<String, List<String>> preferenceLists = is.getPreferenceLists();
    for (List<String> instances : preferenceLists.values()) {
      Assert.assertFalse(instances.contains(disabledInstanceName));
    }
  }
  enableDelayRebalanceInInstance(_gZkClient, CLUSTER_NAME, disabledInstanceName, true);
}
 

@Test(dataProvider = "stateModels")
public void testUserDefinedPreferenceListsInFullAuto(String stateModel, boolean delayEnabled,
    String rebalanceStrateyName) throws Exception {
  String dbName = "Test-DB-" + stateModel;
  createResource(dbName, stateModel, _PARTITIONS, _replica, delayEnabled,
      rebalanceStrateyName);
  try {
    IdealState idealState =
        _gSetupTool.getClusterManagementTool().getResourceIdealState(CLUSTER_NAME, dbName);
    Map<String, List<String>> userDefinedPreferenceLists = idealState.getPreferenceLists();
    List<String> userDefinedPartitions = new ArrayList<>();
    for (String partition : userDefinedPreferenceLists.keySet()) {
      List<String> preferenceList = new ArrayList<>();
      for (int k = _replica; k >= 0; k--) {
        String instance = _participants.get(k).getInstanceName();
        preferenceList.add(instance);
      }
      userDefinedPreferenceLists.put(partition, preferenceList);
      userDefinedPartitions.add(partition);
    }

    ResourceConfig resourceConfig =
        new ResourceConfig.Builder(dbName).setPreferenceLists(userDefinedPreferenceLists).build();
    _configAccessor.setResourceConfig(CLUSTER_NAME, dbName, resourceConfig);

    // TODO remove this sleep after fix https://github.com/apache/helix/issues/526
    Thread.sleep(500);

    Assert.assertTrue(_clusterVerifier.verify(3000));
    verifyUserDefinedPreferenceLists(dbName, userDefinedPreferenceLists,
        userDefinedPartitions);

    while (userDefinedPartitions.size() > 0) {
      IdealState originIS =
          _gSetupTool.getClusterManagementTool().getResourceIdealState(CLUSTER_NAME, dbName);
      Set<String> nonUserDefinedPartitions = new HashSet<>(originIS.getPartitionSet());
      nonUserDefinedPartitions.removeAll(userDefinedPartitions);

      removePartitionFromUserDefinedList(dbName, userDefinedPartitions);
      // TODO: Remove wait once we enable the BestPossibleExternalViewVerifier for the WAGED rebalancer.
      Thread.sleep(1000);
      Assert.assertTrue(_clusterVerifier.verify(3000));
      verifyUserDefinedPreferenceLists(dbName, userDefinedPreferenceLists,
          userDefinedPartitions);
      verifyNonUserDefinedAssignment(dbName, originIS, nonUserDefinedPartitions);
    }
  } finally {
    _gSetupTool.getClusterManagementTool().dropResource(CLUSTER_NAME, dbName);
    _clusterVerifier.verify(5000);
  }
}
 

@Test(dataProvider = "stateModels")
public void testMigrateToFullAutoWhileExpandCluster(
    String stateModel, boolean delayEnabled) throws Exception {
  String db = "Test-DB-" + stateModel;
  if (delayEnabled) {
    createResourceWithDelayedRebalance(CLUSTER_NAME, db, stateModel, _PARTITIONS, _replica,
        _replica - 1, 200000, CrushRebalanceStrategy.class.getName());
  } else {
    createResourceWithDelayedRebalance(CLUSTER_NAME, db, stateModel, _PARTITIONS, _replica,
        _replica, 0, CrushRebalanceStrategy.class.getName());
  }
  IdealState idealState =
      _gSetupTool.getClusterManagementTool().getResourceIdealState(CLUSTER_NAME, db);
  Map<String, List<String>> userDefinedPreferenceLists = idealState.getPreferenceLists();
  List<String> userDefinedPartitions = new ArrayList<>();
  for (String partition : userDefinedPreferenceLists.keySet()) {
    List<String> preferenceList = new ArrayList<>();
    for (int k = _replica; k > 0; k--) {
      String instance = _participants.get(k).getInstanceName();
      preferenceList.add(instance);
    }
    userDefinedPreferenceLists.put(partition, preferenceList);
    userDefinedPartitions.add(partition);
  }

  ResourceConfig resourceConfig =
      new ResourceConfig.Builder(db).setPreferenceLists(userDefinedPreferenceLists).build();
  _configAccessor.setResourceConfig(CLUSTER_NAME, db, resourceConfig);

  // add new instance to the cluster
  int numNodes = _participants.size();
  for (int i = numNodes; i < numNodes + NUM_NODE; i++) {
    String storageNodeName = PARTICIPANT_PREFIX + "_" + (START_PORT + i);
    MockParticipantManager participant = createAndStartParticipant(storageNodeName);
    _participants.add(participant);
    Thread.sleep(50);
  }

  Assert.assertTrue(_clusterVerifier.verifyByPolling());

  _migrationVerifier =
      new MigrationStateVerifier(Collections.singletonMap(db, idealState), _manager);

  _migrationVerifier.reset();
  _migrationVerifier.start();

  while (userDefinedPartitions.size() > 0) {
    removePartitionFromUserDefinedList(db, userDefinedPartitions);
    Thread.sleep(50);
  }


  Assert.assertTrue(_clusterVerifier.verifyByPolling());
  Assert.assertFalse(_migrationVerifier.hasLessReplica());
  Assert.assertFalse(_migrationVerifier.hasMoreReplica());

  _migrationVerifier.stop();
}