Repository: skiptests/astraea
Pull request #340 (2022-05-14 16:51:40)

Enable to fetch producer state
新增方法讓我們擷取producer state


Patch:
diff --git a/app/src/main/java/org/astraea/admin/Admin.java b/app/src/main/java/org/astraea/admin/Admin.java
index 92688549..d6e9dbd2 100644
--- a/app/src/main/java/org/astraea/admin/Admin.java
+++ b/app/src/main/java/org/astraea/admin/Admin.java
@@ -1,6 +1,7 @@
 package org.astraea.admin;
 
 import java.io.Closeable;
+import java.util.Collection;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -102,6 +103,17 @@ public interface Admin extends Closeable {
   /** @return a partition migrator used to move partitions to another broker or folder. */
   Migrator migrator();
 
+  /** @return producer states of all topic partitions */
+  default Map<TopicPartition, Collection<ProducerState>> producerStates() {
+    return producerStates(partitions());
+  }
+
+  /**
+   * @param partitions to search
+   * @return producer states of input topic partitions
+   */
+  Map<TopicPartition, Collection<ProducerState>> producerStates(Set<TopicPartition> partitions);
+
   @Override
   void close();
 }
diff --git a/app/src/main/java/org/astraea/admin/Builder.java b/app/src/main/java/org/astraea/admin/Builder.java
index 8d369591..34fbf7ef 100644
--- a/app/src/main/java/org/astraea/admin/Builder.java
+++ b/app/src/main/java/org/astraea/admin/Builder.java
@@ -85,6 +85,29 @@ public class Builder {
       return new MigratorImpl(admin, this::partitions);
     }
 
+    @Override
+    public Map<TopicPartition, Collection<ProducerState>> producerStates(
+        Set<TopicPartition> partitions) {
+      return Utils.handleException(
+          () ->
+              admin
+                  .describeProducers(
+                      partitions.stream()
+                          .map(TopicPartition::to)
+                          .collect(Collectors.toUnmodifiableList()))
+                  .all()
+                  .get()
+                  .entrySet()
+                  .stream()
+                  .collect(
+                      Collectors.toMap(
+                          e -> TopicPartition.from(e.getKey()),
+                          e ->
+                              e.getValue().activeProducers().stream()
+                                  .map(ProducerState::from)
+                                  .collect(Collectors.toUnmodifiableList()))));
+    }
+
     @Override
     public Set<String> consumerGroupIds() {
       return Utils.handleException(() -> admin.listConsumerGroups().all().get()).stream()
diff --git a/app/src/main/java/org/astraea/admin/ProducerState.java b/app/src/main/java/org/astraea/admin/ProducerState.java
new file mode 100644
index 00000000..966add47
--- /dev/null
+++ b/app/src/main/java/org/astraea/admin/ProducerState.java
@@ -0,0 +1,69 @@
+package org.astraea.admin;
+
+import java.util.Objects;
+
+public final class ProducerState {
+
+  public static ProducerState from(org.apache.kafka.clients.admin.ProducerState state) {
+    return new ProducerState(
+        state.producerId(), state.producerEpoch(), state.lastSequence(), state.lastTimestamp());
+  }
+
+  private final long producerId;
+  private final int producerEpoch;
+  private final int lastSequence;
+  private final long lastTimestamp;
+
+  public ProducerState(long producerId, int producerEpoch, int lastSequence, long lastTimestamp) {
+    this.producerId = producerId;
+    this.producerEpoch = producerEpoch;
+    this.lastSequence = lastSequence;
+    this.lastTimestamp = lastTimestamp;
+  }
+
+  public long producerId() {
+    return producerId;
+  }
+
+  public int producerEpoch() {
+    return producerEpoch;
+  }
+
+  public int lastSequence() {
+    return lastSequence;
+  }
+
+  public long lastTimestamp() {
+    return lastTimestamp;
+  }
+
+  @Override
+  public boolean equals(Object o) {
+    if (this == o) return true;
+    if (o == null || getClass() != o.getClass()) return false;
+    ProducerState that = (ProducerState) o;
+    return producerId == that.producerId
+        && producerEpoch == that.producerEpoch
+        && lastSequence == that.lastSequence
+        && lastTimestamp == that.lastTimestamp;
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(producerId, producerEpoch, lastSequence, lastTimestamp);
+  }
+
+  @Override
+  public String toString() {
+    return "ProducerState{"
+        + "producerId="
+        + producerId
+        + ", producerEpoch="
+        + producerEpoch
+        + ", lastSequence="
+        + lastSequence
+        + ", lastTimestamp="
+        + lastTimestamp
+        + '}';
+  }
+}


Test patch:
diff --git a/app/src/test/java/org/astraea/admin/AdminTest.java b/app/src/test/java/org/astraea/admin/AdminTest.java
index 27accebc..3ef6f46f 100644
--- a/app/src/test/java/org/astraea/admin/AdminTest.java
+++ b/app/src/test/java/org/astraea/admin/AdminTest.java
@@ -393,4 +393,22 @@ public class AdminTest extends RequireBrokerCluster {
               .broker());
     }
   }
+
+  @Test
+  void testProducerStates() throws ExecutionException, InterruptedException {
+    var topic = Utils.randomString(10);
+    try (var producer = Producer.of(bootstrapServers());
+        var admin = Admin.of(bootstrapServers())) {
+      producer.sender().topic(topic).value(new byte[1]).run().toCompletableFuture().get();
+
+      var states = admin.producerStates();
+      Assertions.assertNotEquals(0, states.size());
+      var producerState =
+          states.entrySet().stream()
+              .filter(tp -> tp.getKey().topic().equals(topic))
+              .flatMap(e -> e.getValue().stream())
+              .collect(Collectors.toUnmodifiableList());
+      Assertions.assertEquals(1, producerState.size());
+    }
+  }
 }