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 b66b5641d2..a6f5118213 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 @@ -27,6 +27,7 @@ import java.util.concurrent.ExecutionException; import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicBoolean; import org.apache.kafka.clients.consumer.ConsumerRecord; import org.apache.kafka.clients.producer.ProducerRecord; @@ -35,6 +36,7 @@ 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.context.Lifecycle; import org.springframework.expression.EvaluationContext; import org.springframework.expression.Expression; import org.springframework.integration.MessageTimeoutException; @@ -85,7 +87,8 @@ import org.springframework.util.concurrent.SettableListenableFuture; * * @since 0.5 */ -public class KafkaProducerMessageHandler extends AbstractReplyProducingMessageHandler { +public class KafkaProducerMessageHandler extends AbstractReplyProducingMessageHandler + implements Lifecycle { private static final long DEFAULT_SEND_TIMEOUT = 10000; @@ -97,6 +100,8 @@ public class KafkaProducerMessageHandler extends AbstractReplyProducingMes private final boolean transactional; + private final AtomicBoolean running = new AtomicBoolean(); + private EvaluationContext evaluationContext; private Expression topicExpression; @@ -331,6 +336,23 @@ public class KafkaProducerMessageHandler extends AbstractReplyProducingMes this.evaluationContext = ExpressionUtils.createStandardEvaluationContext(getBeanFactory()); } + @Override + public void start() { + this.running.set(true); + } + + @Override + public void stop() { + if (this.running.compareAndSet(true, false)) { + this.kafkaTemplate.flush(); + } + } + + @Override + public boolean isRunning() { + return this.running.get(); + } + @SuppressWarnings("unchecked") @Override protected Object handleRequestMessage(final Message message) { diff --git a/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/outbound/KafkaProducerMessageHandlerTests.java b/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/outbound/KafkaProducerMessageHandlerTests.java index 1636c12dd9..35fe3a3037 100644 --- a/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/outbound/KafkaProducerMessageHandlerTests.java +++ b/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/outbound/KafkaProducerMessageHandlerTests.java @@ -1,5 +1,5 @@ /* - * Copyright 2016-2018 the original author or authors. + * Copyright 2016-2019 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. @@ -414,11 +414,16 @@ public class KafkaProducerMessageHandlerTests { KafkaTemplate template = new KafkaTemplate(pf); KafkaProducerMessageHandler handler = new KafkaProducerMessageHandler(template); handler.setTopicExpression(new LiteralExpression("bar")); + handler.setBeanFactory(mock(BeanFactory.class)); + handler.afterPropertiesSet(); + handler.start(); handler.handleMessage(new GenericMessage<>("foo")); + handler.stop(); InOrder inOrder = inOrder(producer); inOrder.verify(producer).beginTransaction(); inOrder.verify(producer).send(any(ProducerRecord.class), any(Callback.class)); inOrder.verify(producer).commitTransaction(); + inOrder.verify(producer).flush(); } }