GH-1915: Batch Manual AckMode Improvements

Resolves #1915

- perform acks within the scope of the ack call so users can catch exceptions with sync commits
- when using MANUAL_IMMEDIATE, commit offsets for the batch instead of one-at-a-time

Author: garyrussell

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

@@ -1755,6 +1755,38 @@ private void processAck(ConsumerRecord<K, V> record) {
 			}
 		}
 
+		private void processAcks(ConsumerRecords<K, V> records) {
+			if (!Thread.currentThread().equals(this.consumerThread)) {
+				try {
+					for (ConsumerRecord<K, V> record : records) {
+						this.acks.put(record);
+					}
+					if (this.isManualImmediateAck) {
+						this.consumer.wakeup();
+					}
+				}
+				catch (InterruptedException e) {
+					Thread.currentThread().interrupt();
+					throw new KafkaException("Interrupted while storing ack", e);
+				}
+			}
+			else {
+				if (this.isManualImmediateAck) {
+					try {
+						ackImmediate(records);
+					}
+					catch (@SuppressWarnings(UNUSED) WakeupException e) {
+						// ignore - not polling
+					}
+				}
+				else {
+					for (ConsumerRecord<K, V> record : records) {
+						addOffset(record);
+					}
+				}
+			}
+		}
+
 		private synchronized void ackInOrder(ConsumerRecord<K, V> record) {
 			TopicPartition part = new TopicPartition(record.topic(), record.partition());
 			List<Long> offs = this.offsetsInThisBatch.get(part);
@@ -1806,6 +1838,25 @@ else if (this.syncCommits) {
 			}
 		}
 
+		private void ackImmediate(ConsumerRecords<K, V> records) {
+			Map<TopicPartition, OffsetAndMetadata> commits = new HashMap<>();
+			for (TopicPartition part : records.partitions()) {
+				commits.put(part,
+						new OffsetAndMetadata(records.records(part)
+								.get(records.records(part).size() - 1).offset() + 1));
+			}
+			this.commitLogger.log(() -> "Committing: " + commits);
+			if (this.producer != null) {
+				doSendOffsets(this.producer, commits);
+			}
+			else if (this.syncCommits) {
+				commitSync(commits);
+			}
+			else {
+				commitAsync(commits, 0);
+			}
+		}
+
 		private void commitAsync(Map<TopicPartition, OffsetAndMetadata> commits, int retries) {
 			this.consumer.commitAsync(commits, (offsetsAttempted, exception) -> {
 				if (exception instanceof RetriableCommitFailedException
@@ -2060,9 +2111,6 @@ private void invokeBatchOnMessage(final ConsumerRecords<K, V> records, // NOSONA
 					if (index++ >= this.nackIndex) {
 						toSeek.add(record);
 					}
-					else {
-						this.acks.put(record);
-					}
 				}
 			}
 			if (this.producer != null || (!this.isAnyManualAck && !this.autoCommit)) {
@@ -3026,12 +3074,12 @@ public void acknowledge() {
 				Map<TopicPartition, List<ConsumerRecord<K, V>>> deferred = ListenerConsumer.this.deferredOffsets;
 				if (!this.acked) {
 					for (ConsumerRecord<K, V> record : getHighestOffsetRecords(this.records)) {
-						processAck(record);
 						if (offs != null) {
 							offs.remove(new TopicPartition(record.topic(), record.partition()));
 							deferred.remove(new TopicPartition(record.topic(), record.partition()));
 						}
 					}
+					processAcks(this.records);
 					this.acked = true;
 				}
 			}
@@ -3050,6 +3098,22 @@ public void nack(int index, long sleep) {
 						ListenerConsumer.this.deferredOffsets.forEach((part, recs) -> recs.clear());
 					}
 				}
+				int i = 0;
+				List<ConsumerRecord<K, V>> toAck = new LinkedList<>();
+				for (ConsumerRecord<K, V> record : this.records) {
+					if (i++ < index) {
+						toAck.add(record);
+					}
+					else {
+						break;
+					}
+				}
+				Map<TopicPartition, List<ConsumerRecord<K, V>>> newRecords = new HashMap<>();
+				for (ConsumerRecord<K, V> record : toAck) {
+					newRecords.computeIfAbsent(new TopicPartition(record.topic(), record.partition()),
+							tp -> new LinkedList<>()).add(record);
+				}
+				processAcks(new ConsumerRecords<K, V>(newRecords));
 			}
 
 			@Override

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

@@ -0,0 +1,32 @@
+/*
+ * Copyright 2021 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 org.springframework.kafka.listener.ContainerProperties.AckMode;
+
+/**
+ * @author Gary Russell
+ * @since 2.8
+ *
+ */
+public class ManualImmediateNackBatchTests extends ManualNackBatchTests {
+
+	static {
+		ackMode = AckMode.MANUAL_IMMEDIATE;
+	}
+
+}

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

@@ -73,6 +73,12 @@ public class ManualNackBatchTests {
 
 	private static final String CONTAINER_ID = "container";
 
+	protected static AckMode ackMode;
+
+	static {
+		ackMode = AckMode.MANUAL;
+	}
+
 	@SuppressWarnings("rawtypes")
 	@Autowired
 	private Consumer consumer;
@@ -216,7 +222,7 @@ public ConcurrentKafkaListenerContainerFactory kafkaListenerContainerFactory() {
 			ConcurrentKafkaListenerContainerFactory factory = new ConcurrentKafkaListenerContainerFactory();
 			factory.setConsumerFactory(consumerFactory());
 			factory.setBatchListener(true);
-			factory.getContainerProperties().setAckMode(AckMode.MANUAL);
+			factory.getContainerProperties().setAckMode(ackMode);
 			return factory;
 		}