GH-811: Configurable Reply Headers

Resolves #811

Support copying and/or modifying reply headers.

* Javadoc Polishing

Author: garyrussell

spring-kafka/src/main/java/org/springframework/kafka/annotation/KafkaListenerAnnotationBeanPostProcessor.java

@@ -172,7 +172,7 @@ public void setEndpointRegistry(KafkaListenerEndpointRegistry endpointRegistry)
 
 	/**
 	 * Set the name of the {@link KafkaListenerContainerFactory} to use by default.
-	 * <p>If none is specified, "KafkaListenerContainerFactory" is assumed to be defined.
+	 * <p>If none is specified, "kafkaListenerContainerFactory" is assumed to be defined.
 	 * @param containerFactoryBeanName the {@link KafkaListenerContainerFactory} bean name.
 	 */
 	public void setContainerFactoryBeanName(String containerFactoryBeanName) {

spring-kafka/src/main/java/org/springframework/kafka/config/AbstractKafkaListenerContainerFactory.java

@@ -33,6 +33,7 @@
 import org.springframework.kafka.listener.ErrorHandler;
 import org.springframework.kafka.listener.GenericErrorHandler;
 import org.springframework.kafka.listener.adapter.RecordFilterStrategy;
+import org.springframework.kafka.listener.adapter.ReplyHeadersConfigurer;
 import org.springframework.kafka.support.TopicPartitionInitialOffset;
 import org.springframework.kafka.support.converter.MessageConverter;
 import org.springframework.retry.RecoveryCallback;
@@ -84,6 +85,8 @@ public abstract class AbstractKafkaListenerContainerFactory<C extends AbstractMe
 
 	private AfterRollbackProcessor<K, V> afterRollbackProcessor;
 
+	private ReplyHeadersConfigurer replyHeadersConfigurer;
+
 	/**
 	 * Specify a {@link ConsumerFactory} to use.
 	 * @param consumerFactory The consumer factory.
@@ -230,6 +233,15 @@ public void setAfterRollbackProcessor(AfterRollbackProcessor<K, V> afterRollback
 		this.afterRollbackProcessor = afterRollbackProcessor;
 	}
 
+	/**
+	 * Set a configurer which will be invoked when creating a reply message.
+	 * @param replyHeadersConfigurer the configurer.
+	 * @since 2.2
+	 */
+	public void setReplyHeadersConfigurer(ReplyHeadersConfigurer replyHeadersConfigurer) {
+		this.replyHeadersConfigurer = replyHeadersConfigurer;
+	}
+
 	/**
 	 * Obtain the properties template for this factory - set properties as needed
 	 * and they will be copied to a final properties instance for the endpoint.
@@ -270,6 +282,9 @@ public C createListenerContainer(KafkaListenerEndpoint endpoint) {
 			if (this.replyTemplate != null) {
 				aklEndpoint.setReplyTemplate(this.replyTemplate);
 			}
+			if (this.replyHeadersConfigurer != null) {
+				aklEndpoint.setReplyHeadersConfigurer(this.replyHeadersConfigurer);
+			}
 		}
 
 		endpoint.setupListenerContainer(instance, this.messageConverter);

spring-kafka/src/main/java/org/springframework/kafka/config/AbstractKafkaListenerEndpoint.java

@@ -42,6 +42,7 @@
 import org.springframework.kafka.listener.adapter.FilteringMessageListenerAdapter;
 import org.springframework.kafka.listener.adapter.MessagingMessageListenerAdapter;
 import org.springframework.kafka.listener.adapter.RecordFilterStrategy;
+import org.springframework.kafka.listener.adapter.ReplyHeadersConfigurer;
 import org.springframework.kafka.listener.adapter.RetryingMessageListenerAdapter;
 import org.springframework.kafka.support.TopicPartitionInitialOffset;
 import org.springframework.kafka.support.converter.MessageConverter;
@@ -105,6 +106,8 @@ public abstract class AbstractKafkaListenerEndpoint<K, V>
 
 	private Boolean autoStartup;
 
+	private ReplyHeadersConfigurer replyHeadersConfigurer;
+
 	@Override
 	public void setBeanFactory(BeanFactory beanFactory) throws BeansException {
 		this.beanFactory = beanFactory;
@@ -374,6 +377,15 @@ public void setAutoStartup(Boolean autoStartup) {
 		this.autoStartup = autoStartup;
 	}
 
+	/**
+	 * Set a configurer which will be invoked when creating a reply message.
+	 * @param replyHeadersConfigurer the configurer.
+	 * @since 2.2
+	 */
+	public void setReplyHeadersConfigurer(ReplyHeadersConfigurer replyHeadersConfigurer) {
+		this.replyHeadersConfigurer = replyHeadersConfigurer;
+	}
+
 	@Override
 	public void afterPropertiesSet() {
 		boolean topicsEmpty = getTopics().isEmpty();
@@ -408,7 +420,11 @@ protected abstract MessagingMessageListenerAdapter<K, V> createMessageListener(M
 
 	@SuppressWarnings("unchecked")
 	private void setupMessageListener(MessageListenerContainer container, MessageConverter messageConverter) {
-		Object messageListener = createMessageListener(container, messageConverter);
+		MessagingMessageListenerAdapter<K, V> adapter = createMessageListener(container, messageConverter);
+		if (this.replyHeadersConfigurer != null) {
+			adapter.setReplyHeadersConfigurer(this.replyHeadersConfigurer);
+		}
+		Object messageListener = adapter;
 		Assert.state(messageListener != null, "Endpoint [" + this + "] must provide a non null message listener");
 		if (this.retryTemplate != null) {
 			messageListener = new RetryingMessageListenerAdapter<>((MessageListener<K, V>) messageListener,

spring-kafka/src/main/java/org/springframework/kafka/listener/adapter/MessagingMessageListenerAdapter.java

@@ -25,6 +25,7 @@
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
+import java.util.stream.Collectors;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
@@ -54,11 +55,13 @@
 import org.springframework.kafka.support.converter.RecordMessageConverter;
 import org.springframework.lang.Nullable;
 import org.springframework.messaging.Message;
+import org.springframework.messaging.MessageHeaders;
 import org.springframework.messaging.MessagingException;
 import org.springframework.messaging.converter.MessageConversionException;
 import org.springframework.messaging.handler.annotation.Payload;
 import org.springframework.messaging.support.MessageBuilder;
 import org.springframework.util.Assert;
+import org.springframework.util.ObjectUtils;
 import org.springframework.util.StringUtils;
 
 /**
@@ -108,6 +111,8 @@ public abstract class MessagingMessageListenerAdapter<K, V> implements ConsumerS
 
 	private boolean messageReturnType;
 
+	private ReplyHeadersConfigurer replyHeadersConfigurer;
+
 	public MessagingMessageListenerAdapter(Object bean, Method method) {
 		this.bean = bean;
 		this.inferredType = determineInferredType(method);
@@ -214,6 +219,25 @@ protected boolean isMessageList() {
 		return this.isMessageList;
 	}
 
+	/**
+	 * Return the reply configurer.
+	 * @return the configurer.
+	 * @since 2.2
+	 * @see #setReplyHeadersConfigurer(ReplyHeadersConfigurer)
+	 */
+	protected ReplyHeadersConfigurer getReplyHeadersConfigurer() {
+		return this.replyHeadersConfigurer;
+	}
+
+	/**
+	 * Set a configurer which will be invoked when creating a reply message.
+	 * @param replyHeadersConfigurer the configurer.
+	 * @since 2.2
+	 */
+	public void setReplyHeadersConfigurer(ReplyHeadersConfigurer replyHeadersConfigurer) {
+		this.replyHeadersConfigurer = replyHeadersConfigurer;
+	}
+
 	@Override
 	public void registerSeekCallback(ConsumerSeekCallback callback) {
 		if (this.bean instanceof ConsumerSeekAware) {
@@ -384,6 +408,24 @@ else if (result instanceof Message) {
 				if (sourceIsMessage) {
 					MessageBuilder<Object> builder = MessageBuilder.withPayload(result)
 							.setHeader(KafkaHeaders.TOPIC, topic);
+					if (this.replyHeadersConfigurer != null) {
+						Map<String, Object> headersToCopy = ((Message<?>) source).getHeaders().entrySet().stream()
+							.filter(e -> {
+								String key = e.getKey();
+								return !key.equals(MessageHeaders.ID) && !key.equals(MessageHeaders.TIMESTAMP)
+										&& !key.equals(KafkaHeaders.CORRELATION_ID)
+										&& !key.startsWith(KafkaHeaders.RECEIVED);
+							})
+							.filter(e -> this.replyHeadersConfigurer.shouldCopy(e.getKey(), e.getValue()))
+							.collect(Collectors.toMap(e -> e.getKey(), e -> e.getValue()));
+						if (headersToCopy.size() > 0) {
+							builder.copyHeaders(headersToCopy);
+						}
+						headersToCopy = this.replyHeadersConfigurer.additionalHeaders();
+						if (!ObjectUtils.isEmpty(headersToCopy)) {
+							builder.copyHeaders(headersToCopy);
+						}
+					}
 					if (correlationId != null) {
 						builder.setHeader(KafkaHeaders.CORRELATION_ID, correlationId);
 					}

spring-kafka/src/main/java/org/springframework/kafka/listener/adapter/ReplyHeadersConfigurer.java

@@ -0,0 +1,57 @@
+/*
+ * Copyright 2018 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.kafka.listener.adapter;
+
+import java.util.Map;
+
+import org.springframework.kafka.support.KafkaHeaders;
+import org.springframework.lang.Nullable;
+import org.springframework.messaging.MessageHeaders;
+
+/**
+ * A strategy for configuring which headers, if any, should be set in a reply message.
+ *
+ * @author Gary Russell
+ * @since 2.2
+ *
+ */
+@FunctionalInterface
+public interface ReplyHeadersConfigurer {
+
+	/**
+	 * Return true if the header should be copied to the reply message.
+	 * {@link KafkaHeaders#CORRELATION_ID} will not be offered; it is always copied.
+	 * {@link MessageHeaders#ID} and {@link MessageHeaders#TIMESTAMP} are never copied.
+	 * {@code KafkaHeaders.RECEIVED*} headers are never copied.
+	 * @param headerName the header name.
+	 * @param headerValue the header value.
+	 * @return true to copy.
+	 */
+	boolean shouldCopy(String headerName, Object headerValue);
+
+	/**
+	 * A map of additional headers to add to the reply message.
+	 * IMPORTANT: Any existing headers with the same name will be replaced by those
+	 * returned by this method.
+	 * @return the headers.
+	 */
+	@Nullable
+	default Map<String, Object> additionalHeaders() {
+		return null;
+	}
+
+}

spring-kafka/src/main/java/org/springframework/kafka/support/KafkaHeaders.java

@@ -26,7 +26,15 @@
  */
 public abstract class KafkaHeaders {
 
-	private static final String PREFIX = "kafka_";
+	/**
+	 * The prefix for Kafka headers.
+	 */
+	public static final String PREFIX = "kafka_";
+
+	/**
+	 * The prefix for Kafka headers containing 'received' values.
+	 */
+	public static final String RECEIVED = PREFIX + "received";
 
 	/**
 	 * The header containing the topic when sending data to Kafka.
@@ -72,17 +80,17 @@ public abstract class KafkaHeaders {
 	/**
 	 * The header containing the topic from which the message was received.
 	 */
-	public static final String RECEIVED_TOPIC = PREFIX + "receivedTopic";
+	public static final String RECEIVED_TOPIC = RECEIVED + "Topic";
 
 	/**
 	 * The header containing the message key for the received message.
 	 */
-	public static final String RECEIVED_MESSAGE_KEY = PREFIX + "receivedMessageKey";
+	public static final String RECEIVED_MESSAGE_KEY = RECEIVED + "MessageKey";
 
 	/**
 	 * The header containing the topic partition for the received message.
 	 */
-	public static final String RECEIVED_PARTITION_ID = PREFIX + "receivedPartitionId";
+	public static final String RECEIVED_PARTITION_ID = RECEIVED + "PartitionId";
 
 	/**
 	 * The header for holding the {@link org.apache.kafka.common.record.TimestampType type} of timestamp.

spring-kafka/src/test/java/org/springframework/kafka/requestreply/ReplyingKafkaTemplateTests.java

@@ -20,6 +20,8 @@
 import static org.assertj.core.api.Assertions.fail;
 
 import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
 import java.util.Map;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
@@ -30,7 +32,9 @@
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.producer.ProducerRecord;
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.header.Headers;
 import org.apache.kafka.common.header.internals.RecordHeader;
+import org.apache.kafka.common.header.internals.RecordHeaders;
 import org.junit.ClassRule;
 import org.junit.Rule;
 import org.junit.Test;
@@ -50,6 +54,8 @@
 import org.springframework.kafka.core.KafkaTemplate;
 import org.springframework.kafka.listener.ContainerProperties;
 import org.springframework.kafka.listener.KafkaMessageListenerContainer;
+import org.springframework.kafka.listener.adapter.ReplyHeadersConfigurer;
+import org.springframework.kafka.support.DefaultKafkaHeaderMapper;
 import org.springframework.kafka.support.KafkaHeaders;
 import org.springframework.kafka.support.SimpleKafkaHeaderMapper;
 import org.springframework.kafka.support.TopicPartitionInitialOffset;
@@ -102,11 +108,17 @@ public void testGood() throws Exception {
 		ReplyingKafkaTemplate<Integer, String, String> template = createTemplate(A_REPLY);
 		try {
 			template.setReplyTimeout(30_000);
-			ProducerRecord<Integer, String> record = new ProducerRecord<>(A_REQUEST, "foo");
+			Headers headers = new RecordHeaders();
+			headers.add("baz", "buz".getBytes());
+			ProducerRecord<Integer, String> record = new ProducerRecord<>(A_REQUEST, null, null, null, "foo", headers);
 			RequestReplyFuture<Integer, String, String> future = template.sendAndReceive(record);
 			future.getSendFuture().get(10, TimeUnit.SECONDS); // send ok
 			ConsumerRecord<Integer, String> consumerRecord = future.get(30, TimeUnit.SECONDS);
 			assertThat(consumerRecord.value()).isEqualTo("FOO");
+			Map<String, Object> receivedHeaders = new HashMap<>();
+			new DefaultKafkaHeaderMapper().toHeaders(consumerRecord.headers(), receivedHeaders);
+			assertThat(receivedHeaders).containsKey("baz");
+			assertThat(receivedHeaders).hasSize(2);
 		}
 		finally {
 			template.stop();
@@ -199,12 +211,19 @@ public void testGoodWithSimpleMapper() throws Exception {
 		ReplyingKafkaTemplate<Integer, String, String> template = createTemplate(B_REPLY);
 		try {
 			template.setReplyTimeout(30_000);
-			ProducerRecord<Integer, String> record = new ProducerRecord<>(B_REQUEST, "qux");
+			Headers headers = new RecordHeaders();
+			headers.add("baz", "buz".getBytes());
+			ProducerRecord<Integer, String> record = new ProducerRecord<>(B_REQUEST, null, null, null, "qux", headers);
 			record.headers().add(new RecordHeader(KafkaHeaders.REPLY_TOPIC, B_REPLY.getBytes()));
 			RequestReplyFuture<Integer, String, String> future = template.sendAndReceive(record);
 			future.getSendFuture().get(10, TimeUnit.SECONDS); // send ok
 			ConsumerRecord<Integer, String> consumerRecord = future.get(30, TimeUnit.SECONDS);
 			assertThat(consumerRecord.value()).isEqualTo("qUX");
+			Map<String, Object> receivedHeaders = new HashMap<>();
+			new DefaultKafkaHeaderMapper().toHeaders(consumerRecord.headers(), receivedHeaders);
+			assertThat(receivedHeaders).containsKey("qux");
+			assertThat(receivedHeaders).doesNotContainKey("baz");
+			assertThat(receivedHeaders).hasSize(2);
 		}
 		finally {
 			template.stop();
@@ -290,6 +309,7 @@ public ConcurrentKafkaListenerContainerFactory<Integer, String> kafkaListenerCon
 					new ConcurrentKafkaListenerContainerFactory<>();
 			factory.setConsumerFactory(cf());
 			factory.setReplyTemplate(template());
+			factory.setReplyHeadersConfigurer((k, v) -> k.equals("baz"));
 			return factory;
 		}
 
@@ -302,6 +322,19 @@ public ConcurrentKafkaListenerContainerFactory<Integer, String> simpleMapperFact
 			MessagingMessageConverter messageConverter = new MessagingMessageConverter();
 			messageConverter.setHeaderMapper(new SimpleKafkaHeaderMapper());
 			factory.setMessageConverter(messageConverter);
+			factory.setReplyHeadersConfigurer(new ReplyHeadersConfigurer() {
+
+				@Override
+				public boolean shouldCopy(String headerName, Object headerValue) {
+					return false;
+				}
+
+				@Override
+				public Map<String, Object> additionalHeaders() {
+					return Collections.singletonMap("qux", "fiz");
+				}
+
+			});
 			return factory;
 		}