Repository: skiptests/astraea
Pull request #280 (2022-03-13 10:44:18)

Move cost functions out of partitioner folder
1. 將所有cost相關的程式碼從partitioner移出
2. 將beans的操作整合進`ClusterInfo`


Patch:
diff --git a/app/src/main/java/org/astraea/partitioner/ClusterInfo.java b/app/src/main/java/org/astraea/cost/ClusterInfo.java
similarity index 53%
rename from app/src/main/java/org/astraea/partitioner/ClusterInfo.java
rename to app/src/main/java/org/astraea/cost/ClusterInfo.java
index c3e9f87b..c6bccbe4 100644
--- a/app/src/main/java/org/astraea/partitioner/ClusterInfo.java
+++ b/app/src/main/java/org/astraea/cost/ClusterInfo.java
@@ -1,8 +1,14 @@
-package org.astraea.partitioner;
+package org.astraea.cost;
 
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.NoSuchElementException;
 import java.util.stream.Collectors;
+import org.astraea.metrics.HasBeanObject;
 
 public interface ClusterInfo {
 
@@ -26,6 +32,58 @@ public interface ClusterInfo {
             .map(PartitionInfo::of)
             .collect(Collectors.toUnmodifiableList());
       }
+
+      @Override
+      public Collection<HasBeanObject> beans(int brokerId) {
+        return List.of();
+      }
+
+      @Override
+      public Map<Integer, Collection<HasBeanObject>> allBeans() {
+        return Map.of();
+      }
+    };
+  }
+
+  /**
+   * merge the beans into cluster information
+   *
+   * @param cluster cluster information
+   * @param beans extra beans
+   * @return a new cluster information with extra beans
+   */
+  static ClusterInfo of(ClusterInfo cluster, Map<Integer, Collection<HasBeanObject>> beans) {
+    var all = new HashMap<Integer, List<HasBeanObject>>();
+    cluster
+        .allBeans()
+        .forEach((key, value) -> all.computeIfAbsent(key, k -> new ArrayList<>()).addAll(value));
+    beans.forEach((key, value) -> all.computeIfAbsent(key, k -> new ArrayList<>()).addAll(value));
+    return new ClusterInfo() {
+
+      @Override
+      public List<NodeInfo> nodes() {
+        return cluster.nodes();
+      }
+
+      @Override
+      public List<PartitionInfo> availablePartitions(String topic) {
+        return cluster.availablePartitions(topic);
+      }
+
+      @Override
+      public List<PartitionInfo> partitions(String topic) {
+        return cluster.partitions(topic);
+      }
+
+      @Override
+      public Collection<HasBeanObject> beans(int brokerId) {
+        return all.getOrDefault(brokerId, List.of());
+      }
+
+      @Override
+      public Map<Integer, Collection<HasBeanObject>> allBeans() {
+        return Collections.unmodifiableMap(all);
+      }
     };
   }
 
@@ -77,4 +135,14 @@ public interface ClusterInfo {
    * @return A list of partitions
    */
   List<PartitionInfo> partitions(String topic);
+
+  /**
+   * @param brokerId broker id
+   * @return return the metrics of broker. It returns empty collection if the broker id is
+   *     nonexistent
+   */
+  Collection<HasBeanObject> beans(int brokerId);
+
+  /** @return all beans of all brokers */
+  Map<Integer, Collection<HasBeanObject>> allBeans();
 }
diff --git a/app/src/main/java/org/astraea/partitioner/cost/CostFunction.java b/app/src/main/java/org/astraea/cost/CostFunction.java
similarity index 68%
rename from app/src/main/java/org/astraea/partitioner/cost/CostFunction.java
rename to app/src/main/java/org/astraea/cost/CostFunction.java
index 81475d73..58a57bdc 100644
--- a/app/src/main/java/org/astraea/partitioner/cost/CostFunction.java
+++ b/app/src/main/java/org/astraea/cost/CostFunction.java
@@ -1,10 +1,7 @@
-package org.astraea.partitioner.cost;
+package org.astraea.cost;
 
