Repository: skiptests/astraea
Pull request #314 (2022-05-02 03:15:50)

Configure cost-functions weight
### 說明

這隻PR對`StrictCostDispatcher` 做了修改，讓使用者自訂所要使用的cost-function，並決定各個cost-functions 的權重。
cost-function 的設定方式為，class name作為key；權重作為value，在創建producer 時，放到`Properties`中。這裡用performance tool為範例做設定，
```
./gradlew run --args="performance --bootstrap.servers localhost:9092 --partitioner org.astraea.partitioner.StrictCostDispatcher --configs org.astraea.cost.brokersMetrics.BrokerInputCost=20,org.astraea.cost.brokersMetrics.BrokerOutputCost=1,jmx.port=12345"
```

---

### 結果
現在`StrictCostDispatcher` 對cost-functions 的使用是，
1. 調用所有cost-functions 並取得所有broker 的分數
2. **對個別(broker 的)分數，乘上cost-function所屬的權重**
3. 加總，取得個別broker 的分數
4. 取最小值作為所選擇的broker位置。

實際測試時，發現了一點bug並修改，現在是確實可以使用的。
實際測試後，觀察到確實cost-function 權重設定確實有照預期運作。

---

(以下文長可忽略，是稍微測試的過程)

實際用以下參數，多次發送小資料測試(，每次啟動performance tool 只送100 records, 512 bytes/record)，用小資料是為了方便觀察，先看看程式是否符合預期：
1. BrokerInputCost=20, BrokerOutputCost=1, CpuCost=1, MemoryCost=1
2. BrokerInputCost=1, BrokerOutputCost=20, CpuCost=1, MemoryCost=1
3. BrokerInputCost=1, BrokerOutputCost=1, CpuCost=20, MemoryCost=1

結果與預期：
1. 每次performance tool 的執行都會變換broker 進行發送。
2. 為某一個topic( with one partition) 製造多個consumer group，製造出某broker特別多的consumer pull 資料。而後每次performance tool 的執行都會避開該broker 進行發送。
3. 其中2台broker只使用1個virtual CPU，這兩台的CPU system load (CpuCost 所觀察的指標) 就會很高。而後每次performance tool 的執行都會避開該2台broker 進行發送。

chia7712: > org.astraea.cost.brokersMetrics.BrokerInputCost=1

這邊有機會先考慮之後要如何傳遞參數給各個cost function嗎？
chinghongfang: > 這邊有機會先考慮之後要如何傳遞參數給各個cost function嗎？

請問這個"參數"是指"權重"，意思是要把"權重"丟給cost-function，讓cost-function 處理嘛？
如果是bean metric, partition info... 等參數已經包在`ClusterInfo`裡面了。
chia7712: > 請問這個"參數"是指"權重"，意思是要把"權重"丟給cost-function，讓cost-function 處理嘛？
如果是bean metric, partition info... 等參數已經包在ClusterInfo裡面了。

我指的是假如之後有些cost function需要一些參數才能完成初始化，因此我們會讓cost function多一個介面，例如：
```java
public interface CostFunction {

  default void init(Config config) {
  }
}
```

因此在用strict dispatcher時，我們需要有一個規則把屬於各自cost function的參數餵給他們
chinghongfang: > 我指的是假如之後有些cost function需要一些參數才能完成初始化，因此我們會讓cost function多一個介面

原來如此，是要讓cost-function 也可以依據參數的不同來做調整吧，相當於有兩面的彈性：一個是權重的彈性；一個是cost-function 的調整對吧。

> 因此在用strict dispatcher時，我們需要有一個規則把屬於各自cost function的參數餵給他們

傳入partitioner的格式，改成以下格式
```
org.astraea.cost.brokersMetrics.BrokerInputCost=weight:20;par2:value2;par3:value3
```
傳入cost-function之前，把`weight:20;par2:value2;par3:value3`轉換成`org.astraea.partitioner.Configuration`的格式。
chia7712: > 傳入partitioner的格式，改成以下格式

這個格式看起來不錯。如果要避免此PR太大，可以在另一隻PR再新增參數化cost function的部分

@garyparrot @qoo332001 也麻煩看一下這次要新增的參數化邏輯
chinghongfang: 好的，我在下一隻PR新增上述設定。
還麻煩您再review一次。
chia7712: @chinghongfang #313這隻PR可能會影響到這隻，不過我希望#313先合併這樣方便大家使用新的資料結構`ReplicaInfo`，因此麻煩注意一下，衝突的部分處理完後我再來審視一次

