@Override
public IdealState computeNewIdealState(
String resourceName,
IdealState currentIdealState,
CurrentStateOutput currentStateOutput,
ClusterDataCache clusterData) {
testRebalancerInvoked = true;
for (String partition : currentIdealState.getPartitionSet()) {
String instance = currentIdealState.getPreferenceList(partition).get(0);
currentIdealState.getPreferenceList(partition).clear();
currentIdealState.getPreferenceList(partition).add(instance);
currentIdealState.getInstanceStateMap(partition).clear();
currentIdealState.getInstanceStateMap(partition).put(instance, "MASTER");
}
currentIdealState.setReplicas("1");
return currentIdealState;
}
@Test
public void testCustomizedIdealStateRebalancer() throws InterruptedException {
_setupTool.addResourceToCluster(CLUSTER_NAME, db2, 60, "MasterSlave");
_setupTool.addResourceProperty(
CLUSTER_NAME,
db2,
IdealStateProperty.REBALANCER_CLASS_NAME.toString(),
TestCustomizedIdealStateRebalancer.TestRebalancer.class.getName());
_setupTool.addResourceProperty(
CLUSTER_NAME,
db2,
IdealStateProperty.REBALANCE_MODE.toString(),
RebalanceMode.USER_DEFINED.toString());
_setupTool.rebalanceStorageCluster(CLUSTER_NAME, db2, 3);
boolean result =
ClusterStateVerifier.verifyByZkCallback(
new ExternalViewBalancedVerifier(_gZkClient, CLUSTER_NAME, db2));
Assert.assertTrue(result);
Thread.sleep(1000);
HelixDataAccessor accessor =
new ZKHelixDataAccessor(CLUSTER_NAME, new ZkBaseDataAccessor<ZNRecord>(_gZkClient));
Builder keyBuilder = accessor.keyBuilder();
ExternalView ev = accessor.getProperty(keyBuilder.externalView(db2));
Assert.assertEquals(ev.getPartitionSet().size(), 60);
for (String partition : ev.getPartitionSet()) {
Assert.assertEquals(ev.getStateMap(partition).size(), 1);
}
IdealState is = accessor.getProperty(keyBuilder.idealStates(db2));
for (String partition : is.getPartitionSet()) {
Assert.assertEquals(is.getPreferenceList(partition).size(), 3);
Assert.assertEquals(is.getInstanceStateMap(partition).size(), 3);
}
Assert.assertTrue(testRebalancerCreated);
Assert.assertTrue(testRebalancerInvoked);
}
@Override
public void onCallback(NotificationContext context) {
LOG.info(
"START: MasterSlaveRebalancer.onCallback running at "
+ _context.getHelixManager().getInstanceName());
if (context.getType().equals(NotificationContext.Type.FINALIZE)) {
LOG.info(
"END: MasterSlaveRebalancer.onCallback FINALIZE callback invoked. Likely lost connection to Helix");
return;
}
HelixManager manager = context.getManager();
String clusterName = manager.getClusterName();
HelixAdmin helixAdmin = manager.getClusterManagmentTool();
IdealState idealState =
helixAdmin.getResourceIdealState(clusterName, MySQLConstants.MASTER_SLAVE_RESOURCE_NAME);
if (idealState == null) {
LOG.info(
"END: MasterSlaveRebalancer.onCallback. "
+ MySQLConstants.MASTER_SLAVE_RESOURCE_NAME
+ " is not yet created");
}
PropertyKey.Builder builder = new PropertyKey.Builder(clusterName);
Map<String, LiveInstance> liveInstancesMap =
manager.getHelixDataAccessor().getChildValuesMap(builder.liveInstances());
Map<String, InstanceConfig> instanceConfigs =
manager.getHelixDataAccessor().getChildValuesMap(builder.instanceConfigs());
IdealState newIdealState = new IdealState(idealState.getId());
newIdealState.getRecord().setSimpleFields(idealState.getRecord().getSimpleFields());
newIdealState.getRecord().setListFields(idealState.getRecord().getListFields());
for (String partition : idealState.getPartitionSet()) {
Map<String, String> instanceStateMap = idealState.getInstanceStateMap(partition);
String currMaster = null;
Set<String> slaveSet = new TreeSet<String>();
for (String instance : instanceStateMap.keySet()) {
if ("MASTER".equalsIgnoreCase(instanceStateMap.get(instance))) {
currMaster = instance;
}
if ("SLAVE".equalsIgnoreCase(instanceStateMap.get(instance))) {
slaveSet.add(instance);
}
}
String newMaster = currMaster;
if (!liveInstancesMap.containsKey(currMaster)
|| !instanceConfigs.get(currMaster).getInstanceEnabled()) {
// need to find a new master.
newMaster = findNewMaster(liveInstancesMap, instanceConfigs, currMaster, slaveSet);
}
for (String instance : instanceStateMap.keySet()) {
if (instance.equalsIgnoreCase(newMaster)) {
newIdealState.setPartitionState(partition, instance, "MASTER");
} else {
newIdealState.setPartitionState(partition, instance, "SLAVE");
}
}
}
if (!idealState.equals(newIdealState)) {
LOG.info("New idealstate computed.");
LOG.info(newIdealState.toString());
manager
.getClusterManagmentTool()
.setResourceIdealState(
clusterName, MySQLConstants.MASTER_SLAVE_RESOURCE_NAME, newIdealState);
} else {
LOG.info("No change in IdealState");
}
LOG.info("END: MasterSlaveRebalancer.onCallback");
}
