GH-3616: Defer default topic resolution to the ReplyingKafkaTemplate

Fixes: #3616
Issue link: #3616

The `KafkaProducerMessageHandler` uses an unnecessary logic to determine a default topic/partition.
It is better to push such a logic down to the `ReplyingKafkaTemplate` as more general solution

* Remove `KafkaProducerMessageHandler.getReplyTopic()` logic altogether
* Clean up tests for removed logic
* Add `exclude group: 'ch.qos.logback'` to be able to control logging for SI-Kafka module
* Remove out-dated sentence from `kafka.adoc`

Author: artembilan

build.gradle

@@ -786,7 +786,9 @@ project('spring-integration-kafka') {
     dependencies {
         api 'org.springframework.kafka:spring-kafka'
 
-        testImplementation 'org.springframework.kafka:spring-kafka-test'
+        testImplementation ('org.springframework.kafka:spring-kafka-test') {
+            exclude group: 'ch.qos.logback'
+        }
         testImplementation 'com.fasterxml.jackson.core:jackson-databind'
     }
 }

spring-integration-kafka/src/main/java/org/springframework/integration/kafka/outbound/KafkaProducerMessageHandler.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013-2023 the original author or authors.
+ * Copyright 2013-2024 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.
@@ -19,11 +19,9 @@
 import java.lang.reflect.Type;
 import java.nio.charset.StandardCharsets;
 import java.time.Duration;
-import java.util.Collection;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Set;
-import java.util.TreeSet;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Future;
@@ -33,9 +31,7 @@
 
 import org.apache.kafka.clients.producer.ProducerConfig;
 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.springframework.expression.EvaluationContext;
@@ -423,8 +419,7 @@ public void setUseTemplateConverter(boolean useTemplateConverter) {
 	}
 
 	/**
-	 * Set the time to wait for partition assignment, when used as a gateway, to determine
-	 * the default reply-to topic/partition.
+	 * Set the time to wait for partition assignment, when used as a gateway.
 	 * @param assignmentDuration the assignmentDuration to set.
 	 * @since 6.0
 	 */