Patch:
diff --git a/README.md b/README.md
index e71a5914..c11bbaaf 100644
--- a/README.md
+++ b/README.md
@@ -238,26 +238,29 @@ Run the benchmark from source
 ```
 ### Performance Benchmark Configurations
 1. --bootstrap.servers: the server to connect to
-2. --topic: the topic name. Default: testPerformance-{Time in millis}
-3. --partitions: topic config when creating new topic. Default: 1 
-4. --replicas: topic config when creating new topic. Default: 1
-5. --consumers: the number of consumers (threads). Default: 1
-6. --producers: the number of producers (threads). Default: 1
-7. --run.until: the total number of records sent by the producers or the time for producer to send records.
+2. --compression: the compression algorithm used by producer.
+3. --topic: the topic name. Default: testPerformance-{Time in millis}
+4. --partitions: topic config when creating new topic. Default: 1 
+5. --replicas: topic config when creating new topic. Default: 1
+6. --consumers: the number of consumers (threads). Default: 1
+7. --producers: the number of producers (threads). Default: 1
+8. --run.until: the total number of records sent by the producers or the time for producer to send records.
   The duration formats accepted are (a number) + (a time unit). 
   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. --fixed.size: the flag to let all records have the same size
+9. --record.size: the (bound of) record size in byte. Default: 1 KiB
 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--throughput: the produce rate for all producers. e.g. "--throughput 2MiB". Default: 500 GiB (per second)
-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)
-17--report.path: A path to place the report file. Default: (no report)
-18--report.format: Select output file format. Available format: "csv", "json".
+11. --partitioner: the partitioner to use in producers.
+12. --configs: the configurations pass to partitioner. 
+  The configuration format is "\<key1\>=\<value1\>[,\<key2\>=\<value2\>]*". 
+  eg. "--configs broker.1001.jmx.port=14338,org.astraea.cost.ThroughputCost=1"
+13. --throughput: the produce rate for all producers. e.g. "--throughput 2MiB". Default: 500 GiB (per second)
+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)
+17. --report.path: A path to place the report file. Default: (no report)
+18. --report.format: Select output file format. Available format: "csv", "json". Default: "csv"
+19. --transaction.size: number of records in each transaction. Default: 1
 
 ---
 
diff --git a/app/src/main/java/org/astraea/cost/brokersMetrics/BrokerInputCost.java b/app/src/main/java/org/astraea/cost/brokersMetrics/BrokerInputCost.java
index 92740655..29870e17 100644
--- a/app/src/main/java/org/astraea/cost/brokersMetrics/BrokerInputCost.java
+++ b/app/src/main/java/org/astraea/cost/brokersMetrics/BrokerInputCost.java
@@ -47,12 +47,9 @@ public class BrokerInputCost extends Periodic<Map<Integer, Double>> implements H
                         value.stream()
                             .filter(
                                 hasBeanObject ->
-                                    hasBeanObject
-                                        .beanObject()
-                                        .getProperties()
-                                        .get("name")
+                                    KafkaMetrics.BrokerTopic.BytesInPerSec.metricName()
                                         .equals(
-                                            KafkaMetrics.BrokerTopic.BytesInPerSec.metricName()))
+                                            hasBeanObject.beanObject().getProperties().get("name")))
                             .forEach(
                                 hasBeanObject -> {
                                   var broker = brokersMetric.get(brokerID);
diff --git a/app/src/main/java/org/astraea/cost/brokersMetrics/BrokerOutputCost.java b/app/src/main/java/org/astraea/cost/brokersMetrics/BrokerOutputCost.java
index 012e54e7..7c60c49f 100644
--- a/app/src/main/java/org/astraea/cost/brokersMetrics/BrokerOutputCost.java
+++ b/app/src/main/java/org/astraea/cost/brokersMetrics/BrokerOutputCost.java
@@ -47,12 +47,9 @@ public class BrokerOutputCost extends Periodic<Map<Integer, Double>> implements
                         value.stream()
                             .filter(
                                 hasBeanObject ->
-                                    hasBeanObject
-                                        .beanObject()
-                                        .getProperties()
-                                        .get("name")
+                                    KafkaMetrics.BrokerTopic.BytesOutPerSec.metricName()
                                         .equals(
-                                            KafkaMetrics.BrokerTopic.BytesOutPerSec.metricName()))
+                                            hasBeanObject.beanObject().getProperties().get("name")))
                             .forEach(
                                 hasBeanObject -> {
                                   var broker = brokersMetric.get(brokerID);
diff --git a/app/src/main/java/org/astraea/cost/brokersMetrics/CpuCost.java b/app/src/main/java/org/astraea/cost/brokersMetrics/CpuCost.java
index 8222ba7e..1b33e06a 100644
--- a/app/src/main/java/org/astraea/cost/brokersMetrics/CpuCost.java
+++ b/app/src/main/java/org/astraea/cost/brokersMetrics/CpuCost.java
@@ -76,7 +76,7 @@ public class CpuCost extends Periodic<Map<Integer, Double>> implements HasBroker
 
   @Override
   public Fetcher fetcher() {
-    return client -> List.of(KafkaMetrics.Host.jvmMemory(client));
+    return client -> List.of(KafkaMetrics.Host.operatingSystem(client));
   }
 
   private static class BrokerMetric {
diff --git a/app/src/main/java/org/astraea/partitioner/PartitionerUtils.java b/app/src/main/java/org/astraea/partitioner/PartitionerUtils.java
index e86608bb..48b1a051 100644
--- a/app/src/main/java/org/astraea/partitioner/PartitionerUtils.java
+++ b/app/src/main/java/org/astraea/partitioner/PartitionerUtils.java
@@ -5,6 +5,7 @@ import java.io.IOException;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Properties;
+import java.util.stream.Collectors;
 
 /** do poisson for node's load situation */
 public class PartitionerUtils {
@@ -57,4 +58,25 @@ public class PartitionerUtils {
     }
     return properties;
   }
+
+  /**
+   * This method helps extract the config of form "broker.<brokerID>.jmx.port"="<port>" to a map.
+   * For example, the configuration {"broker.1001.jmx.port"="8000", "broker.1002.jmx.port"="8001"}
+   * will be parsed to map {(1001,8000), (1002,8001)}.
+   *
+   * @param config contains the config of broker id and its jmx port.
+   * @return pairs of broker id and its jmx port
+   */
+  public static Map<Integer, Integer> parseIdJMXPort(Configuration config) {
+    return 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()))
+        .map(e -> Map.entry(Integer.parseInt(e.getKey()), Integer.parseInt(e.getValue())))
+        .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+  }
 }
diff --git a/app/src/main/java/org/astraea/partitioner/StrictCostDispatcher.java b/app/src/main/java/org/astraea/partitioner/StrictCostDispatcher.java
index e93249dc..903502a1 100644
--- a/app/src/main/java/org/astraea/partitioner/StrictCostDispatcher.java
+++ b/app/src/main/java/org/astraea/partitioner/StrictCostDispatcher.java
@@ -5,6 +5,7 @@ import java.util.Collection;
 import java.util.List;
 import java.util.Map;
 import java.util.NoSuchElementException;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.TreeMap;
 import java.util.stream.Collectors;
@@ -26,13 +27,19 @@ import org.astraea.metrics.collector.Receiver;
  * Normally, all brokers use the same JMX port, so you can just define the `jmx.port=12345`. If one
  * of brokers uses different JMX client port, you can define `broker.1000.jmx.port=11111` (`1000` is
  * the broker id) to replace the value of `jmx.port`.
+ *
+ * <p>You can configure the cost functions you want to use. By giving the name of that cost function
+ * and its weight. For example,
+ * `org.astraea.cost.ThroughputCost=1,org.astraea.cost.brokersMetrics.BrokerOutputCost=1`.
  */
 public class StrictCostDispatcher implements Dispatcher {
   public static final String JMX_PORT = "jmx.port";
 
   private final BeanCollector beanCollector =
       BeanCollector.builder().interval(Duration.ofSeconds(4)).build();
-  private final Collection<CostFunction> functions;
+
+  /* The cost-functions we consider and the weight of them. It is visible for test.*/
+  Map<CostFunction, Double> functions;
   private Optional<Integer> jmxPortDefault = Optional.empty();
   private final Map<Integer, Integer> jmxPorts = new TreeMap<>();
   final Map<Integer, Receiver> receivers = new TreeMap<>();
@@ -43,7 +50,10 @@ public class StrictCostDispatcher implements Dispatcher {
 
   // visible for testing
   StrictCostDispatcher(Collection<CostFunction> functions) {
-    this.functions = functions;
+    this.functions =
+        functions.stream()
+            .map(f -> Map.entry(f, 1.0))
+            .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
   }
 
   @Override
@@ -69,16 +79,41 @@ public class StrictCostDispatcher implements Dispatcher {
             .collect(Collectors.toMap(Map.Entry::getKey, e -> e.getValue().current()));
 
     // get scores from all cost functions
-    var scores =
-        functions.stream()
-            .filter(f -> f instanceof HasBrokerCost)
-            .map(f -> (HasBrokerCost) f)
-            .map(f -> f.brokerCost(ClusterInfo.of(clusterInfo, beans)).value())
-            .collect(Collectors.toUnmodifiableList());
+    var scores = computeScore(functions, ClusterInfo.of(clusterInfo, beans));
 
     return bestPartition(partitionLeaders, scores).map(e -> e.getKey().partition()).orElse(0);
   }
 
+  /**
+   * Pass clusterInfo into all cost-functions. The result of each cost-function will multiply on
+   * their corresponding weight.
+   *
+   * @param functions the cost-function objects and their corresponding weight
+   * @return cost-function result multiplied on their corresponding weight
+   */
+  static List<Map<Integer, Double>> computeScore(
+      Map<CostFunction, Double> functions, ClusterInfo clusterInfo) {
+    return functions.entrySet().stream()
+        .filter(e -> e.getKey() instanceof HasBrokerCost)
+        .map(e -> Map.entry((HasBrokerCost) e.getKey(), e.getValue()))
+        .map(
+            functionWeight ->
+                functionWeight
+                    .getKey()
+                    // Execute all cost functions
+                    .brokerCost(clusterInfo)
+                    .value()
+                    .entrySet()
+                    .stream()
+                    // Weight on cost functions result
+                    .map(
+                        IdScore ->
+                            Map.entry(
+                                IdScore.getKey(), IdScore.getValue() * functionWeight.getValue()))
+                    .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)))
+        .collect(Collectors.toUnmodifiableList());
+  }
+
   // visible for testing
   static Optional<Map.Entry<ReplicaInfo, Double>> bestPartition(
       List<ReplicaInfo> partitions, List<Map<Integer, Double>> scores) {
@@ -100,7 +135,7 @@ public class StrictCostDispatcher implements Dispatcher {
         .port(port)
         .fetcher(
             Fetcher.of(
-                functions.stream()
+                functions.keySet().stream()
                     .map(CostFunction::fetcher)
                     .collect(Collectors.toUnmodifiableList())))
         .build();
@@ -111,10 +146,49 @@ public class StrictCostDispatcher implements Dispatcher {
     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))
-        .forEach(e -> jmxPorts.put(Integer.parseInt(e.getKey()), Integer.parseInt(e.getValue())));
+    jmxPorts.putAll(PartitionerUtils.parseIdJMXPort(config));
+
+    functions = parseCostFunctionWeight(config);
+  }
+
+  /**
+   * Helps parse cost-function names and weights. The format of the key and value is "<CostFunction
+   * name>"="<weight>". For example, {"org.astraea.cost.brokersMetrics.BrokerInputCost", "20"} will
+   * be parsed to {(BrokerInputCost object), 20.0}.
+   *
+   * @param config that contains cost-function names and its corresponding weight
+   * @return pairs of cost-function object and its corresponding weight
+   */
+  public static Map<CostFunction, Double> parseCostFunctionWeight(Configuration config) {
+    return config.entrySet().stream()
+        .map(
+            nameAndWeight -> {
+              Class<?> name;
+              double weight;
+              try {
+                name = Class.forName(nameAndWeight.getKey());
+                weight = Double.parseDouble(nameAndWeight.getValue());
+                if (weight < 0.0)
+                  throw new IllegalArgumentException("Cost-function weight should not be negative");
+              } catch (ClassNotFoundException ignore) {
+                /* To delete all config option that is not for configuring cost-function. */
+                return null;
+              }
+              return Map.entry(name, weight);
+            })
+        .filter(Objects::nonNull)
+        .filter(e -> CostFunction.class.isAssignableFrom(e.getKey()))
+        .map(
+            e -> {
+              try {
+                return Map.entry(
+                    (CostFunction) e.getKey().getConstructor().newInstance(), e.getValue());
+              } catch (Exception ex) {
+                ex.printStackTrace();
+                throw new IllegalArgumentException(ex);
+              }
+            })
+        .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
   }
 
   // visible for testing


Test patch:
diff --git a/app/src/test/java/org/astraea/partitioner/PartitionerUtilsTest.java b/app/src/test/java/org/astraea/partitioner/PartitionerUtilsTest.java
index 808ad7fa..e99aad84 100644
--- a/app/src/test/java/org/astraea/partitioner/PartitionerUtilsTest.java
+++ b/app/src/test/java/org/astraea/partitioner/PartitionerUtilsTest.java
@@ -3,6 +3,7 @@ package org.astraea.partitioner;
 import static org.astraea.partitioner.PartitionerUtils.allPoisson;
 import static org.astraea.partitioner.PartitionerUtils.doPoisson;
 import static org.astraea.partitioner.PartitionerUtils.factorial;
+import static org.astraea.partitioner.PartitionerUtils.parseIdJMXPort;
 import static org.astraea.partitioner.PartitionerUtils.weightPoisson;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 
@@ -54,4 +55,22 @@ public class PartitionerUtilsTest {
     Assertions.assertEquals(weightPoisson(0.0, 1.0), 20);
     Assertions.assertEquals(weightPoisson(0.8, 1.0), 1);
   }
+
+  @Test
+  void testParseIdJMXPort() {
+    var config =
+        Configuration.of(Map.of("broker.1001.jmx.port", "8000", "broker.1002.jmx.port", "8001"));
+    var ans = parseIdJMXPort(config);
+    Assertions.assertEquals(2, ans.size());
+    Assertions.assertEquals(8000, ans.get(1001));
+    Assertions.assertEquals(8001, ans.get(1002));
+
+    config = Configuration.of(Map.of("jmx.port", "8000", "broker.1002.jmx.port", "8001"));
+    ans = parseIdJMXPort(config);
+    Assertions.assertEquals(1, ans.size());
+    Assertions.assertEquals(8001, ans.get(1002));
+
+    var config3 = Configuration.of(Map.of("broker.id.jmx.port", "8000"));
+    Assertions.assertThrows(NumberFormatException.class, () -> parseIdJMXPort(config3));
+  }
 }
diff --git a/app/src/test/java/org/astraea/partitioner/StrictCostDispatcherTest.java b/app/src/test/java/org/astraea/partitioner/StrictCostDispatcherTest.java
index 46091e45..b91cae76 100644
--- a/app/src/test/java/org/astraea/partitioner/StrictCostDispatcherTest.java
+++ b/app/src/test/java/org/astraea/partitioner/StrictCostDispatcherTest.java
@@ -11,6 +11,8 @@ import org.astraea.cost.ClusterInfo;
 import org.astraea.cost.HasBrokerCost;
 import org.astraea.cost.NodeInfo;
 import org.astraea.cost.ReplicaInfo;
+import org.astraea.cost.ThroughputCost;
+import org.astraea.cost.brokersMetrics.BrokerInputCost;
 import org.astraea.metrics.collector.Fetcher;
 import org.astraea.metrics.collector.Receiver;
 import org.junit.jupiter.api.Assertions;
@@ -42,6 +44,20 @@ public class StrictCostDispatcherTest {
       dispatcher.configure(Configuration.of(Map.of(StrictCostDispatcher.JMX_PORT, "12345")));
       Assertions.assertEquals(12345, dispatcher.jmxPort(0));
     }
+
+    try (var dispatcher = new StrictCostDispatcher()) {
+      // Test for negative weight.
+      Assertions.assertThrows(
+          IllegalArgumentException.class,
+          () ->
+              dispatcher.configure(Configuration.of(Map.of(ThroughputCost.class.getName(), "-1"))));
+
+      // Test for cost functions configuring
+      dispatcher.configure(
+          Configuration.of(
+              Map.of(ThroughputCost.class.getName(), "0.1", BrokerInputCost.class.getName(), "2")));
+      Assertions.assertEquals(2, dispatcher.functions.size());
+    }
   }
 
   @Test
@@ -71,11 +87,10 @@ public class StrictCostDispatcherTest {
     var receiver = Mockito.mock(Receiver.class);
     Mockito.when(receiver.current()).thenReturn(List.of());
 
-    var costFunction =
+    var costFunction1 =
         new HasBrokerCost() {
           @Override
           public BrokerCost brokerCost(ClusterInfo clusterInfo) {
-            var partitionInfos = clusterInfo.availablePartitions("aa");
             var brokerCost =
                 clusterInfo.allBeans().keySet().stream()
                     .collect(
@@ -84,14 +99,31 @@ public class StrictCostDispatcherTest {
             return () -> brokerCost;
           }
 
+          @Override
+          public Fetcher fetcher() {
+            return client -> List.of();
+          }
+        };
+    var costFunction2 =
+        new HasBrokerCost() {
+          @Override
+          public BrokerCost brokerCost(ClusterInfo clusterInfo) {
+            var brokerCost =
+                clusterInfo.allBeans().keySet().stream()
+                    .collect(
+                        Collectors.toMap(
+                            Function.identity(), id -> id.equals(n0.id()) ? 0.6D : 0.8D));
+            return () -> brokerCost;
+          }
+
           @Override
           public Fetcher fetcher() {
             return client -> List.of();
           }
         };
 
-    var dispatcher =
-        new StrictCostDispatcher(List.of(costFunction)) {
+    try (var dispatcher =
+        new StrictCostDispatcher(List.of(costFunction1, costFunction2)) {
           @Override
           Receiver receiver(String host, int port) {
             return receiver;
@@ -101,20 +133,67 @@ public class StrictCostDispatcherTest {
           int jmxPort(int id) {
             return 0;
           }
-        };
-    var clusterInfo = Mockito.mock(ClusterInfo.class);
+        }) {
+      var clusterInfo = Mockito.mock(ClusterInfo.class);
+
+      // there is no available partition
+      Mockito.when(clusterInfo.availablePartitionLeaders("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 no available partition
-    Mockito.when(clusterInfo.availablePartitionLeaders("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
+      Mockito.when(clusterInfo.availablePartitionLeaders("aa")).thenReturn(List.of(p0));
+      Assertions.assertEquals(1, dispatcher.partition("aa", new byte[0], new byte[0], clusterInfo));
 
-    // there is only one available partition
-    Mockito.when(clusterInfo.availablePartitionLeaders("aa")).thenReturn(List.of(p0));
-    Assertions.assertEquals(1, dispatcher.partition("aa", new byte[0], new byte[0], clusterInfo));
+      // there is no beans, so it just returns the partition.
+      Mockito.when(clusterInfo.availablePartitionLeaders("aa")).thenReturn(List.of(p0, p1));
+      Assertions.assertEquals(2, dispatcher.partition("aa", new byte[0], new byte[0], clusterInfo));
+
+      // cost functions with weight
+      dispatcher.functions = Map.of(costFunction1, 0.1, costFunction2, 0.9);
+      Assertions.assertEquals(1, dispatcher.partition("aa", new byte[0], new byte[0], clusterInfo));
+    }
+  }
+
+  @Test
+  void testParseCostFunctionWeight() {
+    var config =
+        Configuration.of(
+            Map.of(
+                "org.astraea.cost.brokersMetrics.BrokerInputCost",
+                "20",
+                "org.astraea.cost.brokersMetrics.BrokerOutputCost",
+                "1.25"));
+    var ans = StrictCostDispatcher.parseCostFunctionWeight(config);
+    Assertions.assertEquals(2, ans.size());
+    for (var entry : ans.entrySet()) {
+      if (entry
+          .getKey()
+          .getClass()
+          .getName()
+          .equals("org.astraea.cost.brokersMetrics.BrokerInputCost")) {
+        Assertions.assertEquals(20.0, entry.getValue());
+      } else if (entry
+          .getKey()
+          .getClass()
+          .getName()
+          .equals("org.astraea.cost.brokersMetrics.BrokerOutputCost")) {
+        Assertions.assertEquals(1.25, entry.getValue());
+      } else {
+        Assertions.assertEquals(0.0, entry.getValue());
+      }
+    }
 
-    // there is no beans, so it just returns the partition.
-    Mockito.when(clusterInfo.availablePartitionLeaders("aa")).thenReturn(List.of(p0, p1));
-    Assertions.assertEquals(2, dispatcher.partition("aa", new byte[0], new byte[0], clusterInfo));
+    // test negative weight
+    var config2 =
+        Configuration.of(
+            Map.of(
+                "org.astraea.cost.brokersMetrics.BrokerInputCost",
+                "-20",
+                "org.astraea.cost.brokersMetrics.BrokerOutputCost",
+                "1.25"));
+    Assertions.assertThrows(
+        IllegalArgumentException.class,
+        () -> StrictCostDispatcher.parseCostFunctionWeight(config2));
   }
 }