GH-1234: Fix FailedRecordTracker

Fixes #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.

Author: garyrussell

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

@@ -16,9 +16,12 @@
 
 package org.springframework.kafka.listener;
 
+import java.util.HashMap;
+import java.util.Map;
 import java.util.function.BiConsumer;
 
 import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.common.TopicPartition;
 
 import org.springframework.core.log.LogAccessor;
 import org.springframework.lang.Nullable;
@@ -35,7 +38,7 @@
  */
 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;
 
@@ -50,12 +53,18 @@ class FailedRecordTracker {
 
 		Assert.notNull(backOff, "'backOff' cannot be null");
 		if (recoverer == null) {
-			FailedRecord failedRecord = this.failures.get();
-			this.recoverer = (rec, thr) -> logger.error(thr, "Backoff "
-				+ (failedRecord == null
-					? "none"
-					: failedRecord.getBackOffExecution())
-				+ " exhausted for " + rec);
+			this.recoverer = (rec, thr) -> {
+				Map<TopicPartition, FailedRecord> map = this.failures.get();
+				FailedRecord failedRecord = null;
+				if (map != null) {
+					failedRecord = map.get(new TopicPartition(rec.topic(), rec.partition()));
+				}
+				logger.error(thr, "Backoff "
+					+ (failedRecord == null
+						? "none"
+						: failedRecord.getBackOffExecution())
+					+ " exhausted for " + rec);
+			};
 		}
 		else {
 			this.recoverer = recoverer;
@@ -70,10 +79,16 @@ boolean skip(ConsumerRecord<?, ?> record, Exception exception) {
 			recover(record, exception);
 			return true;
 		}
-		FailedRecord failedRecord = this.failures.get();
-		if (failedRecord == null || newFailure(record, failedRecord)) {
-			failedRecord = new FailedRecord(record.topic(), record.partition(), record.offset(), this.backOff.start());
-			this.failures.set(failedRecord);
+		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(), this.backOff.start());
+			map.put(topicPartition, failedRecord);
 		}
 		long nextBackOff = failedRecord.getBackOffExecution().nextBackOff();
 		if (nextBackOff != BackOffExecution.STOP) {
@@ -87,7 +102,10 @@ boolean skip(ConsumerRecord<?, ?> record, Exception exception) {
 		}
 		else {
 			recover(record, exception);
-			this.failures.remove();
+			map.remove(topicPartition);
+			if (map.isEmpty()) {
+				this.failures.remove();
+			}
 			return true;
 		}
 	}
@@ -101,12 +119,6 @@ private void recover(ConsumerRecord<?, ?> record, Exception exception) {
 		}
 	}
 
-	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();
 	}
@@ -117,29 +129,15 @@ void clearThreadState() {
 
 	private static final class FailedRecord {
 
-		private final String topic;
-
-		private final int partition;
-
 		private final long offset;
 
 		private final BackOffExecution backOffExecution;
 
-		FailedRecord(String topic, int partition, long offset, BackOffExecution backOffExecution) {
-			this.topic = topic;
-			this.partition = partition;
+		FailedRecord(long offset, BackOffExecution backOffExecution) {
 			this.offset = offset;
 			this.backOffExecution = backOffExecution;
 		}
 
-		String getTopic() {
-			return this.topic;
-		}
-
-		int getPartition() {
-			return this.partition;
-		}
-
 		long getOffset() {
 			return this.offset;
 		}

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

@@ -19,6 +19,8 @@
 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.kafka.clients.consumer.ConsumerRecord;
@@ -74,4 +76,21 @@ record = new ConsumerRecord<>("bar", 1, 1L, "bar", "baz");
 		assertThat(tracker.skip(record, new RuntimeException())).isTrue();
 	}
 
+	@Test
+	public void testDifferentOrder() {
+		List<ConsumerRecord<?, ?>> records = new ArrayList<>();
+		FailedRecordTracker tracker = new FailedRecordTracker((rec, ex) -> {
+			records.add(rec);
+		}, new FixedBackOff(0L, 2L), mock(LogAccessor.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);
+	}
+
 }