spring-projectsGH-2252: Keep offset metadata in case of batch reprocessing

Author: frosiere

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

@@ -376,13 +376,8 @@ public String getListenerId() {
 		return this.beanName; // the container factory sets the bean name to the id attribute
 	}
 
-	/**
-	 * Get arbitrary static information that will be added to the
-	 * {@link KafkaHeaders#LISTENER_INFO} header of all records.
-	 * @return the info.
-	 * @since 2.8.4
-	 */
 	@Nullable
+	@Override
 	public byte[] getListenerInfo() {
 		return this.listenerInfo != null ? Arrays.copyOf(this.listenerInfo, this.listenerInfo.length) : null;
 	}

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

@@ -139,7 +139,8 @@ && isCommitRecovered() && this.kafkaTemplate.isTransactional()) {
 			ConsumerRecord<K, V> skipped = records.get(0);
 			this.kafkaTemplate.sendOffsetsToTransaction(
 					Collections.singletonMap(new TopicPartition(skipped.topic(), skipped.partition()),
-							new OffsetAndMetadata(skipped.offset() + 1)), consumer.groupMetadata());
+							createOffsetAndMetadata(container, skipped.offset() + 1)
+					), consumer.groupMetadata());
 		}
 
 		if (!recoverable && this.backOff != null) {
@@ -165,4 +166,10 @@ public void clearThreadState() {
 		this.lastIntervals.remove();
 	}
 
+	private OffsetAndMetadata createOffsetAndMetadata(@Nullable MessageListenerContainer container, long offset) {
+		if (container == null) {
+			return new OffsetAndMetadata(offset);
+		}
+		return ListenerUtils.createOffsetAndMetadata(container, offset);
+	}
 }

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

@@ -0,0 +1,51 @@
+/*
+ * Copyright 2017-2022 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.lang.Nullable;
+
+/**
+ * Default implementation for {@link ListenerMetadata}.
+ * @author Francois Rosiere
+ * @since 2.8.6
+ */
+public class DefaultListenerMetadata implements ListenerMetadata {
+
+	private final MessageListenerContainer container;
+
+	public DefaultListenerMetadata(MessageListenerContainer container) {
+		this.container = container;
+	}
+
+	@Override
+	@Nullable
+	public String getListenerId() {
+		return this.container.getListenerId();
+	}
+
+	@Override
+	@Nullable
+	public String getGroupId() {
+		return this.container.getGroupId();
+	}
+
+	@Override
+	@Nullable
+	public byte[] getListenerInfo() {
+		return this.container.getListenerInfo();
+	}
+}

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

@@ -138,7 +138,7 @@ private <K, V> ConsumerRecords<K, V> seekOrRecover(Exception thrownException, @N
 		}
 		Map<TopicPartition, OffsetAndMetadata> offsets = new HashMap<>();
 		toCommit.forEach(rec -> offsets.compute(new TopicPartition(rec.topic(), rec.partition()),
-				(key, val) -> new OffsetAndMetadata(rec.offset() + 1)));
+				(key, val) -> ListenerUtils.createOffsetAndMetadata(container, rec.offset() + 1)));
 		if (offsets.size() > 0) {
 			commit(consumer, container, offsets);
 		}
@@ -149,7 +149,7 @@ private <K, V> ConsumerRecords<K, V> seekOrRecover(Exception thrownException, @N
 				ConsumerRecord<?, ?> recovered = remaining.get(0);
 				commit(consumer, container,
 						Collections.singletonMap(new TopicPartition(recovered.topic(), recovered.partition()),
-								new OffsetAndMetadata(recovered.offset() + 1)));
+								ListenerUtils.createOffsetAndMetadata(container, recovered.offset() + 1)));
 				if (remaining.size() > 1) {
 					throw new KafkaException("Seek to current after exception", getLogLevel(), thrownException);
 				}

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

@@ -560,7 +560,7 @@ private final class ListenerConsumer implements SchedulingAwareRunnable, Consume
 				?  (listenerMetadata, offset) -> new OffsetAndMetadata(offset)
 				: this.containerProperties.getOffsetAndMetadataProvider();
 
-		private final ConsumerAwareListenerMetadata consumerAwareListenerMetadata = new ConsumerAwareListenerMetadata();
+		private final ListenerMetadata listenerMetadata = new DefaultListenerMetadata(KafkaMessageListenerContainer.this);
 
 		private final Consumer<K, V> consumer;
 
@@ -3144,32 +3144,7 @@ public String toString() {
 		}
 
 		private OffsetAndMetadata createOffsetAndMetadata(long offset) {
-			return this.offsetAndMetadataProvider.provide(this.consumerAwareListenerMetadata, offset);
-		}
-
-		private final class ConsumerAwareListenerMetadata implements ListenerMetadata {
-
-			ConsumerAwareListenerMetadata() {
-			}
-
-			@Override
-			@Nullable
-			public String getListenerId() {
-				return getBeanName();
-			}
-
-			@Override
-			@Nullable
-			public String getGroupId() {
-				return ListenerConsumer.this.consumerGroupId;
-			}
-
-			@Override
-			@Nullable
-			public byte[] getListenerInfo() {
-				return ListenerConsumer.this.listenerinfo;
-			}
-
+			return this.offsetAndMetadataProvider.provide(this.listenerMetadata, offset);
 		}
 
 		private final class ConsumerAcknowledgment implements Acknowledgment {

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

@@ -22,6 +22,7 @@
 import java.io.ObjectStreamClass;
 
 import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.common.header.Header;
 import org.apache.kafka.common.header.Headers;
 import org.apache.kafka.common.header.internals.RecordHeaders;
@@ -194,5 +195,21 @@ public static void stoppableSleep(MessageListenerContainer container, long inter
 		while (System.currentTimeMillis() < timeout);
 	}
 
+	/**
+	 * Create a new {@link  OffsetAndMetadata} using the given container and offset.
+	 * @param container a container.
+	 * @param offset an offset.
+	 * @return an offset and metadata.
+	 * @since 2.8.6
+	 */
+	public static OffsetAndMetadata createOffsetAndMetadata(MessageListenerContainer container,
+															long offset) {
+		final OffsetAndMetadataProvider metadataProvider = container.getContainerProperties()
+				.getOffsetAndMetadataProvider();
+		if (metadataProvider != null) {
+			return metadataProvider.provide(new DefaultListenerMetadata(container), offset);
+		}
+		return new OffsetAndMetadata(offset);
+	}
 }
 

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

