GH-2170: Support Custom OffsetAndMetadata

Resolves #2170

Add a way to create a custom OffsetAndMetadata

GH-2170: Remove useless initializations

Author: frosiere

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

@@ -22,6 +22,7 @@
 import java.util.regex.Pattern;
 
 import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
+import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.clients.consumer.OffsetCommitCallback;
 
 import org.springframework.kafka.support.LogIfLevelEnabled;
@@ -89,6 +90,12 @@ public class ConsumerProperties {
 	 */
 	private OffsetCommitCallback commitCallback;
 
+	/**
+	 * A provider for {@link OffsetAndMetadata}; by default, the provider creates an offset and metadata with
+	 * empty metadata. The provider gives a way to customize the metadata.
+	 */
+	private OffsetAndMetadataProvider offsetAndMetadataProvider;
+
 	/**
 	 * Whether or not to call consumer.commitSync() or commitAsync() when the
 	 * container is responsible for commits. Default true.
@@ -275,6 +282,16 @@ public void setCommitCallback(OffsetCommitCallback commitCallback) {
 		this.commitCallback = commitCallback;
 	}
 
+	/**
+	 * Set the offset and metadata provider associated to a commit callback.
+	 * @param offsetAndMetadataProvider an offset and metadata provider.
+	 * @since 2.8.5
+	 * @see #setCommitCallback(OffsetCommitCallback)
+	 */
+	public void setOffsetAndMetadataProvider(OffsetAndMetadataProvider offsetAndMetadataProvider) {
+		this.offsetAndMetadataProvider = offsetAndMetadataProvider;
+	}
+
 	/**
 	 * Return the commit callback.
 	 * @return the callback.
@@ -284,6 +301,15 @@ public OffsetCommitCallback getCommitCallback() {
 		return this.commitCallback;
 	}
 
+	/**
+	 * Return the offset and metadata provider.
+	 * @return the offset and metadata provider.
+	 */
+	@Nullable
+	public OffsetAndMetadataProvider getOffsetAndMetadataProvider() {
+		return this.offsetAndMetadataProvider;
+	}
+
 	/**
 	 * Set whether or not to call consumer.commitSync() or commitAsync() when the
 	 * container is responsible for commits. Default true.
@@ -491,6 +517,7 @@ protected final String renderProperties() {
 						? "\n consumerRebalanceListener=" + this.consumerRebalanceListener
 						: "")
 				+ (this.commitCallback != null ? "\n commitCallback=" + this.commitCallback : "")
+				+ (this.offsetAndMetadataProvider != null ? "\n offsetAndMetadataProvider=" + this.offsetAndMetadataProvider : "")
 				+ "\n syncCommits=" + this.syncCommits
 				+ (this.syncCommitTimeout != null ? "\n syncCommitTimeout=" + this.syncCommitTimeout : "")
 				+ (this.kafkaConsumerProperties.size() > 0 ? "\n properties=" + this.kafkaConsumerProperties : "")

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

@@ -557,6 +557,12 @@ private final class ListenerConsumer implements SchedulingAwareRunnable, Consume
 				? this.containerProperties.getCommitCallback()
 				: new LoggingCommitCallback();
 
+		private final OffsetAndMetadataProvider offsetAndMetadataProvider = this.containerProperties.getOffsetAndMetadataProvider() == null
+				?  (listenerMetadata, offset) -> new OffsetAndMetadata(offset)
+				: this.containerProperties.getOffsetAndMetadataProvider();
+
+		private final ConsumerAwareListenerMetadata consumerAwareListenerMetadata = new ConsumerAwareListenerMetadata();
+
 		private final Consumer<K, V> consumer;
 
 		private final Map<String, Map<Integer, Long>> offsets = new HashMap<>();
@@ -1438,7 +1444,7 @@ private void fixTxOffsetsIfNeeded() {
 							return;
 						}
 						if (position > oamd.offset()) {
-							toFix.put(tp, new OffsetAndMetadata(position));
+							toFix.put(tp, createOffsetAndMetadata(position));
 						}
 					});
 					if (toFix.size() > 0) {
@@ -1910,7 +1916,7 @@ else if (record.offset() < offs.get(0)) {
 		private void ackImmediate(ConsumerRecord<K, V> record) {
 			Map<TopicPartition, OffsetAndMetadata> commits = Collections.singletonMap(
 					new TopicPartition(record.topic(), record.partition()),
-					new OffsetAndMetadata(record.offset() + 1));
+					createOffsetAndMetadata(record.offset() + 1));
 			this.commitLogger.log(() -> COMMITTING + commits);
 			if (this.producer != null) {
 				doSendOffsets(this.producer, commits);
@@ -1926,9 +1932,8 @@ 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));
+				commits.put(part, createOffsetAndMetadata(records.records(part)
+						.get(records.records(part).size() - 1).offset() + 1));
 			}
 			this.commitLogger.log(() -> COMMITTING + commits);
 			if (this.producer != null) {
@@ -2694,7 +2699,7 @@ public void ackCurrent(final ConsumerRecord<K, V> record) {
 			if (this.isRecordAck) {
 				Map<TopicPartition, OffsetAndMetadata> offsetsToCommit =
 						Collections.singletonMap(new TopicPartition(record.topic(), record.partition()),
-								new OffsetAndMetadata(record.offset() + 1));
+								createOffsetAndMetadata(record.offset() + 1));
 				if (this.producer == null) {
 					this.commitLogger.log(() -> COMMITTING + offsetsToCommit);
 					if (this.syncCommits) {
@@ -2996,7 +3001,7 @@ private Map<TopicPartition, OffsetAndMetadata> buildCommits() {
 			for (Entry<String, Map<Integer, Long>> entry : this.offsets.entrySet()) {
 				for (Entry<Integer, Long> offset : entry.getValue().entrySet()) {
 					commits.put(new TopicPartition(entry.getKey(), offset.getKey()),
-							new OffsetAndMetadata(offset.getValue() + 1));
+							createOffsetAndMetadata(offset.getValue() + 1));
 				}
 			}
 			this.offsets.clear();
@@ -3079,6 +3084,29 @@ public String toString() {
 					+ "\n]";
 		}
 
+		private OffsetAndMetadata createOffsetAndMetadata(long offset) {
+			return this.offsetAndMetadataProvider.provide(this.consumerAwareListenerMetadata, offset);
+		}
+
+		private final class ConsumerAwareListenerMetadata implements ListenerMetadata {
+
+			@Override
+			public String getListenerId() {
+				return getBeanName();
+			}
+
+			@Override
+			public String getGroupId() {
+				return ListenerConsumer.this.consumerGroupId;
+			}
+
+			@Override
+			public byte[] getListenerInfo() {
+				return ListenerConsumer.this.listenerinfo;
+			}
+
+		}
+
 		private final class ConsumerAcknowledgment implements Acknowledgment {
 
 			private final ConsumerRecord<K, V> record;
@@ -3272,8 +3300,7 @@ private boolean collectAndCommitIfNecessary(Collection<TopicPartition> partition
 				for (TopicPartition partition : partitions) {
 					try {
 						if (committed.get(partition) == null) { // no existing commit for this group
-							offsetsToCommit.put(partition,
-									new OffsetAndMetadata(ListenerConsumer.this.consumer.position(partition)));
+							offsetsToCommit.put(partition, createOffsetAndMetadata(ListenerConsumer.this.consumer.position(partition)));
 						}
 					}
 					catch (NoOffsetForPartitionException e) {

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

@@ -0,0 +1,45 @@
+/*
+ * Copyright 2016-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;
+
+/**
+ * Metadata associated to a {@link org.springframework.kafka.annotation.KafkaListener}.
+ *
+ * @author Francois Rosiere
+ * @since 2.8.5
+ * @see org.springframework.kafka.annotation.KafkaListener
+ */
+public interface ListenerMetadata {
+
+	/**
+	 * Return the listener id.
+	 * @return the listener id.
+	 */
+	String getListenerId();
+
+	/**
+	 * Return the group id.
+	 * @return the group id.
+	 */
+	String getGroupId();
+
+	/**
+	 * Return the listener info.
+	 * @return the listener info.
+	 */
+	byte[] getListenerInfo();
+}

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

@@ -0,0 +1,40 @@
+/*
+ * Copyright 2016-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.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}.
+ *
+ * @author Francois Rosiere
+ * @since 2.8.5
+ * @see org.apache.kafka.clients.consumer.OffsetCommitCallback
+ */
+public interface OffsetAndMetadataProvider {
+
+	/**
+	 * Provide an offset and metadata object for the given listener metadata and offset.
+	 *
+	 * @param listenerMetadata metadata associated to a listener.
+	 * @param offset an offset.
+	 * @return an offset and metadata.
+	 */
+	OffsetAndMetadata provide(ListenerMetadata listenerMetadata, long offset);
+}

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

@@ -3844,6 +3844,54 @@ public void clearThreadState(Consumer<?, ?> consumer) {
 		container.stop();
 	}
 
+	@Test
+	public void testOffsetAndMetadataWithoutProvider() throws InterruptedException {
+		testOffsetAndMetadata(null, new OffsetAndMetadata(1));
+	}
+
+	@Test
+	public void testOffsetAndMetadataWithProvider() throws InterruptedException {
+		testOffsetAndMetadata((listenerMetadata, offset) ->
+				new OffsetAndMetadata(offset, listenerMetadata.getGroupId()),
+				new OffsetAndMetadata(1, "grp"));
+	}
+
+	@SuppressWarnings("unchecked")
+	private void testOffsetAndMetadata(OffsetAndMetadataProvider provider, OffsetAndMetadata expectedOffsetAndMetadata) throws InterruptedException {
+		final ConsumerFactory<Integer, String> cf = mock(ConsumerFactory.class);
+		final Consumer<Integer, String> consumer = mock(Consumer.class);
+		given(cf.createConsumer(eq("grp"), eq("clientId"), isNull(), any())).willReturn(consumer);
+		given(consumer.poll(any(Duration.class))).willAnswer(i -> new ConsumerRecords<>(
+				Map.of(
+						new TopicPartition("foo", 0),
+						Collections.singletonList(new ConsumerRecord<>("foo", 0, 0L, 1, "foo"))
+				)
+		));
+		final ArgumentCaptor<Map<TopicPartition, OffsetAndMetadata>> offsetsCaptor = ArgumentCaptor.forClass(Map.class);
+		final CountDownLatch latch = new CountDownLatch(1);
+		willAnswer(invocation -> {
+			latch.countDown();
+			return null;
+		}).given(consumer).commitAsync(offsetsCaptor.capture(), any());
+		final ContainerProperties containerProps = new ContainerProperties(new TopicPartitionOffset("foo", 0));
+		containerProps.setGroupId("grp");
+		containerProps.setClientId("clientId");
+		containerProps.setSyncCommits(false);
+		containerProps.setMessageListener((MessageListener<Integer, String>) data -> {
+		});
+		containerProps.setCommitCallback((offsets, exception) -> {
+		});
+		containerProps.setOffsetAndMetadataProvider(provider);
+		final KafkaMessageListenerContainer<Integer, String> container =
+				new KafkaMessageListenerContainer<>(cf, containerProps);
+		container.start();
+		assertThat(latch.await(10, TimeUnit.SECONDS)).isTrue();
+		assertThat(offsetsCaptor.getValue())
+				.hasSize(1)
+				.containsValue(expectedOffsetAndMetadata);
+		container.stop();
+	}
+
 	private Consumer<?, ?> spyOnConsumer(KafkaMessageListenerContainer<Integer, String> container) {
 		Consumer<?, ?> consumer =
 				KafkaTestUtils.getPropertyValue(container, "listenerConsumer.consumer", Consumer.class);