Add seek helper abstract class

- also handle late arriving time stamp seek requests while `processSeeks()` is already running.

* * Properly clean up callbacks for revoked partitions
* clear() wiped them for all threads
* Also remove the TL when revoked

* * Use correct callback with onIdleContainer()
* Add test/doc for external seek calls

Author: garyrussell

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

@@ -0,0 +1,80 @@
+/*
+ * Copyright 2019 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 java.util.Collection;
+import java.util.Collections;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+import org.apache.kafka.common.TopicPartition;
+
+import org.springframework.lang.Nullable;
+
+/**
+ * Manages the {@link ConsumerSeekAware.ConsumerSeekCallback} s for the listener. If the
+ * listener subclasses this class, it can easily seek arbitrary topics/partitions without
+ * having to keep track of the callbacks itself.
+ *
+ * @author Gary Russell
+ * @since 2.3
+ *
+ */
+public abstract class AbstractConsumerSeekAware implements ConsumerSeekAware {
+
+	private final ThreadLocal<ConsumerSeekCallback> callbackForThread = new ThreadLocal<>();
+
+	private final Map<TopicPartition, ConsumerSeekCallback> callbacks = new ConcurrentHashMap<>();
+
+	@Override
+	public void registerSeekCallback(ConsumerSeekCallback callback) {
+		this.callbackForThread.set(callback);
+	}
+
+	@Override
+	public void onPartitionsAssigned(Map<TopicPartition, Long> assignments, ConsumerSeekCallback callback) {
+		ConsumerSeekCallback threadCallback = this.callbackForThread.get();
+		if (threadCallback != null) {
+			assignments.keySet().forEach(tp -> this.callbacks.put(tp, threadCallback));
+		}
+	}
+
+	@Override
+	public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
+		partitions.forEach(tp -> this.callbacks.remove(tp));
+		this.callbackForThread.remove();
+	}
+
+	/**
+	 * Return the callback for the specified topic/partition.
+	 * @param topicPartition the topic/partition.
+	 * @return the callback (or null if there is no assignment).
+	 */
+	@Nullable
+	protected ConsumerSeekCallback getSeekCallbackFor(TopicPartition topicPartition) {
+		return this.callbacks.get(topicPartition);
+	}
+
+	/**
+	 * The map of callbacks for all currently assigned partitions.
+	 * @return the map.
+	 */
+	protected Map<TopicPartition, ConsumerSeekCallback> getSeekCallbacks() {
+		return Collections.unmodifiableMap(this.callbacks);
+	}
+
+}

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

@@ -51,6 +51,17 @@ default void onPartitionsAssigned(Map<TopicPartition, Long> assignments, Consume
 		// do nothing
 	}
 
