Repository: skiptests/astraea
Pull request #303 (2022-04-14 18:36:26)

Score on partition
如https://github.com/skiptests/astraea/issues/289#issuecomment-1087713882 
需要對partition進行打分。更甚者，包括replica也要打分。
更改CostFunction的回傳值，CostFunction將會對`org.apache.kafka.common.TopicPartitionReplica`做打分。

原本有許多cost function都是以broker進行打分的，目前我不做更動，只有讓它們符合現在的介面，本質仍然是對broker打分，以後也許可以做更精準的打分(對partition打分)。

另外，在進行修改的過程中發現，partitioner 會希望 cost function 知道發送到哪個 topic，之後要討論"新增topic 到cost function 參數中"。

chinghongfang: 這裡新增介面，`TopicPartitionReplica`內含有 1) topic name, 2) partition number, 3) broker id，並提供回傳`TopicPartitionReplica`物件的方法。

包含4種建立方法，
1. 使用三個參數topic name, partition number, broker id 建立。
2. 使用`org.apache.kafka.common.TopicPartitionReplica`建立。
3. 使用`PartitionInfo`建立leader `TopicPartitionReplica`。
4. 使用`PartitionInfo`建立leader, followers 的 `TopicPartitionReplica`。


chinghongfang: @chia7712 可以麻煩您再review一次PR嘛？感謝您。

Patch:
diff --git a/app/src/main/java/org/astraea/cost/BrokerCost.java b/app/src/main/java/org/astraea/cost/BrokerCost.java
new file mode 100644
index 00000000..da4b145b
--- /dev/null
+++ b/app/src/main/java/org/astraea/cost/BrokerCost.java
@@ -0,0 +1,9 @@
+package org.astraea.cost;
+
+import java.util.Map;
+
+/** Return type of cost function, `HasBrokerCost`. It returns the score of brokers. */
+public interface BrokerCost {
+  /** @return broker-id and its "cost" */
+  Map<Integer, Double> value();
+}
diff --git a/app/src/main/java/org/astraea/cost/CostFunction.java b/app/src/main/java/org/astraea/cost/CostFunction.java
index b5318f00..61db0f9a 100644
--- a/app/src/main/java/org/astraea/cost/CostFunction.java
+++ b/app/src/main/java/org/astraea/cost/CostFunction.java
@@ -1,22 +1,17 @@
 package org.astraea.cost;
 
-import java.util.Map;
 import org.astraea.metrics.collector.Fetcher;
 
+/**
+ * It is meaningless to implement this interface. Instead, we should implement interfaces like
+ * {@link HasBrokerCost} or {@link HasPartitionCost}.
+ */
 public interface CostFunction {
 
   static CostFunction throughput() {
     return new ThroughputCost();
   }
 
-  /**
-   * score all nodes according to passed beans and cluster information.
-   *
-   * @param clusterInfo cluster information
-   * @return the score of each node. the score range is [0 - 1]
-   */
-  Map<Integer, Double> cost(ClusterInfo clusterInfo);
-
   /** @return the metrics getters. Those getters are used to fetch mbeans. */
   Fetcher fetcher();
 }
