GH-3942: Fix Race in Kafka OB Gateway

Resolves #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.

Author: garyrussell

Diff for: spring-integration-kafka/src/main/java/org/springframework/integration/kafka/dsl/KafkaOutboundGatewaySpec.java

@@ -60,6 +60,18 @@ public S replyMessageConverter(RecordMessageConverter messageConverter) {
 		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 for: spring-integration-kafka/src/main/java/org/springframework/integration/kafka/outbound/KafkaProducerMessageHandler.java

@@ -18,6 +18,7 @@
 
 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 void setUseTemplateConverter(boolean useTemplateConverter) {
 		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 @@ private byte[] getSingleReplyTopic() {
 
 	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 for: spring-integration-kafka/src/test/java/org/springframework/integration/kafka/dsl/KafkaDslTests.java

@@ -400,6 +400,7 @@ public IntegrationFlow sourceFlow() {
 		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))))