From 29c3cd7cddf9b853c57fca2b2118f1b64e5dde30 Mon Sep 17 00:00:00 2001 From: oliverfuehrer Date: Tue, 9 May 2023 10:07:57 +0200 Subject: [PATCH] Dynamic partition improvements in Kafka binder Resolves https://github.com/spring-cloud/spring-cloud-stream/issues/2706 - Update partition count changes in Kafka dynamically during runtime - Checkstyle fixes - PartitionHandler changes - Improved handling of expression with 'payload' in the logic - Removed getter/setter for PartitionHandler and use ReflectionUtils within Test - Renamed property to 'dynamicPartitionUpdatesEnabled' and improved documentation (also one line per sentence) - Improved/renamed test to use embeddedkafka and really test update behaviour - Javadoc --- .../kafka/KafkaMessageChannelBinder.java | 52 +++++++++++++-- .../stream/binder/kafka/KafkaBinderTests.java | 65 +++++++++++++++++++ .../stream/binder/ProducerProperties.java | 20 ++++++ .../main/asciidoc/kafka/kafka_partitions.adoc | 8 +++ 4 files changed, 141 insertions(+), 4 deletions(-) diff --git a/binders/kafka-binder/spring-cloud-stream-binder-kafka/src/main/java/org/springframework/cloud/stream/binder/kafka/KafkaMessageChannelBinder.java b/binders/kafka-binder/spring-cloud-stream-binder-kafka/src/main/java/org/springframework/cloud/stream/binder/kafka/KafkaMessageChannelBinder.java index b7ce90faf..6629039ce 100644 --- a/binders/kafka-binder/spring-cloud-stream-binder-kafka/src/main/java/org/springframework/cloud/stream/binder/kafka/KafkaMessageChannelBinder.java +++ b/binders/kafka-binder/spring-cloud-stream-binder-kafka/src/main/java/org/springframework/cloud/stream/binder/kafka/KafkaMessageChannelBinder.java @@ -53,6 +53,7 @@ import org.apache.kafka.common.header.internals.RecordHeader; import org.apache.kafka.common.header.internals.RecordHeaders; import org.springframework.beans.factory.DisposableBean; +import org.springframework.beans.factory.config.ConfigurableListableBeanFactory; import org.springframework.cloud.stream.binder.AbstractMessageChannelBinder; import org.springframework.cloud.stream.binder.BinderHeaders; import org.springframework.cloud.stream.binder.BinderSpecificPropertiesProvider; @@ -63,6 +64,7 @@ import org.springframework.cloud.stream.binder.ExtendedProducerProperties; import org.springframework.cloud.stream.binder.ExtendedPropertiesBinder; import org.springframework.cloud.stream.binder.HeaderMode; import org.springframework.cloud.stream.binder.MessageValues; +import org.springframework.cloud.stream.binder.PartitionHandler; import org.springframework.cloud.stream.binder.kafka.config.ClientFactoryCustomizer; import org.springframework.cloud.stream.binder.kafka.properties.KafkaBinderConfigurationProperties; import org.springframework.cloud.stream.binder.kafka.properties.KafkaConsumerProperties; @@ -89,6 +91,7 @@ import org.springframework.integration.StaticMessageHeaderAccessor; import org.springframework.integration.acks.AcknowledgmentCallback; import org.springframework.integration.channel.AbstractMessageChannel; import org.springframework.integration.core.MessageProducer; +import org.springframework.integration.expression.ExpressionUtils; import org.springframework.integration.kafka.inbound.KafkaMessageDrivenChannelAdapter; import org.springframework.integration.kafka.inbound.KafkaMessageDrivenChannelAdapter.ListenerMode; import org.springframework.integration.kafka.inbound.KafkaMessageSource; @@ -159,6 +162,7 @@ import org.springframework.util.backoff.FixedBackOff; * @author Yi Liu * @author Chris Bono * @author Byungjun You + * @author Oliver Führer */ public class KafkaMessageChannelBinder extends // @checkstyle:off @@ -458,7 +462,7 @@ public class KafkaMessageChannelBinder extends kafkaTemplate.setAllowNonTransactional(allowNonTransactional); } ProducerConfigurationMessageHandler handler = new ProducerConfigurationMessageHandler( - kafkaTemplate, destination.getName(), producerProperties, producerFB); + kafkaTemplate, destination.getName(), producerProperties, producerFB, getBeanFactory()); if (errorChannel != null) { handler.setSendFailureChannel(errorChannel); } @@ -1478,19 +1482,25 @@ public class KafkaMessageChannelBinder extends return Collections.unmodifiableList(kafkaMessageListenerContainers); } - private final class ProducerConfigurationMessageHandler + final class ProducerConfigurationMessageHandler extends KafkaProducerMessageHandler { private boolean running = true; private final ProducerFactory producerFactory; + PartitionHandler kafkaPartitionHandler = null; + + private String topic; + ProducerConfigurationMessageHandler(KafkaTemplate kafkaTemplate, String topic, ExtendedProducerProperties producerProperties, - ProducerFactory producerFactory) { + ProducerFactory producerFactory, ConfigurableListableBeanFactory beanFactory) { super(kafkaTemplate); + this.topic = topic; + if (producerProperties.getExtension().isUseTopicHeader()) { setTopicExpression(PARSER.parseExpression("headers['" + KafkaHeaders.TOPIC + "'] ?: '" + topic + "'")); } @@ -1516,6 +1526,23 @@ public class KafkaMessageChannelBinder extends setSendTimeoutExpression(producerProperties.getExtension().getSendTimeoutExpression()); } this.producerFactory = producerFactory; + + /* + Activate own instance of a PartitionHandler if necessary/possible to override any other existing + partition calculation (see other usages of PartitionHandler) by using current partition count + (which may have changed at runtime) each time a message is handled. + PartitionKeyExpression 'payload' is not supported here, because of + OutboundContentTypeConvertingInterceptor would have been called before and the payload will be encoded and + not readable for PartitionHandler during handleMessage method. + */ + if (producerProperties.isDynamicPartitionUpdatesEnabled() && + producerProperties.getPartitionKeyExpression() != null && + !(producerProperties.getPartitionKeyExpression().getExpressionString() + .toLowerCase().contains("payload"))) { + kafkaPartitionHandler = + new PartitionHandler(ExpressionUtils.createStandardEvaluationContext(beanFactory), + producerProperties, beanFactory); + } } @Override @@ -1548,6 +1575,24 @@ public class KafkaMessageChannelBinder extends return this.running; } + @Override + public void handleMessage(Message message) { + + // if we use our own partition handler to update partition count we recalculate partition + if (kafkaPartitionHandler != null) { + kafkaPartitionHandler.setPartitionCount(getKafkaTemplate().partitionsFor(this.topic).size()); + int partitionId = kafkaPartitionHandler.determinePartition(message); + + Message newMessage = MessageBuilder + .fromMessage(message) + .setHeader(BinderHeaders.PARTITION_HEADER, partitionId).build(); + + super.handleMessage(newMessage); + } + else { + super.handleMessage(message); + } + } } /** @@ -1661,5 +1706,4 @@ public class KafkaMessageChannelBinder extends } } } - } diff --git a/binders/kafka-binder/spring-cloud-stream-binder-kafka/src/test/java/org/springframework/cloud/stream/binder/kafka/KafkaBinderTests.java b/binders/kafka-binder/spring-cloud-stream-binder-kafka/src/test/java/org/springframework/cloud/stream/binder/kafka/KafkaBinderTests.java index bf5b6575a..766278e5b 100644 --- a/binders/kafka-binder/spring-cloud-stream-binder-kafka/src/test/java/org/springframework/cloud/stream/binder/kafka/KafkaBinderTests.java +++ b/binders/kafka-binder/spring-cloud-stream-binder-kafka/src/test/java/org/springframework/cloud/stream/binder/kafka/KafkaBinderTests.java @@ -17,6 +17,7 @@ package org.springframework.cloud.stream.binder.kafka; import java.io.IOException; +import java.lang.reflect.Field; import java.nio.ByteBuffer; import java.nio.charset.StandardCharsets; import java.time.Duration; @@ -44,6 +45,7 @@ import org.apache.kafka.clients.admin.AdminClient; import org.apache.kafka.clients.admin.AdminClientConfig; import org.apache.kafka.clients.admin.CreateTopicsResult; import org.apache.kafka.clients.admin.DescribeTopicsResult; +import org.apache.kafka.clients.admin.NewPartitions; import org.apache.kafka.clients.admin.NewTopic; import org.apache.kafka.clients.admin.TopicDescription; import org.apache.kafka.clients.consumer.Consumer; @@ -74,6 +76,7 @@ import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.TestInfo; +import org.mockito.ArgumentMatchers; import org.springframework.beans.DirectFieldAccessor; import org.springframework.cloud.stream.binder.Binder; @@ -85,6 +88,7 @@ import org.springframework.cloud.stream.binder.ExtendedConsumerProperties; import org.springframework.cloud.stream.binder.ExtendedProducerProperties; import org.springframework.cloud.stream.binder.HeaderMode; import org.springframework.cloud.stream.binder.PartitionCapableBinderTests; +import org.springframework.cloud.stream.binder.PartitionHandler; import org.springframework.cloud.stream.binder.PartitionTestSupport; import org.springframework.cloud.stream.binder.PollableSource; import org.springframework.cloud.stream.binder.RequeueCurrentMessageException; @@ -153,13 +157,18 @@ import org.springframework.messaging.support.GenericMessage; import org.springframework.messaging.support.MessageBuilder; import org.springframework.util.Assert; import org.springframework.util.MimeTypeUtils; +import org.springframework.util.ReflectionUtils; import org.springframework.util.backoff.FixedBackOff; + import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatExceptionOfType; import static org.assertj.core.api.Assertions.entry; import static org.assertj.core.api.Assertions.fail; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; /** * @author Soby Chacko @@ -167,6 +176,7 @@ import static org.mockito.Mockito.mock; * @author Henryk Konsek * @author Gary Russell * @author Chris Bono + * @Author Oliver Führer */ @EmbeddedKafka(count = 1, controlledShutdown = true, topics = "error.pollableDlq.group-pcWithDlq", brokerProperties = {"transaction.state.log.replication.factor=1", "transaction.state.log.min.isr=1"}) @@ -4005,6 +4015,61 @@ public class KafkaBinderTests extends setupBindingAndAssert("enable-observation.2", binder); } + @Test + void testDynamicPartitionUpdates() throws Exception { + Binder binder = getBinder(); + ExtendedProducerProperties properties = createProducerProperties(); + properties.setPartitionKeyExpression( + spelExpressionParser.parseExpression("headers['partitionKey']")); + properties.setDynamicPartitionUpdatesEnabled(true); + properties.getExtension().getConfiguration().put(ProducerConfig.METADATA_MAX_AGE_CONFIG, "1000"); + + DirectChannel outputChannel = createBindableChannel("output", + createProducerBindingProperties(createProducerProperties())); + + invokeCreateTopic("partitionTopic", 7, 1); + + Binding producerBinding = binder.bindProducer("partitionTopic", + outputChannel, properties); + + KafkaMessageChannelBinder.ProducerConfigurationMessageHandler kafkaProducerMessageHandler = + (KafkaMessageChannelBinder.ProducerConfigurationMessageHandler) TestUtils.getPropertyValue( + producerBinding, "lifecycle", KafkaProducerMessageHandler.class); + + Field kafkaPartitionHandlerField = ReflectionUtils.findField( + KafkaMessageChannelBinder.ProducerConfigurationMessageHandler.class, "kafkaPartitionHandler"); + + PartitionHandler partitionHandler = + (PartitionHandler) kafkaPartitionHandlerField.get(kafkaProducerMessageHandler); + + assertThat(partitionHandler).isNotNull(); + PartitionHandler kafkaPartitionHandlerSpy = spy(partitionHandler); + + kafkaPartitionHandlerField.set(kafkaProducerMessageHandler, kafkaPartitionHandlerSpy); + + // send message with initial partition size + Message message = MessageBuilder + .withPayload("partitionTopic").setHeader("partitionKey", "123").build(); + outputChannel.send(message); + + // change partition size + Map counts = new HashMap<>(); + counts.put("partitionTopic", NewPartitions.increaseTo(11)); + adminClient.createPartitions(counts); + + // wait until metadata is processed in the background + Thread.sleep(2000); + + // send message again with new partition size + Message message2 = MessageBuilder + .withPayload("partitionTopic").setHeader("partitionKey", "456").build(); + outputChannel.send(message2); + + verify(kafkaPartitionHandlerSpy).setPartitionCount(7); + verify(kafkaPartitionHandlerSpy).setPartitionCount(11); + verify(kafkaPartitionHandlerSpy, times(2)).determinePartition(ArgumentMatchers.any()); + } + private void setupBindingAndAssert(String bindingName, AbstractKafkaTestBinder binder) throws Exception { ConfigurableApplicationContext applicationContext = (ConfigurableApplicationContext) binder.getApplicationContext(); TestObservationRegistry observationRegistry = TestObservationRegistry.create(); diff --git a/core/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/binder/ProducerProperties.java b/core/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/binder/ProducerProperties.java index d891c2477..9713d1828 100644 --- a/core/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/binder/ProducerProperties.java +++ b/core/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/binder/ProducerProperties.java @@ -81,6 +81,8 @@ public class ProducerProperties { private PollerProperties poller; + private boolean dynamicPartitionUpdatesEnabled = false; + public String getBindingName() { return bindingName; } @@ -195,6 +197,24 @@ public class ProducerProperties { this.poller = poller; } + /** + * Returns status of property dynamicPartitionUpdatesEnabled. + * @return true if dynamic updates should are enabled otherwise false + */ + public boolean isDynamicPartitionUpdatesEnabled() { + return dynamicPartitionUpdatesEnabled; + } + + /** + * A flag which enables/disables partition count updates during runtime. Disabled by default. + * Depends on binder if supported or not. + * Currently only supported by kafka binder (see 'Partitioning with the Kafka Binder' documentation for details) + * @param enabled true if dynamic updates should be enabled otherwise false + */ + public void setDynamicPartitionUpdatesEnabled(boolean enabled) { + this.dynamicPartitionUpdatesEnabled = enabled; + } + static class ExpressionSerializer extends JsonSerializer { @Override diff --git a/docs/src/main/asciidoc/kafka/kafka_partitions.adoc b/docs/src/main/asciidoc/kafka/kafka_partitions.adoc index 6b53250ff..48ee37888 100644 --- a/docs/src/main/asciidoc/kafka/kafka_partitions.adoc +++ b/docs/src/main/asciidoc/kafka/kafka_partitions.adoc @@ -65,6 +65,14 @@ You can override this default by using the `partitionSelectorExpression` or `par Since partitions are natively handled by Kafka, no special configuration is needed on the consumer side. Kafka allocates partitions across the instances. +NOTE: The partitionCount for a kafka topic may change during runtime (e.g. due to an adminstration task). +The calculated partitions will be different after that (e.g. new partitions will be used then). +Since 4.0.3 of Spring Cloud Stream runtime changes of partition count will be supported. +See also parameter 'spring.kafka.producer.properties.metadata.max.age.ms' to configure update interval. +Due to some limitations it is not possible to use a 'partition-key-expression' which references the 'payload' of a message, the mechanism will be disabled in that case. +The overall behavior is disabled by default and can be enabled using configuration parameter 'producer.dynamicPartitionUpdatesEnabled=true'. + + The following Spring Boot application listens to a Kafka stream and prints (to the console) the partition ID to which each message goes: [source,java]