diff --git a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/simulate/KafkaResource.java b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/simulate/KafkaResource.java index 8ef0c888b7e6..465e4db35d43 100644 --- a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/simulate/KafkaResource.java +++ b/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 props.remove(ProducerConfig.TRANSACTIONAL_ID_CONFIG); try (final KafkaProducer kafkaProducer = new KafkaProducer<>(props)) { - for (ProducerRecord record : records) { - kafkaProducer.send(record); + final List> sendResults = new ArrayList<>(records.size()); + for (final ProducerRecord record : records) { + sendResults.add(kafkaProducer.send(record)); + } + for (final Future sendResult : sendResults) { + sendResult.get(); } } catch (Exception e) { @@ -295,6 +303,18 @@ private KafkaProducer newProducer(Map 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 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 commonClientProperties() { return Map.of("bootstrap.servers", getBootstrapServerUrl()); diff --git a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/simulate/KafkaResourceTest.java b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/simulate/KafkaResourceTest.java index 9b8237417727..a15492ca9929 100644 --- a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/simulate/KafkaResourceTest.java +++ b/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()); }