fix(test): wait for Kafka partitions before publishing

extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/simulate/KafkaResource.java

@@ -27,10 +27,12 @@
 import org.apache.kafka.clients.admin.CreatePartitionsResult;
 import org.apache.kafka.clients.admin.NewPartitions;
 import org.apache.kafka.clients.admin.NewTopic;
+import org.apache.kafka.clients.admin.OffsetSpec;
 import org.apache.kafka.clients.consumer.KafkaConsumer;
 import org.apache.kafka.clients.producer.KafkaProducer;
 import org.apache.kafka.clients.producer.ProducerConfig;
 import org.apache.kafka.clients.producer.ProducerRecord;
+import org.apache.kafka.clients.producer.RecordMetadata;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.serialization.ByteArrayDeserializer;
 import org.apache.kafka.common.serialization.ByteArraySerializer;
@@ -41,6 +43,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.Future;
 import java.util.concurrent.ThreadLocalRandom;
 
 /**
@@ -134,6 +137,7 @@ public void createTopicWithPartitions(String topicName, int numPartitions)
       admin.createTopics(
           List.of(new NewTopic(topicName, numPartitions, (short) 1))
       ).all().get();
+      waitForPartitionsToBeReady(admin, topicName, numPartitions);
     }
     catch (Exception e) {
       throw new RuntimeException(e);
@@ -162,9 +166,9 @@ public void deleteTopic(String topicName)
   }
 
   /**
-   * Increases the number of partitions in the given Kakfa topic. The topic must
-   * already exist. This method waits until the increase in the partition count
-   * has started (but not necessarily finished).
+   * Increases the number of partitions in the given Kafka topic. The topic must
+   * already exist. This method waits until every partition is ready to handle
+   * requests.
    */
   @Override
   public void increasePartitionsInTopic(String topic, int newPartitionCount)
@@ -174,8 +178,8 @@ public void increasePartitionsInTopic(String topic, int newPartitionCount)
           Map.of(topic, NewPartitions.increaseTo(newPartitionCount))
       );
 
-      // Wait for the partitioning to start
       result.values().get(topic).get();
+      waitForPartitionsToBeReady(admin, topic, newPartitionCount);
     }
     catch (Exception e) {
       throw new RuntimeException(e);
@@ -218,8 +222,12 @@ public void produceRecordsWithoutTransaction(List<ProducerRecord<byte[], byte[]>
     props.remove(ProducerConfig.TRANSACTIONAL_ID_CONFIG);
 
     try (final KafkaProducer<byte[], byte[]> kafkaProducer = new KafkaProducer<>(props)) {
-      for (ProducerRecord<byte[], byte[]> record : records) {
-        kafkaProducer.send(record);
+      final List<Future<RecordMetadata>> sendResults = new ArrayList<>(records.size());
+      for (final ProducerRecord<byte[], byte[]> record : records) {
+        sendResults.add(kafkaProducer.send(record));
+      }
+      for (final Future<RecordMetadata> sendResult : sendResults) {
+        sendResult.get();
       }
     }
     catch (Exception e) {
@@ -295,6 +303,18 @@ private KafkaProducer<byte[], byte[]> newProducer(Map<String, Object> extraPrope
     return new KafkaProducer<>(producerProperties);
   }
 
+  private void waitForPartitionsToBeReady(Admin admin, String topic, int partitionCount) throws Exception
+  {
+    // Topic and partition creation may complete before the partition leaders
+    // are ready to handle requests. Verify all partitions through their
+    // leaders before allowing callers to publish records.
+    final Map<TopicPartition, OffsetSpec> partitionOffsets = new HashMap<>();
+    for (int partition = 0; partition < partitionCount; partition++) {
+      partitionOffsets.put(new TopicPartition(topic, partition), OffsetSpec.latest());
+    }
+    admin.listOffsets(partitionOffsets).all().get();
+  }
+
   private Map<String, Object> commonClientProperties()
   {
     return Map.of("bootstrap.servers", getBootstrapServerUrl());

extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/simulate/KafkaResourceTest.java

@@ -20,9 +20,11 @@
 package org.apache.druid.indexing.kafka.simulate;
 
 import org.apache.druid.testing.embedded.EmbeddedDruidCluster;
+import org.apache.kafka.clients.producer.ProducerRecord;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.Timeout;
 
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
 
@@ -60,8 +62,40 @@ public void testKafka()
     final String topicName = "test-topic";
     resource.createTopicWithPartitions(topicName, 3);
     assertEquals(Set.of(topicName), resource.listTopics());
+
+    // Verify that every partition can accept records immediately after creating a topic.
+    resource.produceRecordsWithoutTransaction(
+        List.of(
+            new ProducerRecord<>(topicName, 0, null, new byte[]{1}),
+            new ProducerRecord<>(topicName, 1, null, new byte[]{1}),
+            new ProducerRecord<>(topicName, 2, null, new byte[]{1})
+        )
+    );
+    assertEquals(
+        Map.of("0", 1L, "1", 1L, "2", 1L),
+        resource.getPartitionOffsets(topicName)
+    );
     resource.deleteTopic(topicName);
 
+    final String expandedTopicName = "test-expanded-topic";
+    resource.createTopicWithPartitions(expandedTopicName, 2);
+    resource.increasePartitionsInTopic(expandedTopicName, 4);
+
+    // Verify that every partition can accept records immediately after increasing the partition count.
+    resource.produceRecordsWithoutTransaction(
+        List.of(
+            new ProducerRecord<>(expandedTopicName, 0, null, new byte[]{1}),
+            new ProducerRecord<>(expandedTopicName, 1, null, new byte[]{1}),
+            new ProducerRecord<>(expandedTopicName, 2, null, new byte[]{1}),
+            new ProducerRecord<>(expandedTopicName, 3, null, new byte[]{1})
+        )
+    );
+    assertEquals(
+        Map.of("0", 1L, "1", 1L, "2", 1L, "3", 1L),
+        resource.getPartitionOffsets(expandedTopicName)
+    );
+    resource.deleteTopic(expandedTopicName);
+
     resource.stop();
     assertFalse(resource.isRunning());
   }