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

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

@@ -241,6 +241,7 @@ void resetSequenceNumbers(Consumer<ProducerBatch> resetSequence) {
     private Node consumerGroupCoordinator;
     private boolean coordinatorSupportsBumpingEpoch;
 
+    private volatile State priorState = null;
     private volatile State currentState = State.UNINITIALIZED;
     private volatile RuntimeException lastError = null;
     private volatile ProducerIdAndEpoch producerIdAndEpoch;
@@ -1090,6 +1091,7 @@ private void transitionTo(State target, RuntimeException error) {
         else
             log.debug("Transition from state {} to {}", currentState, target);
 
+        priorState = currentState;
         currentState = target;
     }
 
@@ -1184,15 +1186,19 @@ private TxnOffsetCommitHandler txnOffsetCommitHandler(TransactionalRequestResult
     }
 
     private TransactionalRequestResult handleCachedTransactionRequestResult(
-            Supplier<TransactionalRequestResult> transactionalRequestResultSupplier,
-            State targetState) {
+        Supplier<TransactionalRequestResult> transactionalRequestResultSupplier,
+        State transientState
+    ) {
         ensureTransactional();
 
-        if (pendingResult != null && currentState == targetState) {
+        if (pendingResult != null && !pendingResult.isAcked()) {
             TransactionalRequestResult result = pendingResult;
-            if (result.isCompleted())
+            if (result.isCompleted() && priorState == transientState) {
                 pendingResult = null;
-            return result;
+                return result;
+            } else if (currentState == transientState) {
+                return result;
+            }
         }
 
         pendingResult = transactionalRequestResultSupplier.get();

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

@@ -29,6 +29,7 @@ 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);
@@ -60,6 +61,8 @@ public void await() {
             }
         }
 
+        isAcked = true;
+
         if (!isSuccessful())
             throw error();
     }
@@ -68,10 +71,13 @@ public void await(long timeout, TimeUnit unit) {
         try {
             boolean success = latch.await(timeout, unit);
             if (!isSuccessful()) {
+                isAcked = true;
                 throw error();
             }
             if (!success) {
                 throw new TimeoutException("Timeout expired after " + timeout + " " + unit.name().toLowerCase(Locale.ROOT) + " while awaiting " + operation);
+            } else {
+                isAcked = true;
             }
         } catch (InterruptedException e) {
             throw new InterruptException("Received interrupt while awaiting " + operation, e);
@@ -90,4 +96,8 @@ 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,47 @@ 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, "blah blah");
+
+            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<>();

clients/src/test/java/org/apache/kafka/clients/producer/internals/TransactionManagerTest.java

@@ -1661,6 +1661,10 @@ public void testInvalidProducerEpochConvertToProducerFencedInEndTxn() throws Int
         runUntil(commitResult::isCompleted);
         runUntil(responseFuture::isDone);
 
+        assertThrows(KafkaException.class, commitResult::await);
+        assertFalse(commitResult.isSuccessful());
+        assertTrue(commitResult.isAcked());
+
         // make sure the exception was thrown directly from the follow-up calls.
         assertThrows(KafkaException.class, () -> transactionManager.beginTransaction());
         assertThrows(KafkaException.class, () -> transactionManager.beginCommit());
@@ -3522,13 +3526,17 @@ private void doInitTransactions() {
     }
 
     private void doInitTransactions(long producerId, short epoch) {
-        transactionManager.initializeTransactions();
+        TransactionalRequestResult result = transactionManager.initializeTransactions();
         prepareFindCoordinatorResponse(Errors.NONE, false, CoordinatorType.TRANSACTION, transactionalId);
         runUntil(() -> transactionManager.coordinator(CoordinatorType.TRANSACTION) != null);
         assertEquals(brokerNode, transactionManager.coordinator(CoordinatorType.TRANSACTION));
 
         prepareInitPidResponse(Errors.NONE, false, producerId, epoch);
         runUntil(transactionManager::hasProducerId);
+
+        result.await();
+        assertTrue(result.isSuccessful());
+        assertTrue(result.isAcked());
     }
 
     private void assertAbortableError(Class<? extends RuntimeException> cause) {