From 23bc29958c0124b3077a93ab2710e844cd860f05 Mon Sep 17 00:00:00 2001 From: Gary Russell Date: Fri, 12 Jan 2018 14:08:16 -0500 Subject: [PATCH] GH-184: Polishing - user rebalance listener Needed by the SCST Kafka Binder * GH-184: Polishing - add raw header option --- .../dsl/KafkaInboundChannelAdapterSpec.java | 12 +++++ .../kafka/inbound/KafkaMessageSource.java | 47 ++++++++++++++++++- .../MessageSourceIntegrationTests.java | 19 ++++++++ .../kafka/inbound/MessageSourceTests.java | 3 ++ 4 files changed, 80 insertions(+), 1 deletion(-) diff --git a/spring-integration-kafka/src/main/java/org/springframework/integration/kafka/dsl/KafkaInboundChannelAdapterSpec.java b/spring-integration-kafka/src/main/java/org/springframework/integration/kafka/dsl/KafkaInboundChannelAdapterSpec.java index 0c075d05cd..d742d60017 100644 --- a/spring-integration-kafka/src/main/java/org/springframework/integration/kafka/dsl/KafkaInboundChannelAdapterSpec.java +++ b/spring-integration-kafka/src/main/java/org/springframework/integration/kafka/dsl/KafkaInboundChannelAdapterSpec.java @@ -18,6 +18,8 @@ package org.springframework.integration.kafka.dsl; import java.lang.reflect.Type; +import org.apache.kafka.clients.consumer.ConsumerRebalanceListener; + import org.springframework.integration.dsl.MessageSourceSpec; import org.springframework.integration.kafka.inbound.KafkaMessageSource; import org.springframework.integration.kafka.inbound.KafkaMessageSource.KafkaAckCallbackFactory; @@ -73,4 +75,14 @@ public class KafkaInboundChannelAdapterSpec return this; } + public KafkaInboundChannelAdapterSpec rebalanceListener(ConsumerRebalanceListener rebalanceListener) { + this.target.setRebalanceListener(rebalanceListener); + return this; + } + + public KafkaInboundChannelAdapterSpec rawMessageHeader(boolean rawMessageHeader) { + this.target.setRawMessageHeader(rawMessageHeader); + return this; + } + } diff --git a/spring-integration-kafka/src/main/java/org/springframework/integration/kafka/inbound/KafkaMessageSource.java b/spring-integration-kafka/src/main/java/org/springframework/integration/kafka/inbound/KafkaMessageSource.java index 963c18b088..250b09b8c5 100644 --- a/spring-integration-kafka/src/main/java/org/springframework/integration/kafka/inbound/KafkaMessageSource.java +++ b/spring-integration-kafka/src/main/java/org/springframework/integration/kafka/inbound/KafkaMessageSource.java @@ -43,10 +43,12 @@ import org.apache.kafka.common.errors.WakeupException; import org.springframework.beans.factory.DisposableBean; import org.springframework.integration.IntegrationMessageHeaderAccessor; import org.springframework.integration.endpoint.AbstractMessageSource; +import org.springframework.integration.support.AbstractIntegrationMessageBuilder; import org.springframework.integration.support.AcknowledgmentCallback; import org.springframework.integration.support.AcknowledgmentCallbackFactory; import org.springframework.kafka.core.ConsumerFactory; import org.springframework.kafka.support.Acknowledgment; +import org.springframework.kafka.support.KafkaHeaders; import org.springframework.kafka.support.converter.KafkaMessageHeaders; import org.springframework.kafka.support.converter.MessagingMessageConverter; import org.springframework.kafka.support.converter.RecordMessageConverter; @@ -96,6 +98,10 @@ public class KafkaMessageSource extends AbstractMessageSource private Type payloadType; + private ConsumerRebalanceListener rebalanceListener; + + private boolean rawMessageHeader; + private volatile Consumer consumer; private volatile Collection partitions; @@ -184,11 +190,37 @@ public class KafkaMessageSource extends AbstractMessageSource this.payloadType = payloadType; } + protected ConsumerRebalanceListener getRebalanceListener() { + return this.rebalanceListener; + } + + /** + * Set a rebalance listener. + * @param rebalanceListener the rebalance listener. + */ + public void setRebalanceListener(ConsumerRebalanceListener rebalanceListener) { + this.rebalanceListener = rebalanceListener; + } + @Override public String getComponentType() { return "kafka:message-source"; } + protected boolean isRawMessageHeader() { + return this.rawMessageHeader; + } + + /** + * Set to true to include the raw {@link ConsumerRecord} as a header + * with key {@link KafkaHeaders#RAW_DATA}, + * enabling callers to have access to the record to process errors. + * @param rawMessageHeader true to include the header. + */ + public void setRawMessageHeader(boolean rawMessageHeader) { + this.rawMessageHeader = rawMessageHeader; + } + @Override protected synchronized Object doReceive() { if (this.consumer == null) { @@ -221,11 +253,18 @@ public class KafkaMessageSource extends AbstractMessageSource if (message.getHeaders() instanceof KafkaMessageHeaders) { Map rawHeaders = ((KafkaMessageHeaders) message.getHeaders()).getRawHeaders(); rawHeaders.put(IntegrationMessageHeaderAccessor.ACKNOWLEDGMENT_CALLBACK, ackCallback); + if (this.rawMessageHeader) { + rawHeaders.put(KafkaHeaders.RAW_DATA, record); + } return message; } else { - return getMessageBuilderFactory().fromMessage(message) + AbstractIntegrationMessageBuilder builder = getMessageBuilderFactory().fromMessage(message) .setHeader(IntegrationMessageHeaderAccessor.ACKNOWLEDGMENT_CALLBACK, ackCallback); + if (this.rawMessageHeader) { + builder.setHeader(KafkaHeaders.RAW_DATA, record); + } + return builder; } } @@ -237,11 +276,17 @@ public class KafkaMessageSource extends AbstractMessageSource @Override public void onPartitionsRevoked(Collection partitions) { KafkaMessageSource.this.partitions = Collections.emptyList(); + if (KafkaMessageSource.this.rebalanceListener != null) { + KafkaMessageSource.this.rebalanceListener.onPartitionsRevoked(partitions); + } } @Override public void onPartitionsAssigned(Collection partitions) { KafkaMessageSource.this.partitions = new ArrayList<>(partitions); + if (KafkaMessageSource.this.rebalanceListener != null) { + KafkaMessageSource.this.rebalanceListener.onPartitionsAssigned(partitions); + } } }); diff --git a/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/inbound/MessageSourceIntegrationTests.java b/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/inbound/MessageSourceIntegrationTests.java index 1292783150..c2e4cd5833 100644 --- a/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/inbound/MessageSourceIntegrationTests.java +++ b/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/inbound/MessageSourceIntegrationTests.java @@ -18,9 +18,14 @@ package org.springframework.integration.kafka.inbound; import static org.assertj.core.api.Assertions.assertThat; +import java.util.Collection; import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; import org.apache.kafka.clients.consumer.ConsumerConfig; +import org.apache.kafka.clients.consumer.ConsumerRebalanceListener; +import org.apache.kafka.common.TopicPartition; import org.junit.ClassRule; import org.junit.Test; @@ -50,6 +55,19 @@ public class MessageSourceIntegrationTests { consumerProps.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest"); DefaultKafkaConsumerFactory consumerFactory = new DefaultKafkaConsumerFactory<>(consumerProps); KafkaMessageSource source = new KafkaMessageSource<>(consumerFactory, TOPIC1); + final CountDownLatch assigned = new CountDownLatch(1); + source.setRebalanceListener(new ConsumerRebalanceListener() { + + @Override + public void onPartitionsRevoked(Collection partitions) { + } + + @Override + public void onPartitionsAssigned(Collection partitions) { + assigned.countDown(); + } + + }); Map producerProps = KafkaTestUtils.producerProps(embeddedKafka); DefaultKafkaProducerFactory producerFactory = new DefaultKafkaProducerFactory<>(producerProps); @@ -59,6 +77,7 @@ public class MessageSourceIntegrationTests { template.send(TOPIC1, "baz"); template.send(TOPIC1, "qux"); Message received = source.receive(); + assertThat(assigned.await(10, TimeUnit.SECONDS)).isTrue(); int n = 0; while (n++ < 100 && received == null) { received = source.receive(); diff --git a/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/inbound/MessageSourceTests.java b/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/inbound/MessageSourceTests.java index d0fb191e40..983f99591c 100644 --- a/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/inbound/MessageSourceTests.java +++ b/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/inbound/MessageSourceTests.java @@ -109,8 +109,11 @@ public class MessageSourceTests { ConsumerFactory consumerFactory = mock(ConsumerFactory.class); given(consumerFactory.createConsumer(isNull(), anyString(), isNull())).willReturn(consumer); KafkaMessageSource source = new KafkaMessageSource(consumerFactory, "foo"); + source.setRawMessageHeader(true); Message received = source.receive(); + assertThat(received).isNotNull(); + assertThat(received.getHeaders().get(KafkaHeaders.RAW_DATA)).isInstanceOf(ConsumerRecord.class); StaticMessageHeaderAccessor.getAcknowledgmentCallback(received) .acknowledge(Status.ACCEPT); received = source.receive();