KAFKA-13412; Ensure initTransactions() safe for retry after timeout

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

@@ -964,9 +964,10 @@ private Future<RecordMetadata> doSend(ProducerRecord<K, V> record, Callback call
             // producer callback will make sure to call both 'callback' and interceptor callback
             Callback interceptCallback = new InterceptorCallback<>(callback, this.interceptors, tp);
 
-            if (transactionManager != null && transactionManager.isTransactional()) {
-                transactionManager.failIfNotReadyForSend();
+            if (transactionManager != null) {
+                transactionManager.maybeAddPartition(tp);
             }
+
             RecordAccumulator.RecordAppendResult result = accumulator.append(tp, timestamp, serializedKey,
                     serializedValue, headers, interceptCallback, remainingWaitMs, true, nowMs);
 
@@ -985,9 +986,6 @@ private Future<RecordMetadata> doSend(ProducerRecord<K, V> record, Callback call
                     serializedValue, headers, interceptCallback, remainingWaitMs, false, nowMs);
             }
 
-            if (transactionManager != null && transactionManager.isTransactional())
-                transactionManager.maybeAddPartitionToTransaction(tp);
-
             if (result.batchIsFull || result.newBatchCreated) {
                 log.trace("Waking up the sender since topic {} partition {} is either full or getting a new batch", record.topic(), partition);
                 this.sender.wakeup();

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

@@ -225,7 +225,7 @@ void resetSequenceNumbers(Consumer<ProducerBatch> resetSequence) {
     private final Set<TopicPartition> newPartitionsInTransaction;
     private final Set<TopicPartition> pendingPartitionsInTransaction;
     private final Set<TopicPartition> partitionsInTransaction;
-    private TransactionalRequestResult pendingResult;
+    private PendingStateTransition pendingTransition;
 
     // This is used by the TxnRequestHandlers to control how long to back off before a given request is retried.
     // For instance, this value is lowered by the AddPartitionsToTxnHandler when it receives a CONCURRENT_TRANSACTIONS
@@ -347,11 +347,12 @@ synchronized TransactionalRequestResult initializeTransactions(ProducerIdAndEpoc
                     isEpochBump);
             enqueueRequest(handler);
             return handler.result;
-        }, State.INITIALIZING);
+        }, State.INITIALIZING, "initTransactions");
     }
 
     public synchronized void beginTransaction() {
         ensureTransactional();
+        throwIfPendingState("beginTransaction");
         maybeFailWithError();
         transitionTo(State.IN_TRANSACTION);
     }