@@ -25,6 +25,7 @@
 
 import org.springframework.beans.factory.DisposableBean;
 import org.springframework.context.SmartLifecycle;
+import org.springframework.kafka.support.KafkaHeaders;
 import org.springframework.lang.Nullable;
 
 /**
@@ -194,6 +195,17 @@ default String getListenerId() {
 		throw new UnsupportedOperationException("This container does not support retrieving the listener id");
 	}
 
+	/**
+	 * Get arbitrary static information that will be added to the
+	 * {@link KafkaHeaders#LISTENER_INFO} header of all records.
+	 * @return the info.
+	 * @since 2.8.4
+	 */
+	@Nullable
+	default byte[] getListenerInfo() {
+		throw new UnsupportedOperationException("This container does not support retrieving the listener info");
+	}
+
 	/**
 	 * If this container has child containers, return true if at least one child is running. If there are not
 	 * child containers, returns {@link #isRunning()}.

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

@@ -19,9 +19,8 @@
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 
 /**
- * Provider for {@link OffsetAndMetadata}. In case of async commits of the offsets,
- * the provider can be used in combination with an {@link  org.apache.kafka.clients.consumer.OffsetCommitCallback} to
- * have more granularity in the way to create an {@link OffsetAndMetadata}.
+ * Provider for {@link OffsetAndMetadata}. The provider can be used to have more granularity when creating an
+ * {@link OffsetAndMetadata}. The provider is used for both sync and async commits of the offsets.
  *
  * @author Francois Rosiere
  * @since 2.8.5

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

@@ -211,7 +211,7 @@ public static void seekOrRecover(Exception thrownException, @Nullable List<Consu
 				ConsumerRecord<?, ?> record = records.get(0);
 				Map<TopicPartition, OffsetAndMetadata> offsetToCommit = Collections.singletonMap(
 						new TopicPartition(record.topic(), record.partition()),
-						new OffsetAndMetadata(record.offset() + 1));
+						ListenerUtils.createOffsetAndMetadata(container, record.offset() + 1));
 				if (container.getContainerProperties().isSyncCommits()) {
 					consumer.commitSync(offsetToCommit, container.getContainerProperties().getSyncCommitTimeout());
 				}

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

@@ -75,6 +75,7 @@ void testClassifier() {
 		Consumer<String, String> consumer = mock(Consumer.class);
 		given(consumer.groupMetadata()).willReturn(new ConsumerGroupMetadata("foo"));
 		MessageListenerContainer container = mock(MessageListenerContainer.class);
+		given(container.getContainerProperties()).willReturn(new ContainerProperties("foo"));
 		processor.process(records, consumer, container, illegalState, true, EOSMode.V2);
 		processor.process(records, consumer, container,
 				new DeserializationException("intended", null, false, illegalState), true, EOSMode.V2);

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2021 the original author or authors.
+ * Copyright 2021-2022 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.
@@ -20,6 +20,7 @@
 import static org.mockito.BDDMockito.given;
 import static org.mockito.Mockito.mock;
 
+import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.junit.jupiter.api.Test;
 
 /**
@@ -38,4 +39,24 @@ void stoppableSleep() throws InterruptedException {
 		assertThat(System.currentTimeMillis() - t1).isGreaterThanOrEqualTo(500);
 	}
 
+	@Test
+	void testCreationOfOffsetAndMetadataWithoutProvider() {
+		final MessageListenerContainer container = mock(MessageListenerContainer.class);
+		given(container.getContainerProperties()).willReturn(new ContainerProperties("foo"));
+		final OffsetAndMetadata offsetAndMetadata = ListenerUtils.createOffsetAndMetadata(container, 1L);
+		assertThat(offsetAndMetadata.offset()).isEqualTo(1);
+		assertThat(offsetAndMetadata.metadata()).isEmpty();
+	}
+
+	@Test
+	void testCreationOfOffsetAndMetadataWithProvider() {
+		final MessageListenerContainer container = mock(MessageListenerContainer.class);
+		final ContainerProperties properties = new ContainerProperties("foo");
+		properties.setOffsetAndMetadataProvider((listenerMetadata, offset) -> new OffsetAndMetadata(offset, "my-metadata"));
+		given(container.getContainerProperties()).willReturn(properties);
+		final OffsetAndMetadata offsetAndMetadata = ListenerUtils.createOffsetAndMetadata(container, 1L);
+		assertThat(offsetAndMetadata.offset()).isEqualTo(1);
+		assertThat(offsetAndMetadata.metadata()).isEqualTo("my-metadata");
+	}
 }
+