Add enable2pc to txn manager and add overloaded method initProducerId(bool)

Author: rreddy-22

clients/src/main/java/org/apache/kafka/clients/producer/KafkaProducer.java

@@ -25,7 +25,6 @@
 import org.apache.kafka.clients.consumer.ConsumerRecords;
 import org.apache.kafka.clients.consumer.KafkaConsumer;
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
-import org.apache.kafka.clients.consumer.OffsetCommitCallback;
 import org.apache.kafka.clients.producer.internals.BufferPool;
 import org.apache.kafka.clients.producer.internals.BuiltInPartitioner;
 import org.apache.kafka.clients.producer.internals.KafkaProducerMetrics;
@@ -598,14 +597,17 @@ private TransactionManager configureTransactionState(ProducerConfig config,
 
         if (config.getBoolean(ProducerConfig.ENABLE_IDEMPOTENCE_CONFIG)) {
             final String transactionalId = config.getString(ProducerConfig.TRANSACTIONAL_ID_CONFIG);
+            final boolean enable2PC = config.getBoolean(ProducerConfig.TRANSACTION_TWO_PHASE_COMMIT_ENABLE_CONFIG);
             final int transactionTimeoutMs = config.getInt(ProducerConfig.TRANSACTION_TIMEOUT_CONFIG);
             final long retryBackoffMs = config.getLong(ProducerConfig.RETRY_BACKOFF_MS_CONFIG);
+            
             transactionManager = new TransactionManager(
                 logContext,
                 transactionalId,
                 transactionTimeoutMs,
                 retryBackoffMs,
-                apiVersions
+                apiVersions,
+                enable2PC
             );
 
             if (transactionManager.isTransactional())
@@ -656,6 +658,47 @@ public void initTransactions() {
         transactionManager.maybeUpdateTransactionV2Enabled(true);
     }
 
+    /**
+     * Performs initialization of transactions functionality in this producer instance. This method bootstraps
+     * the producer with a {@code producerId} and also resets the internal state of the producer following a previous
+     * fatal error. Additionally, it allows setting the {@code keepPreparedTxn} flag which, if set to true, puts the producer
+     * into a restricted state that only allows transaction completion operations.
+     * 
+     * <p>
+     * When {@code keepPreparedTxn} is set to {@code true}, the producer will be able to complete in-flight prepared
+     * transactions, but will only allow calling {@link #commitTransaction()}, {@link #abortTransaction()}, or
+     * the to-be-added {@code completeTransaction()} methods. This is to support recovery of prepared transactions 
+     * after a producer restart.
+     *
+     * <p>
+     * Note that this method should only be called once during the lifetime of a producer instance, and must be
+     * called before any other methods which require a {@code transactionalId} to be specified.
+     *
+     * @param keepPreparedTxn whether to keep prepared transactions, restricting the producer to only support completion of
+     *                        prepared transactions. When set to true, the producer will only allow transaction completion
+     *                        operations after initialization.
+     *
+     * @throws IllegalStateException if no {@code transactional.id} has been configured for the producer
+     * @throws org.apache.kafka.common.errors.UnsupportedVersionException fatal error indicating that the broker
+     *         does not support transactions (i.e. if its version is lower than 0.11.0.0). If this is encountered,
+     *         the producer cannot be used for transactional messaging.
+     * @throws org.apache.kafka.common.errors.AuthorizationException fatal error indicating that the configured
+     *         {@code transactional.id} is not authorized. If this is encountered, the producer cannot be used for
+     *         transactional messaging.
+     * @throws KafkaException if the producer has encountered a previous fatal error or for any other unexpected error
+     * @see #initTransactions()
+     */
+    public void initTransactions(boolean keepPreparedTxn) {
+        throwIfNoTransactionManager();
+        throwIfProducerClosed();
+        long now = time.nanoseconds();
+        TransactionalRequestResult result = transactionManager.initializeTransactions(keepPreparedTxn);
+        sender.wakeup();
+        result.await(maxBlockTimeMs, TimeUnit.MILLISECONDS);
+        producerMetrics.recordInit(time.nanoseconds() - now);
+        transactionManager.maybeUpdateTransactionV2Enabled(true);
+    }
+
     /**
      * Should be called before the start of each new transaction. Note that prior to the first invocation
      * of this method, you must invoke {@link #initTransactions()} exactly one time.
@@ -703,7 +746,7 @@ public void beginTransaction() throws ProducerFencedException {
      * <p>
      * Note, that the consumer should have {@code enable.auto.commit=false} and should
      * also not commit offsets manually (via {@link KafkaConsumer#commitSync(Map) sync} or
-     * {@link KafkaConsumer#commitAsync(Map, OffsetCommitCallback) async} commits).
+     * {@link KafkaConsumer#commitAsync()} (Map, OffsetCommitCallback) async} commits).
      * This method will raise {@link TimeoutException} if the producer cannot send offsets before expiration of {@code max.block.ms}.
      * Additionally, it will raise {@link InterruptException} if interrupted.
      *

clients/src/main/java/org/apache/kafka/clients/producer/MockProducer.java

@@ -158,6 +158,23 @@ public void initTransactions() {
         this.sentOffsets = false;
     }
 
+    @Override
+    public void initTransactions(boolean keepPreparedTxn) {
+        verifyNotClosed();
+        verifyNotFenced();
+        if (this.transactionInitialized) {
+            throw new IllegalStateException("MockProducer has already been initialized for transactions.");
+        }
+        if (this.initTransactionException != null) {
+            throw this.initTransactionException;
+        }
+        this.transactionInitialized = true;
+        this.transactionInFlight = false;
+        this.transactionCommitted = false;
+        this.transactionAborted = false;
+        this.sentOffsets = false;
+    }
+
     @Override
     public void beginTransaction() throws ProducerFencedException {
         verifyNotClosed();

clients/src/main/java/org/apache/kafka/clients/producer/Producer.java

@@ -44,6 +44,11 @@ public interface Producer<K, V> extends Closeable {
      */
     void initTransactions();
 
+    /**
+     * See {@link KafkaProducer#initTransactions(boolean)}
+     */
+    void initTransactions(boolean keepPreparedTxn);
+
     /**
      * See {@link KafkaProducer#beginTransaction()}
      */

clients/src/main/java/org/apache/kafka/clients/producer/internals/TransactionManager.java

@@ -138,6 +138,7 @@ public class TransactionManager {
      *
      * <ul>
      *     <li>{@link Producer#initTransactions()} calls {@link #initializeTransactions()}</li>
+     *     <li>{@link Producer#initTransactions(boolean)} calls {@link #initializeTransactions(boolean)}</li>
      *     <li>{@link Producer#beginTransaction()} calls {@link #beginTransaction()}</li>
      *     <li>{@link Producer#commitTransaction()}} calls {@link #beginCommit()}</li>
      *     <li>{@link Producer#abortTransaction()} calls {@link #beginAbort()}
@@ -195,6 +196,7 @@ public class TransactionManager {
     private volatile boolean clientSideEpochBumpRequired = false;
     private volatile long latestFinalizedFeaturesEpoch = -1;
     private volatile boolean isTransactionV2Enabled = false;
+    private final boolean enable2PC;
 
     private enum State {
         UNINITIALIZED,
@@ -255,7 +257,8 @@ public TransactionManager(final LogContext logContext,
                               final String transactionalId,
                               final int transactionTimeoutMs,
                               final long retryBackoffMs,
-                              final ApiVersions apiVersions) {
+                              final ApiVersions apiVersions,
+                              final boolean enable2PC) {
         this.producerIdAndEpoch = ProducerIdAndEpoch.NONE;
         this.transactionalId = transactionalId;
         this.log = logContext.logger(TransactionManager.class);
@@ -273,17 +276,29 @@ public TransactionManager(final LogContext logContext,
         this.retryBackoffMs = retryBackoffMs;
         this.txnPartitionMap = new TxnPartitionMap(logContext);
         this.apiVersions = apiVersions;
+        this.enable2PC = enable2PC;
     }
 
     void setPoisonStateOnInvalidTransition(boolean shouldPoisonState) {
         shouldPoisonStateOnInvalidTransition.set(shouldPoisonState);
     }
 
     public synchronized TransactionalRequestResult initializeTransactions() {
-        return initializeTransactions(ProducerIdAndEpoch.NONE);
+        return initializeTransactions(ProducerIdAndEpoch.NONE, false);
     }
 
     synchronized TransactionalRequestResult initializeTransactions(ProducerIdAndEpoch producerIdAndEpoch) {
+        return initializeTransactions(producerIdAndEpoch, false);
+    }
+
+    public synchronized TransactionalRequestResult initializeTransactions(boolean keepPreparedTxn) {
+        return initializeTransactions(ProducerIdAndEpoch.NONE, keepPreparedTxn);
+    }
+
+    synchronized TransactionalRequestResult initializeTransactions(
+        ProducerIdAndEpoch producerIdAndEpoch,
+        boolean keepPreparedTxn
+    ) {
         maybeFailWithError();
 
         boolean isEpochBump = producerIdAndEpoch != ProducerIdAndEpoch.NONE;
@@ -299,9 +314,12 @@ synchronized TransactionalRequestResult initializeTransactions(ProducerIdAndEpoc
                     .setTransactionalId(transactionalId)
                     .setTransactionTimeoutMs(transactionTimeoutMs)
                     .setProducerId(producerIdAndEpoch.producerId)
-                    .setProducerEpoch(producerIdAndEpoch.epoch);
+                    .setProducerEpoch(producerIdAndEpoch.epoch)
+                    .setEnable2Pc(enable2PC)
+                    .setKeepPreparedTxn(keepPreparedTxn);
+
             InitProducerIdHandler handler = new InitProducerIdHandler(new InitProducerIdRequest.Builder(requestData),
-                    isEpochBump);
+                isEpochBump);
             enqueueRequest(handler);
             return handler.result;
         }, State.INITIALIZING, "initTransactions");

clients/src/test/java/org/apache/kafka/clients/producer/KafkaProducerTest.java

@@ -1289,7 +1289,7 @@ public void testInitTransactionsResponseAfterTimeout() throws Exception {
                     ((FindCoordinatorRequest) request).data().keyType() == FindCoordinatorRequest.CoordinatorType.TRANSACTION.id(),
                 FindCoordinatorResponse.prepareResponse(Errors.NONE, "bad-transaction", NODE));
 
-            Future<?> future = executor.submit(producer::initTransactions);
+            Future<?> future = executor.submit(() -> producer.initTransactions());
             TestUtils.waitForCondition(client::hasInFlightRequests,
                 "Timed out while waiting for expected `InitProducerId` request to be sent");