GH-1171: Custom headers for ReplyingKafkaTemplate

Resolves #1171

Allow the use of custom header names for request/reply semantics.

* Try trusty on travis

Author: garyrussell

.travis.yml

@@ -1,3 +1,4 @@
+dist: trusty
 language: java
 jdk: oraclejdk8
 install: true

spring-kafka/src/main/java/org/springframework/kafka/requestreply/ReplyingKafkaTemplate.java

@@ -20,7 +20,6 @@
 import java.nio.charset.StandardCharsets;
 import java.time.Instant;
 import java.util.Collection;
-import java.util.Iterator;
 import java.util.List;
 import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;
@@ -92,6 +91,12 @@ public class ReplyingKafkaTemplate<K, V, R> extends KafkaTemplate<K, V> implemen
 	private Function<ProducerRecord<K, V>, CorrelationKey> correlationStrategy =
 			ReplyingKafkaTemplate::defaultCorrelationIdStrategy;
 
+	private String correlationHeaderName = KafkaHeaders.CORRELATION_ID;
+
+	private String replyTopicHeaderName = KafkaHeaders.REPLY_TOPIC;
+
+	private String replyPartitionHeaderName = KafkaHeaders.REPLY_PARTITION;
+
 	private volatile boolean running;
 
 	private volatile boolean schedulerInitialized;
@@ -204,6 +209,39 @@ public void setCorrelationIdStrategy(Function<ProducerRecord<K, V>, CorrelationK
 		this.correlationStrategy = correlationStrategy;
 	}
 
