GH-1247: Only wake consumer while polling

Fixes #1247

Previously, stopping the container unconditionally called `consumer.wakeUp()`.
It should only be woken while actually polling, otherwise operations such
as `commitSync()` would fail (after committing the offset but before calling
any callbacks or interceptors).

Add logic to only wake the consumer while it is in the `poll()` method.

**cherry-pick to 2.2.x, 2.1.x**

# Conflicts:
#	spring-kafka/src/main/java/org/springframework/kafka/listener/KafkaMessageListenerContainer.java
#	spring-kafka/src/test/java/org/springframework/kafka/listener/ContainerStoppingErrorHandlerRecordModeTests.java

# Conflicts:
#	spring-kafka/src/main/java/org/springframework/kafka/listener/KafkaMessageListenerContainer.java
#	spring-kafka/src/test/java/org/springframework/kafka/listener/ContainerStoppingBatchErrorHandlerTests.java
#	spring-kafka/src/test/java/org/springframework/kafka/listener/ContainerStoppingErrorHandlerBatchModeTests.java

Author: garyrussell

spring-kafka/src/main/java/org/springframework/kafka/listener/KafkaMessageListenerContainer.java

@@ -32,6 +32,7 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.ScheduledFuture;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.stream.Collectors;
 
 import org.apache.commons.logging.Log;
@@ -292,7 +293,7 @@ public void onSuccess(Object result) {
 				}
 			});
 			setRunning(false);
-			this.listenerConsumer.consumer.wakeup();
+			this.listenerConsumer.wakeIfNecessary();
 		}
 	}
 
@@ -405,6 +406,8 @@ private final class ListenerConsumer implements SchedulingAwareRunnable, Consume
 
 		private volatile Map<TopicPartition, OffsetMetadata> definedPartitions;
 
+		private final AtomicBoolean polling = new AtomicBoolean();
+
 		private volatile Collection<TopicPartition> assignedPartitions;
 
 		private volatile Thread consumerThread;
@@ -727,8 +730,19 @@ public void run() {
 						}
 						publishConsumerPausedEvent(this.consumer.assignment());
 					}
-					ConsumerRecords<K, V> records = this.consumer.poll(this.containerProperties.getPollTimeout());
 					this.lastPoll = System.currentTimeMillis();
