spring-projectsGH-566: Ack Concurrency Issue

Fixes spring-projects#566

If `Acknowledgment.acknowledge()` is called on a foreign thread, there is a concurrency
problem with the `offsets` field; the consumer thread might `clear()` unprocessed acks.

Furthermore, `AckMode.MANUAL_IMMEDIATE` cannot use the `Consumer` object if the ack
is called on a foreign thread - the `Consumer` is not thread-safe.

- Revert `offsets` to simple `HashMap`s
- Only reference `offsets` on the consumer thread
- enqueue foreign acks into the `acks` queue (even "immediate" acks)

**Cherry pick to 2.0.x (and 1.3.x, fixing the lambda in the test and the diamond operators).**

Author: garyrussell

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

@@ -30,7 +30,6 @@
 import java.util.Set;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.ScheduledFuture;
 import java.util.stream.Collectors;
@@ -52,6 +51,7 @@
 import org.apache.kafka.common.errors.WakeupException;
 
 import org.springframework.core.task.SimpleAsyncTaskExecutor;
+import org.springframework.kafka.KafkaException;
 import org.springframework.kafka.core.ConsumerFactory;
 import org.springframework.kafka.core.KafkaResourceHolder;
 import org.springframework.kafka.core.ProducerFactoryUtils;
@@ -346,7 +346,7 @@ private final class ListenerConsumer implements SchedulingAwareRunnable, Consume
 
 		private final Consumer<K, V> consumer;
 
-		private final ConcurrentMap<String, ConcurrentMap<Integer, Long>> offsets = new ConcurrentHashMap<>();
+		private final Map<String, Map<Integer, Long>> offsets = new HashMap<>();
 
 		private final GenericMessageListener<?> genericListener;
 
@@ -406,6 +406,8 @@ private final class ListenerConsumer implements SchedulingAwareRunnable, Consume
 
 		private volatile Collection<TopicPartition> assignedPartitions;
 
+		private volatile Thread consumerThread;
+
 		private int count;
 
 		private long last = System.currentTimeMillis();
@@ -662,6 +664,7 @@ public boolean isLongLived() {
 
 		@Override
 		public void run() {
+			this.consumerThread = Thread.currentThread();
 			if (this.genericListener instanceof ConsumerSeekAware) {
 				((ConsumerSeekAware) this.genericListener).registerSeekCallback(this);
 			}
@@ -818,16 +821,27 @@ record = this.acks.poll();
 		}
 
 		private void processAck(ConsumerRecord<K, V> record) {
-			if (this.isManualImmediateAck) {
+			if (!Thread.currentThread().equals(this.consumerThread)) {
 				try {
-					ackImmediate(record);
+					this.acks.put(record);
 				}
-				catch (WakeupException e) {
-					// ignore - not polling
+				catch (InterruptedException e) {
+					Thread.currentThread().interrupt();
+					throw new KafkaException("Interrupted while storing ack", e);
 				}
 			}
 			else {
-				addOffset(record);
+				if (this.isManualImmediateAck) {
+					try {
+						ackImmediate(record);
+					}
+					catch (WakeupException e) {
+						// ignore - not polling
+					}
+				}
+				else {
+					addOffset(record);
+				}
 			}
 		}
 
@@ -1325,7 +1339,7 @@ private void commitIfNecessary() {
 
 		private Map<TopicPartition, OffsetAndMetadata> buildCommits() {
 			Map<TopicPartition, OffsetAndMetadata> commits = new HashMap<>();
-			for (Entry<String, ConcurrentMap<Integer, Long>> entry : this.offsets.entrySet()) {
+			for (Entry<String, Map<Integer, Long>> entry : this.offsets.entrySet()) {
 				for (Entry<Integer, Long> offset : entry.getValue().entrySet()) {
 					commits.put(new TopicPartition(entry.getKey(), offset.getKey()),
 							new OffsetAndMetadata(offset.getValue() + 1));

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

@@ -551,6 +551,77 @@ public void onMessage(ConsumerRecord<Integer, String> data) {
 		container.stop();
 	}
 
+	@Test
+	public void testRecordAckMockForeignThread() throws Exception {
+		testRecordAckMockForeignThreadGuts(AckMode.MANUAL);
+	}
+
+	@Test
+	public void testRecordAckMockForeignThreadImmediate() throws Exception {
+		testRecordAckMockForeignThreadGuts(AckMode.MANUAL_IMMEDIATE);
+	}
+
+	@SuppressWarnings("unchecked")
+	private void testRecordAckMockForeignThreadGuts(AckMode ackMode) throws Exception {
+		ConsumerFactory<Integer, String> cf = mock(ConsumerFactory.class);
+		Consumer<Integer, String> consumer = mock(Consumer.class);
+		given(cf.createConsumer(isNull(), eq("clientId"), isNull())).willReturn(consumer);
+		final Map<TopicPartition, List<ConsumerRecord<Integer, String>>> records = new HashMap<>();
+		records.put(new TopicPartition("foo", 0), Arrays.asList(
+				new ConsumerRecord<>("foo", 0, 0L, 1, "foo"),
+				new ConsumerRecord<>("foo", 0, 1L, 1, "bar")));
+		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.setAckMode(ackMode);
+		final CountDownLatch latch = new CountDownLatch(2);
+		final List<Acknowledgment> acks = new ArrayList<>();
+		final AtomicReference<Thread> consumerThread = new AtomicReference<>();
+		AcknowledgingMessageListener<Integer, String> messageListener = spy(
+				new AcknowledgingMessageListener<Integer, String>() {
+
+					@Override
+					public void onMessage(ConsumerRecord<Integer, String> data, Acknowledgment acknowledgment) {
+						acks.add(acknowledgment);
+						consumerThread.set(Thread.currentThread());
+						latch.countDown();
+						if (latch.getCount() == 0) {
+							records.clear();
+						}
+					}
+
+				});
+
+		final CountDownLatch commitLatch = new CountDownLatch(1);
+		final AtomicReference<Thread> commitThread = new AtomicReference<>();
+		willAnswer(i -> {
+					commitThread.set(Thread.currentThread());
+					commitLatch.countDown();
+					return null;
+				}
+		).given(consumer).commitSync(any(Map.class));
+
+		containerProps.setMessageListener(messageListener);
+		containerProps.setClientId("clientId");
+		KafkaMessageListenerContainer<Integer, String> container =
+				new KafkaMessageListenerContainer<>(cf, containerProps);
+		container.start();
+		assertThat(latch.await(10, TimeUnit.SECONDS)).isTrue();
+		acks.get(1).acknowledge();
+		assertThat(commitLatch.await(10, TimeUnit.SECONDS)).isTrue();
+		InOrder inOrder = inOrder(messageListener, consumer);
+		inOrder.verify(consumer).poll(1000);
+		inOrder.verify(messageListener, times(2)).onMessage(any(ConsumerRecord.class), any(Acknowledgment.class));
+		inOrder.verify(consumer).commitSync(any(Map.class));
+		container.stop();
+		assertThat(commitThread.get()).isSameAs(consumerThread.get());
+	}
+
 	@SuppressWarnings("unchecked")
 	@Test
 	public void testNonResponsiveConsumerEvent() throws Exception {
@@ -1576,7 +1647,6 @@ public void onPartitionsAssigned(Collection<TopicPartition> partitions) {
 		final CountDownLatch commitLatch = new CountDownLatch(2);
 		willAnswer(invocation -> {
 
-			@SuppressWarnings({ "unchecked" })
 			Map<TopicPartition, OffsetAndMetadata> map = invocation.getArgument(0);
 			try {
 				return invocation.callRealMethod();