diff --git a/app/src/main/java/org/astraea/cost/HasBrokerCost.java b/app/src/main/java/org/astraea/cost/HasBrokerCost.java
new file mode 100644
index 00000000..6da60d10
--- /dev/null
+++ b/app/src/main/java/org/astraea/cost/HasBrokerCost.java
@@ -0,0 +1,11 @@
+package org.astraea.cost;
+
+public interface HasBrokerCost extends CostFunction {
+  /**
+   * score all nodes according to passed beans and cluster information.
+   *
+   * @param clusterInfo cluster information
+   * @return the score of each broker. The score ranges in [0 - 1].
+   */
+  BrokerCost brokerCost(ClusterInfo clusterInfo);
+}
diff --git a/app/src/main/java/org/astraea/cost/HasPartitionCost.java b/app/src/main/java/org/astraea/cost/HasPartitionCost.java
new file mode 100644
index 00000000..66e6556d
--- /dev/null
+++ b/app/src/main/java/org/astraea/cost/HasPartitionCost.java
@@ -0,0 +1,11 @@
+package org.astraea.cost;
+
+public interface HasPartitionCost extends CostFunction {
+  /**
+   * score all nodes according to passed beans and cluster information.
+   *
+   * @param clusterInfo cluster information
+   * @return the score of each partition. The score ranges in [0 - 1].
+   */
+  PartitionCost partitionCost(ClusterInfo clusterInfo);
+}
diff --git a/app/src/main/java/org/astraea/cost/LoadCost.java b/app/src/main/java/org/astraea/cost/LoadCost.java
index 3a10b313..1da9df38 100644
--- a/app/src/main/java/org/astraea/cost/LoadCost.java
+++ b/app/src/main/java/org/astraea/cost/LoadCost.java
@@ -12,7 +12,7 @@ import org.astraea.metrics.kafka.BrokerTopicMetricsResult;
 import org.astraea.metrics.kafka.KafkaMetrics;
 import org.astraea.partitioner.nodeLoadMetric.PartitionerUtils;
 
