Repository: skiptests/astraea
Pull request #185 (2022-01-15 16:02:13)

Specify Broker for performance 
在測試過程中產生的需求，要能夠指定某一個broker（或某幾個brokers）發送資料。吃光該broker的throughput，創造出該broker處於高負載的狀況。

chia7712: 為何不考慮直接指定partition? 也就是我們在performance tool裡面多一個參數 可以指定某一個partition來發送資料，如果覺得指定partition不方便的話，也可以改成指定broker id，然後由performance來查找對應的partition，然後指定partition就好

@chinghongfang 你覺得呢？
chinghongfang: 我也覺得在performance tool直接從producer下手比較好。
wycccccc: 抱歉我沒講清楚大概，目前的設計的確是在performance tool中加入--specify.broker參數指定一個或多個broker。由performance來查找對應的partition，其實我沒找到方法，而partitioner自帶`cluster`參數能夠找到對應的partition，我認為讓他去做這件事是比較合理方便的。
chia7712: > 而partitioner自帶cluster參數能夠找到對應的partition，我認為讓他去做這件事是比較合理方便的。

這樣的設計就綁死要用哪一個partitioner，影響的範圍太大，而且performance要去查broker -> partition應該不是大工程才對，尤其可以用我們的`TopicAdmin`去處理
wycccccc: 目前已經修改成只需要設置`--specify.broker`，不再需要partitioner了
chia7712: 參數的部份是不是要支援多個brokers?我看你做的架構應該是可以直接支援才對
wycccccc: >參數的部份是不是要支援多個brokers?我看你做的架構應該是可以直接支援才對

抱歉，我原來就打算做支持多brokers的，但改着改着忘了。已添加相關功能。

Patch:
diff --git a/app/src/main/java/org/astraea/Utils.java b/app/src/main/java/org/astraea/Utils.java
index cdd34473..b022c9ec 100644
--- a/app/src/main/java/org/astraea/Utils.java
+++ b/app/src/main/java/org/astraea/Utils.java
@@ -7,6 +7,10 @@ import java.net.InetAddress;
 import java.net.UnknownHostException;
 import java.nio.file.Files;
 import java.time.Duration;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Supplier;
 import java.util.stream.Stream;
@@ -111,6 +115,22 @@ public final class Utils {
     }
   }
 
