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 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> replyTopicsAndPartitions = new HashMap<>(); private final KafkaTemplate kafkaTemplate; @@ -162,6 +167,8 @@ public class KafkaProducerMessageHandler extends AbstractReplyProducingMes private boolean useTemplateConverter; + private Duration assignmentDuration = DEFAULT_ASSIGNMENT_TIMEOUT; + private volatile byte[] singleReplyTopic; public KafkaProducerMessageHandler(final KafkaTemplate kafkaTemplate) { @@ -415,6 +422,17 @@ public class KafkaProducerMessageHandler 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 extends AbstractReplyProducingMes private void determineValidReplyTopicsAndPartitions() { ReplyingKafkaTemplate rkt = (ReplyingKafkaTemplate) this.kafkaTemplate; + try { + rkt.waitForAssignment(this.assignmentDuration); + } + catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } Collection replyTopics = rkt.getAssignedReplyTopicPartitions(); Map> 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))))