GH-1244: ReplyingKTemplate reply timeout per send

Resolves #1244

Overloaded `sendAndReceive` with a timeout parameter

Author: garyrussell

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

@@ -16,8 +16,12 @@
 
 package org.springframework.kafka.requestreply;
 
+import java.time.Duration;
+
 import org.apache.kafka.clients.producer.ProducerRecord;
 
+import org.springframework.lang.Nullable;
+
 /**
  * Request/reply operations.
  *
@@ -32,10 +36,19 @@
 public interface ReplyingKafkaOperations<K, V, R> {
 
 	/**
-	 * Send a request and receive a reply.
+	 * Send a request and receive a reply with the default timeout.
 	 * @param record the record to send.
 	 * @return a RequestReplyFuture.
 	 */
 	RequestReplyFuture<K, V, R> sendAndReceive(ProducerRecord<K, V> record);
 
+	/**
+	 * Send a request and receive a reply.
+	 * @param record the record to send.
+	 * @param replyTimeout the reply timeout; if null, the default will be used.
+	 * @return a RequestReplyFuture.
+	 * @since 2.3
+	 */
+	RequestReplyFuture<K, V, R> sendAndReceive(ProducerRecord<K, V> record, @Nullable Duration replyTimeout);
+
 }

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

@@ -18,6 +18,7 @@
 
 import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
+import java.time.Duration;
 import java.time.Instant;
 import java.util.Collection;
 import java.util.List;
@@ -44,6 +45,7 @@
 import org.springframework.kafka.listener.GenericMessageListenerContainer;
 import org.springframework.kafka.support.KafkaHeaders;
 import org.springframework.kafka.support.TopicPartitionOffset;
+import org.springframework.lang.Nullable;
 import org.springframework.scheduling.TaskScheduler;
 import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
 import org.springframework.util.Assert;
@@ -66,7 +68,9 @@ public class ReplyingKafkaTemplate<K, V, R> extends KafkaTemplate<K, V> implemen
 
 	private static final String WITH_CORRELATION_ID = " with correlationId: ";
 
-	private static final long DEFAULT_REPLY_TIMEOUT = 5000L;
+	private static final int FIVE = 5;
+
+	private static final Duration DEFAULT_REPLY_TIMEOUT = Duration.ofSeconds(FIVE);
 
 	private final GenericMessageListenerContainer<K, R> replyContainer;
 
@@ -82,7 +86,7 @@ public class ReplyingKafkaTemplate<K, V, R> extends KafkaTemplate<K, V> implemen
 
 	private boolean autoStartup = true;
 
-	private long replyTimeout = DEFAULT_REPLY_TIMEOUT;
+	private Duration defaultReplyTimeout = DEFAULT_REPLY_TIMEOUT;
 
 	private boolean schedulerSet;
 
@@ -149,13 +153,49 @@ public void setTaskScheduler(TaskScheduler scheduler) {
 		this.schedulerSet = true;
 	}
 
+	/**
+	 * Return the reply timeout used if no replyTimeout is provided in the
+	 * {@link #sendAndReceive(ProducerRecord, Duration)} call.
+	 * @return the timeout.
+	 * @deprecated in favor of {@link #getDefaultReplyTimeout()}.
+	 */
+	@Deprecated
 	protected long getReplyTimeout() {
-		return this.replyTimeout;
+		return this.defaultReplyTimeout.toMillis();
 	}
 
+	/**
+	 * Set the reply timeout used if no replyTimeout is provided in the
+	 * {@link #sendAndReceive(ProducerRecord, Duration)} call.
+	 * @param replyTimeout the timeout.
+	 * @deprecated in favor of {@link #setDefaultReplyTimeout(Duration)}.
+	 */
+	@Deprecated
 	public void setReplyTimeout(long replyTimeout) {
 		Assert.isTrue(replyTimeout >= 0, "'replyTimeout' must be >= 0");
-		this.replyTimeout = replyTimeout;
+		this.defaultReplyTimeout = Duration.ofMillis(replyTimeout);
+	}
+
+	/**
+	 * Return the reply timeout used if no replyTimeout is provided in the
+	 * {@link #sendAndReceive(ProducerRecord, Duration)} call.
+	 * @return the timeout.
+	 * @since 2.3
+	 */
+	protected Duration getDefaultReplyTimeout() {
+		return this.defaultReplyTimeout;
+	}
+
+	/**
+	 * Set the reply timeout used if no replyTimeout is provided in the
+	 * {@link #sendAndReceive(ProducerRecord, Duration)} call.
+	 * @param defaultReplyTimeout the timeout.
+	 * @since 2.3
+	 */
+	public void setDefaultReplyTimeout(Duration defaultReplyTimeout) {
+		Assert.notNull(defaultReplyTimeout, "'defaultReplyTimeout' cannot be null");
+		Assert.isTrue(defaultReplyTimeout.toMillis() >= 0, "'replyTimeout' must be >= 0");
+		this.defaultReplyTimeout = defaultReplyTimeout;
 	}
 
 	@Override
