Repository: skiptests/astraea
Pull request #91 (2021-11-01 15:15:00)

Add consumer assignment listener
1. Select partitioner, jmx server.
2. Add consumer assignment listener.
3. Fix bug in partitioner. (Fix the formula.)

The producer will wait all consumers to get assignments.

chinghongfang: 保留`OffsetPolicy`
* 有些測試是基於`OffsetPolicy`的假設實做的，如PartitionerTest.java，`consumer`要一次抓到20筆資料。
* 有些是保留`OffsetPolicy`可以使測試變簡單，如ProducerTest.java，測試的方法是先傳送一筆資料，再使用`consumer`接收、驗證**確實有傳送資料**。（因為要使用`listener`監聽`consumer`是否得到assignment，`consumer`必須呼叫`poll()`，我們就要使用較複雜的方式，驗證簡單的`testSender()`）
chia7712: @chinghongfang 可否麻煩修一下conflicts

Patch:
diff --git a/app/src/main/java/org/astraea/consumer/Builder.java b/app/src/main/java/org/astraea/consumer/Builder.java
index 43dd1d76..07b47bbb 100644
--- a/app/src/main/java/org/astraea/consumer/Builder.java
+++ b/app/src/main/java/org/astraea/consumer/Builder.java
@@ -20,6 +20,7 @@ public class Builder<Key, Value> {
   private OffsetPolicy offsetPolicy = OffsetPolicy.LATEST;
   private String groupId = "groupId-" + System.currentTimeMillis();
   private final Set<String> topics = new HashSet<>();
+  private ConsumerRebalanceListener listener = ignore -> {};
 
   Builder() {}
 
@@ -61,6 +62,11 @@ public class Builder<Key, Value> {
     return this;
   }
 
+  public Builder<Key, Value> consumerRebalanceListener(ConsumerRebalanceListener listener) {
+    this.listener = Objects.requireNonNull(listener);
+    return this;
+  }
+
   @SuppressWarnings("unchecked")
   public Consumer<Key, Value> build() {
     configs.put(ConsumerConfig.GROUP_ID_CONFIG, Objects.requireNonNull(groupId));
@@ -77,7 +83,7 @@ public class Builder<Key, Value> {
             configs,
             Deserializer.of((Deserializer<Key>) keyDeserializer),
             Deserializer.of((Deserializer<Value>) valueDeserializer));
-    kafkaConsumer.subscribe(topics);
+    kafkaConsumer.subscribe(topics, ConsumerRebalanceListener.of(listener));
     return new Consumer<>() {
       @Override
       public Collection<Record<Key, Value>> poll(Duration timeout) {
diff --git a/app/src/main/java/org/astraea/consumer/ConsumerRebalanceListener.java b/app/src/main/java/org/astraea/consumer/ConsumerRebalanceListener.java
new file mode 100644
index 00000000..12d39f89
--- /dev/null
+++ b/app/src/main/java/org/astraea/consumer/ConsumerRebalanceListener.java
@@ -0,0 +1,31 @@
+package org.astraea.consumer;
+
+import java.util.Collection;
+import java.util.HashSet;
+import java.util.Set;
+import org.apache.kafka.common.TopicPartition;
+
+@FunctionalInterface
+public interface ConsumerRebalanceListener {
+
+  /**
+   * A callback method the user can implement to provide handling of customized offsets on
+   * completion of a successful partition re-assignment. See {@link
+   * org.apache.kafka.clients.consumer.ConsumerRebalanceListener#onPartitionsAssigned(Collection)}
+   * for more detail.
+   */
+  void onPartitionAssigned(Set<TopicPartition> partitions);
+
+  static org.apache.kafka.clients.consumer.ConsumerRebalanceListener of(
+      ConsumerRebalanceListener listener) {
+    return new org.apache.kafka.clients.consumer.ConsumerRebalanceListener() {
+      @Override
+      public void onPartitionsRevoked(Collection<TopicPartition> ignore) {}
+
+      @Override
+      public void onPartitionsAssigned(Collection<TopicPartition> partitions) {
+        listener.onPartitionAssigned(new HashSet<>(partitions));
+      }
+    };
+  }
+}
diff --git a/app/src/main/java/org/astraea/partitioner/partitionerFactory/LinkPartitioner.java b/app/src/main/java/org/astraea/partitioner/partitionerFactory/LinkPartitioner.java
index 0404404b..d3b40786 100644
--- a/app/src/main/java/org/astraea/partitioner/partitionerFactory/LinkPartitioner.java
+++ b/app/src/main/java/org/astraea/partitioner/partitionerFactory/LinkPartitioner.java
@@ -78,6 +78,8 @@ public class LinkPartitioner implements Partitioner {
         if (maxWeightServer == null || item.getValue()[1] > maxWeightServer.getValue()[1]) {
           maxWeightServer = item;
         }
+        currentBrokerHashMap.put(
+            item.getKey(), new int[] {item.getValue()[0], item.getValue()[1] + item.getValue()[0]});
       }
       assert maxWeightServer != null;
       currentBrokerHashMap.put(
diff --git a/app/src/main/java/org/astraea/performance/Performance.java b/app/src/main/java/org/astraea/performance/Performance.java
index 08e62343..8142173d 100644
--- a/app/src/main/java/org/astraea/performance/Performance.java
+++ b/app/src/main/java/org/astraea/performance/Performance.java
@@ -5,15 +5,16 @@ import java.io.IOException;
 import java.time.Duration;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
+import org.apache.kafka.clients.producer.internals.DefaultPartitioner;
 import org.apache.kafka.common.errors.WakeupException;
 import org.astraea.argument.ArgumentUtil;
 import org.astraea.argument.BasicArgumentWithPropFile;
 import org.astraea.concurrent.ThreadPool;
-import org.astraea.consumer.Builder;
 import org.astraea.consumer.Consumer;
 import org.astraea.producer.Producer;
 import org.astraea.topic.TopicAdmin;
@@ -53,7 +54,7 @@ public class Performance {
 
   public static void execute(final Argument param)
       throws InterruptedException, IOException, ExecutionException {
-    try (var topicAdmin = TopicAdmin.of(param.perfProps())) {
+    try (var topicAdmin = TopicAdmin.of(param.props())) {
       topicAdmin
           .creator()
           .numberOfReplicas(param.replicas)
@@ -75,6 +76,7 @@ public class Performance {
     var producerRecords = new AtomicLong(param.records);
     var tracker = new Tracker(producerMetrics, consumerMetrics, param.records);
     var groupId = "groupId-" + System.currentTimeMillis();
+    var getAssignment = new CountDownLatch(param.consumers);
     try (ThreadPool threadPool =
         ThreadPool.builder()
             .executors(
@@ -85,9 +87,9 @@ public class Performance {
                                 Consumer.builder()
                                     .brokers(param.brokers)
                                     .topics(Set.of(param.topic))
-                                    .offsetPolicy(Builder.OffsetPolicy.EARLIEST)
                                     .groupId(groupId)
-                                    .configs(param.perfProps())
+                                    .configs(param.props())
+                                    .consumerRebalanceListener(ignore -> getAssignment.countDown())
                                     .build(),
                                 consumerMetrics.get(i),
                                 consumerRecords))
@@ -97,10 +99,14 @@ public class Performance {
                     .mapToObj(
                         i ->
                             producerExecutor(
-                                Producer.of(param.perfProps()),
+                                Producer.builder()
+                                    .configs(param.props())
+                                    .partitionClassName(param.partitioner)
+                                    .build(),
                                 param,
                                 producerMetrics.get(i),
-                                producerRecords))
+                                producerRecords,
+                                getAssignment))
                     .collect(Collectors.toUnmodifiableList()))
             .executor(tracker)
             .build()) {
@@ -144,11 +150,17 @@ public class Performance {
   }
 
   static ThreadPool.Executor producerExecutor(
-      Producer<byte[], byte[]> producer, Argument param, Metrics metrics, AtomicLong records) {
+      Producer<byte[], byte[]> producer,
+      Argument param,
+      Metrics metrics,
+      AtomicLong records,
+      CountDownLatch getAssignment) {
     byte[] payload = new byte[param.recordSize];
     return new ThreadPool.Executor() {
       @Override
-      public State execute() {
+      public State execute() throws InterruptedException {
+        // Wait for all consumers get assignment.
+        getAssignment.await();
         var currentRecords = records.getAndDecrement();
         if (currentRecords <= 0) return State.DONE;
         long start = System.currentTimeMillis();
@@ -215,8 +227,24 @@ public class Performance {
         validateWith = ArgumentUtil.PositiveLong.class)
     int recordSize = 1024;
 
-    public Map<String, Object> perfProps() {
-      return properties(propFile);
+    @Parameter(
+        names = {"--jmx.servers"},
+        description =
+            "String: server to get jmx metrics <jmx_server>@<broker_id>[,<jmx_server>@<broker_id>]*",
+        validateWith = ArgumentUtil.NotEmptyString.class)
+    String jmxServers = "";
+
+    @Parameter(
+        names = {"--partitioner"},
+        description = "String: the full class name of the desired partitioner",
+        validateWith = ArgumentUtil.NotEmptyString.class)
+    String partitioner = DefaultPartitioner.class.getName();
+
+    @Override
+    public Map<String, Object> props() {
+      Map<String, Object> prop = properties(propFile);
+      if (!this.jmxServers.isEmpty()) prop.put("jmx_servers", this.jmxServers);
+      return prop;
     }
   }
 }
diff --git a/app/src/main/java/org/astraea/performance/Tracker.java b/app/src/main/java/org/astraea/performance/Tracker.java
index 6591637c..3653faa5 100644
--- a/app/src/main/java/org/astraea/performance/Tracker.java
+++ b/app/src/main/java/org/astraea/performance/Tracker.java
@@ -17,22 +17,27 @@ public class Tracker implements ThreadPool.Executor {
 
   @Override
   public State execute() throws InterruptedException {
-    long completed = 0;
-    double bytes = 0;
-    long max = 0;
-    long min = Long.MAX_VALUE;
+    if (logProducers() & logConsumers()) return State.DONE;
+    // 等待1秒，再抓資料輸出
+    Thread.sleep(1000);
+    return State.RUNNING;
+  }
+
+  private boolean logProducers() {
+    var completed = 0L;
+    var bytes = 0D;
+    var max = 0L;
+    var min = Long.MAX_VALUE;
 
-    /* producer */
     for (Metrics data : producerData) {
       completed += data.num();
       bytes += data.avgBytes();
       max = Math.max(max, data.max());
       min = Math.min(min, data.min());
     }
-    if (completed == 0) {
-      Thread.sleep(1000);
-      return State.RUNNING;
-    }
+
+    if (completed == 0) return false;
+
     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");
@@ -43,21 +48,26 @@ public class Tracker implements ThreadPool.Executor {
       System.out.printf(
           "  producer[%d]的發送average latency: %.3fms%n", i, producerData.get(i).avgLatency());
     }
-    if (consumerData.isEmpty()) {
-      if (completed >= records) return State.DONE;
-      else return State.RUNNING;
-    }
-    /* consumer */
-    completed = 0;
-    bytes = 0L;
-    max = 0;
-    min = Long.MAX_VALUE;
+    System.out.println("\n");
+    return completed >= records;
+  }
+
+  private boolean logConsumers() {
+    // there is no consumer, so we just complete this log.
+    if (consumerData.isEmpty()) return true;
+    var completed = 0;
+    var bytes = 0D;
+    var max = 0L;
+    var min = Long.MAX_VALUE;
     for (Metrics data : consumerData) {
       completed += data.num();
       bytes += data.avgBytes();
       max = Math.max(max, data.max());
       min = Math.min(min, data.min());
     }
+
+    if (completed == 0) return false;
+
     System.out.printf("consumer完成度: %.2f%%%n", ((double) completed * 100.0 / (double) records));
     System.out.printf("  輸入%.3fMB/second%n", bytes);
     System.out.println("  端到端max latency: " + max + "ms");
@@ -68,12 +78,7 @@ public class Tracker implements ThreadPool.Executor {
       System.out.printf(
           "  consumer[%d]的端到端average latency: %.3fms%n", i, consumerData.get(i).avgLatency());
     }
-
     System.out.println("\n");
-    // The consumer has consumed the target records
-    if (completed >= records) return State.DONE;
-    // 等待1秒，再抓資料輸出
-    Thread.sleep(1000);
-    return State.RUNNING;
+    return completed >= records;
   }
 }


Test patch:
diff --git a/app/src/test/java/org/astraea/consumer/RebalanceListenerTest.java b/app/src/test/java/org/astraea/consumer/RebalanceListenerTest.java
new file mode 100644
index 00000000..ce50a8d6
--- /dev/null
+++ b/app/src/test/java/org/astraea/consumer/RebalanceListenerTest.java
@@ -0,0 +1,34 @@
+package org.astraea.consumer;
+
+import java.time.Duration;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicInteger;
+import org.astraea.Utils;
+import org.astraea.concurrent.ThreadPool;
+import org.astraea.service.RequireBrokerCluster;
+import org.junit.jupiter.api.Test;
+
+public class RebalanceListenerTest extends RequireBrokerCluster {
+  @Test
+  void testConsumerRebalanceListener() {
+    var getAssignment = new AtomicInteger(0);
+    var topicName = "testRebalanceListener-" + System.currentTimeMillis();
+    try (var consumer =
+        Consumer.builder()
+            .brokers(bootstrapServers())
+            .topics(Set.of(topicName))
+            .consumerRebalanceListener(ignore -> getAssignment.incrementAndGet())
+            .build()) {
+      try (ThreadPool threadPool =
+          ThreadPool.builder()
+              .executor(
+                  () -> {
+                    consumer.poll(Duration.ofSeconds(10));
+                    return ThreadPool.Executor.State.DONE;
+                  })
+              .build()) {
+        Utils.waitFor(() -> getAssignment.get() == 1, Duration.ofSeconds(10));
+      }
+    }
+  }
+}
diff --git a/app/src/test/java/org/astraea/performance/PerformanceTest.java b/app/src/test/java/org/astraea/performance/PerformanceTest.java
index cc684cad..33b9b39c 100644
--- a/app/src/test/java/org/astraea/performance/PerformanceTest.java
+++ b/app/src/test/java/org/astraea/performance/PerformanceTest.java
@@ -1,6 +1,7 @@
 package org.astraea.performance;
 
 import java.util.Set;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.atomic.AtomicLong;
 import org.astraea.Utils;
@@ -23,7 +24,8 @@ public class PerformanceTest extends RequireBrokerCluster {
             Producer.builder().brokers(bootstrapServers()).build(),
             param,
             metrics,
-            new AtomicLong(10))) {
+            new AtomicLong(10),
+            new CountDownLatch(0))) {
       executor.execute();
 
       Utils.waitFor(() -> metrics.num() == 1);