Repository: skiptests/astraea
Pull request #351 (2022-05-17 16:00:38)

A new seek mode: the distance from the latest offset
該方法可以接受一個整數值，並且在建立consuemr時候會先行計算適當的offset，並且將seek到該offset後再將已經設定好的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 c609657b..cb17c699 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 Deserializer<?> valueDeserializer = Deserializer.BYTE_ARRAY;
   private final Set<String> topics = new HashSet<>();
   private ConsumerRebalanceListener listener = ignore -> {};
+  private int distanceFromLatest = -1;
 
   Builder() {}
 
@@ -45,6 +46,18 @@ public class Builder<Key, Value> {
     return config(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "latest");
   }
 
+  /**
+   * set the offset to read from the latest offset. For example, the end offset is 5, and you set
+   * distanceFromLatest to 2, then you will read data from offset: 3
+   *
+   * @param distanceFromLatest the distance from the latest offset
+   * @return this builder
+   */
+  public Builder<Key, Value> distanceFromLatest(int distanceFromLatest) {
+    this.distanceFromLatest = distanceFromLatest;
+    return this;
+  }
+
   public Builder<Key, Value> topics(Set<String> topics) {
     this.topics.addAll(topics);
     return this;
@@ -95,6 +108,18 @@ public class Builder<Key, Value> {
             Deserializer.of((Deserializer<Key>) keyDeserializer),
             Deserializer.of((Deserializer<Value>) valueDeserializer));
     kafkaConsumer.subscribe(topics, ConsumerRebalanceListener.of(listener));
+
+    // this mode is not supported by kafka, so we have to calculate the offset first
+    if (distanceFromLatest > 0) {
+      // 1) poll data until the assignment is completed
+      while (kafkaConsumer.assignment().isEmpty()) kafkaConsumer.poll(Duration.ofMillis(500));
+      var partitions = kafkaConsumer.assignment();
+      // 2) get the end offsets from all subscribed partitions
+      var endOffsets = kafkaConsumer.endOffsets(partitions);
+      // 3) calculate and then seek to the correct offset (end offset - recent offset)
+      endOffsets.forEach(
+          (tp, latest) -> kafkaConsumer.seek(tp, Math.max(0, latest - distanceFromLatest)));
+    }
     return new Consumer<>() {
       @Override
       public Collection<Record<Key, Value>> poll(int recordCount, Duration timeout) {


Test patch:
diff --git a/app/src/test/java/org/astraea/consumer/ConsumerTest.java b/app/src/test/java/org/astraea/consumer/ConsumerTest.java
index 52c66dfc..7f99c882 100644
--- a/app/src/test/java/org/astraea/consumer/ConsumerTest.java
+++ b/app/src/test/java/org/astraea/consumer/ConsumerTest.java
@@ -7,6 +7,7 @@ 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.Utils;
 import org.astraea.producer.Producer;
 import org.astraea.service.RequireBrokerCluster;
 import org.junit.jupiter.api.Assertions;
@@ -117,4 +118,38 @@ public class ConsumerTest extends RequireBrokerCluster {
     // use different group id
     testConsumer.accept("another_group", 0);
   }
+
+  @Test
+  void testDistanceFromLatest() {
+    var count = 10;
+    var topic = Utils.randomString(10);
+    try (var producer = Producer.of(bootstrapServers())) {
+      IntStream.range(0, count)
+          .forEach(
+              i ->
+                  producer
+                      .sender()
+                      .topic(topic)
+                      .value(String.valueOf(count).getBytes(StandardCharsets.UTF_8))
+                      .run());
+      producer.flush();
+    }
+    try (var consumer =
+        Consumer.builder()
+            .bootstrapServers(bootstrapServers())
+            .topics(Set.of(topic))
+            .distanceFromLatest(3)
+            .build()) {
+      Assertions.assertEquals(3, consumer.poll(4, Duration.ofSeconds(5)).size());
+    }
+
+    try (var consumer =
+        Consumer.builder()
+            .bootstrapServers(bootstrapServers())
+            .topics(Set.of(topic))
+            .distanceFromLatest(1000)
+            .build()) {
+      Assertions.assertEquals(10, consumer.poll(11, Duration.ofSeconds(5)).size());
+    }
+  }
 }