From e19e6d4726a6acd52b81820565a22e8020814bda Mon Sep 17 00:00:00 2001
From: Gary Russell <grussell@vmware.com>
Date: Mon, 14 Nov 2022 17:16:50 -0500
Subject: [PATCH] GH-3942: Fix Race in Kafka OB Gateway

Resolves https://github.com/spring-projects/spring-integration/issues/3942

When determining the default reply-to topic/partition, we need to wait for assignment.

Already covered by `KafkaDslTests` (a recent build failure exposed this problem).

**No back-port - 5.5.x uses 2.7.x by default, which does not support this.**

5.5.x users can call `waitForAssignment` on the `ReplyingKafkaTemplate` that is
supplied to the gateways before sending messages.
---
 .../kafka/dsl/KafkaOutboundGatewaySpec.java   | 12 ++++++++++
 .../outbound/KafkaProducerMessageHandler.java | 24 +++++++++++++++++++
 .../integration/kafka/dsl/KafkaDslTests.java  |  1 +
 3 files changed, 37 insertions(+)

diff --git a/spring-integration-kafka/src/main/java/org/springframework/integration/kafka/dsl/KafkaOutboundGatewaySpec.java b/spring-integration-kafka/src/main/java/org/springframework/integration/kafka/dsl/KafkaOutboundGatewaySpec.java
index aba1c68d99..91dad5cf77 100644
--- a/spring-integration-kafka/src/main/java/org/springframework/integration/kafka/dsl/KafkaOutboundGatewaySpec.java
+++ b/spring-integration-kafka/src/main/java/org/springframework/integration/kafka/dsl/KafkaOutboundGatewaySpec.java
@@ -60,6 +60,18 @@ public class KafkaOutboundGatewaySpec<K, V, R, S extends KafkaOutboundGatewaySpe
 		return _this();
 	}
 
+	/**
+	 * Set the time to wait for partition assignment, when used as a gateway, to determine
+	 * the default reply-to topic/partition.
+	 * @param duration the duration.
+	 * @return the spec.
+	 * @since 6.0
+	 */
+	public S assigmentDuration(Duration duration) {
+		this.target.setAssignmentDuration(duration);
+		return _this();
+	}
+
 	/**
 	 * A {@link org.springframework.kafka.core.KafkaTemplate}-based {@link KafkaProducerMessageHandlerSpec} extension.
 	 *
diff --git a/spring-integration-kafka/src/main/java/org/springframework/integration/kafka/outbound/KafkaProducerMessageHandler.java b/spring-integration-kafka/src/main/java/org/springframework/integration/kafka/outbound/KafkaProducerMessageHandler.java
index cf9428f805..7d1a4eb2c2 100644
--- a/spring-integration-kafka/src/main/java/org/springframework/integration/kafka/outbound/KafkaProducerMessageHandler.java
+++ b/spring-integration-kafka/src/main/java/org/springframework/integration/kafka/outbound/KafkaProducerMessageHandler.java
@@ -18,6 +18,7 @@ package org.springframework.integration.kafka.outbound;
 
 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;
@@ -103,6 +104,10 @@ public class KafkaProducerMessageHandler<K, V> extends AbstractReplyProducingMes
 	 */
 	private static final int DEFAULT_TIMEOUT_BUFFER = 5000;
 
+	private static final int TWENTY = 20;
+
+	private static final Duration DEFAULT_ASSIGNMENT_TIMEOUT = Duration.ofSeconds(TWENTY);
+
 	private final Map<String, Set<Integer>> replyTopicsAndPartitions = new HashMap<>();
 
 	private final KafkaTemplate<K, V> kafkaTemplate;
@@ -162,6 +167,8 @@ public class KafkaProducerMessageHandler<K, V> extends AbstractReplyProducingMes
 
 	private boolean useTemplateConverter;
 
+	private Duration assignmentDuration = DEFAULT_ASSIGNMENT_TIMEOUT;
+
 	private volatile byte[] singleReplyTopic;
 
 	public KafkaProducerMessageHandler(final KafkaTemplate<K, V> kafkaTemplate) {
@@ -415,6 +422,17 @@ public class KafkaProducerMessageHandler<K, V> extends AbstractReplyProducingMes
 		this.useTemplateConverter = useTemplateConverter;
 	}
 
+	/**
+	 * Set the time to wait for partition assignment, when used as a gateway, to determine
+	 * the default reply-to topic/partition.
+	 * @param assignmentDuration the assignmentDuration to set.
+	 * @since 6.0
+	 */
+	public void setAssignmentDuration(Duration assignmentDuration) {
+		Assert.notNull(assignmentDuration, "'assignmentDuration' cannot be null");
+		this.assignmentDuration = assignmentDuration;
+	}
+
 	@Override
 	public String getComponentType() {
 		return this.isGateway ? "kafka:outbound-gateway" : "kafka:outbound-channel-adapter";
@@ -647,6 +665,12 @@ public class KafkaProducerMessageHandler<K, V> extends AbstractReplyProducingMes
 
 	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) {
diff --git a/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/dsl/KafkaDslTests.java b/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/dsl/KafkaDslTests.java
index 692ed10975..61810899d8 100644
--- a/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/dsl/KafkaDslTests.java
+++ b/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/dsl/KafkaDslTests.java
@@ -400,6 +400,7 @@ public class KafkaDslTests {
 		public IntegrationFlow outboundGateFlow() {
 			return IntegrationFlow.from(Gate.class)
 					.handle(Kafka.outboundGateway(producerFactory(), replyContainer())
+							.assigmentDuration(Duration.ofSeconds(30))
 							.flushExpression("true")
 							.sync(true)
 							.configureKafkaTemplate(t -> t.defaultReplyTimeout(Duration.ofSeconds(30))))