@@ -361,7 +362,7 @@ public synchronized TransactionalRequestResult beginCommit() {
             maybeFailWithError();
             transitionTo(State.COMMITTING_TRANSACTION);
             return beginCompletingTransaction(TransactionResult.COMMIT);
-        }, State.COMMITTING_TRANSACTION);
+        }, State.COMMITTING_TRANSACTION, "commitTransaction");
     }
 
     public synchronized TransactionalRequestResult beginAbort() {
@@ -373,7 +374,7 @@ public synchronized TransactionalRequestResult beginAbort() {
             // We're aborting the transaction, so there should be no need to add new partitions
             newPartitionsInTransaction.clear();
             return beginCompletingTransaction(TransactionResult.ABORT);
-        }, State.ABORTING_TRANSACTION);
+        }, State.ABORTING_TRANSACTION, "abortTransaction");
     }
 
     private TransactionalRequestResult beginCompletingTransaction(TransactionResult transactionResult) {
@@ -404,10 +405,12 @@ private TransactionalRequestResult beginCompletingTransaction(TransactionResult
     public synchronized TransactionalRequestResult sendOffsetsToTransaction(final Map<TopicPartition, OffsetAndMetadata> offsets,
                                                                             final ConsumerGroupMetadata groupMetadata) {
         ensureTransactional();
+        throwIfPendingState("sendOffsetsToTransaction");
         maybeFailWithError();
-        if (currentState != State.IN_TRANSACTION)
-            throw new KafkaException("Cannot send offsets to transaction either because the producer is not in an " +
-                    "active transaction");
+
+        if (currentState != State.IN_TRANSACTION) {
+            throw new KafkaException("Cannot send offsets to transaction either because in state " + currentState);
+        }
 
         log.debug("Begin adding offsets {} for consumer group {} to transaction", offsets, groupMetadata);
         AddOffsetsToTxnRequest.Builder builder = new AddOffsetsToTxnRequest.Builder(
@@ -423,34 +426,31 @@ public synchronized TransactionalRequestResult sendOffsetsToTransaction(final Ma
         return handler.result;
     }
 
-    public synchronized void maybeAddPartitionToTransaction(TopicPartition topicPartition) {
-        if (isPartitionAdded(topicPartition) || isPartitionPendingAdd(topicPartition))
-            return;
+    public synchronized void maybeAddPartition(TopicPartition topicPartition) {
+        maybeFailWithError();
+        throwIfPendingState("send");
 
-        log.debug("Begin adding new partition {} to transaction", topicPartition);
-        topicPartitionBookkeeper.addPartition(topicPartition);
-        newPartitionsInTransaction.add(topicPartition);
+        if (isTransactional()) {
+            if (!hasProducerId()) {
+                throw new KafkaException("Cannot add partition " + topicPartition +
+                    "to transaction before completing a call to initTransactions");
+            } else if (currentState != State.IN_TRANSACTION) {
+                throw new KafkaException("Cannot add partition " + topicPartition +
+                    " to transaction while in state  " + currentState);
+            } else if (isPartitionAdded(topicPartition) || isPartitionPendingAdd(topicPartition)) {
+                return;
+            } else {
+                log.debug("Begin adding new partition {} to transaction", topicPartition);
+                topicPartitionBookkeeper.addPartition(topicPartition);
+                newPartitionsInTransaction.add(topicPartition);
+            }
+        }
     }
 
     RuntimeException lastError() {
         return lastError;
     }
 
-    public synchronized void failIfNotReadyForSend() {
-        if (hasError())
-            throw new KafkaException("Cannot perform send because at least one previous transactional or " +
-                    "idempotent request has failed with errors.", lastError);
-
-        if (isTransactional()) {
-            if (!hasProducerId())
-                throw new IllegalStateException("Cannot perform a 'send' before completing a call to initTransactions " +
-                        "when transactions are enabled.");
-
-            if (currentState != State.IN_TRANSACTION)
-                throw new IllegalStateException("Cannot call send in state " + currentState);
-        }
-    }
-
     synchronized boolean isSendToPartitionAllowed(TopicPartition tp) {
         if (hasFatalError())
             return false;
@@ -500,8 +500,8 @@ synchronized void transitionToFatalError(RuntimeException exception) {
         log.info("Transiting to fatal error state due to {}", exception.toString());
         transitionTo(State.FATAL_ERROR, exception);
 
-        if (pendingResult != null) {
-            pendingResult.fail(exception);
+        if (pendingTransition != null) {
+            pendingTransition.result.fail(exception);
         }
     }
 
@@ -919,8 +919,8 @@ synchronized void close() {
         KafkaException shutdownException = new KafkaException("The producer closed forcefully");
         pendingRequests.forEach(handler ->
                 handler.fatalError(shutdownException));
-        if (pendingResult != null) {
-            pendingResult.fail(shutdownException);
+        if (pendingTransition != null) {
+            pendingTransition.result.fail(shutdownException);
         }
     }
 
@@ -1183,20 +1183,42 @@ private TxnOffsetCommitHandler txnOffsetCommitHandler(TransactionalRequestResult
         return new TxnOffsetCommitHandler(result, builder);
     }
 
+    private void throwPendingTransitionError(String operation) {
+        throw new KafkaException("Cannot attempt operation `" + operation + "` "
+            + "because the previous call to `" + pendingTransition.operation + "` "
+            + "timed out and must be retried");
+    }
+
+    private void throwIfPendingState(String operation) {
+        if (pendingTransition != null) {
+            if (pendingTransition.result.isAcked()) {
+                pendingTransition = null;
+            } else {
+                throwPendingTransitionError(operation);
+            }
+        }
+    }
+
     private TransactionalRequestResult handleCachedTransactionRequestResult(
-            Supplier<TransactionalRequestResult> transactionalRequestResultSupplier,
-            State targetState) {
+        Supplier<TransactionalRequestResult> transactionalRequestResultSupplier,
+        State nextState,
+        String operation
+    ) {
         ensureTransactional();
 
-        if (pendingResult != null && currentState == targetState) {
-            TransactionalRequestResult result = pendingResult;
-            if (result.isCompleted())
-                pendingResult = null;
-            return result;
+        if (pendingTransition != null) {
+            if (pendingTransition.result.isAcked()) {
+                pendingTransition = null;
+            } else if (nextState != pendingTransition.state) {
+                throwPendingTransitionError(operation);
+            } else {
+                return pendingTransition.result;
+            }
         }
 
-        pendingResult = transactionalRequestResultSupplier.get();
-        return pendingResult;
+        TransactionalRequestResult result = transactionalRequestResultSupplier.get();
+        pendingTransition = new PendingStateTransition(result, nextState, operation);
+        return result;
     }
 
     // package-private for testing
@@ -1762,4 +1784,22 @@ private boolean isFatalException(Errors error) {
                    || error == Errors.PRODUCER_FENCED
                    || error == Errors.UNSUPPORTED_FOR_MESSAGE_FORMAT;
     }
+
+    private static final class PendingStateTransition {
+        private final TransactionalRequestResult result;
+        private final State state;
+        private final String operation;
+
+        private PendingStateTransition(
+            TransactionalRequestResult result,
+            State state,
+            String operation
+        ) {
+            this.result = result;
+            this.state = state;
+            this.operation = operation;
+        }
+    }
+
+
 }

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

@@ -20,15 +20,14 @@
 import org.apache.kafka.common.errors.InterruptException;
 import org.apache.kafka.common.errors.TimeoutException;
 
-import java.util.Locale;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 
 public final class TransactionalRequestResult {
-
     private final CountDownLatch latch;
     private volatile RuntimeException error = null;
     private final String operation;
+    private volatile boolean isAcked = false;
 
     public TransactionalRequestResult(String operation) {
         this(new CountDownLatch(1), operation);
@@ -49,29 +48,20 @@ public void done() {
     }
 
     public void await() {
-        boolean completed = false;
-
-        while (!completed) {
-            try {
-                latch.await();
-                completed = true;
-            } catch (InterruptedException e) {
-                // Keep waiting until done, we have no other option for these transactional requests.
-            }
-        }
-
-        if (!isSuccessful())
-            throw error();
+        this.await(Long.MAX_VALUE, TimeUnit.MILLISECONDS);
     }
 
     public void await(long timeout, TimeUnit unit) {
         try {
             boolean success = latch.await(timeout, unit);
-            if (!isSuccessful()) {
-                throw error();
-            }
             if (!success) {
-                throw new TimeoutException("Timeout expired after " + timeout + " " + unit.name().toLowerCase(Locale.ROOT) + " while awaiting " + operation);
+                throw new TimeoutException("Timeout expired after " + unit.toMillis(timeout) +
+                    "ms while awaiting " + operation);
+            }
+
+            isAcked = true;
+            if (error != null) {
+                throw error;
             }
         } catch (InterruptedException e) {
             throw new InterruptException("Received interrupt while awaiting " + operation, e);
@@ -83,11 +73,15 @@ public RuntimeException error() {
     }
 
     public boolean isSuccessful() {
-        return error == null;
+        return isCompleted() && error == null;
     }
 
     public boolean isCompleted() {
         return latch.getCount() == 0L;
     }
 
+    public boolean isAcked() {
+        return isAcked;
+    }
+
 }

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

@@ -903,6 +903,48 @@ public void testPartitionsForWithNullTopic() {
         }
     }
 
+    @Test
+    public void testInitTransactionsResponseAfterTimeout() throws Exception {
+        int maxBlockMs = 500;
+
+        Map<String, Object> configs = new HashMap<>();
+        configs.put(ProducerConfig.TRANSACTIONAL_ID_CONFIG, "bad-transaction");
+        configs.put(ProducerConfig.MAX_BLOCK_MS_CONFIG, maxBlockMs);
+        configs.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9000");
+
+        Time time = new MockTime();
+        MetadataResponse initialUpdateResponse = RequestTestUtils.metadataUpdateWith(1, singletonMap("topic", 1));
+        ProducerMetadata metadata = newMetadata(0, Long.MAX_VALUE);
+        metadata.updateWithCurrentRequestVersion(initialUpdateResponse, false, time.milliseconds());
+
+        MockClient client = new MockClient(time, metadata);
+
+        ExecutorService executor = Executors.newFixedThreadPool(1);
+
+        Producer<String, String> producer = kafkaProducer(configs, new StringSerializer(),
+            new StringSerializer(), metadata, client, null, time);
+        try {
+            client.prepareResponse(
+                request -> request instanceof FindCoordinatorRequest &&
+                    ((FindCoordinatorRequest) request).data().keyType() == FindCoordinatorRequest.CoordinatorType.TRANSACTION.id(),
+                FindCoordinatorResponse.prepareResponse(Errors.NONE, "bad-transaction", host1));
+
+            Future<?> future = executor.submit(producer::initTransactions);
+            TestUtils.waitForCondition(client::hasInFlightRequests,
+                "Timed out while waiting for expected `InitProducerId` request to be sent");
+
+            time.sleep(maxBlockMs);
+            TestUtils.assertFutureThrows(future, TimeoutException.class);
+
+            client.respond(initProducerIdResponse(1L, (short) 5, Errors.NONE));
+
+            Thread.sleep(1000);
+            producer.initTransactions();
+        } finally {
+            producer.close(Duration.ZERO);
+        }
+    }
+
     @Test
     public void testInitTransactionTimeout() {
         Map<String, Object> configs = new HashMap<>();