@Test
public void nonLeaderTest() throws Exception {
final String tableName = "myTable";
List<String> allTableNames = new ArrayList<String>();
allTableNames.add(tableName);
HelixAdmin helixAdmin;
{
helixAdmin = mock(HelixAdmin.class);
}
{
helixResourceManager = mock(PinotHelixResourceManager.class);
when(helixResourceManager.isLeader()).thenReturn(false);
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();
}
/*
* For tenant creation
*/
@Override
@Post("json")
public Representation post(Representation entity) {
StringRepresentation presentation = null;
try {
PinotResourceManagerResponse response = null;
final Tenant tenant = _objectMapper.readValue(entity.getText(), Tenant.class);
switch (tenant.getTenantRole()) {
case BROKER:
response = _pinotHelixResourceMananger.createBrokerTenant(tenant);
presentation = new StringRepresentation(response.toString());
break;
case SERVER:
response = _pinotHelixResourceMananger.createServerTenant(tenant);
presentation = new StringRepresentation(response.toString());
break;
default:
throw new RuntimeException("Not a valid tenant creation call");
}
} catch (final Exception e) {
presentation = exceptionToStringRepresentation(e);
LOGGER.error("Caught exception while processing put request", e);
setStatus(Status.SERVER_ERROR_INTERNAL);
}
return presentation;
}
@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
@Delete
public Representation delete() {
StringRepresentation presentation = null;
try {
final String tenantName = (String) getRequest().getAttributes().get(TENANT_NAME);
final String type = getReference().getQueryAsForm().getValues("type");
if (type == null) {
presentation =
new StringRepresentation(
"Not specify the type for the tenant name. Please try to append:"
+ "/?type=SERVER or /?type=BROKER ");
} else {
TenantRole tenantRole = TenantRole.valueOf(type.toUpperCase());
PinotResourceManagerResponse res = null;
switch (tenantRole) {
case BROKER:
if (_pinotHelixResourceMananger.isBrokerTenantDeletable(tenantName)) {
res = _pinotHelixResourceMananger.deleteBrokerTenantFor(tenantName);
} else {
res = new PinotResourceManagerResponse();
res.status = STATUS.failure;
res.errorMessage = "Broker Tenant is not null, cannot delete it.";
}
break;
case SERVER:
if (_pinotHelixResourceMananger.isServerTenantDeletable(tenantName)) {
res = _pinotHelixResourceMananger.deleteOfflineServerTenantFor(tenantName);
if (res.isSuccessfull()) {
res = _pinotHelixResourceMananger.deleteRealtimeServerTenantFor(tenantName);
}
} else {
res = new PinotResourceManagerResponse();
res.status = STATUS.failure;
res.errorMessage = "Server Tenant is not null, cannot delete it.";
}
break;
default:
break;
}
presentation = new StringRepresentation(res.toString());
}
} catch (final Exception e) {
presentation = exceptionToStringRepresentation(e);
LOGGER.error("Caught exception while processing delete request", e);
setStatus(Status.SERVER_ERROR_INTERNAL);
}
return presentation;
}
/**
* Helper method to perform idempotent operation to refresh all watches (related to real-time
* segments): - Data change listener for all existing real-time tables. - Child creation listener
* for all existing real-time tables. - Data change listener for all existing real-time segments
*
* @param path
*/
private void refreshWatchers(String path) {
LOGGER.info("Received change notification for path: {}", path);
List<Stat> stats = new ArrayList<>();
List<ZNRecord> tableConfigs =
_pinotHelixResourceManager.getPropertyStore().getChildren(TABLE_CONFIG, stats, 0);
if (tableConfigs == null) {
return;
}
for (ZNRecord tableConfig : tableConfigs) {
try {
AbstractTableConfig abstractTableConfig = AbstractTableConfig.fromZnRecord(tableConfig);
if (abstractTableConfig.isRealTime()) {
String realtimeTable = abstractTableConfig.getTableName();
String realtimeSegmentsPathForTable =
_propertyStorePath + SEGMENTS_PATH + "/" + realtimeTable;
LOGGER.info("Setting data/child changes watch for real-time table '{}'", realtimeTable);
_zkClient.subscribeDataChanges(realtimeSegmentsPathForTable, this);
_zkClient.subscribeChildChanges(realtimeSegmentsPathForTable, this);
List<String> childNames =
_pinotHelixResourceManager
.getPropertyStore()
.getChildNames(SEGMENTS_PATH + "/" + realtimeTable, 0);
if (childNames != null && !childNames.isEmpty()) {
for (String segmentName : childNames) {
String segmentPath = realtimeSegmentsPathForTable + "/" + segmentName;
LOGGER.info("Setting data change watch for real-time segment: {}", segmentPath);
_zkClient.subscribeDataChanges(segmentPath, this);
}
}
}
} catch (JSONException e) {
LOGGER.error("Caught exception while reading table config", e);
} catch (IOException e) {
LOGGER.error(
"Caught exception while setting change listeners for realtime tables/segments", e);
}
}
}
public void start() {
LOGGER.info(
"Starting realtime segments manager, adding a listener on the property store table configs path.");
String zkUrl = _pinotHelixResourceManager.getHelixZkURL();
_zkClient =
new ZkClient(zkUrl, ZkClient.DEFAULT_SESSION_TIMEOUT, ZkClient.DEFAULT_CONNECTION_TIMEOUT);
_zkClient.setZkSerializer(new ZNRecordSerializer());
_zkClient.waitUntilConnected();
// Subscribe to any data/child changes to property
_zkClient.subscribeChildChanges(_tableConfigPath, this);
_zkClient.subscribeDataChanges(_tableConfigPath, this);
// Setup change listeners for already existing tables, if any.
processPropertyStoreChange(_tableConfigPath);
}
/**
* called with optional resourceName if resourceName is not present then it sends back a list of
*
* @return
*/
@Override
@Get
public Representation get() {
StringRepresentation presentation = null;
try {
final String tenantName = (String) getRequest().getAttributes().get(TENANT_NAME);
if (tenantName == null) {
// Return all the tags.
final JSONObject ret = new JSONObject();
final String type = getReference().getQueryAsForm().getValues("type");
if (type == null || type.equals("server")) {
ret.put("SERVER_TENANTS", _pinotHelixResourceMananger.getAllServerTenantNames());
}
if (type == null || type.equals("broker")) {
ret.put("BROKER_TENANTS", _pinotHelixResourceMananger.getAllBrokerTenantNames());
}
presentation = new StringRepresentation(ret.toString(), MediaType.APPLICATION_JSON);
} else {
// Return instances related to given tenant name.
final String type = getReference().getQueryAsForm().getValues("type");
JSONObject resourceGetRet = new JSONObject();
if (type == null) {
resourceGetRet.put(
"ServerInstances",
_pinotHelixResourceMananger.getAllInstancesForServerTenant(tenantName));
resourceGetRet.put(
"BrokerInstances",
_pinotHelixResourceMananger.getAllInstancesForBrokerTenant(tenantName));
} else {
if (type.equals("server")) {
resourceGetRet.put(
"ServerInstances",
_pinotHelixResourceMananger.getAllInstancesForServerTenant(tenantName));
}
if (type.equals("broker")) {
resourceGetRet.put(
"BrokerInstances",
_pinotHelixResourceMananger.getAllInstancesForBrokerTenant(tenantName));
}
}
resourceGetRet.put(TENANT_NAME, tenantName);
presentation =
new StringRepresentation(resourceGetRet.toString(), MediaType.APPLICATION_JSON);
}
} catch (final Exception e) {
presentation = exceptionToStringRepresentation(e);
LOGGER.error("Caught exception while processing get request", e);
setStatus(Status.SERVER_ERROR_INTERNAL);
}
return presentation;
}
private synchronized void assignRealtimeSegmentsToServerInstancesIfNecessary()
throws JSONException, IOException {
// Fetch current ideal state snapshot
Map<String, IdealState> idealStateMap = new HashMap<String, IdealState>();
for (String resource : _pinotHelixResourceManager.getAllRealtimeTables()) {
idealStateMap.put(
resource,
_pinotHelixResourceManager
.getHelixAdmin()
.getResourceIdealState(_pinotHelixResourceManager.getHelixClusterName(), resource));
}
List<String> listOfSegmentsToAdd = new ArrayList<String>();
for (String resource : idealStateMap.keySet()) {
IdealState state = idealStateMap.get(resource);
// Are there any partitions?
if (state.getPartitionSet().size() == 0) {
// No, this is a brand new ideal state, so we will add one new segment to every partition
// and replica
List<String> instancesInResource = new ArrayList<String>();
try {
instancesInResource.addAll(
_pinotHelixResourceManager.getServerInstancesForTable(resource, TableType.REALTIME));
} catch (Exception e) {
LOGGER.error("Caught exception while fetching instances for resource {}", resource, e);
}
// Assign a new segment to all server instances
for (String instanceId : instancesInResource) {
InstanceZKMetadata instanceZKMetadata =
_pinotHelixResourceManager.getInstanceZKMetadata(instanceId);
String groupId = instanceZKMetadata.getGroupId(resource);
String partitionId = instanceZKMetadata.getPartition(resource);
listOfSegmentsToAdd.add(
SegmentNameBuilder.Realtime.build(
resource,
instanceId,
groupId,
partitionId,
String.valueOf(System.currentTimeMillis())));
}
} else {
// Add all server instances to the list of instances for which to assign a realtime segment
Set<String> instancesToAssignRealtimeSegment = new HashSet<String>();
instancesToAssignRealtimeSegment.addAll(
_pinotHelixResourceManager.getServerInstancesForTable(resource, TableType.REALTIME));
// Remove server instances that are currently processing a segment
for (String partition : state.getPartitionSet()) {
RealtimeSegmentZKMetadata realtimeSegmentZKMetadata =
ZKMetadataProvider.getRealtimeSegmentZKMetadata(
_pinotHelixResourceManager.getPropertyStore(),
SegmentNameBuilder.Realtime.extractTableName(partition),
partition);
if (realtimeSegmentZKMetadata.getStatus() == Status.IN_PROGRESS) {
String instanceName = SegmentNameBuilder.Realtime.extractInstanceName(partition);
instancesToAssignRealtimeSegment.remove(instanceName);
}
}
// Assign a new segment to the server instances not currently processing this segment
for (String instanceId : instancesToAssignRealtimeSegment) {
InstanceZKMetadata instanceZKMetadata =
_pinotHelixResourceManager.getInstanceZKMetadata(instanceId);
String groupId = instanceZKMetadata.getGroupId(resource);
String partitionId = instanceZKMetadata.getPartition(resource);
listOfSegmentsToAdd.add(
SegmentNameBuilder.Realtime.build(
resource,
instanceId,
groupId,
partitionId,
String.valueOf(System.currentTimeMillis())));
}
}
}
LOGGER.info(
"Computed list of new segments to add : " + Arrays.toString(listOfSegmentsToAdd.toArray()));
// Add the new segments to the server instances
for (String segmentId : listOfSegmentsToAdd) {
String resourceName = SegmentNameBuilder.Realtime.extractTableName(segmentId);
String instanceName = SegmentNameBuilder.Realtime.extractInstanceName(segmentId);
// Does the ideal state already contain this segment?
if (!idealStateMap.get(resourceName).getPartitionSet().contains(segmentId)) {
// No, add it
// Create the realtime segment metadata
RealtimeSegmentZKMetadata realtimeSegmentMetadataToAdd = new RealtimeSegmentZKMetadata();
realtimeSegmentMetadataToAdd.setTableName(
TableNameBuilder.extractRawTableName(resourceName));
realtimeSegmentMetadataToAdd.setSegmentType(SegmentType.REALTIME);
realtimeSegmentMetadataToAdd.setStatus(Status.IN_PROGRESS);
realtimeSegmentMetadataToAdd.setSegmentName(segmentId);
// Add the new metadata to the property store
ZKMetadataProvider.setRealtimeSegmentZKMetadata(
_pinotHelixResourceManager.getPropertyStore(), realtimeSegmentMetadataToAdd);
// Update the ideal state to add the new realtime segment
HelixHelper.updateIdealState(
_pinotHelixResourceManager.getHelixZkManager(),
resourceName,
idealState ->
PinotTableIdealStateBuilder.addNewRealtimeSegmentToIdealState(
segmentId, idealState, instanceName),
RetryPolicies.exponentialBackoffRetryPolicy(5, 500L, 2.0f));
}
}
}
public void stop() {
LOGGER.info("Stopping realtime segments manager, stopping property store.");
_pinotHelixResourceManager.getPropertyStore().stop();
}
public PinotRealtimeSegmentManager(PinotHelixResourceManager pinotManager) {
_pinotHelixResourceManager = pinotManager;
String clusterName = _pinotHelixResourceManager.getHelixClusterName();
_propertyStorePath = PropertyPathConfig.getPath(PropertyType.PROPERTYSTORE, clusterName);
_tableConfigPath = _propertyStorePath + TABLE_CONFIG;
}
private boolean isLeader() {
return _pinotHelixResourceManager.isLeader();
}