@Test
public void missingEVPartitionTest() throws Exception {
final String tableName = "myTable";
List<String> allTableNames = new ArrayList<String>();
allTableNames.add(tableName);
IdealState idealState = new IdealState(tableName);
idealState.setPartitionState("myTable_0", "pinot1", "ONLINE");
idealState.setPartitionState("myTable_0", "pinot2", "ONLINE");
idealState.setPartitionState("myTable_0", "pinot3", "ONLINE");
idealState.setPartitionState("myTable_1", "pinot1", "ONLINE");
idealState.setPartitionState("myTable_1", "pinot2", "ONLINE");
idealState.setPartitionState("myTable_1", "pinot3", "ONLINE");
idealState.setPartitionState("myTable_2", "pinot3", "OFFLINE");
idealState.setPartitionState("myTable_3", "pinot3", "ONLINE");
idealState.setReplicas("2");
idealState.setRebalanceMode(IdealState.RebalanceMode.CUSTOMIZED);
ExternalView externalView = new ExternalView(tableName);
externalView.setState("myTable_0", "pinot1", "ONLINE");
externalView.setState("myTable_0", "pinot2", "ONLINE");
externalView.setState("myTable_1", "pinot1", "ERROR");
externalView.setState("myTable_1", "pinot2", "ONLINE");
HelixAdmin helixAdmin;
{
helixAdmin = mock(HelixAdmin.class);
when(helixAdmin.getResourceIdealState("StatusChecker", "myTable")).thenReturn(idealState);
when(helixAdmin.getResourceExternalView("StatusChecker", "myTable")).thenReturn(externalView);
}
{
helixResourceManager = mock(PinotHelixResourceManager.class);
when(helixResourceManager.isLeader()).thenReturn(true);
when(helixResourceManager.getAllPinotTableNames()).thenReturn(allTableNames);
when(helixResourceManager.getHelixClusterName()).thenReturn("StatusChecker");
when(helixResourceManager.getHelixAdmin()).thenReturn(helixAdmin);
}
{
config = mock(ControllerConf.class);
when(config.getStatusControllerFrequencyInSeconds()).thenReturn(300);
}
metricsRegistry = new MetricsRegistry();
controllerMetrics = new ControllerMetrics(metricsRegistry);
segmentStatusChecker = new SegmentStatusChecker(helixResourceManager, config);
segmentStatusChecker.setMetricsRegistry(controllerMetrics);
segmentStatusChecker.runSegmentMetrics();
Assert.assertEquals(
controllerMetrics.getValueOfTableGauge(
externalView.getId(), ControllerGauge.SEGMENTS_IN_ERROR_STATE),
1);
Assert.assertEquals(
controllerMetrics.getValueOfTableGauge(
externalView.getId(), ControllerGauge.NUMBER_OF_REPLICAS),
0);
segmentStatusChecker.stop();
}
@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 testSchemataSM() throws Exception {
String className = TestHelper.getTestClassName();
String methodName = TestHelper.getTestMethodName();
String clusterName = className + "_" + methodName;
int n = 5;
MockParticipant[] participants = new MockParticipant[n];
System.out.println("START " + clusterName + " at " + new Date(System.currentTimeMillis()));
TestHelper.setupCluster(
clusterName,
_zkaddr,
12918, // participant start port
"localhost", // participant name prefix
"TestSchemata", // resource name prefix
1, // resources
1, // partitions per resource
n, // number of nodes
0, // replicas
"STORAGE_DEFAULT_SM_SCHEMATA",
false); // don't rebalance
// rebalance ideal-state to use ANY_LIVEINSTANCE for preference list
ZKHelixDataAccessor accessor = new ZKHelixDataAccessor(clusterName, _baseAccessor);
PropertyKey.Builder keyBuilder = accessor.keyBuilder();
PropertyKey key = keyBuilder.idealStates("TestSchemata0");
IdealState idealState = accessor.getProperty(key);
idealState.setReplicas(HelixConstants.StateModelToken.ANY_LIVEINSTANCE.toString());
idealState
.getRecord()
.setListField(
"TestSchemata0_0",
Arrays.asList(HelixConstants.StateModelToken.ANY_LIVEINSTANCE.toString()));
accessor.setProperty(key, idealState);
MockController controller = new MockController(_zkaddr, clusterName, "controller");
controller.syncStart();
// start n-1 participants
for (int i = 1; i < n; i++) {
String instanceName = "localhost_" + (12918 + i);
participants[i] = new MockParticipant(_zkaddr, clusterName, instanceName);
participants[i].syncStart();
}
boolean result =
ClusterStateVerifier.verifyByZkCallback(
new BestPossAndExtViewZkVerifier(_zkaddr, clusterName));
Assert.assertTrue(result);
// start the remaining 1 participant
participants[0] = new MockParticipant(_zkaddr, clusterName, "localhost_12918");
participants[0].syncStart();
// make sure we have all participants in MASTER state
result =
ClusterStateVerifier.verifyByZkCallback(
new BestPossAndExtViewZkVerifier(_zkaddr, clusterName));
Assert.assertTrue(result);
key = keyBuilder.externalView("TestSchemata0");
ExternalView externalView = accessor.getProperty(key);
Map<String, String> stateMap = externalView.getStateMap("TestSchemata0_0");
Assert.assertNotNull(stateMap);
Assert.assertEquals(stateMap.size(), n, "all " + n + " participants should be in Master state");
for (int i = 0; i < n; i++) {
String instanceName = "localhost_" + (12918 + i);
Assert.assertNotNull(stateMap.get(instanceName));
Assert.assertEquals(stateMap.get(instanceName), "MASTER");
}
// clean up
controller.syncStop();
for (int i = 0; i < n; i++) {
participants[i].syncStop();
}
System.out.println("END " + clusterName + " at " + new Date(System.currentTimeMillis()));
}
@Test
public void testZKReconnect() throws Exception {
final AtomicReference<ZkServer> zkServerRef = new AtomicReference<ZkServer>();
final int zkPort = TestHelper.getRandomPort();
final String zkAddr = String.format("localhost:%d", zkPort);
ZkServer zkServer = TestHelper.startZkServer(zkAddr);
zkServerRef.set(zkServer);
String className = TestHelper.getTestClassName();
String methodName = TestHelper.getTestMethodName();
String clusterName = className + "_" + methodName;
// Setup cluster
LOG.info("Setup clusters");
ClusterSetup clusterSetup = new ClusterSetup(zkAddr);
clusterSetup.addCluster(clusterName, true);
// Registers and starts controller
LOG.info("Starts controller");
HelixManager controller =
HelixManagerFactory.getZKHelixManager(clusterName, null, InstanceType.CONTROLLER, zkAddr);
controller.connect();
// Registers and starts participant
LOG.info("Starts participant");
String hostname = "localhost";
String instanceId = String.format("%s_%d", hostname, 1);
clusterSetup.addInstanceToCluster(clusterName, instanceId);
HelixManager participant =
HelixManagerFactory.getZKHelixManager(
clusterName, instanceId, InstanceType.PARTICIPANT, zkAddr);
participant.connect();
LOG.info("Register state machine");
final CountDownLatch latch = new CountDownLatch(1);
participant
.getStateMachineEngine()
.registerStateModelFactory(
"OnlineOffline",
new StateModelFactory<StateModel>() {
@Override
public StateModel createNewStateModel(String stateUnitKey) {
return new SimpleStateModel(latch);
}
},
"test");
String resourceName = "test-resource";
LOG.info("Ideal state assignment");
HelixAdmin helixAdmin = participant.getClusterManagmentTool();
helixAdmin.addResource(
clusterName,
resourceName,
1,
"OnlineOffline",
IdealState.RebalanceMode.CUSTOMIZED.toString());
IdealState idealState = helixAdmin.getResourceIdealState(clusterName, resourceName);
idealState.setReplicas("1");
idealState.setStateModelFactoryName("test");
idealState.setPartitionState(resourceName + "_0", instanceId, "ONLINE");
LOG.info("Shutdown ZK server");
TestHelper.stopZkServer(zkServerRef.get());
Executors.newSingleThreadScheduledExecutor()
.schedule(
new Runnable() {
@Override
public void run() {
try {
LOG.info("Restart ZK server");
// zkServer.set(TestUtils.startZookeeper(zkDir, zkPort));
zkServerRef.set(TestHelper.startZkServer(zkAddr, null, false));
} catch (Exception e) {
LOG.error(e.getMessage(), e);
}
}
},
2L,
TimeUnit.SECONDS);
// future.get();
LOG.info("Before update ideal state");
helixAdmin.setResourceIdealState(clusterName, resourceName, idealState);
LOG.info("After update ideal state");
LOG.info("Wait for OFFLINE->ONLINE state transition");
try {
Assert.assertTrue(latch.await(10, TimeUnit.SECONDS));
// wait until stable state
boolean result =
ClusterStateVerifier.verifyByZkCallback(
new BestPossAndExtViewZkVerifier(zkAddr, clusterName));
Assert.assertTrue(result);
} finally {
participant.disconnect();
zkServerRef.get().shutdown();
}
}
public static IdealState calculateRelayIdealState(
List<String> partitions,
List<String> instances,
String resultRecordName,
int replica,
String firstValue,
String restValue,
String stateModelName) {
Collections.sort(partitions);
Collections.sort(instances);
if (instances.size() % replica != 0) {
throw new HelixException("Instances must be divided by replica");
}
IdealState result = new IdealState(resultRecordName);
result.setNumPartitions(partitions.size());
result.setReplicas("" + replica);
result.setStateModelDefId(StateModelDefinitionId.from(stateModelName));
int groups = instances.size() / replica;
int remainder = instances.size() % replica;
int remainder2 = partitions.size() % groups;
int storageNodeGroupSize = partitions.size() / groups;
for (int i = 0; i < groups; i++) {
int relayStart = 0, relayEnd = 0, storageNodeStart = 0, storageNodeEnd = 0;
if (i < remainder) {
relayStart = (replica + 1) * i;
relayEnd = (replica + 1) * (i + 1);
} else {
relayStart = (replica + 1) * remainder + replica * (i - remainder);
relayEnd = relayStart + replica;
}
// System.out.println("relay start :" + relayStart + " relayEnd:" + relayEnd);
if (i < remainder2) {
storageNodeStart = (storageNodeGroupSize + 1) * i;
storageNodeEnd = (storageNodeGroupSize + 1) * (i + 1);
} else {
storageNodeStart =
(storageNodeGroupSize + 1) * remainder2 + storageNodeGroupSize * (i - remainder2);
storageNodeEnd = storageNodeStart + storageNodeGroupSize;
}
// System.out.println("storageNodeStart :" + storageNodeStart + " storageNodeEnd:" +
// storageNodeEnd);
List<String> snBatch = partitions.subList(storageNodeStart, storageNodeEnd);
List<String> relayBatch = instances.subList(relayStart, relayEnd);
Map<String, List<String>> sublistFields =
calculateSubIdealState(snBatch, relayBatch, replica);
result.getRecord().getListFields().putAll(sublistFields);
}
for (String snName : result.getRecord().getListFields().keySet()) {
Map<String, String> mapField = new TreeMap<String, String>();
List<String> relayCandidates = result.getRecord().getListField(snName);
mapField.put(relayCandidates.get(0), firstValue);
for (int i = 1; i < relayCandidates.size(); i++) {
mapField.put(relayCandidates.get(i), restValue);
}
result.getRecord().getMapFields().put(snName, mapField);
}
System.out.println();
return result;
}