snleee commented on a change in pull request #4504: [Instance Assignment] Add implementations for instance assignment URL: https://github.com/apache/incubator-pinot/pull/4504#discussion_r311773897
########## File path: pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/instance/InstanceTagPoolSelector.java ########## @@ -0,0 +1,135 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.pinot.controller.helix.core.assignment.instance; + +import com.google.common.base.Preconditions; +import java.util.ArrayList; +import java.util.Comparator; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.TreeMap; +import org.apache.helix.model.InstanceConfig; +import org.apache.pinot.common.config.instance.InstanceTagPoolConfig; +import org.apache.pinot.controller.api.pojos.Instance; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + + +/** + * The instance tag/pool selector is responsible for selecting instances based on the tag and pool config. + */ +public class InstanceTagPoolSelector { + private static final Logger LOGGER = LoggerFactory.getLogger(InstanceTagPoolSelector.class); + + private final InstanceTagPoolConfig _tagPoolConfig; + private final String _tableNameWithType; + private final int _tableNameHash; + + public InstanceTagPoolSelector(InstanceTagPoolConfig tagPoolConfig, String tableNameWithType) { + _tagPoolConfig = tagPoolConfig; + _tableNameWithType = tableNameWithType; + _tableNameHash = Math.abs(tableNameWithType.hashCode()); + } + + /** + * Returns a map from pool to instance configs based on the tag and pool config for the given instance configs. + */ + public Map<Integer, List<InstanceConfig>> selectInstances(List<InstanceConfig> instanceConfigs) { + LOGGER.info("Starting instance tag/pool selection for table: {} with hash: {}", _tableNameWithType, _tableNameHash); + + // Filter out the enabled instances with the correct tag + String tag = _tagPoolConfig.getTag(); + Preconditions.checkState(tag != null, "Tag must be configured"); + List<InstanceConfig> candidateInstanceConfigs = new ArrayList<>(); + for (InstanceConfig instanceConfig : instanceConfigs) { + if (instanceConfig.getInstanceEnabled() && instanceConfig.getTags().contains(tag)) { + candidateInstanceConfigs.add(instanceConfig); + } + } + candidateInstanceConfigs.sort(Comparator.comparing(InstanceConfig::getInstanceName)); + int numCandidateInstances = candidateInstanceConfigs.size(); + Preconditions.checkState(numCandidateInstances > 0, "No enabled instance has the tag: %s", tag); + LOGGER.info("{} enabled instances have the tag: {} for table: {}", numCandidateInstances, tag, _tableNameWithType); + + Map<Integer, List<InstanceConfig>> poolToInstanceConfigsMap = new TreeMap<>(); + if (_tagPoolConfig.isPoolBased()) { + // Pool based selection + + // Extract the pool information from the instance configs + for (InstanceConfig instanceConfig : candidateInstanceConfigs) { + Map<String, String> poolMap = instanceConfig.getRecord().getMapField(Instance.POOL_KEY); + if (poolMap != null && poolMap.containsKey(tag)) { + int pool = Integer.parseInt(poolMap.get(tag)); + poolToInstanceConfigsMap.computeIfAbsent(pool, k -> new ArrayList<>()).add(instanceConfig); + } + } + Preconditions + .checkState(!poolToInstanceConfigsMap.isEmpty(), "No enabled instance has the pool configured for the tag: %s", + tag); + Map<Integer, Integer> poolToNumInstancesMap = new TreeMap<>(); + for (Map.Entry<Integer, List<InstanceConfig>> entry : poolToInstanceConfigsMap.entrySet()) { + poolToNumInstancesMap.put(entry.getKey(), entry.getValue().size()); + } + LOGGER.info("Number instances for each pool: {} for table: {}", poolToNumInstancesMap, _tableNameWithType); + + // Calculate the pools to select based on the selection config + Set<Integer> pools = poolToInstanceConfigsMap.keySet(); + List<Integer> poolsToSelect = _tagPoolConfig.getPools(); + if (poolsToSelect != null && !poolsToSelect.isEmpty()) { + Preconditions.checkState(pools.containsAll(poolsToSelect), "Cannot find all instance pools configured: %s", + poolsToSelect); + } else { + int numPools = poolToInstanceConfigsMap.size(); + int numPoolsToSelect = _tagPoolConfig.getNumPools(); + if (numPoolsToSelect > 0) { + Preconditions + .checkState(numPoolsToSelect <= numPools, "Not enough instance pools (%s in the cluster, asked for %s)", + numPools, numPoolsToSelect); + } else { + numPoolsToSelect = numPools; + } + + // Directly return the map if all the pools are selected + if (numPools == numPoolsToSelect) { + LOGGER.info("Selecting all {} pools: {} for table: {}", numPools, pools, _tableNameWithType); + return poolToInstanceConfigsMap; + } + + // Select pools based on the table name hash to evenly distribute the tables + poolsToSelect = new ArrayList<>(numPoolsToSelect); + List<Integer> poolsInCluster = new ArrayList<>(pools); + for (int i = 0; i < numPoolsToSelect; i++) { + poolsToSelect.add(poolsInCluster.get((_tableNameHash + i) % numPools)); + } + } + + // Keep the pools selected + LOGGER.info("Selecting pools: {} for table: {}", poolsToSelect, _tableNameWithType); + pools.retainAll(poolsToSelect); Review comment: ah i see :) ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: [email protected] With regards, Apache Git Services --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
