@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();
}
private void addInstanceTagIfNeeded(String clusterName, String instanceName) {
InstanceConfig instanceConfig = _helixAdmin.getInstanceConfig(clusterName, instanceName);
List<String> instanceTags = instanceConfig.getTags();
if (instanceTags == null || instanceTags.isEmpty()) {
if (ZKMetadataProvider.getClusterTenantIsolationEnabled(
_helixManager.getHelixPropertyStore())) {
_helixAdmin.addInstanceTag(
clusterName,
instanceName,
ControllerTenantNameBuilder.getBrokerTenantNameForTenant(
ControllerTenantNameBuilder.DEFAULT_TENANT_NAME));
} else {
_helixAdmin.addInstanceTag(
clusterName, instanceName, CommonConstants.Helix.UNTAGGED_BROKER_INSTANCE);
}
}
}
public static void deleteResourcePropertyFromHelix(
HelixAdmin admin, String clusterName, String resourceName, String configKey) {
final List<String> keys = new ArrayList<String>();
keys.add(configKey);
final HelixConfigScope scope = getResourceScopeFor(clusterName, resourceName);
admin.removeConfig(scope, keys);
}
@Test()
public void testAdministrator() throws Exception {
System.out.println(
"START " + className + ".testAdministrator() at " + new Date(System.currentTimeMillis()));
final String clusterName = CLUSTER_PREFIX + "_" + className + "_admin";
// basic test
if (_gZkClient.exists("/" + clusterName)) {
_gZkClient.deleteRecursive("/" + clusterName);
}
ZKHelixManager admin =
new ZKHelixManager(clusterName, null, InstanceType.ADMINISTRATOR, ZK_ADDR);
TestHelper.setupEmptyCluster(_gZkClient, clusterName);
admin.connect();
AssertJUnit.assertTrue(admin.isConnected());
HelixAdmin adminTool = admin.getClusterManagmentTool();
// ConfigScope scope = new ConfigScopeBuilder().forCluster(clusterName)
// .forResource("testResource").forPartition("testPartition").build();
HelixConfigScope scope =
new HelixConfigScopeBuilder(ConfigScopeProperty.PARTITION)
.forCluster(clusterName)
.forResource("testResource")
.forPartition("testPartition")
.build();
Map<String, String> properties = new HashMap<String, String>();
properties.put("pKey1", "pValue1");
properties.put("pKey2", "pValue2");
adminTool.setConfig(scope, properties);
properties = adminTool.getConfig(scope, Arrays.asList("pKey1", "pKey2"));
Assert.assertEquals(properties.size(), 2);
Assert.assertEquals(properties.get("pKey1"), "pValue1");
Assert.assertEquals(properties.get("pKey2"), "pValue2");
admin.disconnect();
AssertJUnit.assertFalse(admin.isConnected());
System.out.println(
"END " + className + ".testAdministrator() at " + new Date(System.currentTimeMillis()));
}
@Test
public void missingIdealTest() throws Exception {
final String tableName = "myTable";
List<String> allTableNames = new ArrayList<String>();
allTableNames.add(tableName);
IdealState idealState = null;
HelixAdmin helixAdmin;
{
helixAdmin = mock(HelixAdmin.class);
when(helixAdmin.getResourceIdealState("StatusChecker", "myTable")).thenReturn(null);
when(helixAdmin.getResourceExternalView("StatusChecker", "myTable")).thenReturn(null);
}
{
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(tableName, ControllerGauge.SEGMENTS_IN_ERROR_STATE),
0);
Assert.assertEquals(
controllerMetrics.getValueOfTableGauge(tableName, ControllerGauge.NUMBER_OF_REPLICAS), 0);
segmentStatusChecker.stop();
}
@Test
public void testBounceAll() throws Exception {
// pick numbers that don't divide evenly
final int NUM_PARTICIPANTS = 5;
final int NUM_PARTITIONS = 123;
final int NUM_REPLICAS = 1;
final String RESOURCE_PREFIX = "TestDB";
final String RESOURCE_NAME = RESOURCE_PREFIX + "0";
// create a cluster name based on this test name
String className = TestHelper.getTestClassName();
String methodName = TestHelper.getTestMethodName();
String clusterName = className + "_" + methodName;
System.out.println("START " + clusterName + " at " + new Date(System.currentTimeMillis()));
// Set up cluster
TestHelper.setupCluster(
clusterName,
ZK_ADDR,
12918, // participant port
"localhost", // participant name prefix
"TestDB", // resource name prefix
1, // resources
NUM_PARTITIONS, // partitions per resource
NUM_PARTICIPANTS, // number of nodes
NUM_REPLICAS, // replicas
"OnlineOffline",
RebalanceMode.FULL_AUTO, // use FULL_AUTO mode to test node tagging
true); // do rebalance
// Start the participants
HelixManager[] participants = new HelixManager[NUM_PARTICIPANTS];
for (int i = 0; i < NUM_PARTICIPANTS; i++) {
final String instanceName = "localhost_" + (12918 + i);
participants[i] = createParticipant(clusterName, instanceName);
participants[i].connect();
}
// Start the controller
ClusterControllerManager controller =
new ClusterControllerManager(ZK_ADDR, clusterName, "controller");
controller.syncStart();
// get an admin and accessor
HelixAdmin helixAdmin = new ZKHelixAdmin(_gZkClient);
BaseDataAccessor<ZNRecord> baseAccessor = new ZkBaseDataAccessor<ZNRecord>(_gZkClient);
HelixDataAccessor accessor = new ZKHelixDataAccessor(clusterName, baseAccessor);
PropertyKey.Builder keyBuilder = accessor.keyBuilder();
// do the test
try {
Thread.sleep(1000);
// ensure that the external view coalesces
boolean result =
ClusterStateVerifier.verifyByZkCallback(
new BestPossAndExtViewZkVerifier(ZK_ADDR, clusterName));
Assert.assertTrue(result);
ExternalView stableExternalView =
accessor.getProperty(keyBuilder.externalView(RESOURCE_NAME));
for (HelixManager participant : participants) {
// disable the controller, bounce the node, re-enable the controller, verify assignments
// remained the same
helixAdmin.enableCluster(clusterName, false);
participant.disconnect();
Thread.sleep(1000);
participant = createParticipant(clusterName, participant.getInstanceName());
participant.connect();
Thread.sleep(1000);
helixAdmin.enableCluster(clusterName, true);
Thread.sleep(1000);
result =
ClusterStateVerifier.verifyByZkCallback(
new MatchingExternalViewVerifier(stableExternalView, clusterName));
Assert.assertTrue(result);
}
} finally {
// clean up
controller.syncStop();
for (HelixManager participant : participants) {
participant.disconnect();
}
System.out.println("END " + clusterName + " at " + new Date(System.currentTimeMillis()));
}
}
@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");
}
@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 getBrokerIdealStates(HelixAdmin admin, String clusterName) {
return admin.getResourceIdealState(clusterName, BROKER_RESOURCE);
}
public static ExternalView getExternalViewForResource(
HelixAdmin admin, String clusterName, String resourceName) {
return admin.getResourceExternalView(clusterName, resourceName);
}
public static void updateResourceConfigsFor(
Map<String, String> newConfigs, String resourceName, String clusterName, HelixAdmin admin) {
final HelixConfigScope scope = getResourceScopeFor(clusterName, resourceName);
admin.setConfig(scope, newConfigs);
}
public static Map<String, String> getResourceConfigsFor(
String clusterName, String resourceName, HelixAdmin admin) {
final HelixConfigScope scope = getResourceScopeFor(clusterName, resourceName);
final List<String> keys = admin.getConfigKeys(scope);
return admin.getConfig(scope, keys);
}
public static Map<String, String> getInstanceConfigsMapFor(
String instanceName, String clusterName, HelixAdmin admin) {
final HelixConfigScope scope = getInstanceScopefor(clusterName, instanceName);
final List<String> keys = admin.getConfigKeys(scope);
return admin.getConfig(scope, keys);
}
/**
* Toggle the state of the instance between OFFLINE and ONLINE.
*
* @param instanceName Name of the instance for which to toggle the state.
* @param clusterName Name of the cluster to which the instance belongs.
* @param admin HelixAdmin to access the cluster.
* @param enable Set enable to true for ONLINE and FALSE for OFFLINE.
*/
public static void setInstanceState(
String instanceName, String clusterName, HelixAdmin admin, boolean enable) {
admin.enableInstance(clusterName, instanceName, enable);
}
/**
* Returns all instances for the given cluster.
*
* @param helixAdmin The HelixAdmin object used to interact with the Helix cluster
* @param clusterName Name of the cluster for which to get all the instances for.
* @return Returns a List of strings containing the instance names for the given cluster.
*/
public static List<String> getAllInstances(HelixAdmin helixAdmin, String clusterName) {
return helixAdmin.getInstancesInCluster(clusterName);
}