Repository: skiptests/astraea
Pull request #404 (2022-06-11 17:44:55)

Enable to query transaction state from Admin
這隻PR新增了三個功能：
1. producer現在可以手動指定transaction id
2. admin可以查詢transaction id
3. admin可以查詢transaction state


Patch:
diff --git a/app/src/main/java/org/astraea/app/admin/Admin.java b/app/src/main/java/org/astraea/app/admin/Admin.java
index 28311065..2038d3ef 100644
--- a/app/src/main/java/org/astraea/app/admin/Admin.java
+++ b/app/src/main/java/org/astraea/app/admin/Admin.java
@@ -201,6 +201,21 @@ public interface Admin extends Closeable {
    */
   ClusterInfo clusterInfo(Set<String> topics);
 
+  /** @return all transaction ids */
+  Set<String> transactionIds();
+
+  /** @return all transaction states */
+  default Map<String, Transaction> transactions() {
+    return transactions(transactionIds());
+  }
+  /**
+   * return transaction states associated to input ids
+   *
+   * @param transactionIds to query state
+   * @return transaction states
+   */
+  Map<String, Transaction> transactions(Set<String> transactionIds);
+
   @Override
   void close();
 }
diff --git a/app/src/main/java/org/astraea/app/admin/Builder.java b/app/src/main/java/org/astraea/app/admin/Builder.java
index 01d4c7db..27b91bfc 100644
--- a/app/src/main/java/org/astraea/app/admin/Builder.java
+++ b/app/src/main/java/org/astraea/app/admin/Builder.java
@@ -39,6 +39,7 @@ import org.apache.kafka.clients.admin.MemberDescription;
 import org.apache.kafka.clients.admin.NewPartitionReassignment;
 import org.apache.kafka.clients.admin.NewTopic;
 import org.apache.kafka.clients.admin.OffsetSpec;
+import org.apache.kafka.clients.admin.TransactionListing;
 import org.apache.kafka.common.ElectionType;
 import org.apache.kafka.common.TopicPartitionReplica;
 import org.apache.kafka.common.config.ConfigResource;
@@ -532,6 +533,24 @@ public class Builder {
         }
       };
     }
