[ASSIGNOR] Implement CostAwareAssignor

common/src/main/java/org/astraea/common/assignor/Assignor.java

@@ -36,23 +36,22 @@
 import org.astraea.common.admin.TopicPartition;
 import org.astraea.common.consumer.ConsumerConfigs;
 import org.astraea.common.cost.HasPartitionCost;
-import org.astraea.common.cost.ReplicaLeaderSizeCost;
+import org.astraea.common.cost.NetworkIngressCost;
 import org.astraea.common.metrics.JndiClient;
 import org.astraea.common.metrics.MBeanClient;
 import org.astraea.common.metrics.collector.MetricStore;
 import org.astraea.common.partitioner.PartitionerUtils;
 
 /** Abstract assignor implementation which does some common work (e.g., configuration). */
 public abstract class Assignor implements ConsumerPartitionAssignor, Configurable {
-  private Configuration config;
+  protected Configuration config;
   public static final String COST_PREFIX = "assignor.cost";
   public static final String JMX_PORT = "jmx.port";
   Function<Integer, Integer> jmxPortGetter =
       (id) -> {
         throw new NoSuchElementException("must define either broker.x.jmx.port or jmx.port");
       };
   HasPartitionCost costFunction = HasPartitionCost.EMPTY;
-  // TODO: metric collector may be configured by user in the future.
   // TODO: need to track the performance when using the assignor in large scale consumers, see
   // https://github.com/skiptests/astraea/pull/1162#discussion_r1036285677
   protected MetricStore metricStore = null;
@@ -68,8 +67,6 @@ public abstract class Assignor implements ConsumerPartitionAssignor, Configurabl
    */
   protected abstract Map<String, List<TopicPartition>> assign(
       Map<String, SubscriptionInfo> subscriptions, ClusterInfo clusterInfo);
-  // TODO: replace the topicPartitions by ClusterInfo after Assignor is able to handle Admin
-  // https://github.com/skiptests/astraea/issues/1409
 
   /**
    * Parse config to get JMX port and cost function type.
@@ -171,7 +168,7 @@ public final void configure(Map<String, ?> configs) {
     var defaultJMXPort = config.integer(JMX_PORT);
     this.costFunction =
         costFunctions.isEmpty()
-            ? HasPartitionCost.of(Map.of(new ReplicaLeaderSizeCost(), 1D))
+            ? HasPartitionCost.of(Map.of(new NetworkIngressCost(config), 1D))
             : HasPartitionCost.of(costFunctions);
     this.jmxPortGetter =
         id ->

common/src/main/java/org/astraea/common/assignor/CostAwareAssignor.java

@@ -0,0 +1,88 @@
+/*
+ * 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.astraea.common.assignor;
+
+import java.time.Duration;
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+import org.astraea.common.Configuration;
+import org.astraea.common.admin.ClusterInfo;
+import org.astraea.common.admin.TopicPartition;
+
+/**
+ * This assignor scores the partitions by cost function(s) that user given. Each cost function
+ * evaluate the partitions' cost in each node by metrics depend on which cost function user use. The
+ * default cost function ranks partitions that are in the same node by NetworkIngressCost{@link
+ * org.astraea.common.cost.NetworkIngressCost}
+ *
+ * <p>The important configs are JMX port. Most cost function need the JMX metrics to score
+ * partitions. Normally, all brokers use the same JMX port, so you could just define the
+ * `jmx.port=12345`. If one of brokers uses different JMX client port, you can define
+ * `broker.1001.jmx.port=3456` (`1001` is the broker id) to replace the value of `jmx.port`. If the
+ * jmx port is undefined, only local mbean client is created for each cost function.
+ */
+public class CostAwareAssignor extends Assignor {
+  protected static final String MAX_RETRY_TIME = "max.retry.time";
+  protected static final String SHUFFLE_TIME = "shuffle.time";
+  Duration maxRetryTime = Duration.ofSeconds(30);
+  Duration shuffleTime = Duration.ofSeconds(5);
+
+  @Override
+  protected Map<String, List<TopicPartition>> assign(
+      Map<String, SubscriptionInfo> subscriptions, ClusterInfo clusterInfo) {
+    var subscribedTopics =
+        subscriptions.values().stream()
+            .map(SubscriptionInfo::topics)
+            .flatMap(Collection::stream)
+            .collect(Collectors.toUnmodifiableSet());
+
+    metricStore.wait(
+        (clusterBean) ->
+            costFunction.partitionCost(clusterInfo, clusterBean).value().values().stream()
+                .noneMatch(v -> Double.isNaN(v)),
+        shuffleTime);
+
+    var clusterBean = metricStore.clusterBean();
+    var partitionCost = costFunction.partitionCost(clusterInfo, clusterBean);
+    var cost =
+        partitionCost.value().entrySet().stream()
+            .filter(e -> subscribedTopics.contains(e.getKey().topic()))
+            .collect(Collectors.toUnmodifiableMap(Map.Entry::getKey, Map.Entry::getValue));
+    var incompatiblePartition = partitionCost.incompatibility();
+
+    var shuffler =
+        Shuffler.randomShuffler(
+            subscriptions,
+            cost,
+            incompatiblePartition,
+            config.duration(SHUFFLE_TIME).orElse(shuffleTime).toMillis());
+    return shuffler.shuffle();
+  }
+
+  @Override
+  protected void configure(Configuration config) {
+    config.duration(MAX_RETRY_TIME).ifPresent(v -> this.maxRetryTime = v);
+    config.duration(SHUFFLE_TIME).ifPresent(v -> this.shuffleTime = v);
+  }
+
+  @Override
+  public String name() {
+    return "costAware";
+  }
+}

common/src/main/java/org/astraea/common/assignor/Hint.java

@@ -37,16 +37,28 @@ static Hint of(Set<Hint> hints) {
   static Hint lowCostHint(
       Map<String, SubscriptionInfo> subscriptions, Map<TopicPartition, Double> partitionCost) {
     return (currentAssignment, tp) -> {
-      var candidates =
+      var consumerPerCost =
           currentAssignment.entrySet().stream()
               .filter(e -> subscriptions.get(e.getKey()).topics().contains(tp.topic()))
               .map(
                   e ->
                       Map.entry(
                           e.getKey(), e.getValue().stream().mapToDouble(partitionCost::get).sum()))
-              .sorted(Map.Entry.comparingByValue())
-              .map(Map.Entry::getKey)
-              .toList();
+              .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+      List<String> candidates;
+
+      if (consumerPerCost.containsValue(0.0))
+        candidates =
+            consumerPerCost.entrySet().stream()
+                .filter(e -> e.getValue() == 0.0)
+                .map(Map.Entry::getKey)
+                .toList();
+      else
+        candidates =
+            consumerPerCost.entrySet().stream()
+                .sorted(Map.Entry.comparingByValue())
+                .map(Map.Entry::getKey)
+                .toList();
 
       return candidates.stream().limit((long) Math.ceil(candidates.size() / 2.0)).toList();
     };
@@ -61,7 +73,8 @@ static Hint incompatibleHint(
               .filter(e -> e.getValue().topics().contains(tp.topic()))
               .map(Map.Entry::getKey)
               .toList();
-      if (incompatibilities.get(tp).isEmpty()) return subscriber;
+      if (incompatibilities.containsKey(tp) && incompatibilities.get(tp).isEmpty())
+        return subscriber;
 
       var candidates =
           currentAssignment.entrySet().stream()
@@ -71,6 +84,7 @@ static Hint incompatibleHint(
                       Map.entry(
                           e.getKey(),
                           e.getValue().stream()
+                              .filter(incompatibilities::containsKey)
                               .filter(p -> incompatibilities.get(p).contains(tp))
                               .count()))
               .collect(

common/src/main/java/org/astraea/common/assignor/Limiter.java

@@ -53,29 +53,14 @@ static Limiter incompatibleLimiter(Map<TopicPartition, Set<TopicPartition>> inco
             == 0;
   }
 
-  static Limiter skewCostLimiter(
-      Map<TopicPartition, Double> partitionCost, Map<String, SubscriptionInfo> subscriptions) {
-    var tmpConsumerCost =
-        subscriptions.keySet().stream()
-            .map(c -> Map.entry(c, 0.0))
-            .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
-    partitionCost.entrySet().stream()
-        .sorted(Map.Entry.comparingByValue())
-        .map(Map.Entry::getKey)
-        .forEach(
-            tp -> {
-              var minCostConsumer =
-                  tmpConsumerCost.entrySet().stream().min(Map.Entry.comparingByValue()).get();
-              minCostConsumer.setValue(minCostConsumer.getValue() + partitionCost.get(tp));
-            });
+  static Limiter skewCostLimiter(Map<TopicPartition, Double> partitionCost) {
     var standardDeviation =
         (Function<Collection<Double>, Double>)
             (vs) -> {
               var average = vs.stream().mapToDouble(c -> c).average().getAsDouble();
               return Math.sqrt(
                   vs.stream().mapToDouble(v -> Math.pow(v - average, 2)).average().getAsDouble());
             };
-    var limit = standardDeviation.apply(tmpConsumerCost.values());
 
     return (combinator) -> {
       var sd =
@@ -84,7 +69,7 @@ static Limiter skewCostLimiter(
                   .map(tps -> tps.stream().mapToDouble(partitionCost::get).sum())
                   .collect(Collectors.toSet()));
 
-      return sd < limit;
+      return sd == 0.0;
     };
   }
 }

common/src/main/java/org/astraea/common/assignor/Shuffler.java

@@ -22,7 +22,6 @@
 import java.util.Set;
 import java.util.function.Function;
 import java.util.stream.Collectors;
-import org.astraea.common.Configuration;
 import org.astraea.common.admin.TopicPartition;
 
 public interface Shuffler {
@@ -32,19 +31,17 @@ static Shuffler randomShuffler(
       Map<String, SubscriptionInfo> subscriptions,
       Map<TopicPartition, Double> partitionCost,
       Map<TopicPartition, Set<TopicPartition>> incompatible,
-      Configuration config) {
+      long shuffleTime) {
     var limiters =
         Limiter.of(
             Set.of(
-                Limiter.skewCostLimiter(partitionCost, subscriptions),
-                Limiter.incompatibleLimiter(incompatible)));
+                Limiter.skewCostLimiter(partitionCost), Limiter.incompatibleLimiter(incompatible)));
     var hints =
         Hint.of(
             Set.of(
                 Hint.lowCostHint(subscriptions, partitionCost),
                 Hint.incompatibleHint(subscriptions, incompatible)));
     var generator = Generator.randomGenerator(subscriptions, partitionCost, hints);
-    var shuffleTime = config.duration("shuffle.time").get().toMillis();
     var standardDeviation =
         (Function<Map<String, List<TopicPartition>>, Double>)
             (combinator) -> {
@@ -79,7 +76,6 @@ static Shuffler randomShuffler(
         }
         rejectedCombinators.add(combinator);
       }
-
       return result == null
           ? rejectedCombinators.stream()
               .map(c -> Map.entry(c, standardDeviation.apply(c)))

common/src/test/java/org/astraea/common/assignor/GeneratorTest.java

@@ -0,0 +1,130 @@
+/*
+ * 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.astraea.common.assignor;
+
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
+import java.util.function.BiFunction;
+import org.astraea.common.admin.TopicPartition;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+public class GeneratorTest {
+  @Test
+  void testCombinator() {
+    var checkSubscriptionAndAssignment =
+        (BiFunction<Map<String, List<TopicPartition>>, Map<String, SubscriptionInfo>, Boolean>)
+            (assignment, subscription) ->
+                assignment.entrySet().stream()
+                    .allMatch(
+                        e ->
+                            e.getValue().stream()
+                                .allMatch(
+                                    tp ->
+                                        subscription
+                                            .get(e.getKey())
+                                            .topics()
+                                            .contains(tp.topic())));
+
+    // test same subscription
+    var subscription =
+        Map.of(
+            "c1",
+            new SubscriptionInfo(List.of("t1", "t2", "t3"), null),
+            "c2",
+            new SubscriptionInfo(List.of("t1", "t2", "t3"), null));
+    var cost =
+        Map.of(
+            TopicPartition.of("t1", 0),
+            0.2,
+            TopicPartition.of("t1", 1),
+            0.4,
+            TopicPartition.of("t2", 0),
+            0.2,
+            TopicPartition.of("t3", 0),
+            0.4);
+    var randomGenerator =
+        Generator.randomGenerator(subscription, cost, Hint.lowCostHint(subscription, cost));
+    var result = randomGenerator.get();
+    var partitionsInResult = result.values().stream().flatMap(Collection::stream).toList();
+    var partitionInCluster = cost.keySet();
+    Assertions.assertEquals(partitionInCluster.size(), partitionsInResult.size());
+    Assertions.assertTrue(partitionsInResult.containsAll(partitionInCluster));
+    Assertions.assertTrue(checkSubscriptionAndAssignment.apply(result, subscription));
+    result.forEach((c, tps) -> Assertions.assertFalse(tps.isEmpty()));
+
+    // Validate that unsubscribed partitions are not be assigned to the consumer.
+    subscription =
+        Map.of(
+            "c1",
+            new SubscriptionInfo(List.of("t1"), null),
+            "c2",
+            new SubscriptionInfo(List.of("t1", "t2", "t3"), null));
+    randomGenerator =
+        Generator.randomGenerator(subscription, cost, Hint.lowCostHint(subscription, cost));
+    result = randomGenerator.get();
+    partitionsInResult = result.values().stream().flatMap(Collection::stream).toList();
+
+    Assertions.assertEquals(partitionInCluster.size(), partitionsInResult.size());
+    Assertions.assertTrue(partitionsInResult.containsAll(partitionInCluster));
+    Assertions.assertTrue(checkSubscriptionAndAssignment.apply(result, subscription));
+
+    // Every consumer get at least one partition
+    subscription =
+        Map.of(
+            "c1",
+            new SubscriptionInfo(List.of("t1", "t2", "t3"), null),
+            "c2",
+            new SubscriptionInfo(List.of("t1", "t2", "t3"), null),
+            "c3",
+            new SubscriptionInfo(List.of("t1", "t2", "t3"), null),
+            "c4",
+            new SubscriptionInfo(List.of("t1", "t2", "t3"), null));
+    randomGenerator =
+        Generator.randomGenerator(subscription, cost, Hint.lowCostHint(subscription, cost));
+    result = randomGenerator.get();
+    result.forEach((c, tps) -> Assertions.assertEquals(1, tps.size()));
+
+    // Validate there is no consumer was assigned more than 1 partition
+    cost =
+        Map.of(
+            TopicPartition.of("t1", 0),
+            0.2,
+            TopicPartition.of("t2", 0),
+            0.2,
+            TopicPartition.of("t3", 0),
+            0.4);
+    randomGenerator =
+        Generator.randomGenerator(subscription, cost, Hint.lowCostHint(subscription, cost));
+    result = randomGenerator.get();
+    result.forEach((c, tps) -> Assertions.assertTrue(tps.size() <= 1));
+
+    subscription =
+        Map.of(
+            "c1",
+            new SubscriptionInfo(List.of("t1"), null),
+            "c2",
+            new SubscriptionInfo(List.of("t2"), null),
+            "c3",
+            new SubscriptionInfo(List.of("t3"), null));
+    randomGenerator =
+        Generator.randomGenerator(subscription, cost, Hint.lowCostHint(subscription, cost));
+    result = randomGenerator.get();
+    Assertions.assertTrue(checkSubscriptionAndAssignment.apply(result, subscription));
+  }
+}