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

@@ -28,9 +28,14 @@ import org.apache.kafka.common.requests.{RequestHeader, TransactionResult, Write
 import org.apache.kafka.common.utils.MockTime
 import org.apache.kafka.common.utils.internals.LogContext
 import org.apache.kafka.common.metrics.Metrics
-import org.apache.kafka.common.{Node, TopicPartition}
+import org.apache.kafka.common.{Node, TopicPartition, Uuid}
+import org.apache.kafka.common.metadata.{PartitionChangeRecord, PartitionRecord, RegisterBrokerRecord, TopicRecord}
+import org.apache.kafka.common.metadata.RegisterBrokerRecord.{BrokerEndpoint, BrokerEndpointCollection}
+import org.apache.kafka.common.network.ListenerName
+import org.apache.kafka.common.protocol.ApiMessage
+import org.apache.kafka.common.security.auth.SecurityProtocol
 import org.apache.kafka.coordinator.transaction.{CoordinatorEpochAndTxnMetadata, TransactionMetadata, TransactionState}
-import org.apache.kafka.metadata.MetadataCache
+import org.apache.kafka.metadata.{LeaderConstants, MetadataCache, MetadataCacheFixtures}
 import org.apache.kafka.server.common.{MetadataVersion, TransactionVersion}
 import org.apache.kafka.server.metrics.{KafkaMetricsGroup, KafkaYammerMetrics}
 import org.apache.kafka.server.util.{InterBrokerSendThread, RequestAndCompletionHandler}
@@ -401,6 +406,94 @@ 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())
+    val cache = MetadataCacheFixtures.createCache()
+    val topicId = Uuid.randomUuid()
+    val listenerName = ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT)
+    MetadataCacheFixtures.updateCache(cache, util.List.of[ApiMessage](
+      new TopicRecord().setName(partition1.topic).setTopicId(topicId),
+      new PartitionRecord()
+        .setTopicId(topicId)
+        .setPartitionId(partition1.partition)
+        .setLeader(LeaderConstants.NO_LEADER)
+        .setReplicas(util.List.of[Integer](broker1.id))
+        .setIsr(util.List.of[Integer]())
+    ))
+
+    val manager = new TransactionMarkerChannelManager(
+      KafkaConfig.fromProps(TestUtils.createBrokerConfig(1)), cache, 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())
+
+      MetadataCacheFixtures.updateCache(cache, util.List.of[ApiMessage](
+        new PartitionChangeRecord()
+          .setTopicId(topicId)
+          .setPartitionId(partition1.partition)
+          .setLeader(broker1.id)
+          .setIsr(util.List.of[Integer](broker1.id))
+      ))
+
+      // 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())
+
+      MetadataCacheFixtures.updateCache(cache, util.List.of[ApiMessage](
+        new RegisterBrokerRecord()
+          .setBrokerId(broker1.id)
+          .setRack("rack1")
+          .setFenced(false)
+          .setEndPoints(new BrokerEndpointCollection(util.List.of(
+            new BrokerEndpoint()
+              .setHost(broker1.host)
+              .setPort(broker1.port)
+              .setSecurityProtocol(SecurityProtocol.PLAINTEXT.id)
+              .setName(listenerName.value)
+          )))
+      ))
+
+      val requests = manager.generateRequests().asScala.toList
+      assertEquals(1, requests.size)
+      assertEquals(new Node(broker1.id, broker1.host, broker1.port, "rack1"), 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()

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

@@ -101,9 +101,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);
 
@@ -223,4 +222,4 @@ public interface MetadataCache extends ConfigRepository {
         if (aliveBrokers.isEmpty()) return Optional.empty();
         return Optional.of(aliveBrokers.get(ThreadLocalRandom.current().nextInt(aliveBrokers.size())));
     }
-}
+}

metadata/src/test/java/org/apache/kafka/metadata/MetadataCacheTest.java

@@ -508,6 +508,66 @@ public class MetadataCacheTest {
         }
     }
 
+    @Test
+    public void testGetPartitionLeaderEndpointForKnownPartitionWithoutEndpoint() {
+        MetadataCache cache = createCache();
+        ListenerName listenerName = ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT);
+        String topic = "topic";
+        Uuid topicId = Uuid.randomUuid();
+
+        List<ApiMessage> records = List.of(
+            new RegisterBrokerRecord()
+                .setBrokerId(0)
+                .setBrokerEpoch(BROKER_EPOCH)
+                .setRack(null)
+                .setFenced(false)
+                .setEndPoints(new BrokerEndpointCollection(List.of(
+                    new BrokerEndpoint()
+                        .setHost("foo0")
+                        .setPort(9092)
+                        .setSecurityProtocol(SecurityProtocol.PLAINTEXT.id)
+                        .setName(listenerName.value())
+                ))),
+            new TopicRecord().setName(topic).setTopicId(topicId),
+            new PartitionRecord()
+                .setTopicId(topicId)
+                .setPartitionId(0)
+                .setLeader(0)
+                .setLeaderEpoch(0)
+                .setPartitionEpoch(0)
+                .setIsr(List.of(0))
+                .setReplicas(List.of(0)),
+            new PartitionRecord()
+                .setTopicId(topicId)
+                .setPartitionId(1)
+                .setLeader(1)
+                .setLeaderEpoch(0)
+                .setPartitionEpoch(0)
+                .setIsr(List.of(1))
+                .setReplicas(List.of(1)),
+            new PartitionRecord()
+                .setTopicId(topicId)
+                .setPartitionId(2)
+                .setLeader(LeaderConstants.NO_LEADER)
+                .setLeaderEpoch(0)
+                .setPartitionEpoch(0)
+                .setIsr(List.of())
+                .setReplicas(List.of(0))
+        );
+        updateCache(cache, records);
+
+        assertEquals(Optional.of(new Node(0, "foo0", 9092)), cache.getPartitionLeaderEndpoint(topic, 0, listenerName));
+        assertEquals(Optional.of(Node.noNode()), cache.getPartitionLeaderEndpoint(
+            topic,
+            0,
+            ListenerName.forSecurityProtocol(SecurityProtocol.SSL)
+        ));
+        assertEquals(Optional.of(Node.noNode()), cache.getPartitionLeaderEndpoint(topic, 1, listenerName));
+        assertEquals(Optional.of(Node.noNode()), cache.getPartitionLeaderEndpoint(topic, 2, listenerName));
+        assertEquals(Optional.empty(), cache.getPartitionLeaderEndpoint(topic, 3, listenerName));
+        assertEquals(Optional.empty(), cache.getPartitionLeaderEndpoint("missing-topic", 0, listenerName));
+    }
+
     private void updateCacheWithBrokers(MetadataCache cache, List<Integer> brokerIds,
                                          Uuid topicId, List<ApiMessage> topicRecords) {
         SecurityProtocol securityProtocol = SecurityProtocol.PLAINTEXT;