Repository: skiptests/astraea
Pull request #318 (2022-05-07 11:25:27)

Enable users to limit the returned number of records from Consumer#poll
新增一個方法來限制consumer可以回傳的資料大小


Patch:
diff --git a/app/src/main/java/org/astraea/consumer/Builder.java b/app/src/main/java/org/astraea/consumer/Builder.java
index d62b1ca8..2de5d8e7 100644
--- a/app/src/main/java/org/astraea/consumer/Builder.java
+++ b/app/src/main/java/org/astraea/consumer/Builder.java
@@ -1,14 +1,14 @@
 package org.astraea.consumer;
 
 import java.time.Duration;
+import java.util.ArrayList;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
-import java.util.stream.Collectors;
-import java.util.stream.StreamSupport;
 import org.apache.kafka.clients.consumer.ConsumerConfig;
 import org.apache.kafka.clients.consumer.KafkaConsumer;
 
@@ -36,6 +36,15 @@ public class Builder<Key, Value> {
     return config(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
   }
 
+  /**
+   * make the consumer read data from latest. this is default setting.
+   *
+   * @return this builder
+   */
+  public Builder<Key, Value> fromLatest() {
+    return config(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "latest");
+  }
+
   public Builder<Key, Value> topics(Set<String> topics) {
     this.topics.addAll(topics);
     return this;
@@ -87,11 +96,15 @@ public class Builder<Key, Value> {
     kafkaConsumer.subscribe(topics, ConsumerRebalanceListener.of(listener));
     return new Consumer<>() {
       @Override
-      public Collection<Record<Key, Value>> poll(Duration timeout) {
-        var records = kafkaConsumer.poll(timeout);
-        return StreamSupport.stream(records.spliterator(), false)
-            .map(Record::of)
-            .collect(Collectors.toList());
+      public Collection<Record<Key, Value>> poll(int recordCount, Duration timeout) {
+        var end = System.currentTimeMillis() + timeout.toMillis();
+        var records = new ArrayList<Record<Key, Value>>();
+        while (records.size() < recordCount) {
+          var remaining = end - System.currentTimeMillis();
+          if (remaining <= 0) break;
+          kafkaConsumer.poll(Duration.ofMillis(remaining)).forEach(r -> records.add(Record.of(r)));
+        }
+        return Collections.unmodifiableList(records);
       }
 
       @Override
diff --git a/app/src/main/java/org/astraea/consumer/Consumer.java b/app/src/main/java/org/astraea/consumer/Consumer.java
index 3e936217..5078a83b 100644
--- a/app/src/main/java/org/astraea/consumer/Consumer.java
+++ b/app/src/main/java/org/astraea/consumer/Consumer.java
@@ -7,7 +7,18 @@ import java.util.Map;
 /** An interface for polling records. */
 public interface Consumer<Key, Value> extends AutoCloseable {
 
-  Collection<Record<Key, Value>> poll(Duration timeout);
+  default Collection<Record<Key, Value>> poll(Duration timeout) {
+    return poll(Integer.MAX_VALUE, timeout);
+  }
+
+  /**
+   * try to poll data until there are enough records to return or the timeout is reached.
+   *
+   * @param recordCount max number of returned records.
+   * @param timeout max time to wait data
+   * @return records
+   */
+  Collection<Record<Key, Value>> poll(int recordCount, Duration timeout);
 
   /**
    * Wakeup the consumer. This method is thread-safe and is useful in particular to abort a long


Test patch:
diff --git a/app/src/test/java/org/astraea/consumer/ConsumerTest.java b/app/src/test/java/org/astraea/consumer/ConsumerTest.java
new file mode 100644
index 00000000..2277550e
--- /dev/null
+++ b/app/src/test/java/org/astraea/consumer/ConsumerTest.java
@@ -0,0 +1,112 @@
+package org.astraea.consumer;
+
+import java.nio.charset.StandardCharsets;
+import java.time.Duration;
+import java.util.Set;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.stream.IntStream;
+import org.apache.kafka.common.errors.WakeupException;
+import org.astraea.producer.Producer;
+import org.astraea.service.RequireBrokerCluster;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.Timeout;
+
+public class ConsumerTest extends RequireBrokerCluster {
+
+  private static void produceData(String topic, int size) {
+    try (var producer = Producer.builder().brokers(bootstrapServers()).build()) {
+      IntStream.range(0, size)
+          .forEach(
+              i ->
+                  producer
+                      .sender()
+                      .topic(topic)
+                      .key(String.valueOf(i).getBytes(StandardCharsets.UTF_8))
+                      .run());
+      producer.flush();
+    }
+  }
+
+  @Test
+  void testFromBeginning() {
+    var recordCount = 100;
+    var topic = "testPoll";
+    produceData(topic, recordCount);
+    try (var consumer =
+        Consumer.builder()
+            .topics(Set.of(topic))
+            .brokers(bootstrapServers())
+            .fromBeginning()
+            .build()) {
+
+      Assertions.assertEquals(
+          recordCount, consumer.poll(recordCount, Duration.ofSeconds(10)).size());
+    }
+  }
+
+  @Test
+  void testFromLatest() {
+    var topic = "testFromLatest";
+    produceData(topic, 1);
+    try (var consumer =
+        Consumer.builder().topics(Set.of(topic)).brokers(bootstrapServers()).fromLatest().build()) {
+
+      Assertions.assertEquals(0, consumer.poll(Duration.ofSeconds(3)).size());
+    }
+  }
+
+  @Timeout(7)
+  @Test
+  void testWakeup() throws InterruptedException {
+    var topic = "testWakeup";
+    try (var consumer =
+        Consumer.builder().topics(Set.of(topic)).brokers(bootstrapServers()).fromLatest().build()) {
+      var service = Executors.newSingleThreadExecutor();
+      service.execute(
+          () -> {
+            try {
+              TimeUnit.SECONDS.sleep(3);
+              consumer.wakeup();
+            } catch (InterruptedException ignored) {
+              // swallow
+            }
+          });
+      // this call will be broken after 3 seconds
+      Assertions.assertThrows(WakeupException.class, () -> consumer.poll(Duration.ofSeconds(100)));
+
+      service.shutdownNow();
+      Assertions.assertTrue(service.awaitTermination(3, TimeUnit.SECONDS));
+    }
+  }
+
+  @Test
+  void testGroupId() {
+    var groupId = "testGroupId";
+    var topic = "testGroupId";
+    produceData(topic, 1);
+
+    java.util.function.BiConsumer<String, Integer> testConsumer =
+        (id, expectedSize) -> {
+          try (var consumer =
+              Consumer.builder()
+                  .topics(Set.of(topic))
+                  .brokers(bootstrapServers())
+                  .fromBeginning()
+                  .groupId(id)
+                  .build()) {
+            Assertions.assertEquals(
+                expectedSize, consumer.poll(expectedSize, Duration.ofSeconds(5)).size());
+          }
+        };
+
+    testConsumer.accept(groupId, 1);
+
+    // the data is fetched already, so it should not return any data
+    testConsumer.accept(groupId, 0);
+
+    // use different group id
+    testConsumer.accept("another_group", 0);
+  }
+}