Repository: skiptests/astraea
Pull request #108 (2021-11-15 08:15:02)

印出produce/consume時間
測試時，為了方便紀錄produce/consume所花費的時間，故先新增"印出時間"。


Patch:
diff --git a/README.md b/README.md
index 9a459681..aaa3a4cb 100644
--- a/README.md
+++ b/README.md
@@ -7,6 +7,7 @@ a collection of tools used to balance Kafka data
 - Ching-Hong Fang <fjh7777@gmail.com>
 - Zheng-Xian Li <garyparrottt@gmail.com>
 - Xiang-Jun Sun <sean0651101@gmail.com>
+- Zhi-Mao Teng <zhimao.teng@gmail.com>
 
 # Kafka Tools
 
@@ -18,6 +19,7 @@ This project offers many kafka tools to simplify the life for kafka users.
 4. [Kafka official tool](#kafka-official-tool): run any one specific kafka official tool. All you have to prepare is the docker env.
 5. [Kafka metric client](#kafka-metric-client): utility for accessing kafka Mbean metrics via JMX.
 6. [Replica Collie](#replica-collie): move replicas from brokers to others. You can use this tool to obstruct specific brokers from hosting specific topics.
+7. [Kafka partition score](#Kafka-partition-score): score all broker's partitions. 
 
 [Release page](https://github.com/skiptests/astraea/releases) offers the uber jar including all tools.
 ```shell
@@ -221,3 +223,18 @@ This tool offers an effective way to migrate all replicas from specific brokers
 1. --bootstrap.servers: the server to connect to
 2. --topics: the topics to be moved
 3. --admin.props.file: the file path containing the properties to be passed to kafka admin
+
+## Kafka Partition Score
+
+This tool will score the partition on brokers, the higher score the heavier load.
+
+### Start scoring partitions on broker address "192.168.103.39:9092"
+
+```shell
+./gradlew run --args="score --bootstrap.servers 192.168.103.39:9092"
+```
+
+### Partition Score Configurations
+
+1. --bootstrap.servers: the server to connect to
+
diff --git a/app/src/main/java/org/astraea/App.java b/app/src/main/java/org/astraea/App.java
index a6ef0c43..c00c1520 100644
--- a/app/src/main/java/org/astraea/App.java
+++ b/app/src/main/java/org/astraea/App.java
@@ -9,6 +9,7 @@ import org.astraea.metrics.MetricExplorer;
 import org.astraea.performance.Performance;
 import org.astraea.topic.ReplicaCollie;
 import org.astraea.topic.TopicExplorer;
+import org.astraea.topic.cost.PartitionScore;
 
 public class App {
   private static final Map<String, Class<?>> MAIN_CLASSES =
@@ -16,6 +17,7 @@ public class App {
           "offset", TopicExplorer.class,
           "metrics", MetricExplorer.class,
           "replica", ReplicaCollie.class,
+          "score", PartitionScore.class,
           "performance", Performance.class);
 
   static void execute(Map<String, Class<?>> mains, List<String> args) throws Throwable {
diff --git a/app/src/main/java/org/astraea/performance/Tracker.java b/app/src/main/java/org/astraea/performance/Tracker.java
index 3653faa5..42ba9409 100644
--- a/app/src/main/java/org/astraea/performance/Tracker.java
+++ b/app/src/main/java/org/astraea/performance/Tracker.java
@@ -1,5 +1,6 @@
 package org.astraea.performance;
 
+import java.time.Duration;
 import java.util.List;
 import org.astraea.concurrent.ThreadPool;
 
@@ -8,6 +9,7 @@ public class Tracker implements ThreadPool.Executor {
   private final List<Metrics> producerData;
   private final List<Metrics> consumerData;
   private final long records;
+  private long start = 0L;
 
   public Tracker(List<Metrics> producerData, List<Metrics> consumerData, long records) {
     this.producerData = producerData;
@@ -38,6 +40,16 @@ public class Tracker implements ThreadPool.Executor {
 
     if (completed == 0) return false;
 
+    if (start == 0L) start = System.currentTimeMillis();
+    var duration = Duration.ofMillis(System.currentTimeMillis() - start);
+    System.out.println(
+        "Time: "
+            + duration.toHoursPart()
+            + "hr "
+            + duration.toMinutesPart()
+            + "min "
+            + duration.toSecondsPart()
+            + "sec");
     System.out.printf("producers完成度: %.2f%%%n", ((double) completed * 100.0 / (double) records));
     System.out.printf("  輸出%.3fMB/second%n", bytes);
     System.out.println("  發送max latency: " + max + "ms");
diff --git a/app/src/main/java/org/astraea/topic/cost/CalculateUtils.java b/app/src/main/java/org/astraea/topic/cost/CalculateUtils.java
new file mode 100644
index 00000000..e7204bb9
--- /dev/null
+++ b/app/src/main/java/org/astraea/topic/cost/CalculateUtils.java
@@ -0,0 +1,79 @@
+package org.astraea.topic.cost;
+
+import java.util.HashMap;
+import java.util.Map;
+import org.apache.kafka.common.TopicPartition;
+
+public class CalculateUtils {
+  public static Map<Integer, Map<TopicPartition, Double>> getLoad(
+      Map<Integer, Map<TopicPartition, Integer>> brokerPartitionSize,
+      Map<String, Integer> retentionMillis) {
+    Map<Integer, Map<TopicPartition, Double>> brokerPartitionLoad = new HashMap<>();
+    double load;
+    for (var broker : brokerPartitionSize.keySet()) {
+      Map<TopicPartition, Double> partitionLoad = new HashMap<>();
+      for (var partition : brokerPartitionSize.get(broker).keySet()) {
+        load =
+            (double) brokerPartitionSize.get(broker).get(partition)
+                / retentionMillis.get(partition.topic());
+        partitionLoad.put(partition, load);
+      }
+      brokerPartitionLoad.put(broker, partitionLoad);
+    }
+    return brokerPartitionLoad;
+  }
+
+  public static Map<Integer, Map<TopicPartition, Double>> getScore(
+      Map<Integer, Map<TopicPartition, Double>> load) {
+    Map<Integer, Double> brokerLoad = new HashMap<>();
+    Map<TopicPartition, Double> partitionLoad = new HashMap<>();
+    Map<Integer, Double> partitionSD = new HashMap<>();
+    Map<Integer, Double> partitionMean = new HashMap<>();
+    Map<Integer, Map<TopicPartition, Double>> brokerPartitionScore = new HashMap<>();
+
+    double mean = 0f, LoadSQR = 0f, SD, brokerSD;
+    int partitionNum;
+    for (var broker : load.keySet()) {
+      for (var topicPartition : load.get(broker).keySet()) {
+        mean += load.get(broker).get(topicPartition);
+        partitionLoad.put(topicPartition, load.get(broker).get(topicPartition));
+        LoadSQR += Math.pow(load.get(broker).get(topicPartition), 2);
+      }
+      partitionNum = load.get(broker).keySet().size();
+      brokerLoad.put(broker, mean);
+      mean /= partitionNum;
+      partitionMean.put(broker, mean);
+      SD = Math.pow((LoadSQR - mean * mean * partitionNum) / partitionNum, 0.5);
+      partitionSD.put(broker, SD);
+      mean = 0f;
+      LoadSQR = 0f;
+    }
+
+    for (var i : brokerLoad.keySet()) {
+      mean += brokerLoad.get(i);
+      LoadSQR += Math.pow(brokerLoad.get(i), 2);
+    }
+    mean /= brokerLoad.keySet().size();
+    brokerSD =
+        Math.pow(
+            (LoadSQR - mean * mean * brokerLoad.keySet().size()) / brokerLoad.keySet().size(), 0.5);
+    for (var broker : load.keySet()) {
+      Map<TopicPartition, Double> partitionScore = new HashMap<>();
+      for (var topicPartition : load.get(broker).keySet()) {
+        if (brokerLoad.get(broker) - mean > 0) {
+          partitionScore.put(
+              topicPartition,
+              ((brokerLoad.get(broker) - mean) / brokerSD)
+                  * ((partitionLoad.get(topicPartition) - partitionMean.get(broker))
+                      / partitionSD.get(broker))
+                  * 60.0);
+
+        } else {
+          partitionScore.put(topicPartition, 0.0);
+        }
+        brokerPartitionScore.put(broker, partitionScore);
+      }
+    }
+    return brokerPartitionScore;
+  }
+}
diff --git a/app/src/main/java/org/astraea/topic/cost/GetPartitionInf.java b/app/src/main/java/org/astraea/topic/cost/GetPartitionInf.java
new file mode 100644
index 00000000..f2f79511
--- /dev/null
+++ b/app/src/main/java/org/astraea/topic/cost/GetPartitionInf.java
@@ -0,0 +1,29 @@
+package org.astraea.topic.cost;
+
+import java.util.*;
+import org.apache.kafka.common.TopicPartition;
+import org.astraea.topic.TopicAdmin;
+
+public class GetPartitionInf {
+  static Map<Integer, Map<TopicPartition, Integer>> getSize(TopicAdmin client) {
+    Map<Integer, Map<TopicPartition, Integer>> brokerPartitionSize = new HashMap<>();
+    for (var broker : client.brokerIds()) {
+      Map<TopicPartition, Integer> partitionSize = new HashMap<>();
+      for (var partition : client.replicas(client.topics().keySet()).keySet())
+        if (client.replicas(client.topics().keySet()).get(partition).get(0).broker() == broker)
+          partitionSize.put(
+              partition,
+              (int) client.replicas(client.topics().keySet()).get(partition).get(0).size());
+      brokerPartitionSize.put(broker, partitionSize);
+    }
+    return brokerPartitionSize;
+  }
+
+  static Map<String, Integer> getRetentionMillis(TopicAdmin client) {
+    Map<String, Integer> retentionTime = new HashMap<>();
+    for (var topic : client.topics().keySet()) {
+      retentionTime.put(topic, Integer.parseInt(client.topics().get(topic).get("retention.ms")));
+    }
+    return retentionTime;
+  }
+}
diff --git a/app/src/main/java/org/astraea/topic/cost/PartitionScore.java b/app/src/main/java/org/astraea/topic/cost/PartitionScore.java
new file mode 100644
index 00000000..fcda459c
--- /dev/null
+++ b/app/src/main/java/org/astraea/topic/cost/PartitionScore.java
@@ -0,0 +1,61 @@
+package org.astraea.topic.cost;
+
+import java.util.*;
+import java.util.stream.Collectors;
+import org.apache.kafka.common.TopicPartition;
+import org.astraea.argument.ArgumentUtil;
+import org.astraea.argument.BasicArgument;
+import org.astraea.topic.TopicAdmin;
+
+public class PartitionScore {
+  TopicAdmin admin;
+  Map<Integer, Map<TopicPartition, Integer>> brokerPartitionSize;
+  Map<Integer, Map<TopicPartition, Double>> score;
+  Map<Integer, Map<TopicPartition, Double>> load;
+  Map<String, Integer> retentionMillis;
+
+  public PartitionScore(String address) {
+    admin = TopicAdmin.of(address);
+    brokerPartitionSize = GetPartitionInf.getSize(admin);
+    retentionMillis = GetPartitionInf.getRetentionMillis(admin);
+    load = CalculateUtils.getLoad(brokerPartitionSize, retentionMillis);
+    score = CalculateUtils.getScore(load);
+  }
+
+  public void printScore(Map<Integer, Map<TopicPartition, Double>> score) {
+    Set<TopicPartition> partitionGood = new HashSet<>();
+    Map<Integer, Boolean> BrokerGood = new HashMap<>();
+    for (var broker : score.keySet()) {
+      BrokerGood.put(broker, true);
+      for (var tp : score.get(broker).keySet())
+        if (score.get(broker).get(tp) > 0) BrokerGood.put(broker, false);
+      System.out.println();
+      if (BrokerGood.get(broker)) {
+        System.out.println("broker: " + broker + " is balanced.");
+      } else {
+        System.out.println("broker: " + broker);
+        for (var tp : score.get(broker).keySet()) {
+          if (score.get(broker).get(tp) > 0) {
+            System.out.println(tp + ": " + score.get(broker).get(tp));
+          } else {
+            partitionGood.add(tp);
+          }
+        }
+      }
+    }
+
+    System.out.print(
+        "\nThe following partitions are balanced: "
+            + partitionGood.stream()
+                .map(String::valueOf)
+                .collect(Collectors.joining(", ", "[", "]")));
+  }
+
+  public static void main(String[] args) {
+    var argument = ArgumentUtil.parseArgument(new Argument(), args);
+    var partitionScore = new PartitionScore(argument.brokers);
+    partitionScore.printScore(partitionScore.score);
+  }
+
+  static class Argument extends BasicArgument {}
+}


Test patch:
diff --git a/app/src/test/java/org/astraea/topic/cost/CalculateUtilsTest.java b/app/src/test/java/org/astraea/topic/cost/CalculateUtilsTest.java
new file mode 100644
index 00000000..80dd0748
--- /dev/null
+++ b/app/src/test/java/org/astraea/topic/cost/CalculateUtilsTest.java
@@ -0,0 +1,86 @@
+package org.astraea.topic.cost;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+import java.util.HashMap;
+import java.util.Map;
+import org.apache.kafka.common.TopicPartition;
+import org.astraea.service.RequireBrokerCluster;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+
+public class CalculateUtilsTest extends RequireBrokerCluster {
+  static Map<Integer, Map<TopicPartition, Integer>> fakeBrokerPartitionSize = new HashMap<>();
+  static Map<String, Integer> fakeRetentionMillis = new HashMap<>();
+  static Map<Integer, Map<TopicPartition, Double>> fakePartitionLoad = new HashMap<>();
+
+  @BeforeAll
+  static void setup() {
+    // set partition size and retention time
+    Map<TopicPartition, Integer> fakePartitionSize = new HashMap<>();
+    fakePartitionSize.put(new TopicPartition("test0", 0), 100000000);
+    fakePartitionSize.put(new TopicPartition("test0", 1), 200000000);
+    fakePartitionSize.put(new TopicPartition("test0", 2), 300000000);
+    fakePartitionSize.put(new TopicPartition("test0", 3), 400000000);
+    fakeBrokerPartitionSize.put(0, fakePartitionSize);
+    fakePartitionSize = new HashMap<>();
+    fakePartitionSize.put(new TopicPartition("test1", 0), 500000000);
+    fakePartitionSize.put(new TopicPartition("test1", 1), 600000000);
+    fakePartitionSize.put(new TopicPartition("test1", 2), 700000000);
+    fakePartitionSize.put(new TopicPartition("test1", 3), 800000000);
+    fakeBrokerPartitionSize.put(1, fakePartitionSize);
+    fakeRetentionMillis.put("test0", 604800000);
+    fakeRetentionMillis.put("test1", 604800000);
+
+    // set partition load
+    fakePartitionLoad = new HashMap<>();
+    Map<TopicPartition, Double> fakeTopicLoad = new HashMap<>();
+    fakeTopicLoad.put(new TopicPartition("test0", 0), 0.1);
+    fakeTopicLoad.put(new TopicPartition("test0", 1), 0.2);
+    fakeTopicLoad.put(new TopicPartition("test0", 2), 0.3);
+    fakeTopicLoad.put(new TopicPartition("test0", 3), 0.4);
+    fakePartitionLoad.put(0, fakeTopicLoad);
+    fakeTopicLoad = new HashMap<>();
+    fakeTopicLoad.put(new TopicPartition("test1", 0), 1.5);
+    fakeTopicLoad.put(new TopicPartition("test1", 1), 1.6);
+    fakeTopicLoad.put(new TopicPartition("test1", 2), 1.7);
+    fakeTopicLoad.put(new TopicPartition("test1", 3), 1.8);
+    fakePartitionLoad.put(1, fakeTopicLoad);
+  }
+
+  @Test
+  void testGetLoad() {
+    var Load = CalculateUtils.getLoad(fakeBrokerPartitionSize, fakeRetentionMillis);
+    assertEquals(2, Load.size());
+    assertEquals(4, Load.get(0).size());
+    assertEquals(4, Load.get(1).size());
+    assertEquals(0.165, round(Load.get(0).get(new TopicPartition("test0", 0))));
+    assertEquals(0.331, round(Load.get(0).get(new TopicPartition("test0", 1))));
+    assertEquals(0.496, round(Load.get(0).get(new TopicPartition("test0", 2))));
+    assertEquals(0.661, round(Load.get(0).get(new TopicPartition("test0", 3))));
+    assertEquals(0.827, round(Load.get(1).get(new TopicPartition("test1", 0))));
+    assertEquals(0.992, round(Load.get(1).get(new TopicPartition("test1", 1))));
+    assertEquals(1.157, round(Load.get(1).get(new TopicPartition("test1", 2))));
+    assertEquals(1.323, round(Load.get(1).get(new TopicPartition("test1", 3))));
+  }
+
+  @Test
+  void testGetScore() {
+    var Score = CalculateUtils.getScore(fakePartitionLoad);
+    assertEquals(2, Score.size());
+    assertEquals(4, Score.get(0).size());
+    assertEquals(4, Score.get(1).size());
+    assertEquals(0.0, round(Score.get(0).get(new TopicPartition("test0", 0))));
+    assertEquals(0.0, round(Score.get(0).get(new TopicPartition("test0", 1))));
+    assertEquals(0.0, round(Score.get(0).get(new TopicPartition("test0", 2))));
+    assertEquals(0.0, round(Score.get(0).get(new TopicPartition("test0", 3))));
+    assertEquals(-80.498, round(Score.get(1).get(new TopicPartition("test1", 0))));
+    assertEquals(-26.833, round(Score.get(1).get(new TopicPartition("test1", 1))));
+    assertEquals(26.833, round(Score.get(1).get(new TopicPartition("test1", 2))));
+    assertEquals(80.498, round(Score.get(1).get(new TopicPartition("test1", 3))));
+  }
+
+  double round(double score) {
+    return Math.round(1000 * score) / 1000.0;
+  }
+}
diff --git a/app/src/test/java/org/astraea/topic/cost/GetPartitionInfTest.java b/app/src/test/java/org/astraea/topic/cost/GetPartitionInfTest.java
new file mode 100644
index 00000000..52d70e24
--- /dev/null
+++ b/app/src/test/java/org/astraea/topic/cost/GetPartitionInfTest.java
@@ -0,0 +1,87 @@
+package org.astraea.topic.cost;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import org.astraea.producer.Producer;
+import org.astraea.producer.Serializer;
+import org.astraea.service.RequireBrokerCluster;
+import org.astraea.topic.TopicAdmin;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+
+public class GetPartitionInfTest extends RequireBrokerCluster {
+  static TopicAdmin admin;
+
+  @BeforeAll
+  static void setup() throws ExecutionException, InterruptedException {
+    admin = TopicAdmin.of(bootstrapServers());
+    Map<Integer, String> topicName = new HashMap<>();
+    topicName.put(0, "testPartitionScore0");
+    topicName.put(1, "testPartitionScore1");
+    topicName.put(2, "testPartitionScore2");
+
+    try (var admin = TopicAdmin.of(bootstrapServers())) {
+      admin
+          .creator()
+          .topic(topicName.get(0))
+          .numberOfPartitions(4)
+          .numberOfReplicas((short) 1)
+          .create();
+      admin
+          .creator()
+          .topic(topicName.get(1))
+          .numberOfPartitions(4)
+          .numberOfReplicas((short) 1)
+          .create();
+      admin
+          .creator()
+          .topic(topicName.get(2))
+          .numberOfPartitions(4)
+          .numberOfReplicas((short) 1)
+          .create();
+      // wait for topic
+      TimeUnit.SECONDS.sleep(5);
+    } catch (InterruptedException | IOException e) {
+      e.printStackTrace();
+    }
+    var producer =
+        Producer.builder().brokers(bootstrapServers()).keySerializer(Serializer.STRING).build();
+    int size = 10000;
+    for (int t = 0; t <= 2; t++) {
+      for (int p = 0; p <= 3; p++) {
+        producer
+            .sender()
+            .topic(topicName.get(t))
+            .partition(p)
+            .value(new byte[size])
+            .run()
+            .toCompletableFuture()
+            .get();
+      }
+      size += 10000;
+    }
+  }
+
+  @Test
+  void getSize() {
+
+    var brokerPartitionSize = GetPartitionInf.getSize(admin);
+    assertEquals(3, brokerPartitionSize.size());
+    assertEquals(
+        3 * 4,
+        brokerPartitionSize.get(0).size()
+            + brokerPartitionSize.get(1).size()
+            + brokerPartitionSize.get(2).size());
+  }
+
+  @Test
+  void getRetentionMillis() {
+    var brokerPartitionRetentionMillis = GetPartitionInf.getRetentionMillis(admin);
+    assertEquals(3, brokerPartitionRetentionMillis.size());
+  }
+}
diff --git a/app/src/test/java/org/astraea/topic/cost/PartitionScoreTest.java b/app/src/test/java/org/astraea/topic/cost/PartitionScoreTest.java
new file mode 100644
index 00000000..8b7c5ad3
--- /dev/null
+++ b/app/src/test/java/org/astraea/topic/cost/PartitionScoreTest.java
@@ -0,0 +1,80 @@
+package org.astraea.topic.cost;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import org.astraea.producer.Producer;
+import org.astraea.producer.Serializer;
+import org.astraea.service.RequireBrokerCluster;
+import org.astraea.topic.TopicAdmin;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+
+public class PartitionScoreTest extends RequireBrokerCluster {
+  static TopicAdmin admin;
+
+  @BeforeAll
+  static void setup() throws ExecutionException, InterruptedException {
+    admin = TopicAdmin.of(bootstrapServers());
+    Map<Integer, String> topicName = new HashMap<>();
+    topicName.put(0, "testPartitionScore0");
+    topicName.put(1, "testPartitionScore1");
+    topicName.put(2, "testPartitionScore2");
+    try (var admin = TopicAdmin.of(bootstrapServers())) {
+      admin
+          .creator()
+          .topic(topicName.get(0))
+          .numberOfPartitions(4)
+          .numberOfReplicas((short) 1)
+          .create();
+      admin
+          .creator()
+          .topic(topicName.get(1))
+          .numberOfPartitions(4)
+          .numberOfReplicas((short) 1)
+          .create();
+      admin
+          .creator()
+          .topic(topicName.get(2))
+          .numberOfPartitions(4)
+          .numberOfReplicas((short) 1)
+          .create();
+      // wait for topic
+      TimeUnit.SECONDS.sleep(5);
+    } catch (InterruptedException | IOException e) {
+      e.printStackTrace();
+    }
+    var producer =
+        Producer.builder().brokers(bootstrapServers()).keySerializer(Serializer.STRING).build();
+    int size = 10000;
+    for (int t = 0; t <= 2; t++) {
+      for (int p = 0; p <= 3; p++) {
+        producer
+            .sender()
+            .topic(topicName.get(t))
+            .partition(p)
+            .value(new byte[size])
+            .run()
+            .toCompletableFuture()
+            .get();
+      }
+      size += 10000;
+    }
+    producer.close();
+  }
+
+  @Test
+  void testGetScore() throws ExecutionException, InterruptedException {
+    PartitionScore partitionScore = new PartitionScore(bootstrapServers());
+    assertEquals(3, partitionScore.score.size());
+    assertEquals(
+        3 * 4,
+        partitionScore.score.get(0).size()
+            + partitionScore.score.get(1).size()
+            + partitionScore.score.get(2).size());
+  }
+}