+  /**
+   * Separate host and port
+   *
+   * @param address like 0.0.0.0:9092
+   * @return (0.0.0.0,9092)
+   */
+  public static Map<String, Integer> requireHostPort(List<String> address) {
+    var mapAddress = new HashMap<String, Integer>();
+    address.forEach(
+        str ->
+            mapAddress.put(
+                Arrays.asList(str.split(":")).get(0),
+                Integer.parseInt(Arrays.asList(str.split(":")).get(1))));
+    return mapAddress;
+  }
+
   public static int requirePositive(int value) {
     if (value <= 0)
       throw new IllegalArgumentException("the value: " + value + " must be bigger than zero");
diff --git a/app/src/main/java/org/astraea/performance/Performance.java b/app/src/main/java/org/astraea/performance/Performance.java
index 9e63f9ae..5a4cd928 100644
--- a/app/src/main/java/org/astraea/performance/Performance.java
+++ b/app/src/main/java/org/astraea/performance/Performance.java
@@ -5,10 +5,7 @@ import com.beust.jcommander.Parameter;
 import com.beust.jcommander.ParameterException;
 import java.io.IOException;
 import java.time.Duration;
-import java.util.Collection;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
+import java.util.*;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Future;
@@ -64,6 +61,7 @@ public class Performance {
 
   public static Future<String> execute(final Argument param)
       throws InterruptedException, IOException, ExecutionException {
+    List<Integer> partitions;
     var future = new CompletableFuture<String>();
     try (var topicAdmin = TopicAdmin.of(param.props())) {
       topicAdmin
@@ -72,6 +70,8 @@ public class Performance {
           .numberOfPartitions(param.partitions)
           .topic(param.topic)
           .create();
+
+      partitions = partition(param, topicAdmin);
     }
 
     var consumerMetrics =
@@ -117,6 +117,7 @@ public class Performance {
                                     .build(),
                                 param,
                                 producerMetrics.get(i),
+                                partitions,
                                 manager))
                     .collect(Collectors.toUnmodifiableList()))
             .executor(tracker)
@@ -168,13 +169,14 @@ public class Performance {
       Producer<byte[], byte[]> producer,
       Argument param,
       BiConsumer<Long, Long> observer,
+      List<Integer> partitions,
       Manager manager) {
     return new ThreadPool.Executor() {
       @Override
       public State execute() throws InterruptedException {
         // Wait for all consumers get assignment.
         manager.awaitPartitionAssignment();
-
+        var rand = new Random();
         var payload = manager.payload();
         if (payload.isEmpty()) return State.DONE;
 
@@ -182,6 +184,7 @@ public class Performance {
         producer
             .sender()
             .topic(param.topic)
+            .partition(partitions.get(rand.nextInt(partitions.size())))
             .key(manager.getKey().orElse(null))
             .value(payload.get())
             .timestamp(start)
@@ -203,6 +206,21 @@ public class Performance {
     };
   }
 
+  // visible for test
+  static List<Integer> partition(Argument param, TopicAdmin topicAdmin) {
+    if (positiveSpecifyBroker(param)) {
+      return topicAdmin
+          .partitionsOfBrokers(Set.of(param.topic), new HashSet<>(param.specifyBroker))
+          .stream()
+          .map(topicPartition -> topicPartition.partition())
+          .collect(Collectors.toList());
+    } else return List.of(-1);
+  }
+
+  private static boolean positiveSpecifyBroker(Argument param) {
+    return param.specifyBroker.stream().allMatch(broker -> broker >= 0);
+  }
+
   public static class Argument extends BasicArgumentWithPropFile {
 
     @Parameter(
@@ -295,6 +313,13 @@ public class Performance {
             "String: Distribution name. Available distribution names: \"uniform\", \"zipfian\", \"latest\". Default: (No key)",
         converter = Distribution.DistributionConverter.class)
     Distribution distribution = null;
+
+    @Parameter(
+        names = {"--specify.broker"},
+        description =
+            "String: Used with SpecifyBrokerPartitioner to specify the brokers that partitioner can send.",
+        validateWith = ArgumentUtil.NotEmptyString.class)
+    List<Integer> specifyBroker = List.of(-1);
   }
 
   static class CompressionArgument implements IStringConverter<CompressionType> {
diff --git a/app/src/main/java/org/astraea/producer/Builder.java b/app/src/main/java/org/astraea/producer/Builder.java
index d7a68918..44b9184c 100644
--- a/app/src/main/java/org/astraea/producer/Builder.java
+++ b/app/src/main/java/org/astraea/producer/Builder.java
@@ -86,7 +86,7 @@ public class Builder<Key, Value> {
 
           @Override
           public Sender<Key, Value> partition(int partition) {
-            this.partition = partition;
+            if (partition >= 0) this.partition = partition;
             return this;
           }
 
diff --git a/app/src/main/java/org/astraea/topic/Builder.java b/app/src/main/java/org/astraea/topic/Builder.java
index 24b754f7..d7f143a4 100644
--- a/app/src/main/java/org/astraea/topic/Builder.java
+++ b/app/src/main/java/org/astraea/topic/Builder.java
@@ -211,6 +211,14 @@ public class Builder {
                   .collect(Collectors.toSet()));
     }
 
+    @Override
+    public List<TopicPartition> partitionsOfBrokers(Set<String> topics, Set<Integer> brokersID) {
+      return replicas(topics).entrySet().stream()
+          .filter(e -> e.getValue().stream().anyMatch(r -> brokersID.contains(r.broker())))
+          .map(Map.Entry::getKey)
+          .collect(Collectors.toList());
+    }
+
     @Override
     public Map<TopicPartition, List<Replica>> replicas(Set<String> topics) {
       var replicaInfos =
diff --git a/app/src/main/java/org/astraea/topic/TopicAdmin.java b/app/src/main/java/org/astraea/topic/TopicAdmin.java
index d2cc1ffc..ae5e49a8 100644
--- a/app/src/main/java/org/astraea/topic/TopicAdmin.java
+++ b/app/src/main/java/org/astraea/topic/TopicAdmin.java
@@ -57,6 +57,13 @@ public interface TopicAdmin extends Closeable {
   /** @return all brokers' ids */
   Set<Integer> brokerIds();
 
+  /**
+   * @param topics topic names
+   * @param brokersID brokers ID
+   * @return the partitions of brokers
+   */
+  List<TopicPartition> partitionsOfBrokers(Set<String> topics, Set<Integer> brokersID);
+
   /**
    * @param brokers a Set containing broker's ID
    * @return all log directory


Test patch:
diff --git a/app/src/test/java/org/astraea/performance/PerformanceTest.java b/app/src/test/java/org/astraea/performance/PerformanceTest.java
index a2e50993..260521c4 100644
--- a/app/src/test/java/org/astraea/performance/PerformanceTest.java
+++ b/app/src/test/java/org/astraea/performance/PerformanceTest.java
@@ -1,22 +1,112 @@
 package org.astraea.performance;
 
+import static org.astraea.performance.Performance.partition;
+
 import java.util.List;
 import java.util.Set;
 import java.util.concurrent.ExecutionException;
+import java.util.stream.Collectors;
+import org.apache.kafka.common.TopicPartition;
 import org.astraea.Utils;
 import org.astraea.concurrent.ThreadPool;
 import org.astraea.consumer.Consumer;
 import org.astraea.producer.Producer;
 import org.astraea.service.RequireBrokerCluster;
+import org.astraea.topic.TopicAdmin;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
 
 public class PerformanceTest extends RequireBrokerCluster {
 
+  @Test
+  void testSpecifyBrokerProducerExecutor() {
+    var admin = TopicAdmin.of(bootstrapServers());
+    var topicName = "testConsumerExecutor-" + System.currentTimeMillis();
+    admin.creator().topic(topicName).numberOfPartitions(10).create();
+
+    var metrics = new Metrics();
+    var param = new Performance.Argument();
+    param.brokers = bootstrapServers();
+    param.topic = topicName;
+    param.fixedSize = true;
+    param.exeTime = ExeTime.of("100records");
+    param.specifyBroker = List.of(0);
+    param.consumers = 0;
+    param.partitions = 10;
+    try (ThreadPool.Executor executor =
+        Performance.producerExecutor(
+            Producer.builder().brokers(bootstrapServers()).build(),
+            param,
+            metrics,
+            partition(param, admin),
+            new Manager(param, List.of(), List.of()))) {
+      ThreadPool threadPool = ThreadPool.builder().executor(executor).build();
+      threadPool.waitAll();
+      threadPool.close();
+
+      Utils.waitFor(() -> metrics.num() == 100);
+      var offsets = admin.offsets(Set.of(topicName));
+      var partitions =
+          offsets.entrySet().stream()
+              .filter(entry -> entry.getValue().latest() > 0)
+              .map(entry -> entry.getKey().partition())
+              .collect(Collectors.toList());
+      var partitionsOfBrokers =
+          admin.partitionsOfBrokers(Set.of(topicName), Set.of(0)).stream()
+              .map(TopicPartition::partition)
+              .collect(Collectors.toSet());
+      partitions.forEach(
+          partition -> Assertions.assertTrue(partitionsOfBrokers.contains(partition)));
+    }
+  }
+
+  @Test
+  void testMultipleSpecifyBrokersProducerExecutor() {
+    var admin = TopicAdmin.of(bootstrapServers());
+    var topicName = "testConsumerExecutor-" + System.currentTimeMillis();
+    admin.creator().topic(topicName).numberOfPartitions(10).create();
+
+    var metrics = new Metrics();
+    var param = new Performance.Argument();
+    param.brokers = bootstrapServers();
+    param.topic = topicName;
+    param.fixedSize = true;
+    param.exeTime = ExeTime.of("100records");
+    param.specifyBroker = List.of(0, 1);
+    param.consumers = 0;
+    param.partitions = 10;
+    try (ThreadPool.Executor executor =
+        Performance.producerExecutor(
+            Producer.builder().brokers(bootstrapServers()).build(),
+            param,
+            metrics,
+            partition(param, admin),
+            new Manager(param, List.of(), List.of()))) {
+      ThreadPool threadPool = ThreadPool.builder().executor(executor).build();
+      threadPool.waitAll();
+      threadPool.close();
+
+      Utils.waitFor(() -> metrics.num() == 100);
+      var offsets = admin.offsets(Set.of(topicName));
+      var partitions =
+          offsets.entrySet().stream()
+              .filter(entry -> entry.getValue().latest() > 0)
+              .map(entry -> entry.getKey().partition())
+              .collect(Collectors.toList());
+      var partitionsOfBrokers =
+          admin.partitionsOfBrokers(Set.of(topicName), Set.of(0, 1)).stream()
+              .map(TopicPartition::partition)
+              .collect(Collectors.toSet());
+      partitions.forEach(
+          partition -> Assertions.assertTrue(partitionsOfBrokers.contains(partition)));
+    }
+  }
+
   @Test
   void testProducerExecutor() throws InterruptedException {
     var metrics = new Metrics();
     var param = new Performance.Argument();
+    param.brokers = bootstrapServers();
     param.topic = "testProducerExecutor-" + System.currentTimeMillis();
     param.fixedSize = true;
     param.consumers = 0;
@@ -25,6 +115,7 @@ public class PerformanceTest extends RequireBrokerCluster {
             Producer.builder().brokers(bootstrapServers()).build(),
             param,
             metrics,
+            List.of(-1),
             new Manager(param, List.of(), List.of()))) {
       executor.execute();