+	/**
+	 * When using group management, called when partition assignments are revoked.
+	 * Listeners should discard any callback saved from
+	 * {@link #registerSeekCallback(ConsumerSeekCallback)} on this thread.
+	 * @param partitions the partitions that have been revoked.
+	 * @since 2.3
+	 */
+	default void onPartitionsRevoked(Collection<TopicPartition> partitions) {
+		// do nothing
+	}
+
 	/**
 	 * If the container is configured to emit idle container events, this method is called
 	 * when the container idle event is emitted - allowing a seek operation.

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

@@ -694,6 +694,7 @@ protected ErrorHandler determineErrorHandler(GenericErrorHandler<?> errHandler)
 		}
 
 		private void seekPartitions(Collection<TopicPartition> partitions, boolean idle) {
+			((ConsumerSeekAware) this.genericListener).registerSeekCallback(this);
 			Map<TopicPartition, Long> current = new HashMap<>();
 			for (TopicPartition topicPartition : partitions) {
 				current.put(topicPartition, ListenerConsumer.this.consumer.position(topicPartition));
@@ -847,9 +848,10 @@ private void checkIdle() {
 			}
 		}
 
-		public void wrapUp() {
+		private void wrapUp() {
 			KafkaUtils.clearConsumerGroupId();
 			publishConsumerStoppingEvent(this.consumer);
+			Collection<TopicPartition> partitions = getAssignedPartitions();
 			if (!this.fatalError) {
 				if (this.kafkaTxManager == null) {
 					commitPendingAcks();
@@ -861,7 +863,7 @@ public void wrapUp() {
 					}
 				}
 				else {
-					closeProducers(getAssignedPartitions());
+					closeProducers(partitions);
 				}
 			}
 			else {
@@ -877,6 +879,9 @@ public void wrapUp() {
 			if (this.errorHandler != null) {
 				this.errorHandler.clearThreadState();
 			}
+			if (ListenerConsumer.this.genericListener instanceof ConsumerSeekAware) {
+				((ConsumerSeekAware) ListenerConsumer.this.genericListener).onPartitionsRevoked(partitions);
+			}
 			this.logger.info(() -> getGroupId() + ": Consumer stopped");
 			publishConsumerStoppedEvent();
 		}
@@ -1475,6 +1480,7 @@ private void processSeeks() {
 					if (position == null) {
 						if (offset.isRelativeToCurrent()) {
 							whereTo += this.consumer.position(offset.getTopicPartition());
+							whereTo = Math.max(whereTo, 0);
 						}
 						this.consumer.seek(offset.getTopicPartition(), whereTo);
 					}
@@ -1484,6 +1490,13 @@ else if (position.equals(SeekPosition.BEGINNING)) {
 							this.consumer.seek(offset.getTopicPartition(), whereTo);
 						}
 					}
+					else if (position.equals(SeekPosition.TIMESTAMP)) {
+						// possible late addition since the grouped processing above
+						Map<TopicPartition, OffsetAndTimestamp> offsetsForTimes = this.consumer
+								.offsetsForTimes(
+										Collections.singletonMap(offset.getTopicPartition(), offset.getOffset()));
+						offsetsForTimes.forEach((tp, ot) -> this.consumer.seek(tp, ot.offset()));
+					}
 					else {
 						this.consumer.seekToEnd(Collections.singletonList(offset.getTopicPartition()));
 						if (whereTo != null) {
@@ -1511,6 +1524,7 @@ private void processTimestampSeeks() {
 					}
 					timestampSeeks.put(tpo.getTopicPartition(), tpo.getOffset());
 					seekIterator.remove();
+					traceSeek(tpo);
 				}
 			}
 			if (timestampSeeks != null) {
@@ -1783,6 +1797,9 @@ public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
 						this.consumerAwareListener.onPartitionsRevokedAfterCommit(ListenerConsumer.this.consumer,
 								partitions);
 					}
+					if (ListenerConsumer.this.genericListener instanceof ConsumerSeekAware) {
+						((ConsumerSeekAware) ListenerConsumer.this.genericListener).onPartitionsRevoked(partitions);
+					}
 				}
 				finally {
 					if (ListenerConsumer.this.kafkaTxManager != null) {
@@ -1972,7 +1989,8 @@ private Long computeBackwardWhereTo(long offset, boolean toCurrent, TopicPartiti
 					end = consumerToSeek.position(topicPart);
 				}
 				if (end != null) {
-					return end + offset;
+					long newOffset = end + offset;
+					return newOffset < 0 ? 0 : newOffset;
 				}
 				return null;
 			}

spring-kafka/src/main/java/org/springframework/kafka/listener/adapter/AbstractDelegatingMessageListenerAdapter.java

@@ -16,6 +16,7 @@
 
 package org.springframework.kafka.listener.adapter;
 
+import java.util.Collection;
 import java.util.Map;
 
 import org.apache.commons.logging.LogFactory;
@@ -77,6 +78,13 @@ public void onPartitionsAssigned(Map<TopicPartition, Long> assignments, Consumer
 		}
 	}
 
+	@Override
+	public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
+		if (this.seekAware != null) {
+			this.seekAware.onPartitionsRevoked(partitions);
+		}
+	}
+
 	@Override
 	public void onIdleContainer(Map<TopicPartition, Long> assignments, ConsumerSeekCallback callback) {
 		if (this.seekAware != null) {

spring-kafka/src/main/java/org/springframework/kafka/listener/adapter/MessagingMessageListenerAdapter.java

@@ -253,6 +253,13 @@ public void onPartitionsAssigned(Map<TopicPartition, Long> assignments, Consumer
 		}
 	}
 
+	@Override
+	public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
+		if (this.bean instanceof ConsumerSeekAware) {
+			((ConsumerSeekAware) this.bean).onPartitionsRevoked(partitions);
+		}
+	}
+
 	@Override
 	public void onIdleContainer(Map<TopicPartition, Long> assignments, ConsumerSeekCallback callback) {
 		if (this.bean instanceof ConsumerSeekAware) {

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

@@ -32,6 +32,8 @@
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -76,6 +78,7 @@
 import org.springframework.kafka.core.KafkaTemplate;
 import org.springframework.kafka.core.ProducerFactory;
 import org.springframework.kafka.event.ListenerContainerIdleEvent;
+import org.springframework.kafka.listener.AbstractConsumerSeekAware;
 import org.springframework.kafka.listener.ConcurrentMessageListenerContainer;
 import org.springframework.kafka.listener.ConsumerAwareErrorHandler;
 import org.springframework.kafka.listener.ConsumerAwareListenerErrorHandler;
@@ -152,7 +155,7 @@ public class EnableKafkaIntegrationTests {
 			"annotated22reply", "annotated23", "annotated23reply", "annotated24", "annotated24reply",
 			"annotated25", "annotated25reply1", "annotated25reply2", "annotated26", "annotated27", "annotated28",
 			"annotated29", "annotated30", "annotated30reply", "annotated31", "annotated32", "annotated33",
-			"annotated34", "annotated35", "annotated36", "annotated37", "foo", "manualStart");
+			"annotated34", "annotated35", "annotated36", "annotated37", "foo", "manualStart", "seekOnIdle");
 
 	private static EmbeddedKafkaBroker embeddedKafka = embeddedKafkaRule.getEmbeddedKafka();
 
@@ -201,6 +204,9 @@ public class EnableKafkaIntegrationTests {
 	@Autowired
 	private ConcurrentKafkaListenerContainerFactory<Integer, String> transactionalFactory;
 
+	@Autowired
+	private SeekToLastOnIdleListener seekOnIdleListener;
+
 	@Test
 	public void testAnonymous() {
 		MessageListenerContainer container = this.registry
@@ -757,6 +763,24 @@ public void testProjection() throws InterruptedException {
 		assertThat(this.listener.username).isEqualTo("SomeUsername");
 	}
 
+	@SuppressWarnings("unchecked")
+	@Test
+	public void testSeekToLastOnIdle() throws InterruptedException {
+		this.registry.getListenerContainer("seekOnIdle").start();
+		this.seekOnIdleListener.waitForBalancedAssignment();
+		this.template.send("seekOnIdle", 0, 0, "foo");
+		this.template.send("seekOnIdle", 1, 1, "bar");
+		assertThat(this.seekOnIdleListener.latch1.await(10, TimeUnit.SECONDS)).isTrue();
+		assertThat(this.seekOnIdleListener.latch2.getCount()).isEqualTo(2L);
+		this.seekOnIdleListener.rewindAllOneRecord();
+		assertThat(this.seekOnIdleListener.latch2.await(10, TimeUnit.SECONDS)).isTrue();
+		assertThat(this.seekOnIdleListener.latch3.getCount()).isEqualTo(1L);
+		this.seekOnIdleListener.rewindOnePartitionOneRecord("seekOnIdle", 1);
+		assertThat(this.seekOnIdleListener.latch3.await(10, TimeUnit.SECONDS)).isTrue();
+		this.registry.getListenerContainer("seekOnIdle").stop();
+		assertThat(KafkaTestUtils.getPropertyValue(this.seekOnIdleListener, "callbacks", Map.class)).hasSize(0);
+	}
+
 	@Configuration
 	@EnableKafka
 	@EnableTransactionManagement(proxyTargetClass = true)
@@ -933,6 +957,7 @@ public KafkaListenerContainerFactory<?> batchSpyFactory() {
 			factory.setRecordFilterStrategy(recordFilter());
 			// always send to the same partition so the replies are in order for the test
 			factory.setReplyTemplate(partitionZeroReplyingTemplate());
+			factory.setMissingTopicsFatal(false);
 			return factory;
 		}
 
@@ -968,6 +993,7 @@ public KafkaListenerContainerFactory<?> batchManualFactory2() {
 			ContainerProperties props = factory.getContainerProperties();
 			props.setAckMode(AckMode.MANUAL_IMMEDIATE);
 			props.setIdleEventInterval(100L);
+			props.setPollTimeout(50L);
 			factory.setRecordFilterStrategy(manualFilter());
 			factory.setAckDiscarded(true);
 			factory.setRetryTemplate(new RetryTemplate());
@@ -1050,6 +1076,11 @@ public Listener listener() {
 			return new Listener();
 		}
 
+		@Bean
+		public SeekToLastOnIdleListener seekOnIdle() {
+			return new SeekToLastOnIdleListener();
+		}
+
 		@Bean
 		public IfaceListener<String> ifaceListener() {
 			return new IfaceListenerImpl();
@@ -1698,6 +1729,68 @@ public void registerSeekCallback(ConsumerSeekCallback callback) {
 
 	}
 
+	public static class SeekToLastOnIdleListener extends AbstractConsumerSeekAware {
+
+		private final CountDownLatch latch1 = new CountDownLatch(10);
+
+		private final CountDownLatch latch2 = new CountDownLatch(12);
+
+		private final CountDownLatch latch3 = new CountDownLatch(13);
+
+		private final Set<Thread> consumerThreads = ConcurrentHashMap.newKeySet();
+
+		@KafkaListener(id = "seekOnIdle", topics = "seekOnIdle", autoStartup = "false", concurrency = "2",
+				clientIdPrefix = "seekOnIdle", containerFactory = "kafkaManualAckListenerContainerFactory")
+		public void listen(@SuppressWarnings("unused") String in, Acknowledgment ack) {
+			this.latch1.countDown();
+			this.latch2.countDown();
+			this.latch3.countDown();
+			ack.acknowledge();
+		}
+
+		@Override
+		public void onIdleContainer(Map<org.apache.kafka.common.TopicPartition, Long> assignments,
+				ConsumerSeekCallback callback) {
+
+			if (this.latch1.getCount() > 0) {
+				assignments.keySet().forEach(tp -> callback.seekRelative(tp.topic(), tp.partition(), -1, true));
+			}
+		}
+
+		public void rewindAllOneRecord() {
+			getSeekCallbacks()
+				.forEach((tp, callback) ->
+					callback.seekRelative(tp.topic(), tp.partition(), -1, true));
+		}
+
+		public void rewindOnePartitionOneRecord(String topic, int partition) {
+			getSeekCallbackFor(new org.apache.kafka.common.TopicPartition(topic, partition))
+				.seekRelative(topic, partition, -1, true);
+		}
+
+		@Override
+		public synchronized void onPartitionsAssigned(Map<org.apache.kafka.common.TopicPartition, Long> assignments,
+				ConsumerSeekCallback callback) {
+
+			super.onPartitionsAssigned(assignments, callback);
+			if (assignments.size() > 0) {
+				this.consumerThreads.add(Thread.currentThread());
+				notifyAll();
+			}
+		}
+
+		public synchronized void waitForBalancedAssignment() throws InterruptedException {
+			int n = 0;
+			while (this.consumerThreads.size() < 2) {
+				wait(1000);
+				if (n++ > 20) {
+					throw new IllegalStateException("Balanced distribution did not occur");
+				}
+			}
+		}
+
+	}
+
 	interface IfaceListener<T> {
 
 		void listen(T foo);