You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@pinot.apache.org by ja...@apache.org on 2019/08/08 18:15:56 UTC
[incubator-pinot] branch master updated: [Instance Assignment] Add
implementations for instance assignment (#4504)
This is an automated email from the ASF dual-hosted git repository.
jackie pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-pinot.git
The following commit(s) were added to refs/heads/master by this push:
new 8e02f48 [Instance Assignment] Add implementations for instance assignment (#4504)
8e02f48 is described below
commit 8e02f48f66283c9b4ab4768e8a9d72a86817ded8
Author: Xiaotian (Jackie) Jiang <17...@users.noreply.github.com>
AuthorDate: Thu Aug 8 11:15:49 2019 -0700
[Instance Assignment] Add implementations for instance assignment (#4504)
Generate InstancePartitions based on InstanceConfigs and InstanceAssignmentConfig
Generate default InstanceAssignmentConfig from TableConfig for backward-compatibility
---
.../common/config/SegmentPartitionConfig.java | 2 +-
.../instance/InstanceAssignmentConfigUtils.java | 6 +-
.../instance/InstanceReplicaPartitionConfig.java | 36 +-
.../config/instance/InstanceTagPoolConfig.java | 4 +-
.../pinot/common/config/TableConfigTest.java | 8 +-
.../pinot/controller/api/pojos/Instance.java | 2 +
.../HashBasedRotateInstanceConstraintApplier.java | 60 +++
.../instance/InstanceAssignmentDriver.java | 85 ++++
.../instance/InstanceConstraintApplier.java | 37 ++
.../instance/InstanceReplicaPartitionSelector.java | 189 ++++++++
.../instance/InstanceTagPoolSelector.java | 133 ++++++
...ltimeReplicaGroupSegmentAssignmentStrategy.java | 4 +-
.../instance/InstanceAssignmentTest.java | 526 +++++++++++++++++++++
13 files changed, 1062 insertions(+), 30 deletions(-)
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/config/SegmentPartitionConfig.java b/pinot-common/src/main/java/org/apache/pinot/common/config/SegmentPartitionConfig.java
index f4c6fa8..7220035 100644
--- a/pinot-common/src/main/java/org/apache/pinot/common/config/SegmentPartitionConfig.java
+++ b/pinot-common/src/main/java/org/apache/pinot/common/config/SegmentPartitionConfig.java
@@ -73,7 +73,7 @@ public class SegmentPartitionConfig {
*/
@JsonIgnore
public void setNumPartitions(String column, int numPartitions) {
- ColumnPartitionConfig columnPartitionConfig = _columnPartitionMap.get(column);
+ _columnPartitionMap.get(column).setNumPartitions(numPartitions);
}
/**
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/config/instance/InstanceAssignmentConfigUtils.java b/pinot-common/src/main/java/org/apache/pinot/common/config/instance/InstanceAssignmentConfigUtils.java
index ffeea87..f094527 100644
--- a/pinot-common/src/main/java/org/apache/pinot/common/config/instance/InstanceAssignmentConfigUtils.java
+++ b/pinot-common/src/main/java/org/apache/pinot/common/config/instance/InstanceAssignmentConfigUtils.java
@@ -76,11 +76,11 @@ public class InstanceAssignmentConfigUtils {
Preconditions.checkState(numPartitions > 0, "Number of partitions for column: %s is not properly configured",
partitionColumn);
replicaPartitionConfig.setNumPartitions(numPartitions);
- replicaPartitionConfig.setNumServersPerPartition(replicaGroupStrategyConfig.getNumInstancesPerPartition());
+ replicaPartitionConfig.setNumInstancesPerPartition(replicaGroupStrategyConfig.getNumInstancesPerPartition());
} else {
// If partition column is not configured, use replicaGroupStrategyConfig.getNumInstancesPerPartition() as
- // number of servers per replica for backward-compatibility
- replicaPartitionConfig.setNumServersPerReplica(replicaGroupStrategyConfig.getNumInstancesPerPartition());
+ // number of instances per replica for backward-compatibility
+ replicaPartitionConfig.setNumInstancesPerReplica(replicaGroupStrategyConfig.getNumInstancesPerPartition());
}
} else {
replicaPartitionConfig.setNumReplicas(validationConfig.getReplicasPerPartitionNumber());
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/config/instance/InstanceReplicaPartitionConfig.java b/pinot-common/src/main/java/org/apache/pinot/common/config/instance/InstanceReplicaPartitionConfig.java
index 94d2a87..65d1e49 100644
--- a/pinot-common/src/main/java/org/apache/pinot/common/config/instance/InstanceReplicaPartitionConfig.java
+++ b/pinot-common/src/main/java/org/apache/pinot/common/config/instance/InstanceReplicaPartitionConfig.java
@@ -32,24 +32,24 @@ public class InstanceReplicaPartitionConfig {
private boolean _replicaGroupBased;
@ConfigKey("numServers")
- @ConfigDoc("Number of servers to select for non-replica-group based selection, select all servers if not specified")
- private int _numServers;
+ @ConfigDoc("Number of instances to select for non-replica-group based selection, select all instances if not specified")
+ private int _numInstances;
@ConfigKey("numReplicas")
@ConfigDoc("Number of replicas (replica-groups) for replica-group based selection")
private int _numReplicas;
@ConfigKey("numServersPerReplica")
- @ConfigDoc("Number of servers per replica (replica-group) for replica-group based selection, select all servers if not specified")
- private int _numServersPerReplica;
+ @ConfigDoc("Number of instances per replica for replica-group based selection, select as many instances as possible if not specified")
+ private int _numInstancesPerReplica;
@ConfigKey("numPartitions")
@ConfigDoc("Number of partitions for replica-group based selection, do not partition the replica-group (1 partition) if not specified")
private int _numPartitions;
@ConfigKey("numServersPerPartition")
- @ConfigDoc("Number of servers per partition (within a replica) for replica-group based selection, select all servers if not specified")
- private int _numServersPerPartition;
+ @ConfigDoc("Number of instances per partition (within a replica) for replica-group based selection, select all instances if not specified")
+ private int _numInstancesPerPartition;
@JsonProperty
public boolean isReplicaGroupBased() {
@@ -62,13 +62,13 @@ public class InstanceReplicaPartitionConfig {
}
@JsonProperty
- public int getNumServers() {
- return _numServers;
+ public int getNumInstances() {
+ return _numInstances;
}
@JsonProperty
- public void setNumServers(int numServers) {
- _numServers = numServers;
+ public void setNumInstances(int numInstances) {
+ _numInstances = numInstances;
}
@JsonProperty
@@ -82,13 +82,13 @@ public class InstanceReplicaPartitionConfig {
}
@JsonProperty
- public int getNumServersPerReplica() {
- return _numServersPerReplica;
+ public int getNumInstancesPerReplica() {
+ return _numInstancesPerReplica;
}
@JsonProperty
- public void setNumServersPerReplica(int numServersPerReplica) {
- _numServersPerReplica = numServersPerReplica;
+ public void setNumInstancesPerReplica(int numInstancesPerReplica) {
+ _numInstancesPerReplica = numInstancesPerReplica;
}
@JsonProperty
@@ -102,12 +102,12 @@ public class InstanceReplicaPartitionConfig {
}
@JsonProperty
- public int getNumServersPerPartition() {
- return _numServersPerPartition;
+ public int getNumInstancesPerPartition() {
+ return _numInstancesPerPartition;
}
@JsonProperty
- public void setNumServersPerPartition(int numServersPerPartition) {
- _numServersPerPartition = numServersPerPartition;
+ public void setNumInstancesPerPartition(int numInstancesPerPartition) {
+ _numInstancesPerPartition = numInstancesPerPartition;
}
}
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/config/instance/InstanceTagPoolConfig.java b/pinot-common/src/main/java/org/apache/pinot/common/config/instance/InstanceTagPoolConfig.java
index 9ffe792..25e5139 100644
--- a/pinot-common/src/main/java/org/apache/pinot/common/config/instance/InstanceTagPoolConfig.java
+++ b/pinot-common/src/main/java/org/apache/pinot/common/config/instance/InstanceTagPoolConfig.java
@@ -37,11 +37,11 @@ public class InstanceTagPoolConfig {
private boolean _poolBased;
@ConfigKey("numPools")
- @ConfigDoc("Number of pools to select for pool based selection, contradict to pools")
+ @ConfigDoc("Number of pools to select for pool based selection, contradict to pools, select all pools if neither of them are specified")
private int _numPools;
@ConfigKey("pools")
- @ConfigDoc("Pools to select for pool based selection, contradict to numPools")
+ @ConfigDoc("Pools to select for pool based selection, contradict to numPools, select all pools if neither of them are specified")
private List<Integer> _pools;
@JsonProperty
diff --git a/pinot-common/src/test/java/org/apache/pinot/common/config/TableConfigTest.java b/pinot-common/src/test/java/org/apache/pinot/common/config/TableConfigTest.java
index c94d2c1..b024d16 100644
--- a/pinot-common/src/test/java/org/apache/pinot/common/config/TableConfigTest.java
+++ b/pinot-common/src/test/java/org/apache/pinot/common/config/TableConfigTest.java
@@ -361,7 +361,7 @@ public class TableConfigTest {
InstanceReplicaPartitionConfig replicaPartitionConfig = new InstanceReplicaPartitionConfig();
replicaPartitionConfig.setReplicaGroupBased(true);
replicaPartitionConfig.setNumReplicas(3);
- replicaPartitionConfig.setNumServersPerReplica(5);
+ replicaPartitionConfig.setNumInstancesPerReplica(5);
instanceAssignmentConfig.setReplicaPartitionConfig(replicaPartitionConfig);
TableConfig tableConfig = tableConfigBuilder.setInstanceAssignmentConfigMap(
@@ -462,10 +462,10 @@ public class TableConfigTest {
InstanceReplicaPartitionConfig replicaPartitionConfig = instanceAssignmentConfig.getReplicaPartitionConfig();
assertTrue(replicaPartitionConfig.isReplicaGroupBased());
- assertEquals(replicaPartitionConfig.getNumServers(), 0);
+ assertEquals(replicaPartitionConfig.getNumInstances(), 0);
assertEquals(replicaPartitionConfig.getNumReplicas(), 3);
- assertEquals(replicaPartitionConfig.getNumServersPerReplica(), 5);
+ assertEquals(replicaPartitionConfig.getNumInstancesPerReplica(), 5);
assertEquals(replicaPartitionConfig.getNumPartitions(), 0);
- assertEquals(replicaPartitionConfig.getNumServersPerPartition(), 0);
+ assertEquals(replicaPartitionConfig.getNumInstancesPerPartition(), 0);
}
}
diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/api/pojos/Instance.java b/pinot-controller/src/main/java/org/apache/pinot/controller/api/pojos/Instance.java
index 5200649..7a8ac23 100644
--- a/pinot-controller/src/main/java/org/apache/pinot/controller/api/pojos/Instance.java
+++ b/pinot-controller/src/main/java/org/apache/pinot/controller/api/pojos/Instance.java
@@ -29,6 +29,8 @@ import org.apache.pinot.common.utils.CommonConstants;
*/
//@Example("{\n" + "\t\"host\": \"hostname.example.com\",\n" + "\t\"port\": \"1234\",\n" + "\t\"type\": \"server\"\n" + "}")
public class Instance {
+ public static final String POOL_KEY = "pool";
+
private final String _host;
private final String _port;
private final String _type;
diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/instance/HashBasedRotateInstanceConstraintApplier.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/instance/HashBasedRotateInstanceConstraintApplier.java
new file mode 100644
index 0000000..98934fb
--- /dev/null
+++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/instance/HashBasedRotateInstanceConstraintApplier.java
@@ -0,0 +1,60 @@
+/**
+ * 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.Collections;
+import java.util.List;
+import java.util.Map;
+import org.apache.helix.model.InstanceConfig;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Rotate the instances based on the table name hash to prevent creating hotspot instances. Use this as the default
+ * instance constraint to align with the default instance partitions.
+ */
+public class HashBasedRotateInstanceConstraintApplier implements InstanceConstraintApplier {
+ private static final Logger LOGGER = LoggerFactory.getLogger(HashBasedRotateInstanceConstraintApplier.class);
+
+ private final String _tableNameWithType;
+
+ public HashBasedRotateInstanceConstraintApplier(String tableNameWithType) {
+ _tableNameWithType = tableNameWithType;
+ }
+
+ @Override
+ public Map<Integer, List<InstanceConfig>> applyConstraint(
+ Map<Integer, List<InstanceConfig>> poolToInstanceConfigsMap) {
+ int tableNameHash = Math.abs(_tableNameWithType.hashCode());
+ LOGGER.info("Rotating instances for table: {} with hash: {}", _tableNameWithType, tableNameHash);
+
+ for (Map.Entry<Integer, List<InstanceConfig>> entry : poolToInstanceConfigsMap.entrySet()) {
+ List<InstanceConfig> instanceConfigs = entry.getValue();
+ int numInstanceConfigs = instanceConfigs.size();
+ Preconditions.checkState(numInstanceConfigs > 0, "No instance left in pool: %s", entry.getKey());
+
+ // This should match the default InstancePartitions generated from InstancePartitionsUtils
+ Collections.rotate(instanceConfigs, -(tableNameHash % numInstanceConfigs));
+ }
+
+ return poolToInstanceConfigsMap;
+ }
+}
diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/instance/InstanceAssignmentDriver.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/instance/InstanceAssignmentDriver.java
new file mode 100644
index 0000000..a1fbc40
--- /dev/null
+++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/instance/InstanceAssignmentDriver.java
@@ -0,0 +1,85 @@
+/**
+ * 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.List;
+import java.util.Map;
+import org.apache.helix.model.InstanceConfig;
+import org.apache.pinot.common.config.TableConfig;
+import org.apache.pinot.common.config.instance.InstanceAssignmentConfig;
+import org.apache.pinot.common.config.instance.InstanceAssignmentConfigUtils;
+import org.apache.pinot.common.config.instance.InstanceConstraintConfig;
+import org.apache.pinot.common.config.instance.InstanceReplicaPartitionConfig;
+import org.apache.pinot.common.config.instance.InstanceTagPoolConfig;
+import org.apache.pinot.common.utils.InstancePartitionsType;
+import org.apache.pinot.controller.helix.core.assignment.InstancePartitions;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Driver for the instance assignment.
+ * <p>Instance assignment is performed in 3 steps:
+ * <ul>
+ * <li>Select instances based on the tag/pool configuration</li>
+ * <li>Apply constraints to the instances (optional, multiple constraints can be chained up)</li>
+ * <li>Select instances based on the replica/partition configuration</li>
+ * </ul>
+ */
+public class InstanceAssignmentDriver {
+ private static final Logger LOGGER = LoggerFactory.getLogger(InstanceAssignmentDriver.class);
+
+ private final TableConfig _tableConfig;
+
+ public InstanceAssignmentDriver(TableConfig tableConfig) {
+ _tableConfig = tableConfig;
+ }
+
+ public InstancePartitions assignInstances(InstancePartitionsType instancePartitionsType,
+ List<InstanceConfig> instanceConfigs) {
+ String tableNameWithType = _tableConfig.getTableName();
+ LOGGER.info("Starting {} instance assignment for table: {}", instancePartitionsType, tableNameWithType);
+
+ InstanceAssignmentConfig assignmentConfig =
+ InstanceAssignmentConfigUtils.getInstanceAssignmentConfig(_tableConfig, instancePartitionsType);
+ InstanceTagPoolConfig tagPoolConfig = assignmentConfig.getTagPoolConfig();
+ Preconditions.checkState(tagPoolConfig != null, "Instance tag/pool config is missing");
+ InstanceTagPoolSelector tagPoolSelector = new InstanceTagPoolSelector(tagPoolConfig, tableNameWithType);
+ Map<Integer, List<InstanceConfig>> poolToInstanceConfigsMap = tagPoolSelector.selectInstances(instanceConfigs);
+
+ InstanceConstraintConfig constraintConfig = assignmentConfig.getConstraintConfig();
+ List<InstanceConstraintApplier> constraintAppliers = new ArrayList<>();
+ if (constraintConfig == null) {
+ LOGGER.info("No instance constraint is configured, using default hash-based-rotate instance constraint");
+ constraintAppliers.add(new HashBasedRotateInstanceConstraintApplier(tableNameWithType));
+ }
+ // TODO: support more constraints
+ for (InstanceConstraintApplier constraintApplier : constraintAppliers) {
+ poolToInstanceConfigsMap = constraintApplier.applyConstraint(poolToInstanceConfigsMap);
+ }
+
+ InstanceReplicaPartitionConfig replicaPartitionConfig = assignmentConfig.getReplicaPartitionConfig();
+ Preconditions.checkState(replicaPartitionConfig != null, "Instance replica/partition config is missing");
+ InstanceReplicaPartitionSelector replicaPartitionSelector =
+ new InstanceReplicaPartitionSelector(replicaPartitionConfig, tableNameWithType);
+ return replicaPartitionSelector.selectInstances(poolToInstanceConfigsMap);
+ }
+}
diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/instance/InstanceConstraintApplier.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/instance/InstanceConstraintApplier.java
new file mode 100644
index 0000000..862cedb
--- /dev/null
+++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/instance/InstanceConstraintApplier.java
@@ -0,0 +1,37 @@
+/**
+ * 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 java.util.List;
+import java.util.Map;
+import org.apache.helix.model.InstanceConfig;
+
+
+/**
+ * The instance constraint applier is responsible for filtering out unqualified instances and sorts instances for
+ * picking priority.
+ */
+public interface InstanceConstraintApplier {
+
+ /**
+ * Applies the instance constraint to the given map from pool to instance configs to filter out unqualified instances
+ * and sort instances for picking priority.
+ */
+ Map<Integer, List<InstanceConfig>> applyConstraint(Map<Integer, List<InstanceConfig>> poolToInstanceConfigsMap);
+}
diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/instance/InstanceReplicaPartitionSelector.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/instance/InstanceReplicaPartitionSelector.java
new file mode 100644
index 0000000..8f6976c
--- /dev/null
+++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/instance/InstanceReplicaPartitionSelector.java
@@ -0,0 +1,189 @@
+/**
+ * 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.List;
+import java.util.Map;
+import java.util.TreeMap;
+import org.apache.helix.model.InstanceConfig;
+import org.apache.pinot.common.config.instance.InstanceReplicaPartitionConfig;
+import org.apache.pinot.controller.helix.core.assignment.InstancePartitions;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * The instance replica/partition selector is responsible for selecting the instances for each replica and partition.
+ * <p>NOTE: The replica here refers to the table level replica (replica-group), which is a set of instances that contain
+ * all segments for a table.
+ */
+public class InstanceReplicaPartitionSelector {
+ private static final Logger LOGGER = LoggerFactory.getLogger(InstanceReplicaPartitionSelector.class);
+
+ private final InstanceReplicaPartitionConfig _replicaPartitionConfig;
+ private final String _tableNameWithType;
+
+ public InstanceReplicaPartitionSelector(InstanceReplicaPartitionConfig replicaPartitionConfig,
+ String tableNameWithType) {
+ _replicaPartitionConfig = replicaPartitionConfig;
+ _tableNameWithType = tableNameWithType;
+ }
+
+ /**
+ * Selects instances based on the replica and partition config, and stores the result into the given instance
+ * partitions.
+ */
+ public InstancePartitions selectInstances(Map<Integer, List<InstanceConfig>> poolToInstanceConfigsMap) {
+ int numPools = poolToInstanceConfigsMap.size();
+ Preconditions.checkState(numPools != 0, "No pool qualified for selection");
+
+ int tableNameHash = Math.abs(_tableNameWithType.hashCode());
+ List<Integer> pools = new ArrayList<>(poolToInstanceConfigsMap.keySet());
+ pools.sort(null);
+ LOGGER.info("Starting instance replica/partition selection for table: {} with hash: {} from pools: {}",
+ _tableNameWithType, tableNameHash, pools);
+
+ InstancePartitions instancePartitions = new InstancePartitions(_tableNameWithType);
+
+ if (_replicaPartitionConfig.isReplicaGroupBased()) {
+ // Replica-group based selection
+
+ int numReplicas = _replicaPartitionConfig.getNumReplicas();
+ Preconditions.checkState(numReplicas > 0, "Number of replicas must be positive");
+ Map<Integer, List<Integer>> poolToReplicaIdsMap = new TreeMap<>();
+ for (int replicaId = 0; replicaId < numReplicas; replicaId++) {
+ // Pick one pool for each replica based on the table name hash
+ int pool = pools.get((tableNameHash + replicaId) % numPools);
+ poolToReplicaIdsMap.computeIfAbsent(pool, k -> new ArrayList<>()).add(replicaId);
+ }
+ LOGGER.info("Selecting {} replicas from pool: {} for table: {}", numReplicas, poolToReplicaIdsMap,
+ _tableNameWithType);
+
+ int numInstancesPerReplica = _replicaPartitionConfig.getNumInstancesPerReplica();
+ if (numInstancesPerReplica > 0) {
+ // Check if we have enough instances if number of instances per replica is configured
+ for (Map.Entry<Integer, List<Integer>> entry : poolToReplicaIdsMap.entrySet()) {
+ int pool = entry.getKey();
+ int numInstancesInPool = poolToInstanceConfigsMap.get(pool).size();
+ int numInstancesToSelect = numInstancesPerReplica * entry.getValue().size();
+ Preconditions.checkState(numInstancesToSelect <= numInstancesInPool,
+ "Not enough qualified instances from pool: %s (%s in the pool, asked for %s)", pool, numInstancesInPool,
+ numInstancesToSelect);
+ }
+ } else {
+ // Use as many instances as possible if number of instances per replica is not configured
+ numInstancesPerReplica = Integer.MAX_VALUE;
+ for (Map.Entry<Integer, List<Integer>> entry : poolToReplicaIdsMap.entrySet()) {
+ int pool = entry.getKey();
+ int numReplicasInPool = entry.getValue().size();
+ int numInstancesInPool = poolToInstanceConfigsMap.get(pool).size();
+ Preconditions.checkState(numReplicasInPool <= numInstancesInPool,
+ "Not enough qualified instances from pool: %s, cannot select %s replicas from %s instances", pool,
+ numReplicasInPool, numInstancesInPool);
+ numInstancesPerReplica = Math.min(numInstancesPerReplica, numInstancesInPool / numReplicasInPool);
+ }
+ }
+ LOGGER.info("Selecting {} instances per replica for table: {}", numInstancesPerReplica, _tableNameWithType);
+
+ String[][] replicaIdToInstancesMap = new String[numReplicas][numInstancesPerReplica];
+ for (Map.Entry<Integer, List<Integer>> entry : poolToReplicaIdsMap.entrySet()) {
+ List<InstanceConfig> instanceConfigsInPool = poolToInstanceConfigsMap.get(entry.getKey());
+ List<Integer> replicaIdsInPool = entry.getValue();
+
+ // Use round-robin to assign instances to each replica so that they get instances with similar picking priority
+ int instanceIdInPool = 0;
+ for (int instanceIdInReplica = 0; instanceIdInReplica < numInstancesPerReplica; instanceIdInReplica++) {
+ for (int replicaId : replicaIdsInPool) {
+ replicaIdToInstancesMap[replicaId][instanceIdInReplica] =
+ instanceConfigsInPool.get(instanceIdInPool++).getInstanceName();
+ }
+ }
+ }
+
+ // Assign instances within a replica to one partition if not configured
+ int numPartitions = _replicaPartitionConfig.getNumPartitions();
+ if (numPartitions <= 0) {
+ numPartitions = 1;
+ }
+ // Assign all instances within a replica to each partition if not configured
+ int numInstancesPerPartition = _replicaPartitionConfig.getNumInstancesPerPartition();
+ if (numInstancesPerPartition > 0) {
+ Preconditions.checkState(numInstancesPerPartition <= numInstancesPerReplica,
+ "Number of instances per partition: %s must be smaller or equal to number of instances per replica: %s",
+ numInstancesPerPartition, numInstancesPerReplica);
+ } else {
+ numInstancesPerPartition = numInstancesPerReplica;
+ }
+ LOGGER.info("Selecting {} partitions, {} instances per partition within a replica for table: {}", numPartitions,
+ numInstancesPerPartition, _tableNameWithType);
+
+ // Assign consecutive instances within a replica to each partition.
+ // E.g. (within a replica, 5 instances, 3 partitions, 3 instances per partition)
+ // [i0, i1, i2, i3, i4]
+ // p0 p0 p0 p1 p1
+ // p1 p2 p2 p2
+ for (int replicaId = 0; replicaId < numReplicas; replicaId++) {
+ int instanceIdInReplica = 0;
+ for (int partitionId = 0; partitionId < numPartitions; partitionId++) {
+ List<String> instancesInPartition = new ArrayList<>(numInstancesPerPartition);
+ for (int instanceIdInPartition = 0; instanceIdInPartition < numInstancesPerPartition;
+ instanceIdInPartition++) {
+ instancesInPartition.add(replicaIdToInstancesMap[replicaId][instanceIdInReplica]);
+ instanceIdInReplica = (instanceIdInReplica + 1) % numInstancesPerReplica;
+ }
+ instancesInPartition.sort(null);
+ LOGGER.info("Selecting instances: {} for replica: {}, partition: {} for table: {}", instancesInPartition,
+ replicaId, partitionId, _tableNameWithType);
+ instancePartitions.setInstances(partitionId, replicaId, instancesInPartition);
+ }
+ }
+ } else {
+ // Non-replica-group based selection
+
+ // Pick one pool based on the table name hash
+ int pool = pools.get(tableNameHash % numPools);
+ LOGGER.info("Selecting pool: {} for table: {}", pool, _tableNameWithType);
+ List<InstanceConfig> instanceConfigs = poolToInstanceConfigsMap.get(pool);
+ int numInstanceConfigs = instanceConfigs.size();
+
+ // Assign all instances if not configured
+ int numInstancesToSelect = _replicaPartitionConfig.getNumInstances();
+ if (numInstancesToSelect > 0) {
+ Preconditions.checkState(numInstancesToSelect <= numInstanceConfigs,
+ "Not enough qualified instances from pool: %s (%s in the pool, asked for %s)", pool, numInstanceConfigs,
+ numInstancesToSelect);
+ } else {
+ numInstancesToSelect = numInstanceConfigs;
+ }
+
+ List<String> instancesToSelect = new ArrayList<>(numInstancesToSelect);
+ for (int i = 0; i < numInstancesToSelect; i++) {
+ instancesToSelect.add(instanceConfigs.get(i).getInstanceName());
+ }
+ instancesToSelect.sort(null);
+ LOGGER.info("Selecting instances: {} for table: {}", instancesToSelect, _tableNameWithType);
+ // Set the instances as partition 0 replica 0
+ instancePartitions.setInstances(0, 0, instancesToSelect);
+ }
+
+ return instancePartitions;
+ }
+}
diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/instance/InstanceTagPoolSelector.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/instance/InstanceTagPoolSelector.java
new file mode 100644
index 0000000..5651b89
--- /dev/null
+++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/instance/InstanceTagPoolSelector.java
@@ -0,0 +1,133 @@
+/**
+ * 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;
+
+ public InstanceTagPoolSelector(InstanceTagPoolConfig tagPoolConfig, String tableNameWithType) {
+ _tagPoolConfig = tagPoolConfig;
+ _tableNameWithType = tableNameWithType;
+ }
+
+ /**
+ * 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) {
+ int tableNameHash = Math.abs(_tableNameWithType.hashCode());
+ 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);
+ } else {
+ // Non-pool based selection
+
+ LOGGER.info("Selecting {} instances for table: {}", numCandidateInstances, _tableNameWithType);
+ // Put all instance configs as pool 0
+ poolToInstanceConfigsMap.put(0, candidateInstanceConfigs);
+ }
+ return poolToInstanceConfigsMap;
+ }
+}
diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/RealtimeReplicaGroupSegmentAssignmentStrategy.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/RealtimeReplicaGroupSegmentAssignmentStrategy.java
index bd5126e..05febc8 100644
--- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/RealtimeReplicaGroupSegmentAssignmentStrategy.java
+++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/RealtimeReplicaGroupSegmentAssignmentStrategy.java
@@ -149,8 +149,8 @@ public class RealtimeReplicaGroupSegmentAssignmentStrategy implements SegmentAss
* <p>E.g. (within a replica, 3 servers, 6 partitions)
* <pre>
* "0_0": [i0, i1, i2]
- * p0, p1, p2
- * p3, p4, p5
+ * p0 p1 p2
+ * p3 p4 p5
* </pre>
*/
private List<String> getInstances(InstancePartitions instancePartitions, int partitionId) {
diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/assignment/instance/InstanceAssignmentTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/assignment/instance/InstanceAssignmentTest.java
new file mode 100644
index 0000000..599091d
--- /dev/null
+++ b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/assignment/instance/InstanceAssignmentTest.java
@@ -0,0 +1,526 @@
+/**
+ * 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 java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import org.apache.helix.model.InstanceConfig;
+import org.apache.pinot.common.config.ColumnPartitionConfig;
+import org.apache.pinot.common.config.ReplicaGroupStrategyConfig;
+import org.apache.pinot.common.config.SegmentPartitionConfig;
+import org.apache.pinot.common.config.TableConfig;
+import org.apache.pinot.common.config.TagNameUtils;
+import org.apache.pinot.common.config.TagOverrideConfig;
+import org.apache.pinot.common.config.instance.InstanceAssignmentConfig;
+import org.apache.pinot.common.config.instance.InstanceReplicaPartitionConfig;
+import org.apache.pinot.common.config.instance.InstanceTagPoolConfig;
+import org.apache.pinot.common.utils.CommonConstants.Helix.TableType;
+import org.apache.pinot.common.utils.CommonConstants.Segment.AssignmentStrategy;
+import org.apache.pinot.common.utils.InstancePartitionsType;
+import org.apache.pinot.controller.api.pojos.Instance;
+import org.apache.pinot.controller.helix.core.assignment.InstancePartitions;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.fail;
+
+
+public class InstanceAssignmentTest {
+ private static final String RAW_TABLE_NAME = "myTable";
+ private static final String TENANT_NAME = "tenant";
+ private static final String OFFLINE_TAG = TagNameUtils.getOfflineTagForTenant(TENANT_NAME);
+ private static final String REALTIME_TAG = TagNameUtils.getRealtimeTagForTenant(TENANT_NAME);
+ private static final String SERVER_INSTANCE_ID_PREFIX = "Server_localhost_";
+
+ @Test
+ public void testDefaultOfflineNonReplicaGroup() {
+ TableConfig tableConfig =
+ new TableConfig.Builder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME).setServerTenant(TENANT_NAME).build();
+ InstanceAssignmentDriver driver = new InstanceAssignmentDriver(tableConfig);
+ int numInstances = 10;
+ List<InstanceConfig> instanceConfigs = new ArrayList<>(numInstances);
+ for (int i = 0; i < numInstances; i++) {
+ InstanceConfig instanceConfig = new InstanceConfig(SERVER_INSTANCE_ID_PREFIX + i);
+ instanceConfig.addTag(OFFLINE_TAG);
+ instanceConfigs.add(instanceConfig);
+ }
+
+ // All instances should be assigned as replica 0 partition 0
+ InstancePartitions instancePartitions = driver.assignInstances(InstancePartitionsType.OFFLINE, instanceConfigs);
+ assertEquals(instancePartitions.getNumReplicas(), 1);
+ assertEquals(instancePartitions.getNumPartitions(), 1);
+ List<String> expectedInstances = new ArrayList<>(numInstances);
+ for (int i = 0; i < numInstances; i++) {
+ expectedInstances.add(SERVER_INSTANCE_ID_PREFIX + i);
+ }
+ assertEquals(instancePartitions.getInstances(0, 0), expectedInstances);
+ }
+
+ @Test
+ public void testDefaultOfflineReplicaGroup() {
+ int numReplicas = 3;
+ TableConfig tableConfig =
+ new TableConfig.Builder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME).setServerTenant(TENANT_NAME)
+ .setNumReplicas(numReplicas)
+ .setSegmentAssignmentStrategy(AssignmentStrategy.REPLICA_GROUP_SEGMENT_ASSIGNMENT_STRATEGY).build();
+ int numInstancesPerReplica = 2;
+ ReplicaGroupStrategyConfig replicaGroupStrategyConfig = new ReplicaGroupStrategyConfig();
+ replicaGroupStrategyConfig.setNumInstancesPerPartition(numInstancesPerReplica);
+ tableConfig.getValidationConfig().setReplicaGroupStrategyConfig(replicaGroupStrategyConfig);
+ InstanceAssignmentDriver driver = new InstanceAssignmentDriver(tableConfig);
+ int numInstances = 10;
+ List<InstanceConfig> instanceConfigs = new ArrayList<>(numInstances);
+ for (int i = 0; i < numInstances; i++) {
+ InstanceConfig instanceConfig = new InstanceConfig(SERVER_INSTANCE_ID_PREFIX + i);
+ instanceConfig.addTag(OFFLINE_TAG);
+ instanceConfigs.add(instanceConfig);
+ }
+
+ // Instances should be assigned to 3 replicas with a round-robin fashion, each with 2 instances
+ InstancePartitions instancePartitions = driver.assignInstances(InstancePartitionsType.OFFLINE, instanceConfigs);
+ assertEquals(instancePartitions.getNumReplicas(), numReplicas);
+ assertEquals(instancePartitions.getNumPartitions(), 1);
+ // Instances of index 4 to 7 are not assigned because of the hash-based rotation
+ // Math.abs("myTable_OFFLINE".hashCode()) % 10 = 8
+ // [i8, i9, i0, i1, i2, i3, i4, i5, i6, i7]
+ // r0 r1 r2 r0 r1 r2
+ assertEquals(instancePartitions.getInstances(0, 0),
+ Arrays.asList(SERVER_INSTANCE_ID_PREFIX + 1, SERVER_INSTANCE_ID_PREFIX + 8));
+ assertEquals(instancePartitions.getInstances(0, 1),
+ Arrays.asList(SERVER_INSTANCE_ID_PREFIX + 2, SERVER_INSTANCE_ID_PREFIX + 9));
+ assertEquals(instancePartitions.getInstances(0, 2),
+ Arrays.asList(SERVER_INSTANCE_ID_PREFIX + 0, SERVER_INSTANCE_ID_PREFIX + 3));
+
+ String partitionColumnName = "partition";
+ int numPartitions = 2;
+ replicaGroupStrategyConfig.setPartitionColumn(partitionColumnName);
+ SegmentPartitionConfig segmentPartitionConfig = new SegmentPartitionConfig(
+ Collections.singletonMap(partitionColumnName, new ColumnPartitionConfig("Modulo", numPartitions)));
+ tableConfig.getIndexingConfig().setSegmentPartitionConfig(segmentPartitionConfig);
+
+ // Instances should be assigned to 3 replicas with a round-robin fashion, each with 3 instances, then these 3
+ // instances should be assigned to 2 partitions, each with 2 instances
+ instancePartitions = driver.assignInstances(InstancePartitionsType.OFFLINE, instanceConfigs);
+ assertEquals(instancePartitions.getNumReplicas(), numReplicas);
+ assertEquals(instancePartitions.getNumPartitions(), numPartitions);
+ // Instance of index 7 is not assigned because of the hash-based rotation
+ // Math.abs("myTable_OFFLINE".hashCode()) % 10 = 8
+ // [i8, i9, i0, i1, i2, i3, i4, i5, i6, i7]
+ // r0, r1, r2, r0, r1, r2, r0, r1, r2
+ // r0: [i8, i1, i4]
+ // p0, p0, p1
+ // p1
+ // r1: [i9, i2, i5]
+ // p0, p0, p1
+ // p1
+ // r2: [i0, i3, i6]
+ // p0, p0, p1
+ // p1
+ assertEquals(instancePartitions.getInstances(0, 0),
+ Arrays.asList(SERVER_INSTANCE_ID_PREFIX + 1, SERVER_INSTANCE_ID_PREFIX + 8));
+ assertEquals(instancePartitions.getInstances(1, 0),
+ Arrays.asList(SERVER_INSTANCE_ID_PREFIX + 4, SERVER_INSTANCE_ID_PREFIX + 8));
+ assertEquals(instancePartitions.getInstances(0, 1),
+ Arrays.asList(SERVER_INSTANCE_ID_PREFIX + 2, SERVER_INSTANCE_ID_PREFIX + 9));
+ assertEquals(instancePartitions.getInstances(1, 1),
+ Arrays.asList(SERVER_INSTANCE_ID_PREFIX + 5, SERVER_INSTANCE_ID_PREFIX + 9));
+ assertEquals(instancePartitions.getInstances(0, 2),
+ Arrays.asList(SERVER_INSTANCE_ID_PREFIX + 0, SERVER_INSTANCE_ID_PREFIX + 3));
+ assertEquals(instancePartitions.getInstances(1, 2),
+ Arrays.asList(SERVER_INSTANCE_ID_PREFIX + 0, SERVER_INSTANCE_ID_PREFIX + 6));
+ }
+
+ @Test
+ public void testRealtimeTagOverride() {
+ // Override COMPLETED tag to use OFFLINE tag
+ TagOverrideConfig tagOverrideConfig = new TagOverrideConfig();
+ tagOverrideConfig.setRealtimeConsuming(REALTIME_TAG);
+ tagOverrideConfig.setRealtimeCompleted(OFFLINE_TAG);
+ TableConfig tableConfig =
+ new TableConfig.Builder(TableType.REALTIME).setTableName(RAW_TABLE_NAME).setTagOverrideConfig(tagOverrideConfig)
+ .build();
+ InstanceAssignmentDriver driver = new InstanceAssignmentDriver(tableConfig);
+ int numInstances = 10;
+ List<InstanceConfig> instanceConfigs = new ArrayList<>(numInstances);
+ for (int i = 0; i < numInstances; i++) {
+ InstanceConfig instanceConfig = new InstanceConfig(SERVER_INSTANCE_ID_PREFIX + i);
+ if (i % 2 == 0) {
+ instanceConfig.addTag(REALTIME_TAG);
+ } else {
+ instanceConfig.addTag(OFFLINE_TAG);
+ }
+ instanceConfigs.add(instanceConfig);
+ }
+
+ // All instances with REALTIME tag should be assigned as replica 0 partition 0 for CONSUMING instance partitions
+ InstancePartitions instancePartitions = driver.assignInstances(InstancePartitionsType.CONSUMING, instanceConfigs);
+ assertEquals(instancePartitions.getNumReplicas(), 1);
+ assertEquals(instancePartitions.getNumPartitions(), 1);
+ List<String> expectedInstances = new ArrayList<>(numInstances);
+ for (int i = 0; i < numInstances; i++) {
+ if (i % 2 == 0) {
+ expectedInstances.add(SERVER_INSTANCE_ID_PREFIX + i);
+ }
+ }
+ assertEquals(instancePartitions.getInstances(0, 0), expectedInstances);
+
+ // All instances with OFFLINE tag should be assigned as replica 0 partition 0 for COMPLETED instance partitions
+ instancePartitions = driver.assignInstances(InstancePartitionsType.COMPLETED, instanceConfigs);
+ assertEquals(instancePartitions.getNumReplicas(), 1);
+ assertEquals(instancePartitions.getNumPartitions(), 1);
+ expectedInstances = new ArrayList<>(numInstances);
+ for (int i = 0; i < numInstances; i++) {
+ if (i % 2 != 0) {
+ expectedInstances.add(SERVER_INSTANCE_ID_PREFIX + i);
+ }
+ }
+ assertEquals(instancePartitions.getInstances(0, 0), expectedInstances);
+ }
+
+ @Test
+ public void testPoolBased() {
+ // 10 instances in 2 pools, each with 5 instances
+ int numInstances = 10;
+ int numPools = 2;
+ int numInstancesPerPool = numInstances / numPools;
+ List<InstanceConfig> instanceConfigs = new ArrayList<>(numInstances);
+ for (int i = 0; i < numInstances; i++) {
+ InstanceConfig instanceConfig = new InstanceConfig(SERVER_INSTANCE_ID_PREFIX + i);
+ instanceConfig.addTag(OFFLINE_TAG);
+ int pool = i / numInstancesPerPool;
+ instanceConfig.getRecord()
+ .setMapField(Instance.POOL_KEY, Collections.singletonMap(OFFLINE_TAG, Integer.toString(pool)));
+ instanceConfigs.add(instanceConfig);
+ }
+
+ InstanceAssignmentConfig assignmentConfig = new InstanceAssignmentConfig();
+ // Use all pools
+ InstanceTagPoolConfig tagPoolConfig = new InstanceTagPoolConfig();
+ tagPoolConfig.setTag(OFFLINE_TAG);
+ tagPoolConfig.setPoolBased(true);
+ tagPoolConfig.setNumPools(numPools);
+ assignmentConfig.setTagPoolConfig(tagPoolConfig);
+ // Assign to 2 replicas so that each replica is assigned to one pool
+ int numReplicas = numPools;
+ InstanceReplicaPartitionConfig replicaPartitionConfig = new InstanceReplicaPartitionConfig();
+ replicaPartitionConfig.setReplicaGroupBased(true);
+ replicaPartitionConfig.setNumReplicas(numReplicas);
+ assignmentConfig.setReplicaPartitionConfig(replicaPartitionConfig);
+
+ TableConfig tableConfig = new TableConfig.Builder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME)
+ .setInstanceAssignmentConfigMap(Collections.singletonMap(InstancePartitionsType.OFFLINE, assignmentConfig))
+ .build();
+ InstanceAssignmentDriver driver = new InstanceAssignmentDriver(tableConfig);
+
+ // Math.abs("myTable_OFFLINE".hashCode()) % 2 = 0
+ // All instances in pool 0 should be assigned to replica 0, and all instances in pool 1 should be assigned to
+ // replica 1
+ InstancePartitions instancePartitions = driver.assignInstances(InstancePartitionsType.OFFLINE, instanceConfigs);
+ assertEquals(instancePartitions.getNumReplicas(), numReplicas);
+ assertEquals(instancePartitions.getNumPartitions(), 1);
+ assertEquals(instancePartitions.getInstances(0, 0), Arrays
+ .asList(SERVER_INSTANCE_ID_PREFIX + 0, SERVER_INSTANCE_ID_PREFIX + 1, SERVER_INSTANCE_ID_PREFIX + 2,
+ SERVER_INSTANCE_ID_PREFIX + 3, SERVER_INSTANCE_ID_PREFIX + 4));
+ assertEquals(instancePartitions.getInstances(0, 1), Arrays
+ .asList(SERVER_INSTANCE_ID_PREFIX + 5, SERVER_INSTANCE_ID_PREFIX + 6, SERVER_INSTANCE_ID_PREFIX + 7,
+ SERVER_INSTANCE_ID_PREFIX + 8, SERVER_INSTANCE_ID_PREFIX + 9));
+
+ // Add the third pool with same number of instances
+ numPools = 3;
+ numInstances = numPools * numInstancesPerPool;
+ for (int i = numInstances - numInstancesPerPool; i < numInstances; i++) {
+ InstanceConfig instanceConfig = new InstanceConfig(SERVER_INSTANCE_ID_PREFIX + i);
+ instanceConfig.addTag(OFFLINE_TAG);
+ int pool = numPools - 1;
+ instanceConfig.getRecord()
+ .setMapField(Instance.POOL_KEY, Collections.singletonMap(OFFLINE_TAG, Integer.toString(pool)));
+ instanceConfigs.add(instanceConfig);
+ }
+
+ // Math.abs("myTable_OFFLINE".hashCode()) % 3 = 2
+ // Math.abs("myTable_OFFLINE".hashCode()) % 2 = 0
+ // Pool 0 and 2 will be selected in the pool selection
+ // All instances in pool 0 should be assigned to replica 0, and all instances in pool 2 should be assigned to
+ // replica 1
+ instancePartitions = driver.assignInstances(InstancePartitionsType.OFFLINE, instanceConfigs);
+ assertEquals(instancePartitions.getNumReplicas(), numReplicas);
+ assertEquals(instancePartitions.getNumPartitions(), 1);
+ assertEquals(instancePartitions.getInstances(0, 0), Arrays
+ .asList(SERVER_INSTANCE_ID_PREFIX + 0, SERVER_INSTANCE_ID_PREFIX + 1, SERVER_INSTANCE_ID_PREFIX + 2,
+ SERVER_INSTANCE_ID_PREFIX + 3, SERVER_INSTANCE_ID_PREFIX + 4));
+ assertEquals(instancePartitions.getInstances(0, 1), Arrays
+ .asList(SERVER_INSTANCE_ID_PREFIX + 10, SERVER_INSTANCE_ID_PREFIX + 11, SERVER_INSTANCE_ID_PREFIX + 12,
+ SERVER_INSTANCE_ID_PREFIX + 13, SERVER_INSTANCE_ID_PREFIX + 14));
+
+ // Select all 3 pools in pool selection
+ tagPoolConfig.setNumPools(numPools);
+
+ // Math.abs("myTable_OFFLINE".hashCode()) % 3 = 2
+ // All instances in pool 2 should be assigned to replica 0, and all instances in pool 0 should be assigned to
+ // replica 1
+ instancePartitions = driver.assignInstances(InstancePartitionsType.OFFLINE, instanceConfigs);
+ assertEquals(instancePartitions.getNumReplicas(), numReplicas);
+ assertEquals(instancePartitions.getNumPartitions(), 1);
+ assertEquals(instancePartitions.getInstances(0, 0), Arrays
+ .asList(SERVER_INSTANCE_ID_PREFIX + 10, SERVER_INSTANCE_ID_PREFIX + 11, SERVER_INSTANCE_ID_PREFIX + 12,
+ SERVER_INSTANCE_ID_PREFIX + 13, SERVER_INSTANCE_ID_PREFIX + 14));
+ assertEquals(instancePartitions.getInstances(0, 1), Arrays
+ .asList(SERVER_INSTANCE_ID_PREFIX + 0, SERVER_INSTANCE_ID_PREFIX + 1, SERVER_INSTANCE_ID_PREFIX + 2,
+ SERVER_INSTANCE_ID_PREFIX + 3, SERVER_INSTANCE_ID_PREFIX + 4));
+
+ // Select pool 0 and 1 in pool selection
+ tagPoolConfig.setPools(Arrays.asList(0, 1));
+
+ // Math.abs("myTable_OFFLINE".hashCode()) % 2 = 0
+ // All instances in pool 0 should be assigned to replica 0, and all instances in pool 1 should be assigned to
+ // replica 1
+ instancePartitions = driver.assignInstances(InstancePartitionsType.OFFLINE, instanceConfigs);
+ assertEquals(instancePartitions.getNumReplicas(), numReplicas);
+ assertEquals(instancePartitions.getNumPartitions(), 1);
+ assertEquals(instancePartitions.getInstances(0, 0), Arrays
+ .asList(SERVER_INSTANCE_ID_PREFIX + 0, SERVER_INSTANCE_ID_PREFIX + 1, SERVER_INSTANCE_ID_PREFIX + 2,
+ SERVER_INSTANCE_ID_PREFIX + 3, SERVER_INSTANCE_ID_PREFIX + 4));
+ assertEquals(instancePartitions.getInstances(0, 1), Arrays
+ .asList(SERVER_INSTANCE_ID_PREFIX + 5, SERVER_INSTANCE_ID_PREFIX + 6, SERVER_INSTANCE_ID_PREFIX + 7,
+ SERVER_INSTANCE_ID_PREFIX + 8, SERVER_INSTANCE_ID_PREFIX + 9));
+
+ // Assign instances from 2 pools to 3 replicas
+ numReplicas = numPools;
+ replicaPartitionConfig.setNumReplicas(numReplicas);
+
+ // Math.abs("myTable_OFFLINE".hashCode()) % 2 = 0
+ // [pool0, pool1]
+ // r0 r1
+ // r0
+ // Each replica should have 2 instances assigned
+ // Math.abs("myTable_OFFLINE".hashCode()) % 5 = 3
+ // pool 0: [i3, i4, i0, i1, i2]
+ // r0 r2 r0 r2
+ // pool 1: [i8, i9, i5, i6, i7]
+ // r1 r1
+ instancePartitions = driver.assignInstances(InstancePartitionsType.OFFLINE, instanceConfigs);
+ assertEquals(instancePartitions.getNumReplicas(), numReplicas);
+ assertEquals(instancePartitions.getNumPartitions(), 1);
+ assertEquals(instancePartitions.getInstances(0, 0),
+ Arrays.asList(SERVER_INSTANCE_ID_PREFIX + 0, SERVER_INSTANCE_ID_PREFIX + 3));
+ assertEquals(instancePartitions.getInstances(0, 1),
+ Arrays.asList(SERVER_INSTANCE_ID_PREFIX + 8, SERVER_INSTANCE_ID_PREFIX + 9));
+ assertEquals(instancePartitions.getInstances(0, 2),
+ Arrays.asList(SERVER_INSTANCE_ID_PREFIX + 1, SERVER_INSTANCE_ID_PREFIX + 4));
+ }
+
+ @Test
+ public void testIllegalConfig() {
+ InstanceAssignmentConfig assignmentConfig = new InstanceAssignmentConfig();
+ TableConfig tableConfig = new TableConfig.Builder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME)
+ .setInstanceAssignmentConfigMap(Collections.singletonMap(InstancePartitionsType.OFFLINE, assignmentConfig))
+ .build();
+ InstanceAssignmentDriver driver = new InstanceAssignmentDriver(tableConfig);
+
+ int numInstances = 10;
+ List<InstanceConfig> instanceConfigs = new ArrayList<>(numInstances);
+ for (int i = 0; i < numInstances; i++) {
+ InstanceConfig instanceConfig = new InstanceConfig(SERVER_INSTANCE_ID_PREFIX + i);
+ instanceConfigs.add(instanceConfig);
+ }
+
+ // No instance tag/pool config
+ try {
+ driver.assignInstances(InstancePartitionsType.OFFLINE, instanceConfigs);
+ fail();
+ } catch (IllegalStateException e) {
+ assertEquals(e.getMessage(), "Instance tag/pool config is missing");
+ }
+
+ InstanceTagPoolConfig tagPoolConfig = new InstanceTagPoolConfig();
+ assignmentConfig.setTagPoolConfig(tagPoolConfig);
+
+ // No tag configured
+ try {
+ driver.assignInstances(InstancePartitionsType.OFFLINE, instanceConfigs);
+ fail();
+ } catch (IllegalStateException e) {
+ assertEquals(e.getMessage(), "Tag must be configured");
+ }
+
+ tagPoolConfig.setTag(OFFLINE_TAG);
+
+ // No instance with correct tag
+ try {
+ driver.assignInstances(InstancePartitionsType.OFFLINE, instanceConfigs);
+ fail();
+ } catch (IllegalStateException e) {
+ assertEquals(e.getMessage(), "No enabled instance has the tag: tenant_OFFLINE");
+ }
+
+ for (InstanceConfig instanceConfig : instanceConfigs) {
+ instanceConfig.addTag(OFFLINE_TAG);
+ }
+
+ // No instance replica/partition config
+ try {
+ driver.assignInstances(InstancePartitionsType.OFFLINE, instanceConfigs);
+ fail();
+ } catch (IllegalStateException e) {
+ assertEquals(e.getMessage(), "Instance replica/partition config is missing");
+ }
+
+ InstanceReplicaPartitionConfig replicaPartitionConfig = new InstanceReplicaPartitionConfig();
+ assignmentConfig.setReplicaPartitionConfig(replicaPartitionConfig);
+
+ // All instances should be assigned as replica 0 partition 0
+ InstancePartitions instancePartitions = driver.assignInstances(InstancePartitionsType.OFFLINE, instanceConfigs);
+ assertEquals(instancePartitions.getNumReplicas(), 1);
+ assertEquals(instancePartitions.getNumPartitions(), 1);
+ List<String> expectedInstances = new ArrayList<>(numInstances);
+ for (int i = 0; i < numInstances; i++) {
+ expectedInstances.add(SERVER_INSTANCE_ID_PREFIX + i);
+ }
+ assertEquals(instancePartitions.getInstances(0, 0), expectedInstances);
+
+ // Enable pool
+ tagPoolConfig.setPoolBased(true);
+
+ // No instance has correct pool configured
+ try {
+ driver.assignInstances(InstancePartitionsType.OFFLINE, instanceConfigs);
+ fail();
+ } catch (IllegalStateException e) {
+ assertEquals(e.getMessage(), "No enabled instance has the pool configured for the tag: tenant_OFFLINE");
+ }
+
+ for (int i = 0; i < numInstances; i++) {
+ InstanceConfig instanceConfig = instanceConfigs.get(i);
+ if (i < numInstances / 2) {
+ instanceConfig.getRecord().setMapField(Instance.POOL_KEY, Collections.singletonMap(OFFLINE_TAG, "0"));
+ } else {
+ instanceConfig.getRecord().setMapField(Instance.POOL_KEY, Collections.singletonMap(OFFLINE_TAG, "1"));
+ }
+ }
+
+ // Math.abs("myTable_OFFLINE".hashCode()) % 2 = 0
+ // All instances in pool 0 should be assigned as replica 0 partition 0
+ instancePartitions = driver.assignInstances(InstancePartitionsType.OFFLINE, instanceConfigs);
+ assertEquals(instancePartitions.getNumReplicas(), 1);
+ assertEquals(instancePartitions.getNumPartitions(), 1);
+ expectedInstances.clear();
+ for (int i = 0; i < numInstances / 2; i++) {
+ expectedInstances.add(SERVER_INSTANCE_ID_PREFIX + i);
+ }
+ assertEquals(instancePartitions.getInstances(0, 0), expectedInstances);
+
+ tagPoolConfig.setNumPools(3);
+
+ // Ask for too many pools
+ try {
+ driver.assignInstances(InstancePartitionsType.OFFLINE, instanceConfigs);
+ fail();
+ } catch (IllegalStateException e) {
+ assertEquals(e.getMessage(), "Not enough instance pools (2 in the cluster, asked for 3)");
+ }
+
+ tagPoolConfig.setNumPools(0);
+ tagPoolConfig.setPools(Arrays.asList(0, 2));
+
+ // Ask for pool that does not exist
+ try {
+ driver.assignInstances(InstancePartitionsType.OFFLINE, instanceConfigs);
+ fail();
+ } catch (IllegalStateException e) {
+ assertEquals(e.getMessage(), "Cannot find all instance pools configured: [0, 2]");
+ }
+
+ tagPoolConfig.setPools(null);
+ replicaPartitionConfig.setNumInstances(6);
+
+ // Ask for too many instances
+ try {
+ driver.assignInstances(InstancePartitionsType.OFFLINE, instanceConfigs);
+ fail();
+ } catch (IllegalStateException e) {
+ assertEquals(e.getMessage(), "Not enough qualified instances from pool: 0 (5 in the pool, asked for 6)");
+ }
+
+ replicaPartitionConfig.setNumInstances(0);
+
+ // Enable replica-group
+ replicaPartitionConfig.setReplicaGroupBased(true);
+
+ // Number of replicas must be positive
+ try {
+ driver.assignInstances(InstancePartitionsType.OFFLINE, instanceConfigs);
+ fail();
+ } catch (IllegalStateException e) {
+ assertEquals(e.getMessage(), "Number of replicas must be positive");
+ }
+
+ replicaPartitionConfig.setNumReplicas(11);
+
+ // Ask for too many replicas
+ try {
+ driver.assignInstances(InstancePartitionsType.OFFLINE, instanceConfigs);
+ fail();
+ } catch (IllegalStateException e) {
+ assertEquals(e.getMessage(),
+ "Not enough qualified instances from pool: 0, cannot select 6 replicas from 5 instances");
+ }
+
+ replicaPartitionConfig.setNumReplicas(3);
+ replicaPartitionConfig.setNumInstancesPerReplica(3);
+
+ // Ask for too many instances
+ try {
+ driver.assignInstances(InstancePartitionsType.OFFLINE, instanceConfigs);
+ fail();
+ } catch (IllegalStateException e) {
+ assertEquals(e.getMessage(), "Not enough qualified instances from pool: 0 (5 in the pool, asked for 6)");
+ }
+
+ replicaPartitionConfig.setNumInstancesPerReplica(2);
+ replicaPartitionConfig.setNumInstancesPerPartition(3);
+
+ // Ask for too many instances per partition
+ try {
+ driver.assignInstances(InstancePartitionsType.OFFLINE, instanceConfigs);
+ fail();
+ } catch (IllegalStateException e) {
+ assertEquals(e.getMessage(),
+ "Number of instances per partition: 3 must be smaller or equal to number of instances per replica: 2");
+ }
+
+ replicaPartitionConfig.setNumInstancesPerPartition(0);
+
+ // Math.abs("myTable_OFFLINE".hashCode()) % 5 = 3
+ // pool0: [i3, i4, i0, i1, i2]
+ // r0 r2 r0 r2
+ // pool1: [i8, i9, i5, i6, i7]
+ // r1 r1
+ instancePartitions = driver.assignInstances(InstancePartitionsType.OFFLINE, instanceConfigs);
+ assertEquals(instancePartitions.getNumReplicas(), 3);
+ assertEquals(instancePartitions.getNumPartitions(), 1);
+ assertEquals(instancePartitions.getInstances(0, 0),
+ Arrays.asList(SERVER_INSTANCE_ID_PREFIX + 0, SERVER_INSTANCE_ID_PREFIX + 3));
+ assertEquals(instancePartitions.getInstances(0, 1),
+ Arrays.asList(SERVER_INSTANCE_ID_PREFIX + 8, SERVER_INSTANCE_ID_PREFIX + 9));
+ assertEquals(instancePartitions.getInstances(0, 2),
+ Arrays.asList(SERVER_INSTANCE_ID_PREFIX + 1, SERVER_INSTANCE_ID_PREFIX + 4));
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org