Fix KafkaListenerEndpointRegistry.stop(Runnable)

The `KafkaListenerEndpointRegistry.stop(Runnable)` doesn't change its `running` state

* Removed the callback wrap in favor of setting the state to stop in two places
* Set the running state to false before actually calling stop on the containers

Author: asibross

spring-kafka/src/main/java/org/springframework/kafka/config/KafkaListenerEndpointRegistry.java

@@ -60,6 +60,7 @@
  * @author Juergen Hoeller
  * @author Artem Bilan
  * @author Gary Russell
+ * @author Asi Bross
  *
  * @see KafkaListenerEndpoint
  * @see MessageListenerContainer
@@ -70,8 +71,7 @@ public class KafkaListenerEndpointRegistry implements DisposableBean, SmartLifec
 
 	protected final LogAccessor logger = new LogAccessor(LogFactory.getLog(getClass())); //NOSONAR
 
-	private final Map<String, MessageListenerContainer> listenerContainers =
-			new ConcurrentHashMap<String, MessageListenerContainer>();
+	private final Map<String, MessageListenerContainer> listenerContainers = new ConcurrentHashMap<>();
 
 	private int phase = AbstractMessageListenerContainer.DEFAULT_PHASE;
 
@@ -162,6 +162,7 @@ public void registerListenerContainer(KafkaListenerEndpoint endpoint, KafkaListe
 	@SuppressWarnings("unchecked")
 	public void registerListenerContainer(KafkaListenerEndpoint endpoint, KafkaListenerContainerFactory<?> factory,
 			boolean startImmediately) {
+
 		Assert.notNull(endpoint, "Endpoint must not be null");
 		Assert.notNull(factory, "Factory must not be null");
 
@@ -260,17 +261,19 @@ public void start() {
 
 	@Override
 	public void stop() {
+		this.running = false;
 		for (MessageListenerContainer listenerContainer : getListenerContainers()) {
 			listenerContainer.stop();
 		}
-		this.running = false;
 	}
 
 	@Override
 	public void stop(Runnable callback) {
+		this.running = false;
 		Collection<MessageListenerContainer> listenerContainersToStop = getListenerContainers();
 		if (listenerContainersToStop.size() > 0) {
-			AggregatingCallback aggregatingCallback = new AggregatingCallback(listenerContainersToStop.size(), callback);
+			AggregatingCallback aggregatingCallback = new AggregatingCallback(listenerContainersToStop.size(),
+					callback);
 			for (MessageListenerContainer listenerContainer : listenerContainersToStop) {
 				if (listenerContainer.isRunning()) {
 					listenerContainer.stop(aggregatingCallback);

spring-kafka/src/test/java/org/springframework/kafka/listener/KafkaListenerEndpointRegistryLifecycleTests.java

@@ -0,0 +1,88 @@
+/*
+ * Copyright 2017-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.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.kafka.listener;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.RETURNS_DEEP_STUBS;
+import static org.mockito.Mockito.mock;
+
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+
+import org.junit.jupiter.api.Test;
+
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.kafka.annotation.EnableKafka;
+import org.springframework.kafka.annotation.KafkaListener;
+import org.springframework.kafka.config.ConcurrentKafkaListenerContainerFactory;
+import org.springframework.kafka.config.KafkaListenerEndpointRegistry;
+import org.springframework.kafka.core.ConsumerFactory;
+import org.springframework.test.annotation.DirtiesContext;
+import org.springframework.test.context.junit.jupiter.SpringJUnitConfig;
+
+/**
+ * @author Asi Bross
+ *
+ * @since 2.3.5
+ *
+ */
+@SpringJUnitConfig
+@DirtiesContext
+public class KafkaListenerEndpointRegistryLifecycleTests {
+
+	@Autowired
+	private KafkaListenerEndpointRegistry registry;
+
+	@Test
+	public void lifecycleTest() throws InterruptedException, ExecutionException, TimeoutException {
+		// Registry is started automatically by the application context
+		assertThat(registry.isRunning()).isTrue();
+
+		this.registry.stop();
+		assertThat(registry.isRunning()).isFalse();
+
+		this.registry.start();
+		assertThat(registry.isRunning()).isTrue();
+
+		CompletableFuture<Boolean> isRunning = new CompletableFuture<>();
+		this.registry.stop(() -> isRunning.complete(this.registry.isRunning()));
+		assertThat(isRunning.get(1, TimeUnit.SECONDS)).isFalse();
+	}
+
+	@KafkaListener(topics = "foo", groupId = "bar")
+	public void listen() {
+	}
+
+	@Configuration
+	@EnableKafka
+	public static class Config {
+
+		@SuppressWarnings({ "rawtypes", "unchecked" })
+		@Bean
+		public ConcurrentKafkaListenerContainerFactory kafkaListenerContainerFactory() {
+			ConcurrentKafkaListenerContainerFactory factory = new ConcurrentKafkaListenerContainerFactory();
+			factory.setConsumerFactory(mock(ConsumerFactory.class, RETURNS_DEEP_STUBS));
+			return factory;
+		}
+
+	}
+
+}