From 2490dd5b7ec0b6da151871f66a6ee302a3bd32ff Mon Sep 17 00:00:00 2001
From: Gary Russell <grussell@pivotal.io>
Date: Mon, 16 Sep 2019 16:36:41 -0400
Subject: [PATCH] GH-1234: Fix FailedRecordTracker [2.2.x]

Fixes https://github.com/spring-projects/spring-kafka/issues/1234

Previous implementation assumed the first redelivered record would be
the one that failed; this is not always the case.

Maintain state for each topic/partition/offset.
---
 .../kafka/listener/FailedRecordTracker.java   | 50 ++++++++-----------
 .../listener/FailedRecordTrackerTests.java    | 19 +++++++
 2 files changed, 41 insertions(+), 28 deletions(-)

diff --git a/spring-kafka/src/main/java/org/springframework/kafka/listener/FailedRecordTracker.java b/spring-kafka/src/main/java/org/springframework/kafka/listener/FailedRecordTracker.java
index 5e00e31a..80b33a26 100644
--- a/spring-kafka/src/main/java/org/springframework/kafka/listener/FailedRecordTracker.java
+++ b/spring-kafka/src/main/java/org/springframework/kafka/listener/FailedRecordTracker.java
@@ -17,10 +17,13 @@
 package org.springframework.kafka.listener;
 
 import java.time.temporal.ValueRange;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.function.BiConsumer;
 
 import org.apache.commons.logging.Log;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.common.TopicPartition;
 
 import org.springframework.lang.Nullable;
 
@@ -33,7 +36,7 @@ import org.springframework.lang.Nullable;
  */
 class FailedRecordTracker {
 
-	private final ThreadLocal<FailedRecord> failures = new ThreadLocal<>(); // intentionally not static
+	private final ThreadLocal<Map<TopicPartition, FailedRecord>> failures = new ThreadLocal<>(); // intentionally not static
 
 	private final BiConsumer<ConsumerRecord<?, ?>, Exception> recoverer;
 
@@ -60,14 +63,25 @@ class FailedRecordTracker {
 			recover(record, exception);
 			return true;
 		}
-		FailedRecord failedRecord = this.failures.get();
-		if (this.maxFailures > 0 && (failedRecord == null || newFailure(record, failedRecord))) {
-			this.failures.set(new FailedRecord(record.topic(), record.partition(), record.offset()));
+		Map<TopicPartition, FailedRecord> map = this.failures.get();
+		if (map == null) {
+			this.failures.set(new HashMap<>());
+			map = this.failures.get();
+		}
+		TopicPartition topicPartition = new TopicPartition(record.topic(), record.partition());
+		FailedRecord failedRecord = map.get(topicPartition);
+		if (failedRecord == null || failedRecord.getOffset() != record.offset()) {
+			failedRecord = new FailedRecord(record.offset());
+			map.put(topicPartition, failedRecord);
 			return false;
 		}
 		else if (this.maxFailures > 0 && failedRecord.incrementAndGet() >= this.maxFailures) {
-				recover(record, exception);
-				return true;
+			recover(record, exception);
+			map.remove(topicPartition);
+			if (map.isEmpty()) {
+				this.failures.remove();
+			}
+			return true;
 		}
 		else {
 			return false;
@@ -83,42 +97,22 @@ class FailedRecordTracker {
 		}
 	}
 
-	private boolean newFailure(ConsumerRecord<?, ?> record, FailedRecord failedRecord) {
-		return !failedRecord.getTopic().equals(record.topic())
-				|| failedRecord.getPartition() != record.partition()
-				|| failedRecord.getOffset() != record.offset();
-	}
-
 	void clearThreadState() {
 		this.failures.remove();
 	}
 
 	private static final class FailedRecord {
 
-		private final String topic;
-
-		private final int partition;
-
 		private final long offset;
 
 		private int count;
 
-		FailedRecord(String topic, int partition, long offset) {
-			this.topic = topic;
-			this.partition = partition;
+		FailedRecord(long offset) {
 			this.offset = offset;
 			this.count = 1;
 		}
 
-		private String getTopic() {
-			return this.topic;
-		}
-
-		private int getPartition() {
-			return this.partition;
-		}
-
-		private long getOffset() {
+		long getOffset() {
 			return this.offset;
 		}
 
diff --git a/spring-kafka/src/test/java/org/springframework/kafka/listener/FailedRecordTrackerTests.java b/spring-kafka/src/test/java/org/springframework/kafka/listener/FailedRecordTrackerTests.java
index e85c6492..af0e42b5 100644
--- a/spring-kafka/src/test/java/org/springframework/kafka/listener/FailedRecordTrackerTests.java
+++ b/spring-kafka/src/test/java/org/springframework/kafka/listener/FailedRecordTrackerTests.java
@@ -19,6 +19,8 @@ package org.springframework.kafka.listener;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.Mockito.mock;
 
+import java.util.ArrayList;
+import java.util.List;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 import org.apache.commons.logging.Log;
@@ -57,4 +59,21 @@ public class FailedRecordTrackerTests {
 		assertThat(recovered.get()).isTrue();
 	}
 
+	@Test
+	public void testDifferentOrder() {
+		List<ConsumerRecord<?, ?>> records = new ArrayList<>();
+		FailedRecordTracker tracker = new FailedRecordTracker((rec, ex) -> {
+			records.add(rec);
+		}, 3, mock(Log.class));
+		ConsumerRecord<?, ?> record1 = new ConsumerRecord<>("foo", 0, 0L, "bar", "baz");
+		ConsumerRecord<?, ?> record2 = new ConsumerRecord<>("foo", 1, 0L, "bar", "baz");
+		assertThat(tracker.skip(record1, new RuntimeException())).isFalse();
+		assertThat(tracker.skip(record2, new RuntimeException())).isFalse();
+		assertThat(tracker.skip(record1, new RuntimeException())).isFalse();
+		assertThat(tracker.skip(record2, new RuntimeException())).isFalse();
+		assertThat(tracker.skip(record1, new RuntimeException())).isTrue();
+		assertThat(tracker.skip(record2, new RuntimeException())).isTrue();
+		assertThat(records).hasSize(2);
+	}
+
 }