Repository: skiptests/astraea
Pull request #347 (2022-05-16 10:16:09)

Add topic/partitions to broker for restful APIs
產生下方的統計畫面
```json
    "brokers": [
        {
            "id": 1001,
            "topics": [
                {
                    "topic": "log",
                    "partitionCount": 8
                },
                {
                    "topic": "table",
                    "partitionCount": 10
                }
            ],
```


Patch:
diff --git a/app/src/main/java/org/astraea/admin/Admin.java b/app/src/main/java/org/astraea/admin/Admin.java
index bfb03694..e6106afa 100644
--- a/app/src/main/java/org/astraea/admin/Admin.java
+++ b/app/src/main/java/org/astraea/admin/Admin.java
@@ -95,12 +95,22 @@ public interface Admin extends Closeable {
   /** @return all brokers' ids */
   Set<Integer> brokerIds();
 
+  /**
+   * list all partitions belongs to input brokers
+   *
+   * @param brokerId to search
+   * @return all partition belongs to brokers
+   */
+  default Set<TopicPartition> partitions(int brokerId) {
+    return partitions(topicNames(), Set.of(brokerId)).getOrDefault(brokerId, Set.of());
+  }
+
   /**
    * @param topics topic names
-   * @param brokersID brokers ID
+   * @param brokerIds brokers ID
    * @return the partitions of brokers
    */
-  Set<TopicPartition> partitionsOfBrokers(Set<String> topics, Set<Integer> brokersID);
+  Map<Integer, Set<TopicPartition>> partitions(Set<String> topics, Set<Integer> brokerIds);
 
   /** @return data folders of all broker nodes */
   default Map<Integer, Set<String>> brokerFolders() {
diff --git a/app/src/main/java/org/astraea/admin/Builder.java b/app/src/main/java/org/astraea/admin/Builder.java
index 2471fc66..7eeec61d 100644
--- a/app/src/main/java/org/astraea/admin/Builder.java
+++ b/app/src/main/java/org/astraea/admin/Builder.java
@@ -278,11 +278,19 @@ public class Builder {
     }
 
     @Override
-    public Set<TopicPartition> partitionsOfBrokers(Set<String> topics, Set<Integer> brokersID) {
+    public Map<Integer, Set<TopicPartition>> partitions(
+        Set<String> topics, Set<Integer> brokerIds) {
       return replicas(topics).entrySet().stream()
-          .filter(e -> e.getValue().stream().anyMatch(r -> brokersID.contains(r.broker())))
-          .map(Map.Entry::getKey)
-          .collect(Collectors.toSet());
+          .flatMap(
+              e -> e.getValue().stream().map(replica -> Map.entry(replica.broker(), e.getKey())))
+          .filter(e -> brokerIds.contains(e.getKey()))
+          .collect(Collectors.groupingBy(Map.Entry::getKey))
+          .entrySet()
+          .stream()
+          .collect(
+              Collectors.toMap(
+                  Map.Entry::getKey,
+                  e -> e.getValue().stream().map(Map.Entry::getValue).collect(Collectors.toSet())));
     }
 
     @Override
diff --git a/app/src/main/java/org/astraea/performance/Performance.java b/app/src/main/java/org/astraea/performance/Performance.java
index fddc66fc..a3c16afd 100644
--- a/app/src/main/java/org/astraea/performance/Performance.java
+++ b/app/src/main/java/org/astraea/performance/Performance.java
@@ -4,6 +4,7 @@ import com.beust.jcommander.Parameter;
 import java.io.IOException;
 import java.nio.file.Path;
 import java.time.Duration;
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.HashSet;
 import java.util.List;
@@ -121,7 +122,7 @@ public class Performance {
           .create();
 
       Utils.waitFor(() -> topicAdmin.topicNames().contains(param.topic));
-      partitions = partition(param, topicAdmin);
+      partitions = new ArrayList<>(partition(param, topicAdmin));
     }
 
     var consumerMetrics =
@@ -232,14 +233,16 @@ public class Performance {
   }
 
   // visible for test
-  static List<Integer> partition(Argument param, Admin topicAdmin) {
+  static Set<Integer> partition(Argument param, Admin topicAdmin) {
     if (positiveSpecifyBroker(param)) {
       return topicAdmin
-          .partitionsOfBrokers(Set.of(param.topic), new HashSet<>(param.specifyBroker))
+          .partitions(Set.of(param.topic), new HashSet<>(param.specifyBroker))
+          .values()
           .stream()
+          .flatMap(Collection::stream)
           .map(TopicPartition::partition)
-          .collect(Collectors.toList());
-    } else return List.of(-1);
+          .collect(Collectors.toSet());
+    } else return Set.of(-1);
   }
 
   private static boolean positiveSpecifyBroker(Argument param) {
diff --git a/app/src/main/java/org/astraea/web/BrokerHandler.java b/app/src/main/java/org/astraea/web/BrokerHandler.java
index 51bad7da..b40f7727 100644
--- a/app/src/main/java/org/astraea/web/BrokerHandler.java
+++ b/app/src/main/java/org/astraea/web/BrokerHandler.java
@@ -9,6 +9,7 @@ import java.util.stream.Collectors;
 import java.util.stream.StreamSupport;
 import org.astraea.admin.Admin;
 import org.astraea.admin.Config;
+import org.astraea.admin.TopicPartition;
 
 class BrokerHandler implements Handler {
 
@@ -28,20 +29,39 @@ class BrokerHandler implements Handler {
 
   @Override
   public JsonObject get(Optional<String> target, Map<String, String> queries) {
+
     var brokers =
         admin.brokers(brokers(target)).entrySet().stream()
-            .map(e -> new Broker(e.getKey(), e.getValue()))
+            .map(e -> new Broker(e.getKey(), admin.partitions(e.getKey()), e.getValue()))
             .collect(Collectors.toUnmodifiableList());
     if (target.isPresent() && brokers.size() == 1) return brokers.get(0);
     return new Brokers(brokers);
   }
 
+  static class Topic implements JsonObject {
+    final String topic;
+    final int partitionCount;
+
+    Topic(String topic, int partitionCount) {
+      this.topic = topic;
+      this.partitionCount = partitionCount;
+    }
+  }
+
   static class Broker implements JsonObject {
     final int id;
+    final List<Topic> topics;
     final Map<String, String> configs;
 
-    Broker(int id, Config configs) {
+    Broker(int id, Set<TopicPartition> topicPartitions, Config configs) {
       this.id = id;
+      this.topics =
+          topicPartitions.stream()
+              .collect(Collectors.groupingBy(TopicPartition::topic))
+              .entrySet()
+              .stream()
+              .map(e -> new Topic(e.getKey(), e.getValue().size()))
+              .collect(Collectors.toUnmodifiableList());
       this.configs =
           StreamSupport.stream(configs.spliterator(), false)
               .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));


Test patch:
diff --git a/app/src/test/java/org/astraea/web/BrokerHandlerTest.java b/app/src/test/java/org/astraea/web/BrokerHandlerTest.java
index 7d145f16..155978c5 100644
--- a/app/src/test/java/org/astraea/web/BrokerHandlerTest.java
+++ b/app/src/test/java/org/astraea/web/BrokerHandlerTest.java
@@ -4,6 +4,8 @@ import java.util.Map;
 import java.util.NoSuchElementException;
 import java.util.Optional;
 import java.util.Set;
+import java.util.concurrent.TimeUnit;
+import org.astraea.Utils;
 import org.astraea.admin.Admin;
 import org.astraea.service.RequireBrokerCluster;
 import org.junit.jupiter.api.Assertions;
@@ -12,8 +14,11 @@ import org.junit.jupiter.api.Test;
 public class BrokerHandlerTest extends RequireBrokerCluster {
 
   @Test
-  void testListBrokers() {
+  void testListBrokers() throws InterruptedException {
+    var topic = Utils.randomString(10);
     try (Admin admin = Admin.of(bootstrapServers())) {
+      admin.creator().topic(topic).numberOfPartitions(10).create();
+      TimeUnit.SECONDS.sleep(2);
       var handler = new BrokerHandler(admin);
       var response =
           Assertions.assertInstanceOf(
@@ -24,6 +29,8 @@ public class BrokerHandlerTest extends RequireBrokerCluster {
               id ->
                   Assertions.assertEquals(
                       1, response.brokers.stream().filter(b -> b.id == id).count()));
+      response.brokers.forEach(
+          b -> Assertions.assertTrue(b.topics.stream().anyMatch(t -> t.topic.equals(topic))));
     }
   }
 
@@ -46,9 +53,12 @@ public class BrokerHandlerTest extends RequireBrokerCluster {
   }
 
   @Test
-  void testQuerySingleBroker() {
+  void testQuerySingleBroker() throws InterruptedException {
+    var topic = Utils.randomString(10);
     var brokerId = brokerIds().iterator().next();
     try (Admin admin = Admin.of(bootstrapServers())) {
+      admin.creator().topic(topic).numberOfPartitions(10).create();
+      TimeUnit.SECONDS.sleep(2);
       var handler = new BrokerHandler(admin);
       var broker =
           Assertions.assertInstanceOf(
@@ -56,6 +66,7 @@ public class BrokerHandlerTest extends RequireBrokerCluster {
               handler.get(Optional.of(String.valueOf(brokerId)), Map.of()));
       Assertions.assertEquals(brokerId, broker.id);
       Assertions.assertNotEquals(0, broker.configs.size());
+      Assertions.assertTrue(broker.topics.stream().anyMatch(t -> t.topic.equals(topic)));
     }
   }