/**
* Copyright (C) 2014-2016 LinkedIn Corp. (pinot-core@linkedin.com)
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package com.linkedin.pinot.controller.helix.core.sharding;
import com.linkedin.pinot.common.config.TableNameBuilder;
import com.linkedin.pinot.common.segment.SegmentMetadata;
import com.linkedin.pinot.common.utils.ControllerTenantNameBuilder;
import com.linkedin.pinot.common.utils.Pairs;
import com.linkedin.pinot.common.utils.Pairs.Number2ObjectPair;
import com.linkedin.pinot.common.utils.helix.HelixHelper;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.PriorityQueue;
import org.apache.helix.HelixAdmin;
import org.apache.helix.model.IdealState;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
/**
* Assigns a segment to the instance that has least number of segments.
*/
public class BalanceNumSegmentAssignmentStrategy implements SegmentAssignmentStrategy {
private static final Logger LOGGER = LoggerFactory.getLogger(BalanceNumSegmentAssignmentStrategy.class);
@Override
public List<String> getAssignedInstances(HelixAdmin helixAdmin, String helixClusterName,
SegmentMetadata segmentMetadata, int numReplicas, String tenantName) {
String serverTenantName;
String tableName;
if ("realtime".equalsIgnoreCase(segmentMetadata.getIndexType())) {
tableName = TableNameBuilder.REALTIME.tableNameWithType(segmentMetadata.getTableName());
serverTenantName = ControllerTenantNameBuilder.getRealtimeTenantNameForTenant(tenantName);
} else {
tableName = TableNameBuilder.OFFLINE.tableNameWithType(segmentMetadata.getTableName());
serverTenantName = ControllerTenantNameBuilder.getOfflineTenantNameForTenant(tenantName);
}
List<String> selectedInstances = new ArrayList<String>();
Map<String, Integer> currentNumSegmentsPerInstanceMap = new HashMap<String, Integer>();
List<String> allTaggedInstances = HelixHelper.getEnabledInstancesWithTag(helixAdmin, helixClusterName, serverTenantName);
for (String instance : allTaggedInstances) {
currentNumSegmentsPerInstanceMap.put(instance, 0);
}
// Count number of segments assigned to each instance
IdealState idealState = helixAdmin.getResourceIdealState(helixClusterName, tableName);
if (idealState != null) {
for (String partitionName : idealState.getPartitionSet()) {
Map<String, String> instanceToStateMap = idealState.getInstanceStateMap(partitionName);
if (instanceToStateMap != null) {
for (String instanceName : instanceToStateMap.keySet()) {
if (currentNumSegmentsPerInstanceMap.containsKey(instanceName)) {
currentNumSegmentsPerInstanceMap.put(instanceName, currentNumSegmentsPerInstanceMap.get(instanceName) + 1);
}
// else, ignore. Do not add servers, that are not tagged, to the map
// By this approach, new segments will not be allotted to the server if tags changed
}
}
}
}
// Select up to numReplicas instances with the fewest segments assigned
PriorityQueue<Number2ObjectPair<String>> priorityQueue =
new PriorityQueue<Number2ObjectPair<String>>(numReplicas, Pairs.getDescendingnumber2ObjectPairComparator());
for (String key : currentNumSegmentsPerInstanceMap.keySet()) {
priorityQueue.add(new Number2ObjectPair<String>(currentNumSegmentsPerInstanceMap.get(key), key));
if (priorityQueue.size() > numReplicas) {
priorityQueue.poll();
}
}
while (!priorityQueue.isEmpty()) {
selectedInstances.add(priorityQueue.poll().getB());
}
LOGGER.info("Segment assignment result for : " + segmentMetadata.getName() + ", in resource : "
+ segmentMetadata.getTableName() + ", selected instances: " + Arrays.toString(selectedInstances.toArray()));
return selectedInstances;
}
}