+
+    @Override
+    public Set<String> transactionIds() {
+      return Utils.packException(
+          () ->
+              admin.listTransactions().all().get().stream()
+                  .map(TransactionListing::transactionalId)
+                  .collect(Collectors.toUnmodifiableSet()));
+    }
+
+    @Override
+    public Map<String, Transaction> transactions(Set<String> transactionIds) {
+      return Utils.packException(
+          () ->
+              admin.describeTransactions(transactionIds).all().get().entrySet().stream()
+                  .collect(
+                      Collectors.toMap(Map.Entry::getKey, e -> Transaction.from(e.getValue()))));
+    }
   }
 
   private static class ConfigImpl implements Config {
diff --git a/app/src/main/java/org/astraea/app/admin/Transaction.java b/app/src/main/java/org/astraea/app/admin/Transaction.java
new file mode 100644
index 00000000..db05b625
--- /dev/null
+++ b/app/src/main/java/org/astraea/app/admin/Transaction.java
@@ -0,0 +1,117 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.astraea.app.admin;
+
+import java.util.Objects;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+public class Transaction {
+
+  public static Transaction from(org.apache.kafka.clients.admin.TransactionDescription td) {
+    return new Transaction(
+        td.coordinatorId(),
+        TransactionState.from(td.state()),
+        td.producerId(),
+        td.producerEpoch(),
+        td.transactionTimeoutMs(),
+        td.topicPartitions().stream().map(TopicPartition::from).collect(Collectors.toSet()));
+  }
+
+  private final int coordinatorId;
+  private final TransactionState state;
+  private final long producerId;
+  private final int producerEpoch;
+  private final long transactionTimeoutMs;
+  private final Set<TopicPartition> topicPartitions;
+
+  public Transaction(
+      int coordinatorId,
+      TransactionState state,
+      long producerId,
+      int producerEpoch,
+      long transactionTimeoutMs,
+      Set<TopicPartition> topicPartitions) {
+    this.coordinatorId = coordinatorId;
+    this.state = state;
+    this.producerId = producerId;
+    this.producerEpoch = producerEpoch;
+    this.transactionTimeoutMs = transactionTimeoutMs;
+    this.topicPartitions = topicPartitions;
+  }
+
+  public int coordinatorId() {
+    return coordinatorId;
+  }
+
+  public TransactionState state() {
+    return state;
+  }
+
+  public long producerId() {
+    return producerId;
+  }
+
+  public int producerEpoch() {
+    return producerEpoch;
+  }
+
+  public long transactionTimeoutMs() {
+    return transactionTimeoutMs;
+  }
+
+  public Set<TopicPartition> topicPartitions() {
+    return topicPartitions;
+  }
+
+  @Override
+  public boolean equals(Object o) {
+    if (this == o) return true;
+    if (o == null || getClass() != o.getClass()) return false;
+    Transaction that = (Transaction) o;
+    return coordinatorId == that.coordinatorId
+        && producerId == that.producerId
+        && producerEpoch == that.producerEpoch
+        && transactionTimeoutMs == that.transactionTimeoutMs
+        && state == that.state
+        && Objects.equals(topicPartitions, that.topicPartitions);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(
+        coordinatorId, state, producerId, producerEpoch, transactionTimeoutMs, topicPartitions);
+  }
+
+  @Override
+  public String toString() {
+    return "Transaction{"
+        + "coordinatorId="
+        + coordinatorId
+        + ", state="
+        + state
+        + ", producerId="
+        + producerId
+        + ", producerEpoch="
+        + producerEpoch
+        + ", transactionTimeoutMs="
+        + transactionTimeoutMs
+        + ", topicPartitions="
+        + topicPartitions
+        + '}';
+  }
+}
diff --git a/app/src/main/java/org/astraea/app/admin/TransactionState.java b/app/src/main/java/org/astraea/app/admin/TransactionState.java
new file mode 100644
index 00000000..ef6f2e5d
--- /dev/null
+++ b/app/src/main/java/org/astraea/app/admin/TransactionState.java
@@ -0,0 +1,58 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.astraea.app.admin;
+
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.stream.Collectors;
+
+public enum TransactionState {
+  ONGOING("Ongoing"),
+  PREPARE_ABORT("PrepareAbort"),
+  PREPARE_COMMIT("PrepareCommit"),
+  COMPLETE_ABORT("CompleteAbort"),
+  COMPLETE_COMMIT("CompleteCommit"),
+  EMPTY("Empty"),
+  PREPARE_EPOCH_FENCE("PrepareEpochFence"),
+  UNKNOWN("Unknown");
+
+  private final String name;
+
+  TransactionState(String name) {
+    this.name = name;
+  }
+
+  @Override
+  public String toString() {
+    return name;
+  }
+
+  public static TransactionState from(org.apache.kafka.clients.admin.TransactionState state) {
+    return from(state.toString());
+  }
+
+  public static TransactionState from(String value) {
+    return all().stream()
+        .filter(s -> s.name.equalsIgnoreCase(value))
+        .findFirst()
+        .orElseThrow(() -> new IllegalArgumentException(value + " is not valid transaction state"));
+  }
+
+  public static Collection<TransactionState> all() {
+    return Arrays.stream(TransactionState.values()).collect(Collectors.toUnmodifiableList());
+  }
+}
diff --git a/app/src/main/java/org/astraea/app/producer/Builder.java b/app/src/main/java/org/astraea/app/producer/Builder.java
index 96a3411e..4a8c7cef 100644
--- a/app/src/main/java/org/astraea/app/producer/Builder.java
+++ b/app/src/main/java/org/astraea/app/producer/Builder.java
@@ -21,6 +21,7 @@ import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.Random;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
@@ -77,6 +78,17 @@ public class Builder<Key, Value> {
     return config(ProducerConfig.COMPRESSION_TYPE_CONFIG, compression.nameOfKafka());
   }
 
+  /**
+   * set the transaction id. If you set the transaction id, the builder will always build
+   * transactional producer.
+   *
+   * @param transactionId used to trace by server
+   * @return this builder
+   */
+  public Builder<Key, Value> transactionId(String transactionId) {
+    return config(ProducerConfig.TRANSACTIONAL_ID_CONFIG, transactionId);
+  }
+
   private static <Key, Value> CompletionStage<Metadata> doSend(
       org.apache.kafka.clients.producer.Producer<Key, Value> producer,
       ProducerRecord<Key, Value> record) {
@@ -90,8 +102,14 @@ public class Builder<Key, Value> {
     return completableFuture;
   }
 
+  /**
+   * @return normal producer if there is no transaction id in configs. Otherwise, a transactional
+   *     producer is returned
+   */
   @SuppressWarnings("unchecked")
   public Producer<Key, Value> build() {
+    // if user configs the transaction id, we should build transactional producer
+    if (configs.containsKey(ProducerConfig.TRANSACTIONAL_ID_CONFIG)) return buildTransactional();
     return new NormalProducer<>(
         new KafkaProducer<>(
             configs,
@@ -99,11 +117,18 @@ public class Builder<Key, Value> {
             Serializer.of((Serializer<Value>) valueSerializer)));
   }
 
+  /**
+   * @return always returns a transactional producer. The transaction id is generated automatically
+   *     if it is absent
+   */
   @SuppressWarnings("unchecked")
   public Producer<Key, Value> buildTransactional() {
     var transactionConfigs = new HashMap<>(configs);
-    transactionConfigs.putIfAbsent(
-        ProducerConfig.TRANSACTIONAL_ID_CONFIG, "id" + new Random().nextLong());
+    var transactionId =
+        (String)
+            transactionConfigs.computeIfAbsent(
+                ProducerConfig.TRANSACTIONAL_ID_CONFIG,
+                ignored -> "transaction-id-" + new Random().nextLong());
     // For transactional send
     var transactionProducer =
         new KafkaProducer<>(
@@ -111,7 +136,7 @@ public class Builder<Key, Value> {
             Serializer.of((Serializer<Key>) keySerializer),
             Serializer.of((Serializer<Value>) valueSerializer));
     transactionProducer.initTransactions();
-    return new TransactionalProducer<>(transactionProducer);
+    return new TransactionalProducer<>(transactionProducer, transactionId);
   }
 
   private abstract static class BaseProducer<Key, Value> implements Producer<Key, Value> {
@@ -153,17 +178,20 @@ public class Builder<Key, Value> {
     }
 
     @Override
-    public boolean transactional() {
-      return false;
+    public Optional<String> transactionId() {
+      return Optional.empty();
     }
   }
 
   private static class TransactionalProducer<Key, Value> extends BaseProducer<Key, Value> {
     private final Object lock = new Object();
+    private final String transactionId;
 
     private TransactionalProducer(
-        org.apache.kafka.clients.producer.Producer<Key, Value> kafkaProducer) {
+        org.apache.kafka.clients.producer.Producer<Key, Value> kafkaProducer,
+        String transactionId) {
       super(kafkaProducer);
+      this.transactionId = transactionId;
     }
 
     @Override
@@ -215,5 +243,10 @@ public class Builder<Key, Value> {
     public boolean transactional() {
       return true;
     }
+
+    @Override
+    public Optional<String> transactionId() {
+      return Optional.of(transactionId);
+    }
   }
 }
diff --git a/app/src/main/java/org/astraea/app/producer/Producer.java b/app/src/main/java/org/astraea/app/producer/Producer.java
index b76bc4e1..aa16fcd2 100644
--- a/app/src/main/java/org/astraea/app/producer/Producer.java
+++ b/app/src/main/java/org/astraea/app/producer/Producer.java
@@ -18,6 +18,7 @@ package org.astraea.app.producer;
 
 import java.util.Collection;
 import java.util.Map;
+import java.util.Optional;
 import java.util.concurrent.CompletionStage;
 
 /** An interface for sending records. */
@@ -39,7 +40,12 @@ public interface Producer<Key, Value> extends AutoCloseable {
   void close();
 
   /** @return true if the producer supports transactional. */
-  boolean transactional();
+  default boolean transactional() {
+    return transactionId().isPresent();
+  }
+
+  /** @return the transaction id or empty if the producer does not support transaction. */
+  Optional<String> transactionId();
 
   static Builder<byte[], byte[]> builder() {
     return new Builder<>();


Test patch:
diff --git a/app/src/test/java/org/astraea/app/admin/AdminTest.java b/app/src/test/java/org/astraea/app/admin/AdminTest.java
index f83e3bf8..41a6d319 100644
--- a/app/src/test/java/org/astraea/app/admin/AdminTest.java
+++ b/app/src/test/java/org/astraea/app/admin/AdminTest.java
@@ -791,4 +791,21 @@ public class AdminTest extends RequireBrokerCluster {
       }
     }
   }
+
+  @Test
+  void testTransactionIds() throws ExecutionException, InterruptedException {
+
+    try (var admin = Admin.of(bootstrapServers());
+        var producer =
+            Producer.builder().bootstrapServers(bootstrapServers()).buildTransactional()) {
+      Assertions.assertTrue(producer.transactional());
+      producer.sender().key(new byte[10]).topic("topic").run().toCompletableFuture().get();
+
+      Assertions.assertTrue(admin.transactionIds().contains(producer.transactionId().get()));
+
+      var transaction = admin.transactions().get(producer.transactionId().get());
+      Assertions.assertNotNull(transaction);
+      Assertions.assertEquals(TransactionState.COMPLETE_COMMIT, transaction.state());
+    }
+  }
 }
diff --git a/app/src/test/java/org/astraea/app/admin/TransactionStateTest.java b/app/src/test/java/org/astraea/app/admin/TransactionStateTest.java
new file mode 100644
index 00000000..a8535517
--- /dev/null
+++ b/app/src/test/java/org/astraea/app/admin/TransactionStateTest.java
@@ -0,0 +1,35 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.astraea.app.admin;
+
+import java.util.Arrays;
+import java.util.stream.Collectors;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+public class TransactionStateTest {
+
+  @Test
+  void testConversion() {
+    var r =
+        Arrays.stream(org.apache.kafka.clients.admin.TransactionState.values())
+            .map(TransactionState::from)
+            .collect(Collectors.toSet());
+    Assertions.assertEquals(
+        r.size(), org.apache.kafka.clients.admin.TransactionState.values().length);
+  }
+}
diff --git a/app/src/test/java/org/astraea/app/producer/ProducerTest.java b/app/src/test/java/org/astraea/app/producer/ProducerTest.java
index b6410b67..5c5d84a7 100644
--- a/app/src/test/java/org/astraea/app/producer/ProducerTest.java
+++ b/app/src/test/java/org/astraea/app/producer/ProducerTest.java
@@ -193,4 +193,13 @@ public class ProducerTest extends RequireBrokerCluster {
                 "transactional producer",
                 Producer.builder().bootstrapServers(bootstrapServers()).buildTransactional())));
   }
+
+  @Test
+  void testSetTransactionIdManually() {
+    try (var producer =
+        Producer.builder().bootstrapServers(bootstrapServers()).transactionId("chia").build()) {
+      Assertions.assertTrue(producer.transactional());
+      Assertions.assertTrue(producer.transactionId().isPresent());
+    }
+  }
 }