GH-3178: PartitionOffset support for SeekPosition

* Fixes #3178

* add support for custom SeekPosition via `@PartitionOffset`
* align `@PartitionOffset` to `TopicPartitionOffset`
* add unit test for `@PartitionOffset.SeekPosition`
* add unit test for SpEL partitions to Integer[] and Integer
* address PR review.

Author: Wzy19930507

spring-kafka-docs/src/main/antora/modules/ROOT/pages/kafka/receiving-messages/listener-annotation.adoc

@@ -150,6 +150,25 @@ public void listen(ConsumerRecord<?, ?> record) {
 
 The initial offset will be applied to all 6 partitions.
 
+Since 3.2, `@PartitionOffset` support `SeekPosition.END`, `SeekPosition.BEGINNING`, `SeekPosition.TIMESTAMP`, `seekPosition` match `SeekPosition` enum name:
+
+[source, java]
+----
+@KafkaListener(id = "seekPositionTime", topicPartitions = {
+        @TopicPartition(topic = TOPIC_SEEK_POSITION, partitionOffsets = {
+                @PartitionOffset(partition = "0", initialOffset = "723916800000", seekPosition = "TIMESTAMP"),
+                @PartitionOffset(partition = "1", initialOffset = "0", seekPosition = "BEGINNING"),
+                @PartitionOffset(partition = "2", initialOffset = "0", seekPosition = "END")
+        })
+})
+public void listen(ConsumerRecord<?, ?> record) {
+    ...
+}
+----
+
+If seekPosition set `END` or `BEGINNING` will ignore `initialOffset` and `relativeToCurrent`.
+If seekPosition set `TIMESTAMP`, `initialOffset` means timestamp.
+
 [[manual-acknowledgment]]
 == Manual Acknowledgment
 

spring-kafka-docs/src/main/antora/modules/ROOT/pages/whats-new.adoc

@@ -75,6 +75,11 @@ Provides new class `EndpointHandlerMultiMethod` to handler multi method for retr
 `ConsumerCallback` provides a new API to seek to an offset based on a user-defined function, which takes the current offset in the consumer as an argument.
 See xref:kafka/seek.adoc#seek[Seek API Docs] for more details.
 
+[[x32-annotation-partition-offset-seek-position]]
+=== @PartitionOffset support for SeekPosition
+Adding `seekPosition` property to `@PartitionOffset` support for `TopicPartitionOffset.SeekPosition`.
+See xref:kafka/receiving-messages/listener-annotation.adoc#manual-assignment[manual-assignment] for more details.
+
 [[x32-topic-partition-offset-constructor]]
 === New constructor in TopicPartitionOffset that accepts a function to compute the offset to seek to
 `TopicPartitionOffset` has a new constructor that takes a user-provided function to compute the offset to seek to.

spring-kafka/src/main/java/org/springframework/kafka/annotation/KafkaListenerAnnotationBeanPostProcessor.java

@@ -97,6 +97,7 @@
 import org.springframework.kafka.retrytopic.RetryTopicConfigurer;
 import org.springframework.kafka.retrytopic.RetryTopicSchedulerWrapper;
 import org.springframework.kafka.support.TopicPartitionOffset;
+import org.springframework.kafka.support.TopicPartitionOffset.SeekPosition;
 import org.springframework.lang.Nullable;
 import org.springframework.messaging.converter.GenericMessageConverter;
 import org.springframework.messaging.converter.SmartMessageConverter;
@@ -827,10 +828,8 @@ private String getEndpointGroupId(KafkaListener kafkaListener, @Nullable String
 	private TopicPartitionOffset[] resolveTopicPartitions(KafkaListener kafkaListener) {
 		TopicPartition[] topicPartitions = kafkaListener.topicPartitions();
 		List<TopicPartitionOffset> result = new ArrayList<>();
-		if (topicPartitions.length > 0) {
-			for (TopicPartition topicPartition : topicPartitions) {
-				result.addAll(resolveTopicPartitionsList(topicPartition));
-			}
+		for (TopicPartition topicPartition : topicPartitions) {
+			result.addAll(resolveTopicPartitionsList(topicPartition));
 		}
 		return result.toArray(new TopicPartitionOffset[0]);
 	}
@@ -877,7 +876,7 @@ private List<TopicPartitionOffset> resolveTopicPartitionsList(TopicPartition top
 				() -> "At least one 'partition' or 'partitionOffset' required in @TopicPartition for topic '" + topic + "'");
 		List<TopicPartitionOffset> result = new ArrayList<>();
 		for (String partition : partitions) {
-			resolvePartitionAsInteger((String) topic, resolveExpression(partition), result, null, false, false);
+			resolvePartitionAsInteger((String) topic, resolveExpression(partition), result);
 		}
 		if (partitionOffsets.length == 1 && resolveExpression(partitionOffsets[0].partition()).equals("*")) {
 			result.forEach(tpo -> {
@@ -890,7 +889,8 @@ private List<TopicPartitionOffset> resolveTopicPartitionsList(TopicPartition top
 				Assert.isTrue(!partitionOffset.partition().equals("*"), () ->
 						"Partition wildcard '*' is only allowed in a single @PartitionOffset in " + result);
 				resolvePartitionAsInteger((String) topic, resolveExpression(partitionOffset.partition()), result,
-						resolveInitialOffset(topic, partitionOffset), isRelative(topic, partitionOffset), true);
+						resolveInitialOffset(topic, partitionOffset), isRelative(topic, partitionOffset), true,
+						resolveExpression(partitionOffset.seekPosition()));
 			}
 		}
 		Assert.isTrue(!result.isEmpty(), () -> "At least one partition required for " + topic);
@@ -899,11 +899,11 @@ private List<TopicPartitionOffset> resolveTopicPartitionsList(TopicPartition top
 
 	private Long resolveInitialOffset(Object topic, PartitionOffset partitionOffset) {
 		Object initialOffsetValue = resolveExpression(partitionOffset.initialOffset());
-		Long initialOffset;
+		long initialOffset;
 		if (initialOffsetValue instanceof String str) {
 			Assert.state(StringUtils.hasText(str),
 					() -> "'initialOffset' in @PartitionOffset for topic '" + topic + "' cannot be empty");
-			initialOffset = Long.valueOf(str);
+			initialOffset = Long.parseLong(str);
 		}
 		else if (initialOffsetValue instanceof Long lng) {
 			initialOffset = lng;
@@ -954,20 +954,33 @@ else if (resolvedValue instanceof Iterable) {
 		}
 	}
 
+	private void resolvePartitionAsInteger(String topic, Object resolvedValue, List<TopicPartitionOffset> result) {
+		resolvePartitionAsInteger(topic, resolvedValue, result, null, false, false, null);
+	}
+
 	@SuppressWarnings(UNCHECKED)
-	private void resolvePartitionAsInteger(String topic, Object resolvedValue,
-			List<TopicPartitionOffset> result, @Nullable Long offset, boolean isRelative, boolean checkDups) {
+	private void resolvePartitionAsInteger(String topic, Object resolvedValue, List<TopicPartitionOffset> result,
+			@Nullable Long offset, boolean isRelative, boolean checkDups, @Nullable Object seekPosition) {
 
 		if (resolvedValue instanceof String[] strArr) {
 			for (Object object : strArr) {
-				resolvePartitionAsInteger(topic, object, result, offset, isRelative, checkDups);
+				resolvePartitionAsInteger(topic, object, result, offset, isRelative, checkDups, seekPosition);
 			}
+			return;
 		}
-		else if (resolvedValue instanceof String str) {
+		else if (resolvedValue instanceof Iterable) {
+			for (Object object : (Iterable<Object>) resolvedValue) {
+				resolvePartitionAsInteger(topic, object, result, offset, isRelative, checkDups, seekPosition);
+			}
+			return;
+		}
+
+		TopicPartitionOffset.SeekPosition tpoSp = resloveTopicPartitionOffsetSeekPosition(seekPosition);
+		if (resolvedValue instanceof String str) {
 			Assert.state(StringUtils.hasText(str),
 					() -> "partition in @TopicPartition for topic '" + topic + "' cannot be empty");
 			List<TopicPartitionOffset> collected = parsePartitions(str)
-					.map(part -> new TopicPartitionOffset(topic, part, offset, isRelative))
+					.map(part -> createTopicPartitionOffset(topic, part, offset, isRelative, tpoSp))
 					.toList();
 			if (checkDups) {
 				collected.forEach(tpo -> {
@@ -980,23 +993,47 @@ else if (resolvedValue instanceof String str) {
 		}
 		else if (resolvedValue instanceof Integer[] intArr) {
 			for (Integer partition : intArr) {
-				result.add(new TopicPartitionOffset(topic, partition));
+				result.add(createTopicPartitionOffset(topic, partition, offset, isRelative, tpoSp));
 			}
 		}
 		else if (resolvedValue instanceof Integer intgr) {
-			result.add(new TopicPartitionOffset(topic, intgr));
-		}
-		else if (resolvedValue instanceof Iterable) {
-			for (Object object : (Iterable<Object>) resolvedValue) {
-				resolvePartitionAsInteger(topic, object, result, offset, isRelative, checkDups);
-			}
+			result.add(createTopicPartitionOffset(topic, intgr, offset, isRelative, tpoSp));
 		}
 		else {
 			throw new IllegalArgumentException(String.format(
 					"@KafKaListener for topic '%s' can't resolve '%s' as an Integer or String", topic, resolvedValue));
 		}
 	}
 
+	@Nullable
+	private TopicPartitionOffset.SeekPosition resloveTopicPartitionOffsetSeekPosition(@Nullable Object seekPosition) {
+		TopicPartitionOffset.SeekPosition resloveTpoSp = null;
+		if (seekPosition instanceof String seekPositionName) {
+			String capitalLetterSeekPositionName = seekPositionName.trim().toUpperCase();
+			if (SeekPosition.BEGINNING.name().equals(capitalLetterSeekPositionName)) {
+				resloveTpoSp = SeekPosition.BEGINNING;
+			}
+			else if (SeekPosition.END.name().equals(capitalLetterSeekPositionName)) {
+				resloveTpoSp = SeekPosition.END;
+			}
+			else if (SeekPosition.TIMESTAMP.name().equals(capitalLetterSeekPositionName)) {
+				resloveTpoSp = SeekPosition.TIMESTAMP;
+			}
+		}
+		return resloveTpoSp;
+	}
+
+	private TopicPartitionOffset createTopicPartitionOffset(String topic, int partition, @Nullable Long offset,
+			boolean isRelative, @Nullable SeekPosition seekPosition) {
+
+		if (seekPosition != null) {
+			return new TopicPartitionOffset(topic, partition, offset, seekPosition);
+		}
+		else {
+			return new TopicPartitionOffset(topic, partition, offset, isRelative);
+		}
+	}
+
 	private String resolveExpressionAsString(String value, String attribute) {
 		Object resolved = resolveExpression(value);
 		if (resolved instanceof String str) {

spring-kafka/src/main/java/org/springframework/kafka/annotation/PartitionOffset.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2016-2020 the original author or authors.
+ * Copyright 2016-2024 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,11 +20,14 @@
 import java.lang.annotation.RetentionPolicy;
 import java.lang.annotation.Target;
 
+import org.springframework.kafka.support.TopicPartitionOffset.SeekPosition;
+
 /**
  * Used to add partition/initial offset information to a {@code KafkaListener}.
  *
  * @author Artem Bilan
  * @author Gary Russell
+ * @author Wang Zhiyang
  */
 @Target({})
 @Retention(RetentionPolicy.RUNTIME)
@@ -60,4 +63,17 @@
 	 */
 	String relativeToCurrent() default "false";
 
+	/**
+	 * Position to seek on partition assignment. By default, seek by offset.
+	 * Set {@link SeekPosition} seek position enum name to specify "special"
+	 * seeks, no restrictions on capitalization. If seekPosition set 'BEGINNING'
+	 * or 'END', ignore {@code relativeToCurrent} and {@code initialOffset}.
+	 * If seekPosition set 'TIMESTAMP', initialOffset means time stamp, ignore
+	 * {@code relativeToCurrent}.
+	 * @return special seeks.
+	 * @since 3.2
+	 * @see SeekPosition
+	 */
+	String seekPosition() default "";
+
 }

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

@@ -3103,6 +3103,7 @@ private void initPartitionsIfNeeded() {
 					.filter(e -> SeekPosition.TIMESTAMP.equals(e.getValue().seekPosition))
 					.collect(Collectors.toMap(Entry::getKey, entry -> entry.getValue().offset));
 			if (!times.isEmpty()) {
+				times.forEach((key, value) -> partitions.remove(key));
 				Map<TopicPartition, OffsetAndTimestamp> offsetsForTimes = this.consumer.offsetsForTimes(times);
 				offsetsForTimes.forEach((tp, off) -> {
 					if (off == null) {
@@ -3117,7 +3118,7 @@ private void initPartitionsIfNeeded() {
 			if (this.consumerSeekAwareListener != null) {
 				this.consumerSeekAwareListener.onPartitionsAssigned(this.definedPartitions.keySet().stream()
 							.map(tp -> new SimpleEntry<>(tp, this.consumer.position(tp)))
-							.collect(Collectors.toMap(entry -> entry.getKey(), entry -> entry.getValue())),
+							.collect(Collectors.toMap(SimpleEntry::getKey, SimpleEntry::getValue)),
 						this.seekCallback);
 			}
 		}