protected List<IdealState> setupIdealState(
int nodes, String[] resources, int partitions, int replicas) {
List<IdealState> idealStates = new ArrayList<IdealState>();
List<String> instances = new ArrayList<String>();
for (int i = 0; i < nodes; i++) {
instances.add("localhost_" + i);
}
for (int i = 0; i < resources.length; i++) {
String resourceName = resources[i];
ZNRecord record = new ZNRecord(resourceName);
for (int p = 0; p < partitions; p++) {
List<String> value = new ArrayList<String>();
for (int r = 0; r < replicas; r++) {
value.add("localhost_" + (p + r + 1) % nodes);
}
record.setListField(resourceName + "_" + p, value);
}
IdealState idealState = new IdealState(record);
idealState.setStateModelDefRef("MasterSlave");
idealState.setIdealStateMode(IdealStateModeProperty.AUTO.toString());
idealState.setNumPartitions(partitions);
idealStates.add(idealState);
// System.out.println(idealState);
Builder keyBuilder = accessor.keyBuilder();
accessor.setProperty(keyBuilder.idealStates(resourceName), idealState);
}
return idealStates;
}
@Override
public boolean verify() {
try {
HelixDataAccessor accessor =
new ZKHelixDataAccessor(_clusterName, new ZkBaseDataAccessor<ZNRecord>(_client));
Builder keyBuilder = accessor.keyBuilder();
int numberOfPartitions =
accessor
.getProperty(keyBuilder.idealStates(_resourceName))
.getRecord()
.getListFields()
.size();
ClusterDataCache cache = new ClusterDataCache();
cache.refresh(accessor);
String masterValue =
cache
.getStateModelDef(cache.getIdealState(_resourceName).getStateModelDefRef())
.getStatesPriorityList()
.get(0);
int replicas = Integer.parseInt(cache.getIdealState(_resourceName).getReplicas());
String instanceGroupTag = cache.getIdealState(_resourceName).getInstanceGroupTag();
int instances = 0;
for (String liveInstanceName : cache.getLiveInstances().keySet()) {
if (cache.getInstanceConfigMap().get(liveInstanceName).containsTag(instanceGroupTag)) {
instances++;
}
}
if (instances == 0) {
instances = cache.getLiveInstances().size();
}
return verifyBalanceExternalView(
accessor.getProperty(keyBuilder.externalView(_resourceName)).getRecord(),
numberOfPartitions,
masterValue,
replicas,
instances);
} catch (Exception e) {
return false;
}
}
@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);
}
@Test
public void testBasic() throws Exception {
// Logger.getRootLogger().setLevel(Level.INFO);
String className = TestHelper.getTestClassName();
String methodName = TestHelper.getTestMethodName();
String clusterName = className + "_" + methodName;
final int n = 5;
System.out.println("START " + clusterName + " at " + new Date(System.currentTimeMillis()));
MockParticipantManager[] participants = new MockParticipantManager[n];
TestHelper.setupCluster(
clusterName,
ZK_ADDR,
12918, // participant port
"localhost", // participant name prefix
"TestDB", // resource name prefix
1, // resources
10, // partitions per resource
n, // number of nodes
3, // replicas
"MasterSlave",
true); // do rebalance
ClusterControllerManager controller =
new ClusterControllerManager(ZK_ADDR, clusterName, "controller_0");
controller.syncStart();
// start participants
for (int i = 0; i < n; i++) {
String instanceName = "localhost_" + (12918 + i);
participants[i] = new MockParticipantManager(ZK_ADDR, clusterName, instanceName);
participants[i].syncStart();
}
boolean result =
ClusterStateVerifier.verifyByZkCallback(
new BestPossAndExtViewZkVerifier(ZK_ADDR, clusterName));
Assert.assertTrue(result);
// add a new idealState without registering message handling factory
ClusterSetup setupTool = new ClusterSetup(ZK_ADDR);
setupTool.addResourceToCluster(clusterName, "TestDB1", 16, "MasterSlave");
ZkBaseDataAccessor<ZNRecord> baseAccessor = new ZkBaseDataAccessor<ZNRecord>(_gZkClient);
ZKHelixDataAccessor accessor = new ZKHelixDataAccessor(clusterName, baseAccessor);
Builder keyBuilder = accessor.keyBuilder();
IdealState idealState = accessor.getProperty(keyBuilder.idealStates("TestDB1"));
idealState.setStateModelFactoryName("TestDB1_Factory");
accessor.setProperty(keyBuilder.idealStates("TestDB1"), idealState);
setupTool.rebalanceStorageCluster(clusterName, "TestDB1", 3);
// assert that we have received OFFLINE->SLAVE messages for all partitions
int totalMsgs = 0;
for (int retry = 0; retry < 5; retry++) {
Thread.sleep(100);
totalMsgs = 0;
for (int i = 0; i < n; i++) {
List<Message> msgs =
accessor.getChildValues(keyBuilder.messages(participants[i].getInstanceName()));
totalMsgs += msgs.size();
}
if (totalMsgs == 48) // partition# x replicas
break;
}
Assert.assertEquals(
totalMsgs,
48,
"Should accumulated 48 unprocessed messages (1 O->S per partition per replica) because TestDB1 is added without state-model-factory but was "
+ totalMsgs);
// register "TestDB1_Factory" state model factory
// Logger.getRootLogger().setLevel(Level.INFO);
for (int i = 0; i < n; i++) {
participants[i]
.getStateMachineEngine()
.registerStateModelFactory("MasterSlave", new MockMSModelFactory(), "TestDB1_Factory");
}
result =
ClusterStateVerifier.verifyByZkCallback(
new BestPossAndExtViewZkVerifier(ZK_ADDR, clusterName));
Assert.assertTrue(result);
// clean up
// wait for all zk callbacks done
controller.syncStop();
for (int i = 0; i < 5; i++) {
participants[i].syncStop();
}
System.out.println("END " + clusterName + " at " + new Date(System.currentTimeMillis()));
}
public static IdealState getTableIdealState(HelixManager manager, String resourceName) {
final HelixDataAccessor accessor = manager.getHelixDataAccessor();
final Builder builder = accessor.keyBuilder();
return accessor.getProperty(builder.idealStates(resourceName));
}