-import java.util.List;
 import java.util.Map;
-import org.astraea.metrics.HasBeanObject;
 import org.astraea.metrics.collector.Fetcher;
-import org.astraea.partitioner.ClusterInfo;
 
 public interface CostFunction {
 
@@ -19,7 +16,7 @@ public interface CostFunction {
    * @param clusterInfo cluster information
    * @return the score of each node. the score range is [0 - 1]
    */
-  Map<Integer, Double> cost(Map<Integer, List<HasBeanObject>> beans, ClusterInfo clusterInfo);
+  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/partitioner/NodeInfo.java b/app/src/main/java/org/astraea/cost/NodeInfo.java
similarity index 97%
rename from app/src/main/java/org/astraea/partitioner/NodeInfo.java
rename to app/src/main/java/org/astraea/cost/NodeInfo.java
index f0f8029a..f7633d4a 100644
--- a/app/src/main/java/org/astraea/partitioner/NodeInfo.java
+++ b/app/src/main/java/org/astraea/cost/NodeInfo.java
@@ -1,4 +1,4 @@
-package org.astraea.partitioner;
+package org.astraea.cost;
 
 import java.util.Objects;
 
diff --git a/app/src/main/java/org/astraea/partitioner/PartitionInfo.java b/app/src/main/java/org/astraea/cost/PartitionInfo.java
similarity index 95%
rename from app/src/main/java/org/astraea/partitioner/PartitionInfo.java
rename to app/src/main/java/org/astraea/cost/PartitionInfo.java
index 197b77b5..50d3ffca 100644
--- a/app/src/main/java/org/astraea/partitioner/PartitionInfo.java
+++ b/app/src/main/java/org/astraea/cost/PartitionInfo.java
@@ -1,4 +1,4 @@
-package org.astraea.partitioner;
+package org.astraea.cost;
 
 public interface PartitionInfo {
 
diff --git a/app/src/main/java/org/astraea/partitioner/cost/ThroughputCost.java b/app/src/main/java/org/astraea/cost/ThroughputCost.java
similarity index 79%
rename from app/src/main/java/org/astraea/partitioner/cost/ThroughputCost.java
rename to app/src/main/java/org/astraea/cost/ThroughputCost.java
index 7464f92b..92670ae5 100644
--- a/app/src/main/java/org/astraea/partitioner/cost/ThroughputCost.java
+++ b/app/src/main/java/org/astraea/cost/ThroughputCost.java
@@ -1,5 +1,6 @@
-package org.astraea.partitioner.cost;
+package org.astraea.cost;
 
+import java.util.Collection;
 import java.util.List;
 import java.util.Map;
 import java.util.stream.Collectors;
@@ -7,15 +8,12 @@ import org.astraea.metrics.HasBeanObject;
 import org.astraea.metrics.collector.Fetcher;
 import org.astraea.metrics.kafka.BrokerTopicMetricsResult;
 import org.astraea.metrics.kafka.KafkaMetrics;
-import org.astraea.partitioner.ClusterInfo;
-import org.astraea.partitioner.NodeInfo;
 
 public class ThroughputCost implements CostFunction {
 
   @Override
-  public Map<Integer, Double> cost(
-      Map<Integer, List<HasBeanObject>> beans, ClusterInfo clusterInfo) {
-    var score = score(beans);
+  public Map<Integer, Double> cost(ClusterInfo clusterInfo) {
+    var score = score(clusterInfo.allBeans());
 
     var max = score.values().stream().mapToDouble(v -> v).max().orElse(1);
 
@@ -24,7 +22,7 @@ public class ThroughputCost implements CostFunction {
         .collect(Collectors.toMap(n -> n, n -> score.getOrDefault(n, 0.0D) / max));
   }
 
-  Map<Integer, Double> score(Map<Integer, List<HasBeanObject>> beans) {
+  Map<Integer, Double> score(Map<Integer, Collection<HasBeanObject>> beans) {
     // TODO: this implementation only consider the oneMinuteRate ...
     return beans.entrySet().stream()
         .collect(
diff --git a/app/src/main/java/org/astraea/metrics/collector/Receiver.java b/app/src/main/java/org/astraea/metrics/collector/Receiver.java
index ccaa103e..95e4eecc 100644
--- a/app/src/main/java/org/astraea/metrics/collector/Receiver.java
+++ b/app/src/main/java/org/astraea/metrics/collector/Receiver.java
@@ -1,6 +1,6 @@
 package org.astraea.metrics.collector;
 
-import java.util.List;
+import java.util.Collection;
 import org.astraea.metrics.HasBeanObject;
 
 /**
@@ -20,7 +20,7 @@ public interface Receiver extends AutoCloseable {
    *
    * @return current mbeans.
    */
-  List<HasBeanObject> current();
+  Collection<HasBeanObject> current();
 
   @Override
   void close();
diff --git a/app/src/main/java/org/astraea/partitioner/Dispatcher.java b/app/src/main/java/org/astraea/partitioner/Dispatcher.java
index 772605c9..199b0a25 100644
--- a/app/src/main/java/org/astraea/partitioner/Dispatcher.java
+++ b/app/src/main/java/org/astraea/partitioner/Dispatcher.java
@@ -4,6 +4,7 @@ import java.util.Map;
 import java.util.stream.Collectors;
 import org.apache.kafka.clients.producer.Partitioner;
 import org.apache.kafka.common.Cluster;
+import org.astraea.cost.ClusterInfo;
 
 public interface Dispatcher extends Partitioner {
   /**
diff --git a/app/src/main/java/org/astraea/partitioner/StrictCostDispatcher.java b/app/src/main/java/org/astraea/partitioner/StrictCostDispatcher.java
index d23be8d1..55128ee5 100644
--- a/app/src/main/java/org/astraea/partitioner/StrictCostDispatcher.java
+++ b/app/src/main/java/org/astraea/partitioner/StrictCostDispatcher.java
@@ -9,10 +9,12 @@ import java.util.Optional;
 import java.util.TreeMap;
 import java.util.stream.Collectors;
 import org.astraea.Utils;
+import org.astraea.cost.ClusterInfo;
+import org.astraea.cost.CostFunction;
+import org.astraea.cost.PartitionInfo;
 import org.astraea.metrics.collector.BeanCollector;
 import org.astraea.metrics.collector.Fetcher;
 import org.astraea.metrics.collector.Receiver;
-import org.astraea.partitioner.cost.CostFunction;
 
 /**
  * this dispatcher scores the nodes by multiples cost functions. Each function evaluate the target
@@ -68,7 +70,7 @@ public class StrictCostDispatcher implements Dispatcher {
     // get scores from all cost functions
     var scores =
         functions.stream()
-            .map(f -> f.cost(beans, clusterInfo))
+            .map(f -> f.cost(ClusterInfo.of(clusterInfo, beans)))
             .collect(Collectors.toUnmodifiableList());
 
     return bestPartition(partitions, scores).map(e -> e.getKey().partition()).orElse(0);
diff --git a/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/NodeLoadClient.java b/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/NodeLoadClient.java
index 861d309f..202890e4 100644
--- a/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/NodeLoadClient.java
+++ b/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/NodeLoadClient.java
@@ -18,11 +18,11 @@ import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantLock;
 import java.util.stream.Collectors;
 import org.astraea.Utils;
+import org.astraea.cost.ClusterInfo;
+import org.astraea.cost.NodeInfo;
 import org.astraea.metrics.collector.Receiver;
 import org.astraea.metrics.java.HasJvmMemory;
 import org.astraea.metrics.kafka.KafkaMetrics;
-import org.astraea.partitioner.ClusterInfo;
-import org.astraea.partitioner.NodeInfo;
 
 /**
  * this clas is responsible for obtaining jmx metrics from BeanCollector and calculating the
@@ -158,7 +158,8 @@ public class NodeLoadClient {
                       .metrics
                       .get(KafkaMetrics.BrokerTopic.BytesInPerSec.name())
                       .current()
-                      .get(0)
+                      .iterator()
+                      .next()
                       .beanObject()
                       .getAttributes()
                       .get("Count")
@@ -169,7 +170,8 @@ public class NodeLoadClient {
                       .metrics
                       .get(KafkaMetrics.BrokerTopic.BytesOutPerSec.name())
                       .current()
-                      .get(0)
+                      .iterator()
+                      .next()
                       .beanObject()
                       .getAttributes()
                       .get("Count")
diff --git a/app/src/main/java/org/astraea/partitioner/smoothPartitioner/SmoothWeightPartitioner.java b/app/src/main/java/org/astraea/partitioner/smoothPartitioner/SmoothWeightPartitioner.java
index 5d94d7a1..b9c39017 100644
--- a/app/src/main/java/org/astraea/partitioner/smoothPartitioner/SmoothWeightPartitioner.java
+++ b/app/src/main/java/org/astraea/partitioner/smoothPartitioner/SmoothWeightPartitioner.java
@@ -17,7 +17,7 @@ import java.util.concurrent.locks.ReentrantLock;
 import java.util.stream.Collectors;
 import org.apache.kafka.clients.producer.Partitioner;
 import org.apache.kafka.common.Cluster;
-import org.astraea.partitioner.ClusterInfo;
+import org.astraea.cost.ClusterInfo;
 import org.astraea.partitioner.Configuration;
 import org.astraea.partitioner.nodeLoadMetric.NodeLoadClient;
 
diff --git a/gradle/wrapper/gradle-wrapper.properties b/gradle/wrapper/gradle-wrapper.properties
index 976bbe0a..7086e5cb 100644
--- a/gradle/wrapper/gradle-wrapper.properties
+++ b/gradle/wrapper/gradle-wrapper.properties
@@ -1,6 +1,6 @@
 distributionBase=GRADLE_USER_HOME
 distributionPath=wrapper/dists
 # please update the gradle version in gradlew too
-distributionUrl=https\://services.gradle.org/distributions/gradle-7.4-bin.zip
+distributionUrl=https\://services.gradle.org/distributions/gradle-7.4.1-bin.zip
 zipStoreBase=GRADLE_USER_HOME
 zipStorePath=wrapper/dists
diff --git a/gradlew b/gradlew
index 017e6341..5abce60e 100755
--- a/gradlew
+++ b/gradlew
@@ -116,7 +116,7 @@ esac
 
 # gradle-wrapper.jar is not in git repo anymore, so we have to download it now.
 if [ ! -e "$APP_HOME/gradle/wrapper/gradle-wrapper.jar" ]; then
-  if ! curl -s -S --retry 3 -L -o "$APP_HOME/gradle/wrapper/gradle-wrapper.jar" "https://raw.githubusercontent.com/gradle/gradle/v7.4.0/gradle/wrapper/gradle-wrapper.jar"; then
+  if ! curl -s -S --retry 3 -L -o "$APP_HOME/gradle/wrapper/gradle-wrapper.jar" "https://raw.githubusercontent.com/gradle/gradle/v7.4.1/gradle/wrapper/gradle-wrapper.jar"; then
     rm -f "$APP_HOME/gradle/wrapper/gradle-wrapper.jar"
     echo "failed to download gradle-wrapper.jar"
     exit 1


Test patch:
diff --git a/app/src/test/java/org/astraea/partitioner/ClusterInfoTest.java b/app/src/test/java/org/astraea/cost/ClusterInfoTest.java
similarity index 61%
rename from app/src/test/java/org/astraea/partitioner/ClusterInfoTest.java
rename to app/src/test/java/org/astraea/cost/ClusterInfoTest.java
index 5d2d6434..de532ed5 100644
--- a/app/src/test/java/org/astraea/partitioner/ClusterInfoTest.java
+++ b/app/src/test/java/org/astraea/cost/ClusterInfoTest.java
@@ -1,7 +1,10 @@
-package org.astraea.partitioner;
+package org.astraea.cost;
 
+import java.util.Collection;
 import java.util.List;
+import java.util.Map;
 import org.apache.kafka.common.Cluster;
+import org.astraea.metrics.HasBeanObject;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
 import org.mockito.Mockito;
@@ -30,4 +33,22 @@ public class ClusterInfoTest {
     Assertions.assertEquals(
         NodeInfo.of(node), clusterInfo.partitions(partition.topic()).get(0).leader());
   }
+
+  @Test
+  void testEmptyBeans() {
+    var clusterInfo = ClusterInfo.of(Mockito.mock(org.apache.kafka.common.Cluster.class));
+    Assertions.assertEquals(0, clusterInfo.allBeans().size());
+    Assertions.assertEquals(0, clusterInfo.beans(19).size());
+  }
+
+  @Test
+  void testBeans() {
+    var beans = Map.of(1, (Collection<HasBeanObject>) List.of(Mockito.mock(HasBeanObject.class)));
+    var origin = Mockito.mock(ClusterInfo.class);
+    Mockito.when(origin.allBeans()).thenReturn(Map.of());
+    var clusterInfo = ClusterInfo.of(origin, beans);
+    Assertions.assertEquals(1, clusterInfo.allBeans().size());
+    Assertions.assertEquals(0, clusterInfo.beans(19).size());
+    Assertions.assertEquals(1, clusterInfo.beans(1).size());
+  }
 }
diff --git a/app/src/test/java/org/astraea/partitioner/NodeInfoTest.java b/app/src/test/java/org/astraea/cost/NodeInfoTest.java
similarity index 93%
rename from app/src/test/java/org/astraea/partitioner/NodeInfoTest.java
rename to app/src/test/java/org/astraea/cost/NodeInfoTest.java
index d44e110f..61b6434e 100644
--- a/app/src/test/java/org/astraea/partitioner/NodeInfoTest.java
+++ b/app/src/test/java/org/astraea/cost/NodeInfoTest.java
@@ -1,4 +1,4 @@
-package org.astraea.partitioner;
+package org.astraea.cost;
 
 import org.apache.kafka.common.Node;
 import org.junit.jupiter.api.Assertions;
diff --git a/app/src/test/java/org/astraea/partitioner/PartitionInfoTest.java b/app/src/test/java/org/astraea/cost/PartitionInfoTest.java
similarity index 95%
rename from app/src/test/java/org/astraea/partitioner/PartitionInfoTest.java
rename to app/src/test/java/org/astraea/cost/PartitionInfoTest.java
index 9a6159f1..3a3aac74 100644
--- a/app/src/test/java/org/astraea/partitioner/PartitionInfoTest.java
+++ b/app/src/test/java/org/astraea/cost/PartitionInfoTest.java
@@ -1,4 +1,4 @@
-package org.astraea.partitioner;
+package org.astraea.cost;
 
 import org.apache.kafka.common.Node;
 import org.junit.jupiter.api.Assertions;
diff --git a/app/src/test/java/org/astraea/partitioner/cost/ThroughputCostTest.java b/app/src/test/java/org/astraea/cost/ThroughputCostTest.java
similarity index 84%
rename from app/src/test/java/org/astraea/partitioner/cost/ThroughputCostTest.java
rename to app/src/test/java/org/astraea/cost/ThroughputCostTest.java
index 48cb0d9d..820ded3d 100644
--- a/app/src/test/java/org/astraea/partitioner/cost/ThroughputCostTest.java
+++ b/app/src/test/java/org/astraea/cost/ThroughputCostTest.java
@@ -1,10 +1,8 @@
-package org.astraea.partitioner.cost;
+package org.astraea.cost;
 
 import java.util.List;
 import java.util.Map;
 import org.astraea.metrics.kafka.BrokerTopicMetricsResult;
-import org.astraea.partitioner.ClusterInfo;
-import org.astraea.partitioner.NodeInfo;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
 import org.mockito.Mockito;
@@ -34,8 +32,9 @@ public class ThroughputCostTest {
 
     var cluster = Mockito.mock(ClusterInfo.class);
     Mockito.when(cluster.nodes()).thenReturn(List.of(node));
+    Mockito.when(cluster.allBeans()).thenReturn(Map.of());
 
-    var cost = throughputCost.cost(Map.of(10, List.of(bean)), cluster);
+    var cost = throughputCost.cost(ClusterInfo.of(cluster, Map.of(10, List.of(bean))));
     Assertions.assertEquals(1, cost.size());
     Assertions.assertEquals(1, cost.get(10));
   }
diff --git a/app/src/test/java/org/astraea/metrics/collector/BeanCollectorTest.java b/app/src/test/java/org/astraea/metrics/collector/BeanCollectorTest.java
index 522383a3..df7f3fd9 100644
--- a/app/src/test/java/org/astraea/metrics/collector/BeanCollectorTest.java
+++ b/app/src/test/java/org/astraea/metrics/collector/BeanCollectorTest.java
@@ -75,7 +75,7 @@ public class BeanCollectorTest {
 
     var c0 = receiver.current();
     Assertions.assertEquals(1, c0.size());
-    var firstObject = c0.get(0);
+    var firstObject = c0.iterator().next();
     sleep(1);
 
     var c1 = receiver.current();
diff --git a/app/src/test/java/org/astraea/partitioner/DispatcherTest.java b/app/src/test/java/org/astraea/partitioner/DispatcherTest.java
index 342d44a5..d15917f4 100644
--- a/app/src/test/java/org/astraea/partitioner/DispatcherTest.java
+++ b/app/src/test/java/org/astraea/partitioner/DispatcherTest.java
@@ -3,6 +3,7 @@ package org.astraea.partitioner;
 import java.util.Map;
 import java.util.Objects;
 import java.util.concurrent.atomic.AtomicInteger;
+import org.astraea.cost.ClusterInfo;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
 
diff --git a/app/src/test/java/org/astraea/partitioner/StrictCostDispatcherTest.java b/app/src/test/java/org/astraea/partitioner/StrictCostDispatcherTest.java
index f194629a..f471e71f 100644
--- a/app/src/test/java/org/astraea/partitioner/StrictCostDispatcherTest.java
+++ b/app/src/test/java/org/astraea/partitioner/StrictCostDispatcherTest.java
@@ -5,10 +5,12 @@ import java.util.Map;
 import java.util.NoSuchElementException;
 import java.util.function.Function;
 import java.util.stream.Collectors;
-import org.astraea.metrics.HasBeanObject;
+import org.astraea.cost.ClusterInfo;
+import org.astraea.cost.CostFunction;
+import org.astraea.cost.NodeInfo;
+import org.astraea.cost.PartitionInfo;
 import org.astraea.metrics.collector.Fetcher;
 import org.astraea.metrics.collector.Receiver;
-import org.astraea.partitioner.cost.CostFunction;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
 import org.mockito.Mockito;
@@ -65,9 +67,8 @@ public class StrictCostDispatcherTest {
     var costFunction =
         new CostFunction() {
           @Override
-          public Map<Integer, Double> cost(
-              Map<Integer, List<HasBeanObject>> beans, ClusterInfo clusterInfo) {
-            return beans.keySet().stream()
+          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));
           }
diff --git a/app/src/test/java/org/astraea/partitioner/nodeLoadMetric/NodeLoadClientTest.java b/app/src/test/java/org/astraea/partitioner/nodeLoadMetric/NodeLoadClientTest.java
index 4b85e90a..fd25eb11 100644
--- a/app/src/test/java/org/astraea/partitioner/nodeLoadMetric/NodeLoadClientTest.java
+++ b/app/src/test/java/org/astraea/partitioner/nodeLoadMetric/NodeLoadClientTest.java
@@ -8,8 +8,8 @@ import java.util.HashMap;
 import java.util.List;
 import java.util.concurrent.atomic.AtomicInteger;
 import org.astraea.Utils;
-import org.astraea.partitioner.ClusterInfo;
-import org.astraea.partitioner.NodeInfo;
+import org.astraea.cost.ClusterInfo;
+import org.astraea.cost.NodeInfo;
 import org.astraea.service.RequireBrokerCluster;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.Assertions;