+					this.polling.set(true);
+					ConsumerRecords<K, V> records = this.consumer.poll(this.containerProperties.getPollTimeout());
+					if (!this.polling.compareAndSet(true, false)) {
+						/*
+						 * There is a small race condition where wakeIfNecessary was called between
+						 * exiting the poll and before we reset the boolean.
+						 */
+						if (records.count() > 0 && this.logger.isDebugEnabled()) {
+							this.logger.debug("Discarding polled records, container stopped: " + records.count());
+						}
+						return;
+					}
 					if (this.consumerPaused && !isPaused()) {
 						if (this.logger.isDebugEnabled()) {
 							this.logger.debug("Resuming consumption from: " + this.consumer.paused());
@@ -808,6 +822,12 @@ public void run() {
 			this.logger.info("Consumer stopped");
 		}
 
+		void wakeIfNecessary() {
+			if (this.polling.getAndSet(false)) {
+				this.consumer.wakeup();
+			}
+		}
+
 		/**
 		 * Handle exceptions thrown by the consumer outside of message listener
 		 * invocation (e.g. commit exceptions).

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2017 the original author or authors.
+ * 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.
@@ -91,7 +91,6 @@ public void stopContainerAfterException() throws Exception {
 		InOrder inOrder = inOrder(this.consumer);
 		inOrder.verify(this.consumer).subscribe(any(Collection.class), any(ConsumerRebalanceListener.class));
 		inOrder.verify(this.consumer).poll(1000);
-		inOrder.verify(this.consumer).wakeup();
 		inOrder.verify(this.consumer).unsubscribe();
 		inOrder.verify(this.consumer).close();
 		inOrder.verifyNoMoreInteractions();

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2017 the original author or authors.
+ * 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.
@@ -93,7 +93,6 @@ public void stopContainerAfterException() throws Exception {
 		InOrder inOrder = inOrder(this.consumer);
 		inOrder.verify(this.consumer).subscribe(any(Collection.class), any(ConsumerRebalanceListener.class));
 		inOrder.verify(this.consumer).poll(1000);
-		inOrder.verify(this.consumer).wakeup();
 		inOrder.verify(this.consumer).unsubscribe();
 		inOrder.verify(this.consumer).close();
 		inOrder.verifyNoMoreInteractions();

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

@@ -101,7 +101,6 @@ public void stopContainerAfterException() throws Exception {
 				Collections.singletonMap(new TopicPartition("foo", 0), new OffsetAndMetadata(2L)));
 		inOrder.verify(this.consumer).commitSync(
 				Collections.singletonMap(new TopicPartition("foo", 1), new OffsetAndMetadata(1L)));
-		inOrder.verify(this.consumer).wakeup();
 		inOrder.verify(this.consumer).unsubscribe();
 		inOrder.verify(this.consumer).close();
 		inOrder.verifyNoMoreInteractions();

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

@@ -25,6 +25,7 @@
 import static org.mockito.BDDMockito.willAnswer;
 import static org.mockito.Mockito.inOrder;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
@@ -550,6 +551,68 @@ public void onMessage(ConsumerRecord<Integer, String> data) {
 		inOrder.verify(messageListener).onMessage(any(ConsumerRecord.class));
 		inOrder.verify(consumer).commitSync(any(Map.class));
 		container.stop();
+		verify(consumer).wakeup();
+	}
+
+	@SuppressWarnings("unchecked")
+	@Test
+	public void testRecordAckAfterStop() throws Exception {
+		ConsumerFactory<Integer, String> cf = mock(ConsumerFactory.class);
+		Consumer<Integer, String> consumer = mock(Consumer.class);
+		given(cf.createConsumer(eq("grp"), eq("clientId"), isNull())).willReturn(consumer);
+		final Map<TopicPartition, List<ConsumerRecord<Integer, String>>> records = new HashMap<>();
+		records.put(new TopicPartition("foo", 0), Collections.singletonList(
+				new ConsumerRecord<>("foo", 0, 0L, 1, "foo")));
+		ConsumerRecords<Integer, String> consumerRecords = new ConsumerRecords<>(records);
+		given(consumer.poll(anyLong())).willAnswer(i -> {
+			Thread.sleep(50);
+			return consumerRecords;
+		});
+		TopicPartitionInitialOffset[] topicPartition = new TopicPartitionInitialOffset[] {
+				new TopicPartitionInitialOffset("foo", 0) };
+		ContainerProperties containerProps = new ContainerProperties(topicPartition);
+		containerProps.setGroupId("grp");
+		containerProps.setAckMode(AckMode.RECORD);
+		final CountDownLatch latch1 = new CountDownLatch(1);
+		final CountDownLatch latch2 = new CountDownLatch(1);
+		MessageListener<Integer, String> messageListener = spy(
+				new MessageListener<Integer, String>() { // Cannot be lambda: Mockito doesn't mock final classes
+
+					@Override
+					public void onMessage(ConsumerRecord<Integer, String> data) {
+						latch1.countDown();
+						try {
+							latch2.await(10, TimeUnit.SECONDS);
+						}
+						catch (InterruptedException e) {
+							Thread.currentThread().interrupt();
+						}
+					}
+
+				});
+
+		final CountDownLatch commitLatch = new CountDownLatch(1);
+		willAnswer(i -> {
+					commitLatch.countDown();
+					return null;
+				}
+		).given(consumer).commitSync(any());
+
+		containerProps.setMessageListener(messageListener);
+		containerProps.setClientId("clientId");
+		containerProps.setShutdownTimeout(5L);
+		KafkaMessageListenerContainer<Integer, String> container =
+				new KafkaMessageListenerContainer<>(cf, containerProps);
+		container.start();
+		assertThat(latch1.await(10, TimeUnit.SECONDS)).isTrue();
+		container.stop();
+		latch2.countDown();
+		assertThat(commitLatch.await(10, TimeUnit.SECONDS)).isTrue();
+		InOrder inOrder = inOrder(messageListener, consumer);
+		inOrder.verify(consumer).poll(ContainerProperties.DEFAULT_POLL_TIMEOUT);
+		inOrder.verify(messageListener).onMessage(any(ConsumerRecord.class));
+		inOrder.verify(consumer).commitSync(any());
+		verify(consumer, never()).wakeup();
 	}
 
 	@Test