KAFKA-20734: Return noNode for known partitions without endpoints (#22680)

In `KRaftMetadataCache`, `getPartitionLeaderEndpoint` returns empty when
the previous leader is offline after the refactor.

With this PR, the method now returns `noNode` when the leader partition
exists, aligning with the original Scala implementation.

Changes are covered by unit tests. Please feel free to review.

Reviewers: Kevin Wu <kwu@confluent.io>, Justine Olshan
 <jolshan@confluent.io>

core/src/test/scala/unit/kafka/coordinator/transaction/TransactionMarkerChannelManagerTest.scala

@@ -400,6 +400,70 @@ class TransactionMarkerChannelManagerTest {
     assertEquals(Map(broker1 -> expectedBroker1Request), secondDrainedRequests)
   }
 
+  @ParameterizedTest
+  @EnumSource(classOf[TransactionResult])
+  def shouldRetryMarkersUntilLeaderEndpointIsAvailable(result: TransactionResult): Unit = {
+    mockCache()
+    txnMetadata2.state(if (result == TransactionResult.COMMIT) TransactionState.PREPARE_COMMIT else TransactionState.PREPARE_ABORT)
+    val expectedTransition = txnMetadata2.prepareComplete(time.milliseconds())
+    // Return noNode for addTxnMarkersToSend, and the first two generateRequests()
+    when(metadataCache.getPartitionLeaderEndpoint(
+      ArgumentMatchers.eq(partition1.topic),
+      ArgumentMatchers.eq(partition1.partition),
+      any())
+    ).thenReturn(Optional.of(Node.noNode))
+      .thenReturn(Optional.of(Node.noNode))
+      .thenReturn(Optional.of(Node.noNode))
+      .thenReturn(Optional.of(broker1)
+    )
+
+    val manager = new TransactionMarkerChannelManager(
+      KafkaConfig.fromProps(TestUtils.createBrokerConfig(1)), metadataCache, networkClient, txnStateManager, time)
+    try {
+      manager.addTxnMarkersToSend(coordinatorEpoch, result, txnMetadata2, expectedTransition)
+
+      // A known partition without a leader must retain its marker and remain pending.
+      assertTrue(manager.generateRequests().isEmpty)
+      assertEquals(1, manager.queueForUnknownBroker.totalNumMarkers)
+      assertEquals(1, manager.numTxnsWithPendingMarkers)
+      assertTrue(txnMetadata2.topicPartitions.contains(partition1))
+      verify(txnStateManager, times(0)).appendTransactionToLog(any(), any(), any(), any(), any(), any())
+
+      // Knowing the leader ID is insufficient while its broker registration is missing.
+      assertTrue(manager.generateRequests().isEmpty)
+      assertEquals(1, manager.queueForUnknownBroker.totalNumMarkers)
+      assertEquals(1, manager.numTxnsWithPendingMarkers)
+      assertTrue(txnMetadata2.topicPartitions.contains(partition1))
+      verify(txnStateManager, times(0)).appendTransactionToLog(any(), any(), any(), any(), any(), any())
+
+      val requests = manager.generateRequests().asScala.toList
+      assertEquals(1, requests.size)
+      assertEquals(broker1, requests.head.destination)
+      val request = requests.head.request.asInstanceOf[WriteTxnMarkersRequest.Builder].build()
+      assertEquals(util.List.of(new WriteTxnMarkersRequest.TxnMarkerEntry(
+        producerId2, producerEpoch, coordinatorEpoch, result, util.List.of(partition1), TransactionVersion.TV_2.featureLevel()
+      )), request.markers())
+      assertEquals(0, manager.queueForUnknownBroker.totalNumMarkers)
+      assertEquals(1, manager.numTxnsWithPendingMarkers)
+      verify(txnStateManager, times(0)).appendTransactionToLog(any(), any(), any(), any(), any(), any())
+
+      val response = new WriteTxnMarkersResponse(
+        util.Map.of(producerId2: java.lang.Long, util.Map.of(partition1, Errors.NONE)))
+      requests.head.handler.onComplete(new ClientResponse(
+        new RequestHeader(ApiKeys.WRITE_TXN_MARKERS, 0, "client", 1), null, null,
+        time.milliseconds(), time.milliseconds(), false, null, null, response))
+
+      assertTrue(txnMetadata2.topicPartitions.isEmpty)
+      assertEquals(0, manager.numTxnsWithPendingMarkers)
+      verify(txnStateManager).appendTransactionToLog(
+        ArgumentMatchers.eq(transactionalId2), ArgumentMatchers.eq(coordinatorEpoch),
+        ArgumentMatchers.eq(expectedTransition), any(), any(), any())
+      assertTrue(manager.generateRequests().isEmpty)
+    } finally {
+      manager.shutdown()
+    }
+  }
+
   @Test
   def shouldRemoveMarkersForTxnPartitionWhenPartitionEmigrated(): Unit = {
     mockCache()

core/src/test/scala/unit/kafka/server/MetadataCacheTest.scala

@@ -23,9 +23,9 @@ import org.apache.kafka.common.network.ListenerName
 import org.apache.kafka.common.protocol.{ApiMessage, Errors}
 import org.apache.kafka.common.record.internal.RecordBatch
 import org.apache.kafka.common.security.auth.SecurityProtocol
-import org.apache.kafka.common.{DirectoryId, TopicPartition, Uuid}
+import org.apache.kafka.common.{DirectoryId, Node, TopicPartition, Uuid}
 import org.apache.kafka.image.{MetadataDelta, MetadataImage, MetadataProvenance}
-import org.apache.kafka.metadata.{KRaftMetadataCache, LeaderRecoveryState, MetadataCache}
+import org.apache.kafka.metadata.{KRaftMetadataCache, LeaderConstants, LeaderRecoveryState, MetadataCache}
 import org.apache.kafka.server.common.KRaftVersion
 import org.junit.jupiter.api.Assertions._
 import org.junit.jupiter.api.Test
@@ -34,7 +34,7 @@ import org.junit.jupiter.params.provider.MethodSource
 
 import java.util
 import java.util.Arrays.asList
-import java.util.Collections
+import java.util.{Collections, Optional}
 import java.util.stream.Collectors
 import scala.collection.{Seq, mutable}
 import scala.jdk.CollectionConverters._
@@ -982,6 +982,54 @@ class MetadataCacheTest {
     assertTrue(fencedBroker.isFenced)
   }
 
+  @Test def testGetPartitionLeaderEndpointForKnownPartitionWithoutEndpoint(): Unit = {
+    val metadataCache = new KRaftMetadataCache(0, () => KRaftVersion.KRAFT_VERSION_0)
+    val listenerName = ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT)
+    val topic = "topic"
+    val topicId = Uuid.randomUuid
+    val records = Seq(new RegisterBrokerRecord()
+      .setBrokerId(0)
+      .setBrokerEpoch(0L)
+      .setRack(null)
+      .setFenced(false)
+      .setEndPoints(new BrokerEndpointCollection(Seq(new BrokerEndpoint()
+        .setHost("foo0")
+        .setPort(9092)
+        .setSecurityProtocol(SecurityProtocol.PLAINTEXT.id)
+        .setName(listenerName.value)).iterator.asJava)),
+      new TopicRecord()
+        .setName(topic)
+        .setTopicId(topicId),
+      new PartitionRecord()
+        .setTopicId(topicId)
+        .setPartitionId(0)
+        .setLeader(0)
+        .setLeaderEpoch(0)
+        .setPartitionEpoch(0)
+        .setIsr(util.List.of(0)).setReplicas(util.List.of(0)),
+      new PartitionRecord()
+        .setTopicId(topicId)
+        .setPartitionId(1)
+        .setLeader(1)
+        .setLeaderEpoch(0)
+        .setPartitionEpoch(0)
+        .setIsr(util.List.of(1))
+        .setReplicas(util.List.of(1)),
+      new PartitionRecord()
+        .setTopicId(topicId)
+        .setPartitionId(2)
+        .setLeader(LeaderConstants.NO_LEADER)
+        .setLeaderEpoch(0).setPartitionEpoch(0)
+        .setIsr(util.List.of).setReplicas(util.List.of(0)))
+    MetadataCacheTest.updateCache(metadataCache, records)
+    assertEquals(Optional.of(new Node(0, "foo0", 9092)), metadataCache.getPartitionLeaderEndpoint(topic, 0, listenerName))
+    assertEquals(Optional.of(Node.noNode), metadataCache.getPartitionLeaderEndpoint(topic, 0, ListenerName.forSecurityProtocol(SecurityProtocol.SSL)))
+    assertEquals(Optional.of(Node.noNode), metadataCache.getPartitionLeaderEndpoint(topic, 1, listenerName))
+    assertEquals(Optional.of(Node.noNode), metadataCache.getPartitionLeaderEndpoint(topic, 2, listenerName))
+    assertEquals(Optional.empty, metadataCache.getPartitionLeaderEndpoint(topic, 3, listenerName))
+    assertEquals(Optional.empty, metadataCache.getPartitionLeaderEndpoint("missing-topic", 0, listenerName))
+  }
+
   val oldRequestControllerEpoch: Int = 122
   val newRequestControllerEpoch: Int = 123
 

metadata/src/main/java/org/apache/kafka/metadata/KRaftMetadataCache.java

@@ -417,17 +417,18 @@ public class KRaftMetadataCache implements MetadataCache {
     }
 
     /**
-     * If the leader is not known, return None;
-     * If the leader is known and corresponding node is available, return Some(node)
-     * If the leader is known but corresponding node with the listener name is not available, return Some(NO_NODE)
+     * If the topic or partition is not known, return None.
+     * If the leader endpoint is available, return Some(node).
+     * Otherwise, return Some(NO_NODE).
      */
     @Override
     public Optional<Node> getPartitionLeaderEndpoint(String topicName, int partitionId, ListenerName listenerName) {
         MetadataImage image = currentImage;
         return Optional.ofNullable(image.topics().getTopic(topicName))
             .flatMap(topic -> Optional.ofNullable(topic.partitions().get(partitionId)))
-            .flatMap(partition -> Optional.ofNullable(image.cluster().broker(partition.leader))
-                .map(broker -> broker.node(listenerName.value()).orElse(Node.noNode())));
+            .map(partition -> Optional.ofNullable(image.cluster().broker(partition.leader))
+            .flatMap(broker -> broker.node(listenerName.value()))
+            .orElse(Node.noNode()));
     }
 
     @Override

metadata/src/main/java/org/apache/kafka/metadata/MetadataCache.java

@@ -102,9 +102,8 @@ public interface MetadataCache extends ConfigRepository {
     /**
      * Get a partition leader's endpoint
      *
-     * @return  If the leader is known, and the listener name is available, return Some(node). If the leader is known,
-     *          but the listener is unavailable, return Some(Node.NO_NODE). Otherwise, if the leader is not known,
-     *          return None
+     * @return  If the topic or partition is not known, return None. If the leader endpoint is available, return
+     *          Some(node). Otherwise, return Some(Node.NO_NODE).
      */
     Optional<Node> getPartitionLeaderEndpoint(String topic, int partitionId, ListenerName listenerName);
 
@@ -224,4 +223,4 @@ public interface MetadataCache extends ConfigRepository {
         if (aliveBrokers.isEmpty()) return Optional.empty();
         return Optional.of(aliveBrokers.get(ThreadLocalRandom.current().nextInt(aliveBrokers.size())));
     }
-}
+}