Repository: skiptests/astraea
Pull request #291 (2022-03-30 08:44:52)

Rewrite `SmoothWeightPartitioner`
此PR重構了原來的`SmoothWeightPartitioner`，使用`CostFunction`實做。並做了些許的變動：

1. 省去部份初始判斷(`lastTime==-1...`)
2. 把時間判斷集中(原本是`NodeLoadClient`, `SmoothWeightPartitioner`分別做"1秒1更新")
3. 省去還未實做的功能(如：`throughputAbility`, `standardDeviation`...等功能)
4. 算法切分不同，所以unit test另外寫了一份。

另外改寫的過程中發現了很多問題：

1. 現在算法結果難以使用0~1表示
2. 許多型別轉換(是否必要)
3. 參數重複參考(是否必要)
4. `CostFunction`回傳值的 "意義"

@wycccccc 看看算法是否須要再做調整？


chia7712: 我想先討論一下框架的部分：

> 現在算法結果難以使用0~1表示

可否多描述一下困難點？

> CostFunction回傳值的 "意義"

1. 給每個節點評分，數值越高代表越好
2. 正規化到0 - 1之間，如此在最後量化時才方便加權取得總分
chinghongfang: > 現在算法結果難以使用0~1表示

主要是看到原來的Smooth運算要使用int，不確定換成double來計算還會不會符合原意。
另外，現在的算法十分的複雜，有時需要用int運算，有時要轉回double，有時會出現負值。如果這些轉換不是必要的話，希望可以都使用double？

> CostFunction回傳值的 "意義"

現在有想到一個可能的問題，因為值都被正規化，大家的分數都會被壓到0~1之間，也看不清原來值得大小。
例如：
這是在我們現在的算法發現的，但以後也可能會發生
broker1:
BytesInPersec 10Bytes
BytesOutPersec 5GBytes
broker2:
BytesInPersec 1Bytes
BytesOutPersec 10GBytes
那麼，兩者被正規化以後，只能看出彼此間的倍數關係。所以雖然現在可以明顯看出誰的工作量較大(broker2)，但正規化後反而會認為broker1較大，如果兩項指標的權重相同的話。
chia7712: > 那麼，兩者被正規化以後，只能看出彼此間的倍數關係。

這邊正規指得是講彼此的評分擺到0-1這個尺度上，至於是線性、等差、指數則是cost function自己決定。所以不一定只能呈現倍數關係。

> 但正規化後反而會認為broker1較大，如果兩項指標的權重相同的話。

要用哪些cost以及哪些cost比較重要（權重）會是之後調校的一個重點，如果像你說的權重弄得一樣，導致天秤失衡，那就是代表權重不能一樣，而不是正規化有問題（而且如何正規化是cost function可自行決定的部分）
chia7712: > 另外，現在的算法十分的複雜，有時需要用int運算，有時要轉回double，有時會出現負值。如果這些轉換不是必要的話，希望可以都使用double？

計算會否失真可能要請 @wycccccc 幫忙確認，或是可以留下一隻PR
chinghongfang: > 這邊正規指得是講彼此的評分擺到0-1這個尺度上，至於是線性、等差、指數則是cost function自己決定。所以不一定只能呈現倍數關係。

瞭解，所以是不固定回傳值，把這個意義留給cost function決定。
cost function就只有上面提到的兩點(0~1 和 越高越好)要求。
所以可見未來可能會有多種cost function，取相同資訊(e.g. BytesInPerSec)，但使用不同的正規化方法，對吧？

> 要用哪些cost以及哪些cost比較重要（權重）會是之後調校的一個重點，如果像你說的權重弄得一樣，導致天秤失衡，那就是代表權重不能一樣，而不是正規化有問題（而且如何正規化是cost function可自行決定的部分）

是的，不過上面的例子就算調整權重，也沒辦法解決，因為數據反過來發生，狀況還是會產生。(除非權重是動態改變的)
所以我想應該要從cost function的回傳值下手，就上面的說法，也許要使用其他正規化方法，不要讓它隱藏重要的資料。

chia7712: > 不過上面的例子就算調整權重，也沒辦法解決，因為數據反過來發生，狀況還是會產生。(除非權重是動態改變的)

不好意思，我有點捕捉不到意思。在你給的例子中一個進資料較多、一個出資料較多。你是想討論達到什麼樣的平衡？如果目標是進資料的平衡，那進資料的權重會往上拉，反之則拉高出資料的權重。

如果你是想平衡「吞吐量」（不管讀寫），那就該考慮使用只考慮吞吐量的cost function

如果我誤會你的意思還請補充
chinghongfang: > 在你給的例子中一個進資料較多、一個出資料較多。你是想討論達到什麼樣的平衡？

確實，應該要先想我們想要達成什麼樣的平衡。
我再思考看看。十分感謝。
chia7712: > 現在有想到一個可能的問題，因為值都被正規化，大家的分數都會被壓到0~1之間，也看不清原來值得大小。

這個部分還請做最後確認，現在的方式會不會有所謂"看不清原來值得大小"的狀況

Patch:
diff --git a/README.md b/README.md
index 7371e018..1e97b1c1 100644
--- a/README.md
+++ b/README.md
@@ -248,14 +248,15 @@ Run the benchmark from source
   The time units can be "days", "day", "h", "m", "s", "ms", "us", "ns".
   e.g. "--run.until 1m" or "--run.until 89242records" Default: 1000records
 8. --record.size: the (bound of) record size in byte. Default: 1 KiB
-9. --prop.file: the path to property file.
-10. --partitioner: the partitioner to use in producers
-11. --createCSV: put the metrics into a csv file if this flag is set. Default: false
-12. --compression: the compression algorithm used by producer. Available algorithm are none, gzip, snappy, lz4, and zstd. Default: (NONE)
-13. --jmx.servers: the jmx server addresses of the brokers 
-14. --key.distribution: name of the distribution on key. Available distribution names: "uniform", "zipfian", "latest", "fixed". Default: (No key)
-15. --size.distribution: name of the distribution on value size. Available distribution names: "uniform", "zipfian", "latest", "fixed". Default: "uniform"
-16. --specify.broker: list of broker IDs to produce records to. Default: (Do Not Specify)
+9. --fixed.size: the flag to let all records have the same size
+10. --prop.file: the path to property file.
+11. --partitioner: the partitioner to use in producers
+12. --jmx.servers: the jmx server addresses of the brokers 
+13. --key.distribution: name of the distribution on key. Available distribution names: "uniform", "zipfian", "latest", "fixed". Default: (No key)
+14. --size.distribution: name of the distribution on value size. Available distribution names: "uniform", "zipfian", "latest", "fixed". Default: "uniform"
+15. --specify.broker: list of broker IDs to produce records to. Default: (Do Not Specify)
+16. --report.path: A path to place the report file. Default: (no report)
+17. --report.format: Select output file format. Available format: "csv", "json".
 
 ---
 
