spring-projectsGH-3108: Unsafe access to consumer in seek API

Fixes: spring-projects#3108

* When the seek API that uses the user-provided function to compute the offset to seek,
  there is a concurrency issue in which the Kafka consumer is used unsafely.

See this for details: spring-projects#3078 (comment)

Author: sobychacko

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

@@ -57,3 +57,9 @@ See xref:retrytopic/topic-naming.adoc#single-topic-maxinterval-delay[Single Topi
 `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-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.
+When this constructor is used, the framework calls the function with the input argument of the current consumer offset position.
+See xref:kafka/seek.adoc#seek[Seek API Docs] for more details.
+

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

@@ -3005,11 +3005,15 @@ private void processSeeks() {
 					SeekPosition position = offset.getPosition();
 					TopicPartition topicPartition = offset.getTopicPartition();
 					Long whereTo = offset.getOffset();
+					Function<Long, Long> offsetComputeFunction = offset.getOffsetComputeFunction();
 					if (position == null) {
 						if (offset.isRelativeToCurrent()) {
 							whereTo += this.consumer.position(topicPartition);
 							whereTo = Math.max(whereTo, 0);
 						}
+						else if (offsetComputeFunction != null) {
+							whereTo = offsetComputeFunction.apply(this.consumer.position(topicPartition));
+						}
 						this.consumer.seek(topicPartition, whereTo);
 					}
 					else if (SeekPosition.TIMESTAMP.equals(position)) {
@@ -3262,8 +3266,7 @@ public void seek(String topic, int partition, long offset) {
 
 		@Override
 		public void seek(String topic, int partition, Function<Long, Long> offsetComputeFunction) {
-			this.seeks.add(new TopicPartitionOffset(topic, partition, offsetComputeFunction.apply(
-					this.consumer.position(new TopicPartition(topic, partition)))));
+			this.seeks.add(new TopicPartitionOffset(topic, partition, offsetComputeFunction));
 		}
 
 		@Override

spring-kafka/src/main/java/org/springframework/kafka/support/TopicPartitionOffset.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2019-2022 the original author or authors.
+ * Copyright 2019-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.
@@ -17,6 +17,7 @@
 package org.springframework.kafka.support;
 
 import java.util.Objects;
+import java.util.function.Function;
 
 import org.apache.kafka.common.TopicPartition;
 
@@ -41,6 +42,7 @@
  *
  * @author Artem Bilan
  * @author Gary Russell
+ * @author Soby Chacko
  *
  * @since 2.3
  */
@@ -77,6 +79,8 @@ public enum SeekPosition {
 
 	private boolean relativeToCurrent;
 
+	private Function<Long, Long> offsetComputeFunction;
+
 	/**
 	 * Construct an instance with no initial offset management.
 	 * @param topic the topic.
@@ -98,6 +102,19 @@ public TopicPartitionOffset(String topic, int partition, Long offset) {
 		this(topic, partition, offset, false);
 	}
 
+	/**
+	 * Construct an instance with the provided function to compute the offset.
+	 * @param topic the topic.
+	 * @param partition the partition.
+	 * @param offsetComputeFunction function to compute the offset.
+	 * @since 3.2.0
+	 */
+	public TopicPartitionOffset(String topic, int partition, Function<Long, Long> offsetComputeFunction) {
+		this.topicPartition = new TopicPartition(topic, partition);
+		this.offsetComputeFunction = offsetComputeFunction;
+		this.position = null;
+	}
+
 	/**
 	 * Construct an instance with the provided initial offset.
 	 * @param topic the topic.
@@ -198,6 +215,10 @@ public SeekPosition getPosition() {
 		return this.position;
 	}
 
+	public Function<Long, Long> getOffsetComputeFunction() {
+		return offsetComputeFunction;
+	}
+
 	@Override
 	public boolean equals(Object o) {
 		if (this == o) {

spring-kafka/src/test/java/org/springframework/kafka/annotation/EnableKafkaIntegrationTests.java

@@ -50,6 +50,7 @@
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Collectors;
+import java.util.stream.IntStream;
 
 import org.aopalliance.intercept.MethodInterceptor;
 import org.apache.commons.logging.Log;
@@ -189,7 +190,7 @@
 		"annotated29", "annotated30", "annotated30reply", "annotated31", "annotated32", "annotated33",
 		"annotated34", "annotated35", "annotated36", "annotated37", "foo", "manualStart", "seekOnIdle",
 		"annotated38", "annotated38reply", "annotated39", "annotated40", "annotated41", "annotated42",
-		"annotated43", "annotated43reply"})
+		"annotated43", "annotated43reply", "seekToComputeFn"})
 @TestPropertySource(properties = "spel.props=fetch.min.bytes=420000,max.poll.records=10")
 public class EnableKafkaIntegrationTests {
 
@@ -248,6 +249,9 @@ public class EnableKafkaIntegrationTests {
 	@Autowired
 	private SeekToLastOnIdleListener seekOnIdleListener;
 
+	@Autowired
+	private SeekToOffsetFromComputeFunction seekToOffsetFromComputeFunction;
+
 	@Autowired
 	private MeterRegistry meterRegistry;
 
@@ -1069,6 +1073,13 @@ void classLevelTwoInstancesSameClass() {
 		assertThat(this.registry.getListenerContainer("multiTwoTwo")).isNotNull();
 	}
 
+	@Test
+	void seekToOffsetComputedFromFunction() throws InterruptedException {
+		IntStream.range(0, 10).forEach(i -> template.send("seekToComputeFn", 0, i, "my-data"));
+		assertThat(this.seekToOffsetFromComputeFunction.latch1.await(10, TimeUnit.SECONDS)).isTrue();
+		assertThat(this.seekToOffsetFromComputeFunction.latch2.await(10, TimeUnit.SECONDS)).isTrue();
+	}
+
 	@Configuration
 	@EnableKafka
 	@EnableTransactionManagement(proxyTargetClass = true)
@@ -1425,6 +1436,10 @@ public SeekToLastOnIdleListener seekOnIdle() {
 			return new SeekToLastOnIdleListener();
 		}
 
+		@Bean
+		public SeekToOffsetFromComputeFunction seekToOffsetFromComputeFunction() {
+			return new SeekToOffsetFromComputeFunction();
+		}
 		@Bean
 		public IfaceListener<String> ifaceListener() {
 			return new IfaceListenerImpl();
@@ -2333,6 +2348,31 @@ public Object postProcessBeforeInitialization(Object bean, String beanName) thro
 		}
 	}
 
+	public static class SeekToOffsetFromComputeFunction extends AbstractConsumerSeekAware {
+
+		CountDownLatch latch1 = new CountDownLatch(10);
+		CountDownLatch latch2 = new CountDownLatch(1);
+
+		@KafkaListener(id = "seekToComputeFn", topics = "seekToComputeFn")
+		public void listen(String in) throws InterruptedException {
+			if (latch2.getCount() > 0) { // if latch2 is zero, the test condition is met
+				if (latch1.getCount() == 0) { // Seek happened on the consumer
+					latch2.countDown();
+				}
+				if (latch1.getCount() > 0) {
+					latch1.countDown();
+					if (latch1.getCount() == 0) {
+						ConsumerSeekCallback seekToComputeFn = getSeekCallbackFor(
+								new org.apache.kafka.common.TopicPartition("seekToComputeFn", 0));
+						assertThat(seekToComputeFn).isNotNull();
+						seekToComputeFn.
+								seek("seekToComputeFn", 0, current -> 0L);
+					}
+				}
+			}
+		}
+	}
+
 	public static class SeekToLastOnIdleListener extends AbstractConsumerSeekAware {
 
 		final CountDownLatch latch1 = new CountDownLatch(10);