/**
* Remove a resource (offline/realtime table) from the Broker's ideal state.
*
* @param helixManager The HelixManager object for accessing helix cluster.
* @param resourceTag Name of the resource that needs to be removed from Broker ideal state.
*/
public static void removeResourceFromBrokerIdealState(
HelixManager helixManager, final String resourceTag) {
Function<IdealState, IdealState> updater =
new Function<IdealState, IdealState>() {
@Override
public IdealState apply(IdealState idealState) {
if (idealState.getPartitionSet().contains(resourceTag)) {
idealState.getPartitionSet().remove(resourceTag);
return idealState;
} else {
return null;
}
}
};
// Removing partitions from ideal state
LOGGER.info("Trying to remove resource from idealstats");
HelixHelper.updateIdealState(
helixManager,
CommonConstants.Helix.BROKER_RESOURCE_INSTANCE,
updater,
DEFAULT_RETRY_POLICY);
}
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 static void updateBrokerConfig(
Map<String, String> brokerResourceConfig, HelixAdmin admin, String clusterName) {
HelixHelper.updateResourceConfigsFor(brokerResourceConfig, BROKER_RESOURCE, clusterName, admin);
}
public static Map<String, String> getBrokerResourceConfig(HelixAdmin admin, String clusterName) {
return HelixHelper.getResourceConfigsFor(clusterName, BROKER_RESOURCE, admin);
}