+	/**
+	 * Set a custom header name for the correlation id. Default
+	 * {@link KafkaHeaders#CORRELATION_ID}.
+	 * @param correlationHeaderName the header name.
+	 * @since 2.3
+	 */
+	public void setCorrelationHeaderName(String correlationHeaderName) {
+		Assert.notNull(correlationHeaderName, "'correlationHeaderName' cannot be null");
+		this.correlationHeaderName = correlationHeaderName;
+	}
+
+	/**
+	 * Set a custom header name for the reply topic. Default
+	 * {@link KafkaHeaders#REPLY_TOPIC}.
+	 * @param replyTopicHeaderName the header name.
+	 * @since 2.3
+	 */
+	public void setReplyTopicHeaderName(String replyTopicHeaderName) {
+		Assert.notNull(replyTopicHeaderName, "'replyTopicHeaderName' cannot be null");
+		this.replyTopicHeaderName = replyTopicHeaderName;
+	}
+
+	/**
+	 * Set a custom header name for the reply partition. Default
+	 * {@link KafkaHeaders#REPLY_PARTITION}.
+	 * @param replyPartitionHeaderName the reply partition header name.
+	 * @since 2.3
+	 */
+	public void setReplyPartitionHeaderName(String replyPartitionHeaderName) {
+		Assert.notNull(replyPartitionHeaderName, "'replyPartitionHeaderName' cannot be null");
+		this.replyPartitionHeaderName = replyPartitionHeaderName;
+	}
+
 	@Override
 	public void afterPropertiesSet() {
 		if (!this.schedulerSet && !this.schedulerInitialized) {
@@ -249,12 +287,12 @@ public RequestReplyFuture<K, V, R> sendAndReceive(ProducerRecord<K, V> record) {
 		Headers headers = record.headers();
 		boolean hasReplyTopic = headers.lastHeader(KafkaHeaders.REPLY_TOPIC) != null;
 		if (!hasReplyTopic && this.replyTopic != null) {
-			headers.add(new RecordHeader(KafkaHeaders.REPLY_TOPIC, this.replyTopic));
+			headers.add(new RecordHeader(this.replyTopicHeaderName, this.replyTopic));
 			if (this.replyPartition != null) {
-				headers.add(new RecordHeader(KafkaHeaders.REPLY_PARTITION, this.replyPartition));
+				headers.add(new RecordHeader(this.replyPartitionHeaderName, this.replyPartition));
 			}
 		}
-		headers.add(new RecordHeader(KafkaHeaders.CORRELATION_ID, correlationId.getCorrelationId()));
+		headers.add(new RecordHeader(this.correlationHeaderName, correlationId.getCorrelationId()));
 		this.logger.debug(() -> "Sending: " + record + WITH_CORRELATION_ID + correlationId);
 		RequestReplyFuture<K, V, R> future = new RequestReplyFuture<>();
 		this.futures.put(correlationId, future);
@@ -338,18 +376,15 @@ private static <K, V> CorrelationKey defaultCorrelationIdStrategy(
 	@Override
 	public void onMessage(List<ConsumerRecord<K, R>> data) {
 		data.forEach(record -> {
-			Iterator<Header> iterator = record.headers().iterator();
+			Header correlationHeader = record.headers().lastHeader(this.correlationHeaderName);
 			CorrelationKey correlationId = null;
-			while (correlationId == null && iterator.hasNext()) {
-				Header next = iterator.next();
-				if (next.key().equals(KafkaHeaders.CORRELATION_ID)) {
-					correlationId = new CorrelationKey(next.value());
-				}
+			if (correlationHeader != null) {
+				correlationId = new CorrelationKey(correlationHeader.value());
 			}
 			if (correlationId == null) {
 				this.logger.error(() -> "No correlationId found in reply: " + record
 						+ " - to use request/reply semantics, the responding server must return the correlation id "
-						+ " in the '" + KafkaHeaders.CORRELATION_ID + "' header");
+						+ " in the '" + this.correlationHeaderName + "' header");
 			}
 			else {
 				RequestReplyFuture<K, V, R> future = this.futures.remove(correlationId);

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

@@ -25,6 +25,7 @@
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
 
+import java.nio.ByteBuffer;
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Collection;
@@ -121,9 +122,14 @@ public class ReplyingKafkaTemplateTests {
 
 	private static final String F_REQUEST = "fRequest";
 
+	private static final String G_REPLY = "gReply";
+
+	private static final String G_REQUEST = "gRequest";
+
 	@ClassRule
 	public static EmbeddedKafkaRule embeddedKafkaRule = new EmbeddedKafkaRule(1, true, 5, A_REQUEST, A_REPLY,
-			B_REQUEST, B_REPLY, C_REQUEST, C_REPLY, D_REQUEST, D_REPLY, E_REQUEST, E_REPLY, F_REQUEST, F_REPLY);
+			B_REQUEST, B_REPLY, C_REQUEST, C_REPLY, D_REQUEST, D_REPLY, E_REQUEST, E_REPLY, F_REQUEST, F_REPLY,
+			G_REQUEST, G_REPLY);
 
 	private static EmbeddedKafkaBroker embeddedKafka = embeddedKafkaRule.getEmbeddedKafka();
 
@@ -152,10 +158,8 @@ public void testGood() throws Exception {
 			new DefaultKafkaHeaderMapper().toHeaders(consumerRecord.headers(), receivedHeaders);
 			assertThat(receivedHeaders).containsKey("baz");
 			assertThat(receivedHeaders).hasSize(2);
-			assertThat(KafkaTestUtils.getPropertyValue(
-					this.registry.getListenerContainer(A_REQUEST), "containerProperties.missingTopicsFatal",
-					Boolean.class))
-				.isFalse();
+			assertThat(this.registry.getListenerContainer(A_REQUEST).getContainerProperties().isMissingTopicsFatal())
+					.isFalse();
 		}
 		finally {
 			template.stop();
@@ -481,6 +485,28 @@ public AggregatingReplyingKafkaTemplate<Integer, String, String> aggregatingTemp
 		return template;
 	}
 
+	@Test
+	public void withCustomHeaders() throws Exception {
+		ReplyingKafkaTemplate<Integer, String, String> template = createTemplate(new TopicPartitionOffset(G_REPLY, 1));
+		template.setCorrelationHeaderName("custom.correlation.id");
+		template.setReplyTopicHeaderName("custom.reply.to");
+		template.setReplyPartitionHeaderName("custom.reply.partition");
+		try {
+			template.setReplyTimeout(30_000);
+			Headers headers = new RecordHeaders();
+			ProducerRecord<Integer, String> record = new ProducerRecord<>(G_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("fooWithCustomHeaders");
+			assertThat(consumerRecord.partition()).isEqualTo(1);
+		}
+		finally {
+			template.stop();
+			template.destroy();
+		}
+	}
+
 	@Configuration
 	@EnableKafka
 	public static class Config {
@@ -568,6 +594,16 @@ public String dListener2(String in) {
 			return in.substring(0, 1) + in.substring(1).toUpperCase();
 		}
 
+		@KafkaListener(id = G_REQUEST, topics = G_REQUEST)
+		public void gListener(Message<String> in) {
+			String replyTopic = new String(in.getHeaders().get("custom.reply.to",  byte[].class));
+			int replyPart = ByteBuffer.wrap(in.getHeaders().get("custom.reply.partition", byte[].class)).getInt();
+			ProducerRecord<Integer, String> record = new ProducerRecord<>(replyTopic, replyPart, null,
+					in.getPayload() + "WithCustomHeaders");
+			record.headers().add(new RecordHeader("custom.correlation.id",
+					in.getHeaders().get("custom.correlation.id", byte[].class)));
+			template().send(record);
+		}
 	}
 
 	@KafkaListener(topics = C_REQUEST, groupId = C_REQUEST)

src/reference/asciidoc/kafka.adoc

@@ -544,7 +544,7 @@ public class KRequestingApplication {
 
 Note that we can use Boot's auto-configured container factory to create the reply container.
 
-The template sets a header called `KafkaHeaders.CORRELATION_ID`, which must be echoed back by the server side.
+The template sets a header (named `KafkaHeaders.CORRELATION_ID` by default), which must be echoed back by the server side.
 
 In this case, the following `@KafkaListener` application responds:
 
@@ -617,6 +617,17 @@ In this case, though, the reply container must not use Kafka's group management
 NOTE: The `DefaultKafkaHeaderMapper` requires Jackson to be on the classpath (for the `@KafkaListener`).
 If it is not available, the message converter has no header mapper, so you must configure a `MessagingMessageConverter` with a `SimpleKafkaHeaderMapper`, as shown earlier.
 
+By default, 3 headers are used:
+
+* `KafkaHeaders.CORRELATION_ID` - used to correlate the reply to a request
+* `KafkaHeaders.REPLY_TOPIC` - used to tell the server where to reply
+* `KafkaHeaders.REPLY_PARTITION` - (optional) used to tell the server which partition to reply to
+
+These header names are used by the `@KafkaListener` infrastructure to route the reply.
+
+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`).
+
 [[aggregating-request-reply]]
 ===== Aggregating Multiple Replies
 

src/reference/asciidoc/whats-new.adoc

@@ -129,3 +129,8 @@ This enables the code to access to extra information that was missing in the old
 
 You can now override the default broker list property name in the annotation.
 See <<kafka-testing-embeddedkafka-annotation>> for more information.
+
+==== ReplyingKafkaTemplate Changes
+
+You can now customize the header names for correlation, reply topic and reply partition.
+See <<replying-template>> for more information.