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

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<byte[], byte[]> {
 
 		private boolean running = true;
 
 		private final ProducerFactory<byte[], byte[]> producerFactory;
 
+		PartitionHandler kafkaPartitionHandler = null;
+
+		private String topic;
+
 		ProducerConfigurationMessageHandler(KafkaTemplate<byte[], byte[]> kafkaTemplate,
 				String topic,
 				ExtendedProducerProperties<KafkaProducerProperties> producerProperties,
-				ProducerFactory<byte[], byte[]> producerFactory) {
+				ProducerFactory<byte[], byte[]> 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
 			}
 		}
 	}
-
 }

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<KafkaProducerProperties> 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<MessageChannel> 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<String, NewPartitions> 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();

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<Expression> {
 
 		@Override

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]