NUM_REPLICAS, // replicas
"OnlineOffline", RebalanceMode.FULL_AUTO, // use FULL_AUTO mode to test node tagging
true); // do rebalance
// Tag the resource
final HelixAdmin helixAdmin = new ZKHelixAdmin(_zkclient);
IdealState idealState = helixAdmin.getResourceIdealState(clusterName, RESOURCE_NAME);
idealState.setInstanceGroupTag(TAG);
helixAdmin.setResourceIdealState(clusterName, RESOURCE_NAME, idealState);
// Get a data accessor
final HelixDataAccessor accessor = new ZKHelixDataAccessor(clusterName, _baseAccessor);
final PropertyKey.Builder keyBuilder = accessor.keyBuilder();
// Tag the participants
for (int i = 0; i < NUM_PARTICIPANTS; i++) {
final String instanceName = "localhost_" + (12918 + i);
helixAdmin.addInstanceTag(clusterName, instanceName, TAG);
}
// Start controller
MockController controller =
new MockController(_zkaddr, clusterName, "controller");
controller.syncStart();
// Start participants
MockParticipant[] participants = new MockParticipant[NUM_PARTICIPANTS];
for (int i = 0; i < NUM_PARTICIPANTS; i++) {
final String instanceName = "localhost_" + (12918 + i);
participants[i] = new MockParticipant(_zkaddr, clusterName, instanceName);
participants[i].syncStart();
}
// Verify that there are NUM_PARTITIONS partitions in the external view, each having
// NUM_REPLICAS replicas, where all assigned replicas are to tagged nodes, and they are all
// ONLINE.
Verifier v = new Verifier() {
@Override
public boolean verify() throws Exception {
ExternalView externalView =
HelixTestUtil.pollForProperty(ExternalView.class, accessor,
keyBuilder.externalView(RESOURCE_NAME), true);
if (externalView == null) {
return false;
}
Set<String> taggedInstances =
Sets.newHashSet(helixAdmin.getInstancesInClusterWithTag(clusterName, TAG));
Set<String> partitionSet = externalView.getPartitionSet();
if (partitionSet.size() != NUM_PARTITIONS) {
return false;
}
for (String partitionName : partitionSet) {
Map<String, String> stateMap = externalView.getStateMap(partitionName);
if (stateMap.size() != NUM_REPLICAS) {
return false;
}
for (String participantName : stateMap.keySet()) {
if (!taggedInstances.contains(participantName)) {
return false;
}
String state = stateMap.get(participantName);
if (!state.equalsIgnoreCase("ONLINE")) {
return false;
}
}
}
return true;
}
};
// Run the verifier for both nodes tagged
boolean initialResult = TestHelper.verify(v, 10 * 1000);
Assert.assertTrue(initialResult);
// Untag a node
helixAdmin.removeInstanceTag(clusterName, "localhost_12918", TAG);
// Verify again
boolean finalResult = TestHelper.verify(v, 10 * 1000);
Assert.assertTrue(finalResult);