spring-projectsGH-2348: Custom Correlation Consumer Side

Resolves spring-projects#2348

The replying template supports a custom header for correlation for cases
when the consumer side is not a Spring app and uses a different header.

Support a custom header name on the consumer side, for cases where the
client side is not Spring and uses a different header.

Author: garyrussell

spring-kafka-docs/src/main/asciidoc/kafka.adoc

@@ -726,6 +726,9 @@ These header names are used by the `@KafkaListener` infrastructure to route the
 Starting with version 2.3, you can customize the header names - the template has 3 properties `correlationHeaderName`, `replyTopicHeaderName`, and `replyPartitionHeaderName`.
 This is useful if your server is not a Spring application (or does not use the `@KafkaListener`).
 
+NOTE: Conversely, if the requesting application is not a spring application and puts correlation information in a different header, starting with version 3.0, you can configure a custom `correlationHeaderName` on the listener container factory and that header will be echoed back.
+Previously, the listener had to echo custom correlation headers.
+
 [[exchanging-messages]]
 ====== Request/Reply with `Message<?>` s
 

spring-kafka-docs/src/main/asciidoc/whats-new.adoc

@@ -52,3 +52,9 @@ See <<kafka-template>>.
 
 The futures returned by this class are now `CompletableFuture` s instead of `ListenableFuture` s.
 See <<replying-template>> and <<exchanging-messages>>.
+
+[[x30-listener]]
+==== `@KafkaListener` Changes
+
+You can now use a custom correlation header which will be echoed in any reply message.
+See the note at the end of <<replying-template>> for more information.

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

@@ -44,6 +44,7 @@
 import org.springframework.kafka.listener.adapter.ReplyHeadersConfigurer;
 import org.springframework.kafka.requestreply.ReplyingKafkaOperations;
 import org.springframework.kafka.support.JavaUtils;
+import org.springframework.kafka.support.KafkaHeaders;
 import org.springframework.kafka.support.TopicPartitionOffset;
 import org.springframework.kafka.support.converter.MessageConverter;
 import org.springframework.util.Assert;
@@ -108,6 +109,8 @@ public abstract class AbstractKafkaListenerContainerFactory<C extends AbstractMe
 
 	private ContainerCustomizer<K, V, C> containerCustomizer;
 
+	private String correlationHeaderName;
+
 	@Override
 	public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
 		this.applicationContext = applicationContext;
@@ -321,6 +324,17 @@ public void setContainerCustomizer(ContainerCustomizer<K, V, C> containerCustomi
 		this.containerCustomizer = containerCustomizer;
 	}
 