@@ -500,8 +495,7 @@ protected Object handleRequestMessage(final Message<?> message) {
 		final ProducerRecord<K, V> producerRecord;
 		boolean flush =
 				Boolean.TRUE.equals(this.flushExpression.getValue(this.evaluationContext, message, Boolean.class));
-		boolean preBuilt = message.getPayload() instanceof ProducerRecord;
-		if (preBuilt) {
+		if (message.getPayload() instanceof ProducerRecord) {
 			producerRecord = (ProducerRecord<K, V>) message.getPayload();
 		}
 		else {
@@ -517,11 +511,11 @@ protected Object handleRequestMessage(final Message<?> message) {
 		CompletableFuture<SendResult<K, V>> sendFuture;
 		RequestReplyFuture<K, V, Object> gatewayFuture = null;
 		try {
-			if (this.isGateway
-					&& (!preBuilt || producerRecord.headers().lastHeader(KafkaHeaders.REPLY_TOPIC) == null)) {
-				producerRecord.headers().add(new RecordHeader(KafkaHeaders.REPLY_TOPIC, getReplyTopic(message)));
-				gatewayFuture = ((ReplyingKafkaTemplate<K, V, Object>) this.kafkaTemplate)
-						.sendAndReceive(producerRecord);
+			if (this.isGateway) {
+				waitForAssignment();
+				addReplyTopicIfAny(message.getHeaders(), producerRecord.headers());
+				gatewayFuture =
+						((ReplyingKafkaTemplate<K, V, Object>) this.kafkaTemplate).sendAndReceive(producerRecord);
 				sendFuture = gatewayFuture.getSendFuture();
 			}
 			else {
@@ -554,24 +548,6 @@ protected Object handleRequestMessage(final Message<?> message) {
 		return processReplyFuture(gatewayFuture);
 	}
 
-	private void sendFutureIfRequested(CompletableFuture<SendResult<K, V>> sendFuture, Object futureToken) {
-
-		if (futureToken != null) {
-			MessageChannel futures = getFuturesChannel();
-			if (futures != null) {
-				try {
-					futures.send(getMessageBuilderFactory()
-							.withPayload(sendFuture)
-							.setHeader(KafkaIntegrationHeaders.FUTURE_TOKEN, futureToken)
-							.build());
-				}
-				catch (Exception e) {
-					this.logger.error(e, "Failed to send sendFuture");
-				}
-			}
-		}
-	}
-
 	@SuppressWarnings("unchecked")
 	private ProducerRecord<K, V> createProducerRecord(final Message<?> message) {
 		MessageHeaders messageHeaders = message.getHeaders();
@@ -604,91 +580,56 @@ private ProducerRecord<K, V> createProducerRecord(final Message<?> message) {
 			payload = null;
 		}
 
-		Headers headers = null;
+		Headers headers = new RecordHeaders();
 		if (this.headerMapper != null) {
-			headers = new RecordHeaders();
 			this.headerMapper.fromHeaders(messageHeaders, headers);
 		}
+
 		return this.producerRecordCreator.create(message, topic, partitionId, timestamp, (K) messageKey, payload,
 				headers);
 	}
 
-	private byte[] getReplyTopic(Message<?> message) { // NOSONAR
-		if (this.replyTopicsAndPartitions.isEmpty()) {
-			determineValidReplyTopicsAndPartitions();
-		}
-		Object replyHeader = message.getHeaders().get(KafkaHeaders.REPLY_TOPIC);
-		byte[] replyTopic = null;
-		String topicToCheck = null;
-		if (replyHeader instanceof String) {
-			replyTopic = ((String) replyHeader).getBytes(StandardCharsets.UTF_8);
-			topicToCheck = (String) replyHeader;
-		}
-		else if (replyHeader instanceof byte[]) {
-			replyTopic = (byte[]) replyHeader;
-		}
-		else if (replyHeader != null) {
-			throw new IllegalStateException(KafkaHeaders.REPLY_TOPIC + " must be String or byte[]");
+	private void waitForAssignment() {
+		ReplyingKafkaTemplate<?, ?, ?> rkt = (ReplyingKafkaTemplate<?, ?, ?>) this.kafkaTemplate;
+		try {
+			rkt.waitForAssignment(this.assignmentDuration);
 		}
-		if (replyTopic == null) {
-			if (this.replyTopicsAndPartitions.size() == 1) {
-				replyTopic = getSingleReplyTopic();
-			}
-			else {
-				throw new IllegalStateException("No reply topic header and no default reply topic can be determined; "
-						+ "container's assigned partitions: " + this.replyTopicsAndPartitions);
-			}
+		catch (InterruptedException e) {
+			Thread.currentThread().interrupt();
 		}
-		else {
-			if (topicToCheck == null) {
-				topicToCheck = new String(replyTopic, StandardCharsets.UTF_8);
-			}
-			if (!this.replyTopicsAndPartitions.containsKey(topicToCheck)) {
-				throw new IllegalStateException("The reply topic header ["
-						+ topicToCheck +
-						"] does not match any reply container topic: " + this.replyTopicsAndPartitions.keySet());
+	}
+
+	@Nullable
+	private void addReplyTopicIfAny(MessageHeaders messageHeaders, Headers headers) {
+		if (this.isGateway) {
+			Object replyHeader = messageHeaders.get(KafkaHeaders.REPLY_TOPIC);
+			if (replyHeader instanceof String topicString) {
+				headers.add(KafkaHeaders.REPLY_TOPIC, topicString.getBytes(StandardCharsets.UTF_8));
 			}
-		}
-		Integer replyPartition = message.getHeaders().get(KafkaHeaders.REPLY_PARTITION, Integer.class);
-		if (replyPartition != null) {
-			if (topicToCheck == null) {
-				topicToCheck = new String(replyTopic, StandardCharsets.UTF_8);
+			else if (replyHeader instanceof byte[] topicBytes) {
+				headers.add(KafkaHeaders.REPLY_TOPIC, topicBytes);
 			}
-			if (!this.replyTopicsAndPartitions.get(topicToCheck).contains(replyPartition)) {
-				throw new IllegalStateException("The reply partition header ["
-						+ replyPartition + "] does not match any reply container partition for topic ["
-						+ topicToCheck + "]: " + this.replyTopicsAndPartitions.get(topicToCheck));
+			else if (replyHeader != null) {
+				throw new IllegalStateException(KafkaHeaders.REPLY_TOPIC + " must be String or byte[]");
 			}
 		}
-		return replyTopic;
 	}
 
-	private byte[] getSingleReplyTopic() {
-		if (this.singleReplyTopic == null) {
-			this.singleReplyTopic = this.replyTopicsAndPartitions.keySet()
-					.iterator()
-					.next()
-					.getBytes(StandardCharsets.UTF_8);
-		}
-		return this.singleReplyTopic;
-	}
+	private void sendFutureIfRequested(CompletableFuture<SendResult<K, V>> sendFuture, Object futureToken) {
 
-	private void determineValidReplyTopicsAndPartitions() {
-		ReplyingKafkaTemplate<?, ?, ?> rkt = (ReplyingKafkaTemplate<?, ?, ?>) this.kafkaTemplate;
-		try {
-			rkt.waitForAssignment(this.assignmentDuration);
-		}
-		catch (InterruptedException e) {
-			Thread.currentThread().interrupt();
-		}
-		Collection<TopicPartition> replyTopics = rkt.getAssignedReplyTopicPartitions();
-		Map<String, Set<Integer>> topicsAndPartitions = new HashMap<>();
-		if (replyTopics != null) {
-			replyTopics.forEach(tp -> {
-				topicsAndPartitions.computeIfAbsent(tp.topic(), (k) -> new TreeSet<>());
-				topicsAndPartitions.get(tp.topic()).add(tp.partition());
-			});
-			this.replyTopicsAndPartitions.putAll(topicsAndPartitions);
+		if (futureToken != null) {
+			MessageChannel futures = getFuturesChannel();
+			if (futures != null) {
+				try {
+					futures.send(getMessageBuilderFactory()
+							.withPayload(sendFuture)
+							.setHeader(KafkaIntegrationHeaders.FUTURE_TOKEN, futureToken)
+							.build());
+				}
+				catch (Exception e) {
+					this.logger.error(e, "Failed to send sendFuture");
+				}
+			}
 		}
 	}
 
@@ -804,7 +745,7 @@ private Message<?> dontLeakHeaders(Message<?> message) {
 	 * @param <K> the key type.
 	 * @param <V> the value type.
 	 *
-	 * @since 3.2.1
+	 * @since 5.4
 	 *
 	 */
 	@FunctionalInterface

spring-integration-kafka/src/test/java/org/springframework/integration/kafka/outbound/KafkaProducerMessageHandlerTests.java

@@ -121,17 +121,17 @@
  */
 class KafkaProducerMessageHandlerTests {
 
-	private static String topic1 = "testTopic1out";
+	private static final String topic1 = "testTopic1out";
 
-	private static String topic2 = "testTopic2out";
+	private static final String topic2 = "testTopic2out";
 
-	private static String topic3 = "testTopic3out";
+	private static final String topic3 = "testTopic3out";
 
-	private static String topic4 = "testTopic4out";
+	private static final String topic4 = "testTopic4out";
 
-	private static String topic5 = "testTopic5out";
+	private static final String topic5 = "testTopic5out";
 
-	private static String topic6 = "testTopic6in";
+	private static final String topic6 = "testTopic6in";
 
 	private static EmbeddedKafkaBroker embeddedKafka;
 
@@ -467,31 +467,6 @@ public void onPartitionsAssigned(Collection<TopicPartition> partitions) {
 		assertThat(reply.getHeaders().get(KafkaHeaders.TOPIC)).isNull();
 		assertThat(reply.getHeaders().get(KafkaHeaders.CORRELATION_ID)).isNull();
 
-		final Message<?> messageToHandle1 = MessageBuilder.withPayload("foo")
-				.setHeader(KafkaHeaders.TOPIC, topic5)
-				.setHeader(KafkaHeaders.KEY, 2)
-				.setHeader(KafkaHeaders.PARTITION, 1)
-				.setHeader(KafkaHeaders.REPLY_TOPIC, "bad")
-				.build();
-
-		assertThatExceptionOfType(MessageHandlingException.class)
-				.isThrownBy(() -> handler.handleMessage(messageToHandle1))
-				.withStackTraceContaining("The reply topic header [bad] does not match any reply container topic: "
-						+ "[" + topic6 + "]");
-
-		final Message<?> messageToHandle2 = MessageBuilder.withPayload("foo")
-				.setHeader(KafkaHeaders.TOPIC, topic5)
-				.setHeader(KafkaHeaders.KEY, 2)
-				.setHeader(KafkaHeaders.PARTITION, 1)
-				.setHeader(KafkaHeaders.REPLY_PARTITION, 999)
-				.build();
-
-		assertThatExceptionOfType(MessageHandlingException.class)
-				.isThrownBy(() -> handler.handleMessage(messageToHandle2))
-				.withStackTraceContaining("The reply partition header [999] " +
-						"does not match any reply container partition for topic ["
-						+ topic6 + "]: [0, 1]");
-
 		template.stop();
 		// discard from the test consumer
 		KafkaTestUtils.getSingleRecord(consumer, topic6);

src/reference/antora/modules/ROOT/pages/kafka.adoc

@@ -468,9 +468,6 @@ The outbound gateway is for request/reply operations.
 It differs from most Spring Integration gateways in that the sending thread does not block in the gateway, and the reply is processed on the reply listener container thread.
 If your code invokes the gateway behind a synchronous https://docs.spring.io/spring-integration/reference/html/messaging-endpoints-chapter.html#gateway[Messaging Gateway], the user thread blocks there until the reply is received (or a timeout occurs).
 
-IMPORTANT: The gateway does not accept requests until the reply container has been assigned its topics and partitions.
-It is suggested that you add a `ConsumerRebalanceListener` to the template's reply container properties and wait for the `onPartitionsAssigned` call before sending messages to the gateway.
-
 The `KafkaProducerMessageHandler` `sendTimeoutExpression` default is `delivery.timeout.ms` Kafka producer property `+ 5000` so that the actual Kafka error after a timeout is propagated to the application, instead of a timeout generated by this framework.
 This has been changed for consistency because you may get unexpected behavior (Spring may time out the `send()`, while it is actually, eventually, successful).
 IMPORTANT: That timeout is 120 seconds by default, so you may wish to reduce it to get more timely failures.