@@ -282,6 +322,11 @@ public void stop(Runnable callback) {
 
 	@Override
 	public RequestReplyFuture<K, V, R> sendAndReceive(ProducerRecord<K, V> record) {
+		return sendAndReceive(record, null);
+	}
+
+	@Override
+	public RequestReplyFuture<K, V, R> sendAndReceive(ProducerRecord<K, V> record, @Nullable Duration replyTimeout) {
 		Assert.state(this.running, "Template has not been start()ed"); // NOSONAR (sync)
 		CorrelationKey correlationId = this.correlationStrategy.apply(record);
 		Assert.notNull(correlationId, "the created 'correlationId' cannot be null");
@@ -304,11 +349,11 @@ public RequestReplyFuture<K, V, R> sendAndReceive(ProducerRecord<K, V> record) {
 			this.futures.remove(correlationId);
 			throw new KafkaException("Send failed", e);
 		}
-		scheduleTimeout(record, correlationId);
+		scheduleTimeout(record, correlationId, replyTimeout == null ? this.defaultReplyTimeout : replyTimeout);
 		return future;
 	}
 
-	private void scheduleTimeout(ProducerRecord<K, V> record, CorrelationKey correlationId) {
+	private void scheduleTimeout(ProducerRecord<K, V> record, CorrelationKey correlationId, Duration replyTimeout) {
 		this.scheduler.schedule(() -> {
 			RequestReplyFuture<K, V, R> removed = this.futures.remove(correlationId);
 			if (removed != null) {
@@ -317,7 +362,7 @@ private void scheduleTimeout(ProducerRecord<K, V> record, CorrelationKey correla
 					removed.setException(new KafkaReplyTimeoutException("Reply timed out"));
 				}
 			}
-		}, Instant.now().plusMillis(this.replyTimeout));
+		}, Instant.now().plus(replyTimeout));
 	}
 
 	/**

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

@@ -17,6 +17,7 @@
 package org.springframework.kafka.requestreply;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
 import static org.assertj.core.api.Assertions.fail;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.isNull;
@@ -151,7 +152,7 @@ public void captureTestName(TestInfo info) {
 	public void testGood() throws Exception {
 		ReplyingKafkaTemplate<Integer, String, String> template = createTemplate(A_REPLY);
 		try {
-			template.setReplyTimeout(30_000);
+			template.setDefaultReplyTimeout(Duration.ofSeconds(30));
 			Headers headers = new RecordHeaders();
 			headers.add("baz", "buz".getBytes());
 			ProducerRecord<Integer, String> record = new ProducerRecord<>(A_REQUEST, null, null, null, "foo", headers);
@@ -165,6 +166,11 @@ public void testGood() throws Exception {
 			assertThat(receivedHeaders).hasSize(2);
 			assertThat(this.registry.getListenerContainer(A_REQUEST).getContainerProperties().isMissingTopicsFatal())
 					.isFalse();
+			ProducerRecord<Integer, String> record2 =
+					new ProducerRecord<>(A_REQUEST, null, null, null, "slow", headers);
+			assertThatExceptionOfType(ExecutionException.class)
+					.isThrownBy(() -> template.sendAndReceive(record2, Duration.ZERO).get(10, TimeUnit.SECONDS))
+					.withCauseExactlyInstanceOf(KafkaReplyTimeoutException.class);
 		}
 		finally {
 			template.stop();
@@ -176,7 +182,7 @@ public void testGood() throws Exception {
 	public void testMultiListenerMessageReturn() throws Exception {
 		ReplyingKafkaTemplate<Integer, String, String> template = createTemplate(C_REPLY);
 		try {
-			template.setReplyTimeout(30_000);
+			template.setDefaultReplyTimeout(Duration.ofSeconds(30));
 			ProducerRecord<Integer, String> record = new ProducerRecord<>(C_REQUEST, "foo");
 			record.headers().add(new RecordHeader(KafkaHeaders.REPLY_TOPIC, C_REPLY.getBytes()));
 			RequestReplyFuture<Integer, String, String> future = template.sendAndReceive(record);
@@ -195,7 +201,7 @@ public void testGoodDefaultReplyHeaders() throws Exception {
 		ReplyingKafkaTemplate<Integer, String, String> template = createTemplate(
 				new TopicPartitionOffset(A_REPLY, 3));
 		try {
-			template.setReplyTimeout(30_000);
+			template.setDefaultReplyTimeout(Duration.ofSeconds(30));
 			ProducerRecord<Integer, String> record = new ProducerRecord<>(A_REQUEST, "bar");
 			RequestReplyFuture<Integer, String, String> future = template.sendAndReceive(record);
 			future.getSendFuture().get(10, TimeUnit.SECONDS); // send ok
@@ -213,7 +219,7 @@ public void testGoodDefaultReplyHeaders() throws Exception {
 	public void testGoodSamePartition() throws Exception {
 		ReplyingKafkaTemplate<Integer, String, String> template = createTemplate(A_REPLY);
 		try {
-			template.setReplyTimeout(30_000);
+			template.setDefaultReplyTimeout(Duration.ofSeconds(30));
 			ProducerRecord<Integer, String> record = new ProducerRecord<>(A_REQUEST, 2, null, "baz");
 			record.headers().add(new RecordHeader(KafkaHeaders.REPLY_TOPIC, A_REPLY.getBytes()));
 			record.headers()
@@ -235,7 +241,7 @@ public void testGoodSamePartition() throws Exception {
 	public void testTimeout() throws Exception {
 		ReplyingKafkaTemplate<Integer, String, String> template = createTemplate(A_REPLY);
 		try {
-			template.setReplyTimeout(1);
+			template.setDefaultReplyTimeout(Duration.ofMillis(1));
 			ProducerRecord<Integer, String> record = new ProducerRecord<>(A_REQUEST, "fiz");
 			record.headers().add(new RecordHeader(KafkaHeaders.REPLY_TOPIC, A_REPLY.getBytes()));
 			RequestReplyFuture<Integer, String, String> future = template.sendAndReceive(record);
@@ -265,7 +271,7 @@ public void testTimeout() throws Exception {
 	public void testGoodWithSimpleMapper() throws Exception {
 		ReplyingKafkaTemplate<Integer, String, String> template = createTemplate(B_REPLY);
 		try {
-			template.setReplyTimeout(30_000);
+			template.setDefaultReplyTimeout(Duration.ofSeconds(30));
 			Headers headers = new RecordHeaders();
 			headers.add("baz", "buz".getBytes());
 			ProducerRecord<Integer, String> record = new ProducerRecord<>(B_REQUEST, null, null, null, "qux", headers);
@@ -291,7 +297,7 @@ public void testAggregateNormal() throws Exception {
 		AggregatingReplyingKafkaTemplate<Integer, String, String> template = aggregatingTemplate(
 				new TopicPartitionOffset(D_REPLY, 0), 2);
 		try {
-			template.setReplyTimeout(30_000);
+			template.setDefaultReplyTimeout(Duration.ofSeconds(30));
 			ProducerRecord<Integer, String> record = new ProducerRecord<>(D_REQUEST, null, null, null, "foo");
 			RequestReplyFuture<Integer, String, Collection<ConsumerRecord<Integer, String>>> future =
 					template.sendAndReceive(record);
@@ -320,7 +326,7 @@ public void testAggregateTimeout() throws Exception {
 		AggregatingReplyingKafkaTemplate<Integer, String, String> template = aggregatingTemplate(
 				new TopicPartitionOffset(E_REPLY, 0), 3);
 		try {
-			template.setReplyTimeout(5_000);
+			template.setDefaultReplyTimeout(Duration.ofSeconds(5));
 			ProducerRecord<Integer, String> record = new ProducerRecord<>(E_REQUEST, null, null, null, "foo");
 			RequestReplyFuture<Integer, String, Collection<ConsumerRecord<Integer, String>>> future =
 					template.sendAndReceive(record);
@@ -355,7 +361,7 @@ public void testAggregateTimeoutPartial() throws Exception {
 				new TopicPartitionOffset(F_REPLY, 0), 3);
 		template.setReturnPartialOnTimeout(true);
 		try {
-			template.setReplyTimeout(5_000);
+			template.setDefaultReplyTimeout(Duration.ofSeconds(5));
 			ProducerRecord<Integer, String> record = new ProducerRecord<>(F_REQUEST, null, null, null, "foo");
 			RequestReplyFuture<Integer, String, Collection<ConsumerRecord<Integer, String>>> future =
 					template.sendAndReceive(record);
@@ -395,7 +401,7 @@ public void testAggregateOrphansNotStored() throws Exception {
 			return null;
 		}).given(producer).send(any(), any());
 		AggregatingReplyingKafkaTemplate template = new AggregatingReplyingKafkaTemplate(pf, container, coll -> true);
-		template.setReplyTimeout(30_000);
+		template.setDefaultReplyTimeout(Duration.ofSeconds(30));
 		template.start();
 		List<ConsumerRecord> records = new ArrayList<>();
 		ConsumerRecord record = new ConsumerRecord("two", 0, 0L, null, "test1");
@@ -497,7 +503,7 @@ public void withCustomHeaders() throws Exception {
 		template.setReplyTopicHeaderName("custom.reply.to");
 		template.setReplyPartitionHeaderName("custom.reply.partition");
 		try {
-			template.setReplyTimeout(30_000);
+			template.setDefaultReplyTimeout(Duration.ofSeconds(30));
 			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);
@@ -575,7 +581,10 @@ public Map<String, Object> additionalHeaders() {
 
 		@KafkaListener(id = A_REQUEST, topics = A_REQUEST)
 		@SendTo  // default REPLY_TOPIC header
-		public String handleA(String in) {
+		public String handleA(String in) throws InterruptedException {
+			if (in.equals("slow")) {
+				Thread.sleep(50);
+			}
 			return in.toUpperCase();
 		}
 

src/reference/asciidoc/kafka.adoc

@@ -329,19 +329,24 @@ public KafkaTemplate<Integer, CustomValue> kafkaTemplate() {
 
 Version 2.1.3 introduced a subclass of `KafkaTemplate` to provide request/reply semantics.
 The class is named `ReplyingKafkaTemplate` and has one method (in addition to those in the superclass).
-The following listing shows the method's signature:
+The following listing shows the method signatures:
 
 ====
 [source, java]
 ----
 RequestReplyFuture<K, V, R> sendAndReceive(ProducerRecord<K, V> record);
+
+RequestReplyFuture<K, V, R> sendAndReceive(ProducerRecord<K, V> record,
+    Duration replyTimeout);
 ----
 ====
 
 The result is a `ListenableFuture` that is asynchronously populated with the result (or an exception, for a timeout).
 The result also has a `sendFuture` property, which is the result of calling `KafkaTemplate.send()`.
 You can use this future to determine the result of the send operation.
 
+If the first method is used, or the `replyTimeout` argument is `null`, the template's `defaultReplyTimeout` property is used (5 seconds by default).
+
 The following Spring Boot application shows an example of how to use the feature:
 
 ====
@@ -359,9 +364,9 @@ public class KRequestingApplication {
         return args -> {
             ProducerRecord<String, String> record = new ProducerRecord<>("kRequests", "foo");
             RequestReplyFuture<String, String, String> replyFuture = template.sendAndReceive(record);
-            SendResult<String, String> sendResult = replyFuture.getSendFuture().get();
+            SendResult<String, String> sendResult = replyFuture.getSendFuture().get(10, TimeUnit.SECONDS);
             System.out.println("Sent ok: " + sendResult.getRecordMetadata());
-            ConsumerRecord<String, String> consumerRecord = replyFuture.get();
+            ConsumerRecord<String, String> consumerRecord = replyFuture.get(10, TimeUnit.SECONDS);
             System.out.println("Return value: " + consumerRecord.value());
         };
     }

src/reference/asciidoc/whats-new.adoc

@@ -112,6 +112,8 @@ See the javadocs, <<serdes>>, and <<serde>> for more informaion.
 
 When a reply times out, the future is completed exceptionally with a `KafkaReplyTimeoutException` instead of a `KafkaException`.
 
+Also, an overloaded `sendAndReceive` method is now provided that allows specifying the reply timeout on a per message basis.
+
 ==== AggregatingReplyingKafkaTemplate
 
 Extends the `ReplyingKafkaTemplate` by aggregating replies from multiple receivers.