+	/**
+	 * Set a custom header name for the correlation id. Default
+	 * {@link KafkaHeaders#CORRELATION_ID}. This header will be echoed back in any reply
+	 * message.
+	 * @param correlationHeaderName the header name.
+	 * @since 3.0
+	 */
+	public void setCorrelationHeaderName(String correlationHeaderName) {
+		this.correlationHeaderName = correlationHeaderName;
+	}
+
 	@SuppressWarnings("deprecation")
 	@Override
 	public void afterPropertiesSet() {
@@ -363,7 +377,8 @@ private void configureEndpoint(AbstractKafkaListenerEndpoint<K, V> aklEndpoint)
 				.acceptIfNotNull(this.ackDiscarded, aklEndpoint::setAckDiscarded)
 				.acceptIfNotNull(this.replyTemplate, aklEndpoint::setReplyTemplate)
 				.acceptIfNotNull(this.replyHeadersConfigurer, aklEndpoint::setReplyHeadersConfigurer)
-				.acceptIfNotNull(this.batchToRecordAdapter, aklEndpoint::setBatchToRecordAdapter);
+				.acceptIfNotNull(this.batchToRecordAdapter, aklEndpoint::setBatchToRecordAdapter)
+				.acceptIfNotNull(this.correlationHeaderName, aklEndpoint::setCorrelationHeaderName);
 		if (aklEndpoint.getBatchListener() == null) {
 			JavaUtils.INSTANCE
 					.acceptIfNotNull(this.batchListener, aklEndpoint::setBatchListener);

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

@@ -45,6 +45,8 @@
 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.support.JavaUtils;
+import org.springframework.kafka.support.KafkaHeaders;
 import org.springframework.kafka.support.TopicPartitionOffset;
 import org.springframework.kafka.support.converter.MessageConverter;
 import org.springframework.lang.Nullable;
@@ -117,6 +119,8 @@ public abstract class AbstractKafkaListenerEndpoint<K, V>
 
 	private byte[] listenerInfo;
 
+	private String correlationHeaderName;
+
 	@Override
 	public void setBeanFactory(BeanFactory beanFactory) throws BeansException {
 		this.beanFactory = beanFactory;
@@ -445,6 +449,16 @@ public void setBatchToRecordAdapter(BatchToRecordAdapter<K, V> batchToRecordAdap
 		this.batchToRecordAdapter = batchToRecordAdapter;
 	}
 
+	/**
+	 * Set a custom header name for the correlation id. Default
+	 * {@link KafkaHeaders#CORRELATION_ID}. This header will be echoed back in any reply
+	 * message.
+	 * @param correlationHeaderName the header name.
+	 * @since 3.0
+	 */
+	public void setCorrelationHeaderName(String correlationHeaderName) {
+		this.correlationHeaderName = correlationHeaderName;
+	}
 
 	@Override
 	public void afterPropertiesSet() {
@@ -485,9 +499,9 @@ private void setupMessageListener(MessageListenerContainer container,
 			@Nullable MessageConverter messageConverter) {
 
 		MessagingMessageListenerAdapter<K, V> adapter = createMessageListener(container, messageConverter);
-		if (this.replyHeadersConfigurer != null) {
-			adapter.setReplyHeadersConfigurer(this.replyHeadersConfigurer);
-		}
+		JavaUtils.INSTANCE
+				.acceptIfNotNull(this.replyHeadersConfigurer, adapter::setReplyHeadersConfigurer)
+				.acceptIfNotNull(this.correlationHeaderName, adapter::setCorrelationHeaderName);
 		adapter.setSplitIterables(this.splitIterables);
 		Object messageListener = adapter;
 		boolean isBatchListener = isBatchListener();

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

@@ -128,11 +128,30 @@ public abstract class MessagingMessageListenerAdapter<K, V> implements ConsumerS
 
 	private boolean splitIterables = true;
 
+	private String correlationHeaderName = KafkaHeaders.CORRELATION_ID;
+
+	/**
+	 * Create an instance with the provided bean and method.
+	 * @param bean the bean.
+	 * @param method the method.
+	 */
 	public MessagingMessageListenerAdapter(Object bean, Method method) {
 		this.bean = bean;
 		this.inferredType = determineInferredType(method); // NOSONAR = intentionally not final
 	}
 
+	/**
+	 * Set a custom header name for the correlation id. Default
+	 * {@link KafkaHeaders#CORRELATION_ID}. This header will be echoed back in any reply
+	 * message.
+	 * @param correlationHeaderName the header name.
+	 * @since 3.0
+	 */
+	public void setCorrelationHeaderName(String correlationHeaderName) {
+		Assert.notNull(correlationHeaderName, "'correlationHeaderName' cannot be null");
+		this.correlationHeaderName = correlationHeaderName;
+	}
+
 	/**
 	 * Set the MessageConverter.
 	 * @param messageConverter the converter.
@@ -478,7 +497,7 @@ private Message<?> checkHeaders(Object result, String topic, Object source) { //
 		MessageHeaders headers = reply.getHeaders();
 		boolean needsTopic = headers.get(KafkaHeaders.TOPIC) == null;
 		boolean sourceIsMessage = source instanceof Message;
-		boolean needsCorrelation = headers.get(KafkaHeaders.CORRELATION_ID) == null && sourceIsMessage;
+		boolean needsCorrelation = headers.get(this.correlationHeaderName) == null && sourceIsMessage;
 		boolean needsPartition = headers.get(KafkaHeaders.PARTITION) == null && sourceIsMessage
 				&& getReplyPartition((Message<?>) source) != null;
 		if (needsTopic || needsCorrelation || needsPartition) {
@@ -487,8 +506,8 @@ private Message<?> checkHeaders(Object result, String topic, Object source) { //
 				builder.setHeader(KafkaHeaders.TOPIC, topic);
 			}
 			if (needsCorrelation && sourceIsMessage) {
-				builder.setHeader(KafkaHeaders.CORRELATION_ID,
-						((Message<?>) source).getHeaders().get(KafkaHeaders.CORRELATION_ID));
+				builder.setHeader(this.correlationHeaderName,
+						((Message<?>) source).getHeaders().get(this.correlationHeaderName));
 			}
 			if (sourceIsMessage && reply.getHeaders().get(KafkaHeaders.REPLY_PARTITION) == null) {
 				setPartition(builder, (Message<?>) source);
@@ -503,8 +522,8 @@ private void sendSingleResult(Object result, String topic, @Nullable Object sour
 		byte[] correlationId = null;
 		boolean sourceIsMessage = source instanceof Message;
 		if (sourceIsMessage
-				&& ((Message<?>) source).getHeaders().get(KafkaHeaders.CORRELATION_ID) != null) {
-			correlationId = ((Message<?>) source).getHeaders().get(KafkaHeaders.CORRELATION_ID, byte[].class);
+				&& ((Message<?>) source).getHeaders().get(this.correlationHeaderName) != null) {
+			correlationId = ((Message<?>) source).getHeaders().get(this.correlationHeaderName, byte[].class);
 		}
 		if (sourceIsMessage) {
 			sendReplyForMessageSource(result, topic, source, correlationId);
@@ -515,15 +534,15 @@ private void sendSingleResult(Object result, String topic, @Nullable Object sour
 	}
 
 	@SuppressWarnings("unchecked")
-	private void sendReplyForMessageSource(Object result, String topic, Object source, byte[] correlationId) {
+	private void sendReplyForMessageSource(Object result, String topic, Object source, @Nullable byte[] correlationId) {
 		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.equals(this.correlationHeaderName)
 								&& !key.startsWith(KafkaHeaders.RECEIVED);
 					})
 					.filter(e -> this.replyHeadersConfigurer.shouldCopy(e.getKey(), e.getValue()))
@@ -537,7 +556,7 @@ private void sendReplyForMessageSource(Object result, String topic, Object sourc
 			}
 		}
 		if (correlationId != null) {
-			builder.setHeader(KafkaHeaders.CORRELATION_ID, correlationId);
+			builder.setHeader(this.correlationHeaderName, correlationId);
 		}
 		setPartition(builder, ((Message<?>) source));
 		this.replyTemplate.send(builder.build());