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
This commit is contained in:
oliverfuehrer
2023-05-09 10:07:57 +02:00
committed by Soby Chacko
parent e19bdd4381
commit 29c3cd7cdd
4 changed files with 141 additions and 4 deletions

View File

@@ -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
}
}
}
}

View File

@@ -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();