refactor: use ConsumerRecords#nextOffsets in fixTxOffsets

Signed-off-by: Su Ko <rhtn1128@gmail.com>

Author: bandalgomsu

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

@@ -172,6 +172,7 @@
  * @author Christian Fredriksson
  * @author Timofey Barabanov
  * @author Janek Lasocki-Biczysko
+ * @author Su Ko
  */
 public class KafkaMessageListenerContainer<K, V> // NOSONAR line count
 		extends AbstractMessageListenerContainer<K, V> implements ConsumerPauseResumeEventPublisher {
@@ -843,6 +844,8 @@ private final class ListenerConsumer implements SchedulingAwareRunnable, Consume
 
 		private @Nullable ConsumerRecords<K, V> remainingRecords;
 
+		private @Nullable ConsumerRecords<K, V> lastRecords;
+
 		private boolean pauseForPending;
 
 		private boolean firstPoll;
@@ -1532,7 +1535,7 @@ private void doProcessCommits() {
 		private void invokeIfHaveRecords(@Nullable ConsumerRecords<K, V> records) {
 			if (records != null && records.count() > 0) {
 				this.receivedSome = true;
-				savePositionsIfNeeded(records);
+				saveLastRecordsIfNeeded(records);
 				notIdle();
 				notIdlePartitions(records.partitions());
 				invokeListener(records);
@@ -1604,50 +1607,59 @@ private void notIdle() {
 			}
 		}
 
-		private void savePositionsIfNeeded(ConsumerRecords<K, V> records) {
+		private void saveLastRecordsIfNeeded(ConsumerRecords<K, V> records) {
 			if (this.fixTxOffsets) {
-				this.savedPositions.clear();
-				records.partitions().forEach(tp -> this.savedPositions.put(tp, this.consumer.position(tp)));
+				this.lastRecords = records;
 			}
 		}
 
 		private void fixTxOffsetsIfNeeded() {
-			if (this.fixTxOffsets) {
-				try {
-					Map<TopicPartition, OffsetAndMetadata> toFix = new HashMap<>();
-					this.lastCommits.forEach((tp, oamd) -> {
-						long position = this.consumer.position(tp);
-						Long saved = this.savedPositions.get(tp);
-						if (saved != null && saved != position) {
-							this.logger.debug(() -> "Skipping TX offset correction - seek(s) have been performed; "
-									+ "saved: " + this.savedPositions + ", "
-									+ "committed: " + oamd + ", "
-									+ "current: " + tp + "@" + position);
-							return;
-						}
-						if (position > oamd.offset()) {
-							toFix.put(tp, createOffsetAndMetadata(position));
-						}
-					});
-					if (!toFix.isEmpty()) {
-						this.logger.debug(() -> "Fixing TX offsets: " + toFix);
-						if (this.kafkaTxManager == null) {
-							commitOffsets(toFix);
-						}
-						else {
-							Objects.requireNonNull(this.transactionTemplate).executeWithoutResult(
-									status -> doSendOffsets(getTxProducer(), toFix));
-						}
-					}
-				}
-				catch (Exception e) {
-					this.logger.error(e, () -> "Failed to correct transactional offset(s): "
-							+ ListenerConsumer.this.lastCommits);
+			if (!this.fixTxOffsets) {
+				return;
+			}
+
+			try {
+				if (this.lastRecords == null) {
+					this.logger.trace(() -> "No previous records available for TX offset fix.");
+					return;
 				}
-				finally {
-					ListenerConsumer.this.lastCommits.clear();
+
+				Map<TopicPartition, OffsetAndMetadata> nextOffsets = this.lastRecords.nextOffsets();
+				Map<TopicPartition, OffsetAndMetadata> toFix = new HashMap<>();
+
+				nextOffsets.forEach((tp, nextOffset) -> {
+					OffsetAndMetadata committed = this.lastCommits.get(tp);
+
+					if (committed == null) {
+						this.logger.debug(() -> "No committed offset for " + tp + "; skipping TX offset fix.");
+						return;
+					}
+
+					if (!this.lastRecords.records(tp).isEmpty() && nextOffset.offset() > committed.offset()) {
+						toFix.put(tp, nextOffset);
+					}
+				});
+
+				if (!toFix.isEmpty()) {
+					this.logger.debug(() ->
+							String.format("Fixing TX offsets for %d partitions: %s", toFix.size(), toFix));
+					if (this.kafkaTxManager == null) {
+						commitOffsets(toFix);
+					}
+					else {
+						Objects.requireNonNull(this.transactionTemplate)
+								.executeWithoutResult(status -> doSendOffsets(getTxProducer(), toFix));
+					}
 				}
 			}
+			catch (Exception e) {
+				this.logger.error(e, () -> "Failed to correct transactional offset(s): "
+						+ ListenerConsumer.this.lastCommits);
+			}
+			finally {
+				ListenerConsumer.this.lastCommits.clear();
+				this.lastRecords = null;
+			}
 		}
 
 		private @Nullable ConsumerRecords<K, V> doPoll() {