diff --git a/app/src/main/java/org/astraea/argument/PathField.java b/app/src/main/java/org/astraea/argument/PathField.java
new file mode 100644
index 00000000..1cced1e6
--- /dev/null
+++ b/app/src/main/java/org/astraea/argument/PathField.java
@@ -0,0 +1,11 @@
+package org.astraea.argument;
+
+import java.nio.file.FileSystems;
+import java.nio.file.Path;
+
+public class PathField extends Field<Path> {
+  @Override
+  public Path convert(String value) {
+    return FileSystems.getDefault().getPath(value);
+  }
+}
diff --git a/app/src/main/java/org/astraea/cost/LoadCost.java b/app/src/main/java/org/astraea/cost/LoadCost.java
new file mode 100644
index 00000000..3a10b313
--- /dev/null
+++ b/app/src/main/java/org/astraea/cost/LoadCost.java
@@ -0,0 +1,180 @@
+package org.astraea.cost;
+
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+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.nodeLoadMetric.PartitionerUtils;
+
+public class LoadCost implements CostFunction {
+  private final Map<Integer, BrokerMetric> brokersMetric = new HashMap<>();
+  private final Map<String, Double> metricNameAndWeight =
+      Map.of(
+          KafkaMetrics.BrokerTopic.BytesInPerSec.metricName(),
+          0.5,
+          KafkaMetrics.BrokerTopic.BytesOutPerSec.metricName(),
+          0.5);
+
+  /** Do "Poisson" and "weightPoisson" calculation on "load". And change output to double. */
+  @Override
+  public Map<Integer, Double> cost(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()));
+  }
+
+  /**
+   * The result is computed by "BytesInPerSec.count" and "BytesOutPerSec.count".
+   *
+   * <ol>
+   *   <li>We normalize the two metric (by divide sum of each metric).
+   *   <li>We compute the sum on the two metric with a specific weight.
+   *   <li>Compare the weighted sum with two boundary(0.5*avg and 1.5*avg) to get the "load" {0,1,2}
+   *   <li>Sum up the historical "load" and divide by 20 the get the score.
+   * </ol>
+   *
+   * <p>e.g. We have 3 brokers with information:<br>
+   * broker1: BytesInPerSec.count=50, BytesOutPerSec.count=15 (<1,50000,21000>)<br>
+   * broker2: BytesInPerSec.count=100, BytesOutPerSec.count=2 (<2,100000,2000>)<br>
+   * broker3: BytesInPerSec.count=200, BytesOutPerSec.count=1 (<3,200000,1000>)<br>
+   *
+   * <ol>
+   *   <li>Normalize: about <1, 1/7, 7/8> <2, 2/7, 1/12> <3, 4/7, 1/24>
+   *   <li>WeightedSum: about <1, 57/112> <2, 31/168> <3, 103/336>
+   *   <li>Load:<1,2><2,1><3,1>
+   * </ol>
+   *
+   * broker1 score: 2<br>
+   * broker2 score: 1<br>
+   * broker3 score: 1<br>
+   *
+   * @return brokerID with "(sum of load)".
+   */
+  Map<Integer, Integer> computeLoad(Map<Integer, Collection<HasBeanObject>> allBeans) {
+    // Store mbean data into local stat("brokersMetric").
+    allBeans.forEach(
+        (brokerID, value) -> {
+          if (!brokersMetric.containsKey(brokerID))
+            brokersMetric.put(brokerID, new BrokerMetric(brokerID));
+          value.stream()
+              .filter(hasBeanObject -> hasBeanObject instanceof BrokerTopicMetricsResult)
+              .map(hasBeanObject -> (BrokerTopicMetricsResult) hasBeanObject)
+              .forEach(
+                  result ->
+                      brokersMetric
+                          .get(brokerID)
+                          .updateCount(
+                              result.beanObject().getProperties().get("name"), result.count()));
+        });
+
+    // Sum all count with the same mbean name
+    var total =
+        brokersMetric.values().stream()
+            .map(brokerMetric -> brokerMetric.currentCount)
+            .reduce(
+                new HashMap<>(),
+                (accumulate, currentCount) -> {
+                  currentCount.forEach(
+                      (name, count) ->
+                          accumulate.put(name, accumulate.getOrDefault(name, 0L) + count));
+                  return accumulate;
+                });
+
+    // Reduce all count for all brokers' mbean current count by "divide total" (current/total). And
+    // then get the weighted sum according to predefined wights "metricNameAndWeight".
+    // It is called "brokerSituation" in original name.
+    var weightedSum =
+        brokersMetric.entrySet().stream()
+            .collect(
+                Collectors.toMap(
+                    Map.Entry::getKey,
+                    entry ->
+                        entry.getValue().currentCount.entrySet().stream()
+                            .mapToDouble(
+                                nameAndCount ->
+                                    (double) (nameAndCount.getValue() + 1)
+                                        / (total.getOrDefault(nameAndCount.getKey(), 1L) + 1)
+                                        * metricNameAndWeight.getOrDefault(
+                                            nameAndCount.getKey(), 0.0))
+                            .sum()));
+
+    var avgWeightedSum =
+        weightedSum.values().stream().mapToDouble(d -> d).sum() / brokersMetric.size();
+
+    // Record the "load" into local stat("brokersMetric").
+    weightedSum.forEach(
+        (brokerID, theWeightedSum) -> {
+          if (theWeightedSum < avgWeightedSum * 0.5) brokersMetric.get(brokerID).updateLoad(0);
+          else if (theWeightedSum < avgWeightedSum * 1.5) brokersMetric.get(brokerID).updateLoad(1);
+          else brokersMetric.get(brokerID).updateLoad(2);
+        });
+
+    // The sum of historical and current load.
+    return brokersMetric.entrySet().stream()
+        .collect(
+            Collectors.toMap(
+                Map.Entry::getKey, entry -> entry.getValue().load.stream().mapToInt(i -> i).sum()));
+  }
+
+  /** @return the metrics getters. Those getters are used to fetch mbeans. */
+  @Override
+  public Fetcher fetcher() {
+    return client ->
+        List.of(
+            KafkaMetrics.BrokerTopic.BytesInPerSec.fetch(client),
+            KafkaMetrics.BrokerTopic.BytesOutPerSec.fetch(client));
+  }
+
+  private static class BrokerMetric {
+    private final int brokerID;
+
+    // mbean data. They are:
+    // ("BytesInPerSec", BytesInPerSec.count), ("BytesOutPerSec", ByteOutPerSec.count)
+    private final Map<String, Long> accumulateCount = new HashMap<>();
+    private final Map<String, Long> currentCount = new HashMap<>();
+
+    // Record the latest 10 numbers only.
+    private final List<Integer> load =
+        IntStream.range(0, 10).mapToObj(i -> 0).collect(Collectors.toList());
+    private int loadIndex = 0;
+
+    BrokerMetric(int brokerID) {
+      this.brokerID = brokerID;
+    }
+
+    /**
+     * This method records the difference between last update and current given "count" e.g.
+     *
+     * <p>At time 1: updateCount("ByteInPerSec", 100);<br>
+     * At time 3: updateCount("ByteOutPerSec", 20);<br>
+     * At time 20: updateCount("ByteInPerSec", 150);<br>
+     * At time 22: updateCount("ByteOutPerSec", 90);<br>
+     * At time 35: updateCount("ByteInPerSec", 170)<br>
+     * Then in time [20, 35), currentCount.get("ByteInPerSec") is 50
+     *
+     * <p>
+     */
+    void updateCount(String domainName, Long count) {
+      currentCount.put(domainName, count - accumulateCount.getOrDefault(domainName, 0L));
+      accumulateCount.put(domainName, count);
+    }
+
+    /**
+     * This method record input data into a list. This list contains the latest 10 record. Each time
+     * it is called, the current index, "loadIndex", is increased by 1.
+     */
+    void updateLoad(Integer load) {
+      this.load.set(loadIndex, load);
+      loadIndex = (loadIndex + 1) % 10;
+    }
+  }
+}
diff --git a/app/src/main/java/org/astraea/cost/MemoryWarningCost.java b/app/src/main/java/org/astraea/cost/MemoryWarningCost.java
new file mode 100644
index 00000000..1f0e23a4
--- /dev/null
+++ b/app/src/main/java/org/astraea/cost/MemoryWarningCost.java
@@ -0,0 +1,46 @@
+package org.astraea.cost;
+
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+import org.astraea.Utils;
+import org.astraea.metrics.collector.Fetcher;
+import org.astraea.metrics.java.HasJvmMemory;
+import org.astraea.metrics.kafka.KafkaMetrics;
+
+/** Warning on heap usage >= 80%. */
+public class MemoryWarningCost implements CostFunction {
+  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));
+  }
+
+  @Override
+  public Fetcher fetcher() {
+    return client -> List.of(KafkaMetrics.Host.jvmMemory(client));
+  }
+}
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 202890e4..add80c70 100644
--- a/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/NodeLoadClient.java
+++ b/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/NodeLoadClient.java
@@ -184,7 +184,7 @@ public class NodeLoadClient {
           broker.output =
               Integer.parseInt(String.valueOf((currentOutput - broker.lastOutPut) / 1024));
           broker.lastOutPut = currentOutput;
-          broker.input = Integer.parseInt(String.valueOf((currentInput - broker.lastInput) / 1024));
+          broker.input = (int) ((currentInput - broker.lastInput) / 1024);
           broker.lastInput = currentInput;
           Objects.requireNonNull(broker.metrics.get("Memory"));
           broker.jvm =
diff --git a/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/PartitionerUtils.java b/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/PartitionerUtils.java
index b557cd74..06e7bc5c 100644
--- a/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/PartitionerUtils.java
+++ b/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/PartitionerUtils.java
@@ -10,7 +10,7 @@ import java.util.Properties;
 public class PartitionerUtils {
   private PartitionerUtils() {}
 
-  public static HashMap<Integer, Double> allPoisson(Map<Integer, Integer> overLoadCount) {
+  public static Map<Integer, Double> allPoisson(Map<Integer, Integer> overLoadCount) {
     var poissonMap = new HashMap<Integer, Double>();
     var lambda = avgLoadCount(overLoadCount);
     overLoadCount.forEach((nodeID, count) -> poissonMap.put(nodeID, doPoisson(lambda, count)));
diff --git a/app/src/main/java/org/astraea/partitioner/smoothPartitioner/SmoothWeight.java b/app/src/main/java/org/astraea/partitioner/smoothPartitioner/SmoothWeight.java
new file mode 100644
index 00000000..b28aaffa
--- /dev/null
+++ b/app/src/main/java/org/astraea/partitioner/smoothPartitioner/SmoothWeight.java
@@ -0,0 +1,47 @@
+package org.astraea.partitioner.smoothPartitioner;
+
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Objects;
+
+/**
+ * Given initial key-score pair, it will output a preferred key with higher score. The score of the
+ * chosen key will decrease. And the other score will increment by its original score. It may result
+ * in "higher score with higher chosen rate".
+ */
+public final class SmoothWeight {
+  private Map<Integer, Integer> load;
+  private int loadSum;
+  public Map<Integer, Integer> weight;
+
+  public SmoothWeight() {}
+
+  public SmoothWeight(Map<Integer, Integer> load) {
+    init(load);
+  }
+
+  public synchronized void init(Map<Integer, Integer> load) {
+    this.load = new HashMap<>(load);
+    this.weight = new HashMap<>(load);
+    this.loadSum = load.values().stream().mapToInt(d -> d).sum();
+  }
+
+  /**
+   * Get the preferred ID, and update the state.
+   *
+   * @return the preferred ID
+   */
+  public int getAndChoose() {
+    Objects.requireNonNull(this.weight);
+
+    this.weight.replaceAll((ID, value) -> value + this.load.get(ID));
+    var maxID =
+        this.weight.entrySet().stream()
+            .max(Comparator.comparingDouble(Map.Entry::getValue))
+            .orElseGet(() -> Map.entry(0, 0))
+            .getKey();
+    this.weight.computeIfPresent(maxID, (ID, value) -> value - loadSum);
+    return maxID;
+  }
+}
diff --git a/app/src/main/java/org/astraea/partitioner/smoothPartitioner/SmoothWeightDispatcher.java b/app/src/main/java/org/astraea/partitioner/smoothPartitioner/SmoothWeightDispatcher.java
new file mode 100644
index 00000000..eaddabaa
--- /dev/null
+++ b/app/src/main/java/org/astraea/partitioner/smoothPartitioner/SmoothWeightDispatcher.java
@@ -0,0 +1,155 @@
+package org.astraea.partitioner.smoothPartitioner;
+
+import java.time.Duration;
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.NoSuchElementException;
+import java.util.Optional;
+import java.util.Random;
+import java.util.TreeMap;
+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.CostFunction;
+import org.astraea.cost.LoadCost;
+import org.astraea.cost.MemoryWarningCost;
+import org.astraea.metrics.collector.BeanCollector;
+import org.astraea.metrics.collector.Fetcher;
+import org.astraea.metrics.collector.Receiver;
+import org.astraea.partitioner.Configuration;
+import org.astraea.partitioner.Dispatcher;
+
+/**
+ * Based on the jmx metrics obtained from Kafka, it records the load status of the node over a
+ * period of time. Predict the future status of each node through the poisson of the load status.
+ * Finally, the result of poisson is used as the weight to perform smooth weighted RoundRobin.
+ */
+public class SmoothWeightDispatcher implements Dispatcher {
+  public static final String JMX_PORT = "jmx.port";
+  private final BeanCollector beanCollector =
+      BeanCollector.builder().interval(Duration.ofSeconds(1)).build();
+  private Optional<Integer> jmxPortDefault = Optional.empty();
+  private final Map<Integer, Integer> jmxPorts = new TreeMap<>();
+  private final Map<Integer, Receiver> receivers = new TreeMap<>();
+
+  private final Collection<CostFunction> functions =
+      List.of(new LoadCost(), new MemoryWarningCost());
+  // Memory warning can dominate the score.
+  private final Map<Object, Double> costFraction =
+      Map.of(LoadCost.class, 1.0, MemoryWarningCost.class, -100.0);
+
+  private final SmoothWeight smoothWeightCal = new SmoothWeight();
+  private final Random random = new Random();
+  // Fetch data every n seconds
+  private long lastFetchTime = 0L;
+  private final Lock lock = new ReentrantLock();
+
+  @Override
+  public int partition(String topic, byte[] key, byte[] value, ClusterInfo clusterInfo) {
+    var partitions = clusterInfo.availablePartitions(topic);
+    // just return first partition if there is no available partitions
+    if (partitions.isEmpty()) return 0;
+
+    // just return the only one available partition
+    if (partitions.size() == 1) return partitions.iterator().next().partition();
+
+    if (Utils.overSecond(lastFetchTime, 1) && lock.tryLock()) {
+      try {
+        // add new receivers for new brokers
+        partitions.stream()
+            .filter(p -> !receivers.containsKey(p.leader().id()))
+            .forEach(
+                p ->
+                    receivers.put(
+                        p.leader().id(), receiver(p.leader().host(), jmxPort(p.leader().id()))));
+
+        // fetch the latest beans for each node
+        var beans =
+            receivers.entrySet().stream()
+                .collect(Collectors.toMap(Map.Entry::getKey, e -> e.getValue().current()));
+
+        var compoundScore =
+            functions.stream()
+                .map(
+                    f -> {
+                      var map = f.cost(ClusterInfo.of(clusterInfo, beans));
+                      map.replaceAll((k, v) -> v * costFraction.get(f.getClass()));
+                      return map;
+                    })
+                .reduce(new HashMap<>(), SmoothWeightDispatcher::costCompound)
+                .entrySet()
+                .stream()
+                .map(e -> Map.entry(e.getKey(), (int) (double) e.getValue()))
+                .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+
+        smoothWeightCal.init(compoundScore);
+      } finally {
+        lastFetchTime = System.currentTimeMillis();
+        lock.unlock();
+      }
+    }
+
+    // Make "smooth weight choosing" on the score.
+    var targetBroker = smoothWeightCal.getAndChoose();
+    var targetPartitions =
+        partitions.stream()
+            .filter(p -> p.leader().id() == targetBroker)
+            .collect(Collectors.toList());
+
+    return targetPartitions.get(random.nextInt(targetPartitions.size())).partition();
+  }
+
+  // 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(identity::putIfAbsent);
+    return identity;
+  }
+
+  Receiver receiver(String host, int port) {
+    return beanCollector
+        .register()
+        .host(host)
+        .port(port)
+        .fetcher(
+            Fetcher.of(
+                functions.stream()
+                    .map(CostFunction::fetcher)
+                    .collect(Collectors.toUnmodifiableList())))
+        .build();
+  }
+
+  @Override
+  public void configure(Configuration config) {
+    jmxPortDefault = config.integer(JMX_PORT);
+
+    // seeks for custom jmx ports.
+    config.entrySet().stream()
+        .filter(e -> e.getKey().startsWith("broker."))
+        .filter(e -> e.getKey().endsWith(JMX_PORT))
+        .map(
+            e ->
+                Map.entry(
+                    e.getKey().replaceAll("^broker[.]", "").replaceAll("[.]" + JMX_PORT + "$", ""),
+                    e.getValue()))
+        .forEach(e -> jmxPorts.put(Integer.parseInt(e.getKey()), Integer.parseInt(e.getValue())));
+  }
+
+  // visible for testing
+  int jmxPort(int id) {
+    if (jmxPorts.containsKey(id)) return jmxPorts.get(id);
+    return jmxPortDefault.orElseThrow(
+        () -> new NoSuchElementException("broker: " + id + " does not have jmx port"));
+  }
+
+  @Override
+  public void close() {
+    receivers.values().forEach(Utils::close);
+    receivers.clear();
+  }
+}
diff --git a/app/src/main/java/org/astraea/performance/FileWriter.java b/app/src/main/java/org/astraea/performance/FileWriter.java
deleted file mode 100644
index b1ca1ed1..00000000
--- a/app/src/main/java/org/astraea/performance/FileWriter.java
+++ /dev/null
@@ -1,119 +0,0 @@
-package org.astraea.performance;
-
-import java.io.BufferedWriter;
-import java.io.IOException;
-import java.text.SimpleDateFormat;
-import java.util.Date;
-import java.util.stream.IntStream;
-import org.astraea.concurrent.Executor;
-import org.astraea.concurrent.State;
-import org.astraea.utils.DataUnit;
-
-public class FileWriter implements Executor {
-  private final Manager manager;
-  private final Tracker tracker;
-  private BufferedWriter writer;
-  private String CSVName;
-
-  public FileWriter(Manager manager, Tracker tracker) {
-    this.manager = manager;
-    this.tracker = tracker;
-    initFileWriter();
-  }
-
-  @Override
-  public State execute() throws InterruptedException {
-    if (logToCSV()) return State.DONE;
-    Thread.sleep(1000);
-    return State.RUNNING;
-  }
-
-  public String CSVName() {
-    return CSVName;
-  }
-
-  private void initFileWriter() {
-    CSVName = "Performance" + new SimpleDateFormat("yyyyMMddHHmmss").format(new Date()) + ".csv";
-    try {
-      writer = new BufferedWriter(new java.io.FileWriter(CSVName));
-      writer.write(
-          "Time \\ Name, Consumed/Produced, Output throughput (/sec), Input throughput (/sec), "
-              + "Publish max latency (ms), Publish min latency (ms), "
-              + "End-to-end max latency (ms), End-to-end min latency (ms)");
-      IntStream.range(0, tracker.producerResult().bytes.size())
-          .forEach(
-              i -> {
-                try {
-                  writer.write(
-                      ",Producer["
-                          + i
-                          + "] current throughput (/sec), Producer["
-                          + i
-                          + "] current publish latency (ms)");
-                } catch (IOException ignore) {
-                }
-              });
-      IntStream.range(0, tracker.consumerResult().bytes.size())
-          .forEach(
-              i -> {
-                try {
-                  writer.write(
-                      ",Consumer["
-                          + i
-                          + "] current throughput (/sec), Consumer["
-                          + i
-                          + "] current ene-to-end latency (ms)");
-                } catch (IOException ignore) {
-                }
-              });
-      writer.newLine();
-    } catch (IOException ignore) {
-      writer = null;
-    }
-  }
-
-  private boolean logToCSV() {
-    var producerResult = tracker.producerResult();
-    var consumerResult = tracker.consumerResult();
-    if (writer == null || producerResult.completedRecords == 0L) return false;
-    var duration = tracker.duration();
-    var producerPercentage =
-        Math.min(
-            100D,
-            manager.exeTime().percentage(producerResult.completedRecords, duration.toMillis()));
-    var consumerPercentage = consumerResult.completedRecords * 100D / manager.producedRecords();
-    try {
-      writer.write(
-          duration.toHoursPart()
-              + "h"
-              + duration.toMinutesPart()
-              + "m"
-              + duration.toSecondsPart()
-              + "s");
-      writer.write(String.format(",%.2f%% / %.2f%%", consumerPercentage, producerPercentage));
-      writer.write("," + DataUnit.Byte.of(producerResult.totalCurrentBytes()));
-      writer.write("," + DataUnit.Byte.of(consumerResult.totalCurrentBytes()));
-      writer.write("," + producerResult.maxLatency + "," + producerResult.minLatency);
-      writer.write("," + consumerResult.maxLatency + "," + consumerResult.minLatency);
-      for (int i = 0; i < producerResult.bytes.size(); ++i) {
-        writer.write("," + DataUnit.Byte.of(producerResult.currentBytes.get(i)));
-        writer.write("," + producerResult.averageLatencies.get(i));
-      }
-      for (int i = 0; i < consumerResult.bytes.size(); ++i) {
-        writer.write("," + DataUnit.Byte.of(consumerResult.currentBytes.get(i)));
-        writer.write("," + consumerResult.averageLatencies.get(i));
-      }
-      writer.newLine();
-    } catch (IOException ignore) {
-    }
-    return manager.producedDone() && manager.consumedDone();
-  }
-
-  @Override
-  public void close() {
-    try {
-      writer.close();
-    } catch (IOException ignore) {
-    }
-  }
-}
diff --git a/app/src/main/java/org/astraea/performance/Performance.java b/app/src/main/java/org/astraea/performance/Performance.java
index 02e89d22..0a4b7ceb 100644
--- a/app/src/main/java/org/astraea/performance/Performance.java
+++ b/app/src/main/java/org/astraea/performance/Performance.java
@@ -2,7 +2,9 @@ package org.astraea.performance;
 
 import com.beust.jcommander.Parameter;
 import java.io.IOException;
+import java.nio.file.Path;
 import java.time.Duration;
+import java.util.Collection;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
@@ -23,6 +25,7 @@ import org.astraea.Utils;
 import org.astraea.argument.CompressionField;
 import org.astraea.argument.NonEmptyStringField;
 import org.astraea.argument.NonNegativeShortField;
+import org.astraea.argument.PathField;
 import org.astraea.argument.PositiveLongField;
 import org.astraea.argument.PositiveShortField;
 import org.astraea.argument.StringMapField;
@@ -95,6 +98,11 @@ public class Performance {
 
     var manager = new Manager(param, producerMetrics, consumerMetrics);
     var tracker = new Tracker(producerMetrics, consumerMetrics, manager);
+    Collection<Executor> fileWriter =
+        (param.CSVPath != null)
+            ? List.of(
+                ReportFormat.createFileWriter(param.reportFormat, param.CSVPath, manager, tracker))
+            : List.of();
     var groupId = "groupId-" + System.currentTimeMillis();
     try (var threadPool =
         ThreadPool.builder()
@@ -133,7 +141,7 @@ public class Performance {
                                 manager))
                     .collect(Collectors.toUnmodifiableList()))
             .executor(tracker)
-            .executor((param.createCSV) ? new FileWriter(manager, tracker) : () -> State.DONE)
+            .executors(fileWriter)
             .build()) {
       threadPool.waitAll();
       return new Result(param.topic);
@@ -341,11 +349,6 @@ public class Performance {
       return props;
     }
 
-    @Parameter(
-        names = {"--createCSV"},
-        description = "put the metrics into a csv file if this flag is set")
-    boolean createCSV = false;
-
     @Parameter(
         names = {"--compression"},
         description =
@@ -383,6 +386,18 @@ public class Performance {
             "String: Used with SpecifyBrokerPartitioner to specify the brokers that partitioner can send.",
         validateWith = NonEmptyStringField.class)
     List<Integer> specifyBroker = List.of(-1);
+
+    @Parameter(
+        names = {"--report.path"},
+        description = "String: A path to place the report. Default: (no report)",
+        converter = PathField.class)
+    Path CSVPath = null;
+
+    @Parameter(
+        names = {"--report.format"},
+        description = "Output format for the report",
+        converter = ReportFormat.ReportFormatConverter.class)
+    ReportFormat reportFormat = ReportFormat.CSV;
   }
 
   public static class Result {
diff --git a/app/src/main/java/org/astraea/performance/ReportFormat.java b/app/src/main/java/org/astraea/performance/ReportFormat.java
new file mode 100644
index 00000000..f03782b1
--- /dev/null
+++ b/app/src/main/java/org/astraea/performance/ReportFormat.java
@@ -0,0 +1,245 @@
+package org.astraea.performance;
+
+import com.beust.jcommander.IStringConverter;
+import com.beust.jcommander.ParameterException;
+import java.io.BufferedWriter;
+import java.io.FileWriter;
+import java.io.IOException;
+import java.nio.file.FileSystems;
+import java.nio.file.Path;
+import java.text.SimpleDateFormat;
+import java.time.Duration;
+import java.util.Date;
+import java.util.stream.IntStream;
+import org.astraea.Utils;
+import org.astraea.concurrent.Executor;
+import org.astraea.concurrent.State;
+import org.astraea.utils.DataUnit;
+
+public enum ReportFormat {
+  CSV("csv"),
+  JSON("json");
+
+  private final String name;
+
+  ReportFormat(String name) {
+    this.name = name;
+  }
+
+  public static class ReportFormatConverter implements IStringConverter<ReportFormat> {
+    @Override
+    public ReportFormat convert(String value) {
+      switch (value.toLowerCase()) {
+        case "csv":
+          return ReportFormat.CSV;
+        case "json":
+          return ReportFormat.JSON;
+        default:
+          throw new ParameterException("Invalid file format. Use \"csv\" or \"json\"");
+      }
+    }
+  }
+
+  public static Executor createFileWriter(
+      ReportFormat reportFormat, Path path, Manager manager, Tracker tracker) throws IOException {
+    var filePath =
+        FileSystems.getDefault()
+            .getPath(
+                path.toString(),
+                "Performance"
+                    + new SimpleDateFormat("yyyyMMddHHmmss").format(new Date())
+                    + reportFormat);
+    var writer = new BufferedWriter(new FileWriter(filePath.toFile()));
+    switch (reportFormat) {
+      case CSV:
+        initCSVFormat(
+            writer, tracker.producerResult().bytes.size(), tracker.consumerResult().bytes.size());
+        return new Executor() {
+          @Override
+          public State execute() throws InterruptedException {
+            if (logToCSV(writer, manager, tracker)) return State.DONE;
+            Thread.sleep(1000);
+            return State.RUNNING;
+          }
+
+          @Override
+          public void close() {
+            Utils.close(writer);
+          }
+        };
+      case JSON:
+        writer.write("{");
+        return new Executor() {
+          @Override
+          public State execute() throws InterruptedException {
+            if (logToJSON(writer, manager, tracker)) return State.DONE;
+            Thread.sleep(1000);
+            return State.RUNNING;
+          }
+
+          @Override
+          public void close() {
+            try {
+              writer.write("}");
+              Utils.close(writer);
+            } catch (IOException ignore) {
+            }
+          }
+        };
+      default:
+        throw new IllegalArgumentException("Invalid format.");
+    }
+  }
+
+  private static void initCSVFormat(BufferedWriter writer, int producerCounts, int consumerCounts)
+      throws IOException {
+    writer.write(
+        "Time \\ Name, Consumed/Produced, Output throughput (/sec), Input throughput (/sec), "
+            + "Publish max latency (ms), Publish min latency (ms), "
+            + "End-to-end max latency (ms), End-to-end min latency (ms)");
+    IntStream.range(0, producerCounts)
+        .forEach(
+            i -> {
+              try {
+                writer.write(
+                    ",Producer["
+                        + i
+                        + "] current throughput (/sec), Producer["
+                        + i
+                        + "] current publish latency (ms)");
+              } catch (IOException ignore) {
+              }
+            });
+    IntStream.range(0, consumerCounts)
+        .forEach(
+            i -> {
+              try {
+                writer.write(
+                    ",Consumer["
+                        + i
+                        + "] current throughput (/sec), Consumer["
+                        + i
+                        + "] current ene-to-end latency (ms)");
+              } catch (IOException ignore) {
+              }
+            });
+    writer.newLine();
+  }
+
+  private static boolean logToCSV(BufferedWriter writer, Manager manager, Tracker tracker) {
+    var result = processResult(manager, tracker);
+    if (result.producerResult.completedRecords == 0) return false;
+    try {
+      writer.write(
+          result.duration.toHoursPart()
+              + "h"
+              + result.duration.toMinutesPart()
+              + "m"
+              + result.duration.toSecondsPart()
+              + "s");
+      writer.write(
+          String.format(",%.2f%% / %.2f%%", result.consumerPercentage, result.producerPercentage));
+      writer.write("," + DataUnit.Byte.of(result.producerResult.totalCurrentBytes()));
+      writer.write("," + DataUnit.Byte.of(result.consumerResult.totalCurrentBytes()));
+      writer.write("," + result.producerResult.maxLatency + "," + result.producerResult.minLatency);
+      writer.write("," + result.consumerResult.maxLatency + "," + result.consumerResult.minLatency);
+      for (int i = 0; i < result.producerResult.bytes.size(); ++i) {
+        writer.write("," + DataUnit.Byte.of(result.producerResult.currentBytes.get(i)));
+        writer.write("," + result.producerResult.averageLatencies.get(i));
+      }
+      for (int i = 0; i < result.consumerResult.bytes.size(); ++i) {
+        writer.write("," + DataUnit.Byte.of(result.consumerResult.currentBytes.get(i)));
+        writer.write("," + result.consumerResult.averageLatencies.get(i));
+      }
+      writer.newLine();
+    } catch (IOException ignore) {
+    }
+    return manager.producedDone() && manager.consumedDone();
+  }
+
+  /** Write to writer. Output: "(timestamp)": { (many metrics ...) } */
+  private static boolean logToJSON(BufferedWriter writer, Manager manager, Tracker tracker) {
+    var result = processResult(manager, tracker);
+    if (result.producerResult.completedRecords == 0) return false;
+    try {
+      writer.write(
+          String.format(
+              "\"%dh%dm%ds\": {",
+              result.duration.toHoursPart(),
+              result.duration.toMinutesPart(),
+              result.duration.toSecondsPart()));
+      writer.write(
+          String.format(
+              "\"consumerPercentage\": %.2f%%, \"Producer Percentage\": %.2f%%",
+              result.consumerPercentage, result.producerPercentage));
+      writer.write(
+          ", \"outputThroughput\": " + result.producerResult.averageBytes(result.duration));
+      writer.write(", \"inputThroughput\": " + result.consumerResult.averageBytes(result.duration));
+      writer.write(", \"publishMaxLatency\": " + result.producerResult.maxLatency);
+      writer.write(", \"publishMinLatency\": " + result.producerResult.minLatency);
+      writer.write(", \"E2EMaxLatency\": " + result.consumerResult.maxLatency);
+      writer.write(", \"E2EMinLatency\": " + result.consumerResult.minLatency);
+
+      writer.write(", \"producerThroughput\": [");
+      for (int i = 0; i < result.producerResult.bytes.size(); ++i) {
+        writer.write(Tracker.avg(result.duration, result.producerResult.bytes.get(i)) + ", ");
+      }
+      writer.write("], \"producerLatency\": [");
+      for (int i = 0; i < result.producerResult.bytes.size(); ++i) {
+        writer.write(result.producerResult.averageLatencies.get(i) + ", ");
+      }
+      writer.write("], \"consumerThroughput\": [");
+      for (int i = 0; i < result.consumerResult.bytes.size(); ++i) {
+        writer.write(Tracker.avg(result.duration, result.consumerResult.bytes.get(i)) + ", ");
+      }
+      writer.write("], \"consumerLatency\": [");
+      for (int i = 0; i < result.consumerResult.bytes.size(); ++i) {
+        writer.write(result.consumerResult.averageLatencies.get(i) + ", ");
+      }
+      writer.write("]}");
+      writer.newLine();
+    } catch (IOException ignore) {
+    }
+    return manager.producedDone() && manager.consumedDone();
+  }
+
+  private static ProcessedResult processResult(Manager manager, Tracker tracker) {
+    var producerResult = tracker.producerResult();
+    var consumerResult = tracker.consumerResult();
+    var duration = tracker.duration();
+    return new ProcessedResult(
+        tracker.consumerResult(),
+        tracker.producerResult(),
+        duration,
+        Math.min(
+            100D,
+            manager.exeTime().percentage(producerResult.completedRecords, duration.toMillis())),
+        consumerResult.completedRecords * 100D / manager.producedRecords());
+  }
+
+  static class ProcessedResult {
+    public final Tracker.Result consumerResult;
+    public final Tracker.Result producerResult;
+    public final Duration duration;
+    public final double consumerPercentage;
+    public final double producerPercentage;
+
+    ProcessedResult(
+        Tracker.Result consumerResult,
+        Tracker.Result producerResult,
+        Duration duration,
+        double consumerPercentage,
+        double producerPercentage) {
+      this.consumerResult = consumerResult;
+      this.producerResult = producerResult;
+      this.duration = duration;
+      this.consumerPercentage = consumerPercentage;
+      this.producerPercentage = producerPercentage;
+    }
+  }
+
+  @Override
+  public String toString() {
+    return name;
+  }
+}


Test patch:
diff --git a/app/src/test/java/org/astraea/cost/FakeClusterInfo.java b/app/src/test/java/org/astraea/cost/FakeClusterInfo.java
new file mode 100644
index 00000000..7f6d13df
--- /dev/null
+++ b/app/src/test/java/org/astraea/cost/FakeClusterInfo.java
@@ -0,0 +1,33 @@
+package org.astraea.cost;
+
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
+import org.astraea.metrics.HasBeanObject;
+
+public class FakeClusterInfo implements ClusterInfo {
+  @Override
+  public List<NodeInfo> nodes() {
+    return List.of();
+  }
+
+  @Override
+  public List<PartitionInfo> availablePartitions(String topic) {
+    return List.of();
+  }
+
+  @Override
+  public List<PartitionInfo> partitions(String topic) {
+    return List.of();
+  }
+
+  @Override
+  public Collection<HasBeanObject> beans(int brokerId) {
+    return List.of();
+  }
+
+  @Override
+  public Map<Integer, Collection<HasBeanObject>> allBeans() {
+    return Map.of();
+  }
+}
diff --git a/app/src/test/java/org/astraea/cost/LoadCostTest.java b/app/src/test/java/org/astraea/cost/LoadCostTest.java
new file mode 100644
index 00000000..e520b208
--- /dev/null
+++ b/app/src/test/java/org/astraea/cost/LoadCostTest.java
@@ -0,0 +1,60 @@
+package org.astraea.cost;
+
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
+import org.astraea.metrics.HasBeanObject;
+import org.astraea.metrics.jmx.BeanObject;
+import org.astraea.metrics.kafka.BrokerTopicMetricsResult;
+import org.astraea.metrics.kafka.KafkaMetrics;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+import org.mockito.Mockito;
+
+public class LoadCostTest {
+  @Test
+  void testComputeLoad() {
+    var loadCostFunction = new LoadCost();
+    var allBeans = exampleClusterInfo().allBeans();
+    var load = loadCostFunction.computeLoad(allBeans);
+
+    Assertions.assertEquals(2, load.get(1));
+    Assertions.assertEquals(1, load.get(2));
+    Assertions.assertEquals(1, load.get(3));
+
+    load = loadCostFunction.computeLoad(allBeans);
+
+    // count does not change so all broker get one more score
+    Assertions.assertEquals(3, load.get(1));
+    Assertions.assertEquals(2, load.get(2));
+    Assertions.assertEquals(2, load.get(3));
+  }
+
+  private ClusterInfo exampleClusterInfo() {
+    var BytesInPerSec1 = mockResult(KafkaMetrics.BrokerTopic.BytesInPerSec.metricName(), 50000L);
+    var BytesInPerSec2 = mockResult(KafkaMetrics.BrokerTopic.BytesInPerSec.metricName(), 100000L);
+    var BytesInPerSec3 = mockResult(KafkaMetrics.BrokerTopic.BytesInPerSec.metricName(), 200000L);
+    var BytesOutPerSec1 = mockResult(KafkaMetrics.BrokerTopic.BytesOutPerSec.metricName(), 210L);
+    var BytesOutPerSec2 = mockResult(KafkaMetrics.BrokerTopic.BytesOutPerSec.metricName(), 20L);
+    var BytesOutPerSec3 = mockResult(KafkaMetrics.BrokerTopic.BytesOutPerSec.metricName(), 10L);
+
+    Collection<HasBeanObject> broker1 = List.of(BytesInPerSec1, BytesOutPerSec1);
+    Collection<HasBeanObject> broker2 = List.of(BytesInPerSec2, BytesOutPerSec2);
+    Collection<HasBeanObject> broker3 = List.of(BytesInPerSec3, BytesOutPerSec3);
+    return new FakeClusterInfo() {
+      @Override
+      public Map<Integer, Collection<HasBeanObject>> allBeans() {
+        return Map.of(1, broker1, 2, broker2, 3, broker3);
+      }
+    };
+  }
+
+  private BrokerTopicMetricsResult mockResult(String name, long count) {
+    var result = Mockito.mock(BrokerTopicMetricsResult.class);
+    var bean = Mockito.mock(BeanObject.class);
+    Mockito.when(result.beanObject()).thenReturn(bean);
+    Mockito.when(bean.getProperties()).thenReturn(Map.of("name", name));
+    Mockito.when(result.count()).thenReturn(count);
+    return result;
+  }
+}
diff --git a/app/src/test/java/org/astraea/cost/MemoryWarningCostTest.java b/app/src/test/java/org/astraea/cost/MemoryWarningCostTest.java
new file mode 100644
index 00000000..5485de1f
--- /dev/null
+++ b/app/src/test/java/org/astraea/cost/MemoryWarningCostTest.java
@@ -0,0 +1,53 @@
+package org.astraea.cost;
+
+import java.lang.management.MemoryUsage;
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
+import org.astraea.metrics.HasBeanObject;
+import org.astraea.metrics.java.HasJvmMemory;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+import org.mockito.Mockito;
+
+public class MemoryWarningCostTest {
+  @Test
+  void testCost() throws InterruptedException {
+    var jvmMemory1 = mockResult(50L, 100L);
+    var jvmMemory2 = mockResult(81L, 100L);
+    var jvmMemory3 = mockResult(80L, 150L);
+
+    Collection<HasBeanObject> broker1 = List.of(jvmMemory1);
+    Collection<HasBeanObject> broker2 = List.of(jvmMemory2);
+    Collection<HasBeanObject> broker3 = List.of(jvmMemory3);
+    ClusterInfo clusterInfo =
+        new FakeClusterInfo() {
+          @Override
+          public Map<Integer, Collection<HasBeanObject>> allBeans() {
+            return Map.of(1, broker1, 2, broker2, 3, broker3);
+          }
+        };
+
+    var memoryWarning = new MemoryWarningCost();
+    var scores = memoryWarning.cost(clusterInfo);
+    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);
+    Assertions.assertEquals(0.0, scores.get(1));
+    Assertions.assertEquals(1.0, scores.get(2));
+    Assertions.assertEquals(0.0, scores.get(3));
+  }
+
+  private HasJvmMemory mockResult(long used, long max) {
+    var jvmMemory = Mockito.mock(HasJvmMemory.class);
+    var memoryUsage = Mockito.mock(MemoryUsage.class);
+    Mockito.when(jvmMemory.heapMemoryUsage()).thenReturn(memoryUsage);
+    Mockito.when(memoryUsage.getUsed()).thenReturn(used);
+    Mockito.when(memoryUsage.getMax()).thenReturn(max);
+    return jvmMemory;
+  }
+}
diff --git a/app/src/test/java/org/astraea/partitioner/nodeLoadMetric/PartitionerUtilsTest.java b/app/src/test/java/org/astraea/partitioner/nodeLoadMetric/PartitionerUtilsTest.java
index 0b82ba67..e243c392 100644
--- a/app/src/test/java/org/astraea/partitioner/nodeLoadMetric/PartitionerUtilsTest.java
+++ b/app/src/test/java/org/astraea/partitioner/nodeLoadMetric/PartitionerUtilsTest.java
@@ -7,6 +7,7 @@ import static org.astraea.partitioner.nodeLoadMetric.PartitionerUtils.weightPois
 import static org.junit.jupiter.api.Assertions.assertEquals;
 
 import java.util.HashMap;
+import java.util.Map;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
 
@@ -33,7 +34,7 @@ public class PartitionerUtilsTest {
     testNodesLoadCount.put(1, 5);
     testNodesLoadCount.put(2, 8);
 
-    HashMap<Integer, Double> poissonMap;
+    Map<Integer, Double> poissonMap;
     HashMap<Integer, Double> testPoissonMap = new HashMap<>();
     testPoissonMap.put(0, 0.12465201948308113);
     testPoissonMap.put(1, 0.6159606548330632);
diff --git a/app/src/test/java/org/astraea/partitioner/smoothPartitioner/SmoothWeightTest.java b/app/src/test/java/org/astraea/partitioner/smoothPartitioner/SmoothWeightTest.java
new file mode 100644
index 00000000..56434137
--- /dev/null
+++ b/app/src/test/java/org/astraea/partitioner/smoothPartitioner/SmoothWeightTest.java
@@ -0,0 +1,20 @@
+package org.astraea.partitioner.smoothPartitioner;
+
+import java.util.Map;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+public class SmoothWeightTest {
+  @Test
+  public void testGetAndChoose() {
+    var smoothWeight = new SmoothWeight(Map.of(1, 5, 2, 3, 3, 1));
+
+    Assertions.assertEquals(1, smoothWeight.getAndChoose());
+    Assertions.assertEquals(2, smoothWeight.getAndChoose());
+    Assertions.assertEquals(1, smoothWeight.getAndChoose());
+    Assertions.assertEquals(1, smoothWeight.getAndChoose());
+    Assertions.assertEquals(2, smoothWeight.getAndChoose());
+    Assertions.assertEquals(1, smoothWeight.getAndChoose());
+    Assertions.assertEquals(3, smoothWeight.getAndChoose());
+  }
+}