-public class LoadCost implements CostFunction {
+public class LoadCost implements HasBrokerCost {
   private final Map<Integer, BrokerMetric> brokersMetric = new HashMap<>();
   private final Map<String, Double> metricNameAndWeight =
       Map.of(
@@ -23,13 +23,15 @@ public class LoadCost implements CostFunction {
 
   /** Do "Poisson" and "weightPoisson" calculation on "load". And change output to double. */
   @Override
-  public Map<Integer, Double> cost(ClusterInfo clusterInfo) {
+  public BrokerCost brokerCost(ClusterInfo clusterInfo) {
     var load = computeLoad(clusterInfo.allBeans());
 
     // Poisson calculation (-> Poisson -> throughputAbility -> to double)
-    return PartitionerUtils.allPoisson(load).entrySet().stream()
-        .map(e -> Map.entry(e.getKey(), PartitionerUtils.weightPoisson(e.getValue(), 1.0)))
-        .collect(Collectors.toMap(Map.Entry::getKey, e -> (double) e.getValue()));
+    var brokerScore =
+        PartitionerUtils.allPoisson(load).entrySet().stream()
+            .map(e -> Map.entry(e.getKey(), PartitionerUtils.weightPoisson(e.getValue(), 1.0)))
+            .collect(Collectors.toMap(Map.Entry::getKey, e -> (double) e.getValue()));
+    return () -> brokerScore;
   }
 
   /**
diff --git a/app/src/main/java/org/astraea/cost/MemoryWarningCost.java b/app/src/main/java/org/astraea/cost/MemoryWarningCost.java
index 1f0e23a4..51521315 100644
--- a/app/src/main/java/org/astraea/cost/MemoryWarningCost.java
+++ b/app/src/main/java/org/astraea/cost/MemoryWarningCost.java
@@ -9,34 +9,36 @@ import org.astraea.metrics.java.HasJvmMemory;
 import org.astraea.metrics.kafka.KafkaMetrics;
 
 /** Warning on heap usage >= 80%. */
-public class MemoryWarningCost implements CostFunction {
+public class MemoryWarningCost implements HasBrokerCost {
   private final long constructTime = System.currentTimeMillis();
 
   /** @return 1 if the heap usage >= 80%, 0 otherwise. */
   @Override
-  public Map<Integer, Double> cost(ClusterInfo clusterInfo) {
-    return clusterInfo.allBeans().entrySet().stream()
-        .map(
-            e ->
-                Map.entry(
-                    e.getKey(),
-                    e.getValue().stream()
-                        .filter(beanObject -> beanObject instanceof HasJvmMemory)
-                        .map(beanObject -> (HasJvmMemory) beanObject)
-                        .map(
-                            hasJvmMemory -> {
-                              if (Utils.overSecond(constructTime, 10)
-                                  && (hasJvmMemory.heapMemoryUsage().getUsed() + 0.0)
-                                          / (hasJvmMemory.heapMemoryUsage().getMax() + 1)
-                                      >= 0.8) {
-                                return 1.0;
-                              } else {
-                                return 0.0;
-                              }
-                            })
-                        .findAny()
-                        .orElseThrow()))
-        .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+  public BrokerCost brokerCost(ClusterInfo clusterInfo) {
+    var brokerScore =
+        clusterInfo.allBeans().entrySet().stream()
+            .map(
+                e ->
+                    Map.entry(
+                        e.getKey(),
+                        e.getValue().stream()
+                            .filter(beanObject -> beanObject instanceof HasJvmMemory)
+                            .map(beanObject -> (HasJvmMemory) beanObject)
+                            .map(
+                                hasJvmMemory -> {
+                                  if (Utils.overSecond(constructTime, 10)
+                                      && (hasJvmMemory.heapMemoryUsage().getUsed() + 0.0)
+                                              / (hasJvmMemory.heapMemoryUsage().getMax() + 1)
+                                          >= 0.8) {
+                                    return 1.0;
+                                  } else {
+                                    return 0.0;
+                                  }
+                                })
+                            .findAny()
+                            .orElseThrow()))
+            .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+    return () -> brokerScore;
   }
 
   @Override
diff --git a/app/src/main/java/org/astraea/cost/PartitionCost.java b/app/src/main/java/org/astraea/cost/PartitionCost.java
new file mode 100644
index 00000000..e90f7e55
--- /dev/null
+++ b/app/src/main/java/org/astraea/cost/PartitionCost.java
@@ -0,0 +1,22 @@
+package org.astraea.cost;
+
+import java.util.Map;
+
+/** Return type of cost function, `HasPartitionCost`. It returns the score of partitions. */
+public interface PartitionCost {
+  /**
+   * Get the cost of all leader partitions with the given topic name.
+   *
+   * @param topic the topic name we want to query for.
+   * @return the cost of all leader partitions, with respect to the given topic.
+   */
+  Map<TopicPartition, Double> value(String topic);
+
+  /**
+   * Get the cost of all partitions (leader/followers) with the given broker ID.
+   *
+   * @param brokerId the broker we want to query for.
+   * @return the cost of all partitions (leader/followers), with respect to the given broker.
+   */
+  Map<TopicPartition, Double> value(int brokerId);
+}
diff --git a/app/src/main/java/org/astraea/cost/ThroughputCost.java b/app/src/main/java/org/astraea/cost/ThroughputCost.java
index 92670ae5..73c3950f 100644
--- a/app/src/main/java/org/astraea/cost/ThroughputCost.java
+++ b/app/src/main/java/org/astraea/cost/ThroughputCost.java
@@ -9,19 +9,21 @@ import org.astraea.metrics.collector.Fetcher;
 import org.astraea.metrics.kafka.BrokerTopicMetricsResult;
 import org.astraea.metrics.kafka.KafkaMetrics;
 
-public class ThroughputCost implements CostFunction {
+public class ThroughputCost implements HasBrokerCost {
 
   @Override
-  public Map<Integer, Double> cost(ClusterInfo clusterInfo) {
+  public BrokerCost brokerCost(ClusterInfo clusterInfo) {
     var score = score(clusterInfo.allBeans());
 
     var max = score.values().stream().mapToDouble(v -> v).max().orElse(1);
 
-    return clusterInfo.nodes().stream()
-        .map(NodeInfo::id)
-        .collect(Collectors.toMap(n -> n, n -> score.getOrDefault(n, 0.0D) / max));
+    return () ->
+        clusterInfo.nodes().stream()
+            .map(NodeInfo::id)
+            .collect(Collectors.toMap(n -> n, n -> score.getOrDefault(n, 0.0D) / max));
   }
 
+  /* Score by broker. */
   Map<Integer, Double> score(Map<Integer, Collection<HasBeanObject>> beans) {
     // TODO: this implementation only consider the oneMinuteRate ...
     return beans.entrySet().stream()
diff --git a/app/src/main/java/org/astraea/cost/TopicPartition.java b/app/src/main/java/org/astraea/cost/TopicPartition.java
new file mode 100644
index 00000000..abff3f4b
--- /dev/null
+++ b/app/src/main/java/org/astraea/cost/TopicPartition.java
@@ -0,0 +1,50 @@
+package org.astraea.cost;
+
+public interface TopicPartition extends Comparable<TopicPartition> {
+  static TopicPartition of(String topic, int partition) {
+    final int prime = 31;
+    final int hash = (prime + topic.hashCode()) * prime + partition;
+    return new TopicPartition() {
+      @Override
+      public String topic() {
+        return topic;
+      }
+
+      @Override
+      public int partition() {
+        return partition;
+      }
+
+      @Override
+      public int hashCode() {
+        return hash;
+      }
+
+      @Override
+      public boolean equals(Object o) {
+        if (!(o instanceof TopicPartition)) return false;
+        var topicPartition = (TopicPartition) o;
+        return topicPartition.topic().equals(topic()) && topicPartition.partition() == partition();
+      }
+
+      @Override
+      public int compareTo(TopicPartition topicPartition) {
+        if (!this.topic().equals(topicPartition.topic())) {
+          return this.topic().compareTo(topicPartition.topic());
+        } else if (this.partition() != topicPartition.partition()) {
+          return this.partition() - topicPartition.partition();
+        } else {
+          return 0;
+        }
+      }
+    };
+  }
+
+  static TopicPartition of(org.apache.kafka.common.TopicPartition topicPartition) {
+    return TopicPartition.of(topicPartition.topic(), topicPartition.partition());
+  }
+
+  String topic();
+
+  int partition();
+}
diff --git a/app/src/main/java/org/astraea/partitioner/StrictCostDispatcher.java b/app/src/main/java/org/astraea/partitioner/StrictCostDispatcher.java
index 55128ee5..f25426b6 100644
--- a/app/src/main/java/org/astraea/partitioner/StrictCostDispatcher.java
+++ b/app/src/main/java/org/astraea/partitioner/StrictCostDispatcher.java
@@ -11,6 +11,7 @@ import java.util.stream.Collectors;
 import org.astraea.Utils;
 import org.astraea.cost.ClusterInfo;
 import org.astraea.cost.CostFunction;
+import org.astraea.cost.HasBrokerCost;
 import org.astraea.cost.PartitionInfo;
 import org.astraea.metrics.collector.BeanCollector;
 import org.astraea.metrics.collector.Fetcher;
@@ -70,7 +71,9 @@ public class StrictCostDispatcher implements Dispatcher {
     // get scores from all cost functions
     var scores =
         functions.stream()
-            .map(f -> f.cost(ClusterInfo.of(clusterInfo, beans)))
+            .filter(f -> f instanceof HasBrokerCost)
+            .map(f -> (HasBrokerCost) f)
+            .map(f -> f.brokerCost(ClusterInfo.of(clusterInfo, beans)).value())
             .collect(Collectors.toUnmodifiableList());
 
     return bestPartition(partitions, scores).map(e -> e.getKey().partition()).orElse(0);
diff --git a/app/src/main/java/org/astraea/partitioner/smoothPartitioner/SmoothWeightDispatcher.java b/app/src/main/java/org/astraea/partitioner/smoothPartitioner/SmoothWeightDispatcher.java
index eaddabaa..48c03a2f 100644
--- a/app/src/main/java/org/astraea/partitioner/smoothPartitioner/SmoothWeightDispatcher.java
+++ b/app/src/main/java/org/astraea/partitioner/smoothPartitioner/SmoothWeightDispatcher.java
@@ -15,6 +15,7 @@ import java.util.stream.Collectors;
 import org.astraea.Utils;
 import org.astraea.cost.ClusterInfo;
 import org.astraea.cost.CostFunction;
+import org.astraea.cost.HasBrokerCost;
 import org.astraea.cost.LoadCost;
 import org.astraea.cost.MemoryWarningCost;
 import org.astraea.metrics.collector.BeanCollector;
@@ -74,9 +75,11 @@ public class SmoothWeightDispatcher implements Dispatcher {
 
         var compoundScore =
             functions.stream()
+                .filter(f -> f instanceof HasBrokerCost)
+                .map(f -> (HasBrokerCost) f)
                 .map(
                     f -> {
-                      var map = f.cost(ClusterInfo.of(clusterInfo, beans));
+                      var map = f.brokerCost(ClusterInfo.of(clusterInfo, beans)).value();
                       map.replaceAll((k, v) -> v * costFraction.get(f.getClass()));
                       return map;
                     })
@@ -106,7 +109,7 @@ public class SmoothWeightDispatcher implements Dispatcher {
   // Just add all cost function score together
   static Map<Integer, Double> costCompound(
       Map<Integer, Double> identity, Map<Integer, Double> cost) {
-    cost.forEach((ID, value) -> identity.computeIfPresent(ID, (k, v) -> v + value));
+    cost.forEach((id, value) -> identity.computeIfPresent(id, (k, v) -> v + value));
     cost.forEach(identity::putIfAbsent);
     return identity;
   }


Test patch:
diff --git a/app/src/test/java/org/astraea/cost/MemoryWarningCostTest.java b/app/src/test/java/org/astraea/cost/MemoryWarningCostTest.java
index 5485de1f..c38dca8b 100644
--- a/app/src/test/java/org/astraea/cost/MemoryWarningCostTest.java
+++ b/app/src/test/java/org/astraea/cost/MemoryWarningCostTest.java
@@ -4,6 +4,7 @@ import java.lang.management.MemoryUsage;
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import org.astraea.metrics.HasBeanObject;
 import org.astraea.metrics.java.HasJvmMemory;
 import org.junit.jupiter.api.Assertions;
@@ -26,17 +27,33 @@ public class MemoryWarningCostTest {
           public Map<Integer, Collection<HasBeanObject>> allBeans() {
             return Map.of(1, broker1, 2, broker2, 3, broker3);
           }
+
+          @Override
+          public Set<String> topics() {
+            return Set.of("t");
+          }
+
+          @Override
+          public List<PartitionInfo> availablePartitions(String topic) {
+            return List.of(
+                PartitionInfo.of(
+                    "t", 0, NodeInfo.of(1, "host1", 9092), List.of(), List.of(), List.of()),
+                PartitionInfo.of(
+                    "t", 0, NodeInfo.of(2, "host2", 9092), List.of(), List.of(), List.of()),
+                PartitionInfo.of(
+                    "t", 0, NodeInfo.of(3, "host3", 9092), List.of(), List.of(), List.of()));
+          }
         };
 
     var memoryWarning = new MemoryWarningCost();
-    var scores = memoryWarning.cost(clusterInfo);
+    var scores = memoryWarning.brokerCost(clusterInfo).value();
     Assertions.assertEquals(0.0, scores.get(1));
     Assertions.assertEquals(0.0, scores.get(2));
     Assertions.assertEquals(0.0, scores.get(3));
 
     Thread.sleep(10000);
 
-    scores = memoryWarning.cost(clusterInfo);
+    scores = memoryWarning.brokerCost(clusterInfo).value();
     Assertions.assertEquals(0.0, scores.get(1));
     Assertions.assertEquals(1.0, scores.get(2));
     Assertions.assertEquals(0.0, scores.get(3));
diff --git a/app/src/test/java/org/astraea/cost/ThroughputCostTest.java b/app/src/test/java/org/astraea/cost/ThroughputCostTest.java
index 820ded3d..e7982cf6 100644
--- a/app/src/test/java/org/astraea/cost/ThroughputCostTest.java
+++ b/app/src/test/java/org/astraea/cost/ThroughputCostTest.java
@@ -2,6 +2,7 @@ package org.astraea.cost;
 
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import org.astraea.metrics.kafka.BrokerTopicMetricsResult;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
@@ -33,8 +34,12 @@ public class ThroughputCostTest {
     var cluster = Mockito.mock(ClusterInfo.class);
     Mockito.when(cluster.nodes()).thenReturn(List.of(node));
     Mockito.when(cluster.allBeans()).thenReturn(Map.of());
+    Mockito.when(cluster.topics()).thenReturn(Set.of("t"));
+    Mockito.when(cluster.availablePartitions("t"))
+        .thenReturn(List.of(PartitionInfo.of("t", 0, node, List.of(), List.of(), List.of())));
 
-    var cost = throughputCost.cost(ClusterInfo.of(cluster, Map.of(10, List.of(bean))));
+    var cost =
+        throughputCost.brokerCost(ClusterInfo.of(cluster, Map.of(10, List.of(bean)))).value();
     Assertions.assertEquals(1, cost.size());
     Assertions.assertEquals(1, cost.get(10));
   }
diff --git a/app/src/test/java/org/astraea/partitioner/StrictCostDispatcherTest.java b/app/src/test/java/org/astraea/partitioner/StrictCostDispatcherTest.java
index b14a6d57..405daeac 100644
--- a/app/src/test/java/org/astraea/partitioner/StrictCostDispatcherTest.java
+++ b/app/src/test/java/org/astraea/partitioner/StrictCostDispatcherTest.java
@@ -3,10 +3,12 @@ package org.astraea.partitioner;
 import java.util.List;
 import java.util.Map;
 import java.util.NoSuchElementException;
+import java.util.Set;
 import java.util.function.Function;
 import java.util.stream.Collectors;
+import org.astraea.cost.BrokerCost;
 import org.astraea.cost.ClusterInfo;
-import org.astraea.cost.CostFunction;
+import org.astraea.cost.HasBrokerCost;
 import org.astraea.cost.NodeInfo;
 import org.astraea.cost.PartitionInfo;
 import org.astraea.metrics.collector.Fetcher;
@@ -100,12 +102,16 @@ public class StrictCostDispatcherTest {
     Mockito.when(receiver.current()).thenReturn(List.of());
 
     var costFunction =
-        new CostFunction() {
+        new HasBrokerCost() {
           @Override
-          public Map<Integer, Double> cost(ClusterInfo clusterInfo) {
-            return clusterInfo.allBeans().keySet().stream()
-                .collect(
-                    Collectors.toMap(Function.identity(), id -> id.equals(n0.id()) ? 0.9D : 0.5D));
+          public BrokerCost brokerCost(ClusterInfo clusterInfo) {
+            var partitionInfos = clusterInfo.availablePartitions("aa");
+            var brokerCost =
+                clusterInfo.allBeans().keySet().stream()
+                    .collect(
+                        Collectors.toMap(
+                            Function.identity(), id -> id.equals(n0.id()) ? 0.9D : 0.5D));
+            return () -> brokerCost;
           }
 
           @Override
@@ -130,6 +136,7 @@ public class StrictCostDispatcherTest {
 
     // there is no available partition
     Mockito.when(clusterInfo.availablePartitions("aa")).thenReturn(List.of());
+    Mockito.when(clusterInfo.topics()).thenReturn(Set.of("aa"));
     Assertions.assertEquals(0, dispatcher.partition("aa", new byte[0], new byte[0], clusterInfo));
 
     // there is only one available partition