GH-184: Polishing - user rebalance listener

Needed by the SCST Kafka Binder

* GH-184: Polishing - add raw header option

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<K, V>
 		return this;
 	}
 
+	public KafkaInboundChannelAdapterSpec<K, V> rebalanceListener(ConsumerRebalanceListener rebalanceListener) {
+		this.target.setRebalanceListener(rebalanceListener);
+		return this;
+	}
+
+	public KafkaInboundChannelAdapterSpec<K, V> rawMessageHeader(boolean rawMessageHeader) {
+		this.target.setRawMessageHeader(rawMessageHeader);
+		return this;
+	}
+
 }

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<K, V> extends AbstractMessageSource<Object>
 
 	private Type payloadType;
 
+	private ConsumerRebalanceListener rebalanceListener;
+
+	private boolean rawMessageHeader;
+
 	private volatile Consumer<K, V> consumer;
 
 	private volatile Collection<TopicPartition> partitions;
@@ -184,11 +190,37 @@ public class KafkaMessageSource<K, V> extends AbstractMessageSource<Object>
 		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<K, V> extends AbstractMessageSource<Object>
 		if (message.getHeaders() instanceof KafkaMessageHeaders) {
 			Map<String, Object> 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<K, V> extends AbstractMessageSource<Object>
 				@Override
 				public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
 					KafkaMessageSource.this.partitions = Collections.emptyList();
+					if (KafkaMessageSource.this.rebalanceListener != null) {
+						KafkaMessageSource.this.rebalanceListener.onPartitionsRevoked(partitions);
+					}
 				}
 
 				@Override
 				public void onPartitionsAssigned(Collection<TopicPartition> partitions) {
 					KafkaMessageSource.this.partitions = new ArrayList<>(partitions);
+					if (KafkaMessageSource.this.rebalanceListener != null) {
+						KafkaMessageSource.this.rebalanceListener.onPartitionsAssigned(partitions);
+					}
 				}
 
 			});

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<Integer, String> consumerFactory = new DefaultKafkaConsumerFactory<>(consumerProps);
 		KafkaMessageSource<Integer, String> source = new KafkaMessageSource<>(consumerFactory, TOPIC1);
+		final CountDownLatch assigned = new CountDownLatch(1);
+		source.setRebalanceListener(new ConsumerRebalanceListener() {
+
+			@Override
+			public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
+			}
+
+			@Override
+			public void onPartitionsAssigned(Collection<TopicPartition> partitions) {
+				assigned.countDown();
+			}
+
+		});
 
 		Map<String, Object> producerProps = KafkaTestUtils.producerProps(embeddedKafka);
 		DefaultKafkaProducerFactory<Object, Object> producerFactory = new DefaultKafkaProducerFactory<>(producerProps);
@@ -59,6 +77,7 @@ public class MessageSourceIntegrationTests {
 		template.send(TOPIC1, "baz");
 		template.send(TOPIC1, "qux");
 		Message<Object> received = source.receive();
+		assertThat(assigned.await(10, TimeUnit.SECONDS)).isTrue();
 		int n = 0;
 		while (n++ < 100 && received == null) {
 			received = source.receive();

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