Seek Improvements

- add convenience methods to `AbstractConsumerSeekAware`

- compile tests with Java 11

Author: garyrussell

.travis.yml

@@ -1,6 +1,6 @@
-dist: trusty
+dist: bionic
 language: java
-jdk: oraclejdk8
+jdk: openjdk11
 install: true
 before_cache:
   - rm -f $HOME/.gradle/caches/modules-2/modules-2.lock

build.gradle

@@ -133,6 +133,10 @@ subprojects { subproject ->
 		targetCompatibility = 1.8
 	}
 
+	compileTestJava {
+		sourceCompatibility = 11
+	}
+
 	compileTestKotlin {
 		kotlinOptions {
 			jvmTarget = '1.8'

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2019 the original author or authors.
+ * Copyright 2019-2020 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.
@@ -18,6 +18,8 @@
 
 import java.util.Collection;
 import java.util.Collections;
+import java.util.LinkedList;
+import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 
@@ -40,6 +42,8 @@ public abstract class AbstractConsumerSeekAware implements ConsumerSeekAware {
 
 	private final Map<TopicPartition, ConsumerSeekCallback> callbacks = new ConcurrentHashMap<>();
 
+	private final Map<ConsumerSeekCallback, List<TopicPartition>> callbacksToTopic = new ConcurrentHashMap<>();
+
 	@Override
 	public void registerSeekCallback(ConsumerSeekCallback callback) {
 		this.callbackForThread.set(callback);
@@ -49,13 +53,27 @@ public void registerSeekCallback(ConsumerSeekCallback callback) {
 	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));
+			assignments.keySet().forEach(tp -> {
+				this.callbacks.put(tp, threadCallback);
+				this.callbacksToTopic.computeIfAbsent(threadCallback, key -> new LinkedList<>()).add(tp);
+			});
 		}
 	}
 
 	@Override
 	public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
-		partitions.forEach(tp -> this.callbacks.remove(tp));
+		partitions.forEach(tp -> {
+			ConsumerSeekCallback removed = this.callbacks.remove(tp);
+			if (removed != null) {
+				List<TopicPartition> topics = this.callbacksToTopic.get(removed);
+				if (topics != null) {
+					topics.remove(tp);
+					if (topics.size() == 0) {
+						this.callbacksToTopic.remove(removed);
+					}
+				}
+			}
+		});
 	}
 
 	@Override
@@ -81,4 +99,38 @@ protected Map<TopicPartition, ConsumerSeekCallback> getSeekCallbacks() {
 		return Collections.unmodifiableMap(this.callbacks);
 	}
 
+	/**
+	 * Return the currently registered callbacks and their associated {@link TopicPartition}(s).
+	 * @return the map of callbacks and partitions.
+	 * @since 2.6
+	 */
+	protected Map<ConsumerSeekCallback, List<TopicPartition>> getCallbacksAndTopics() {
+		return Collections.unmodifiableMap(this.callbacksToTopic);
+	}
+
+	/**
+	 * Seek all assigned partitions to the beginning.
+	 * @since 2.6
+	 */
+	public void seekToBeginning() {
+		getCallbacksAndTopics().forEach((cb, topics) -> cb.seekToBeginning(topics));
+	}
+
+	/**
+	 * Seek all assigned partitions to the end.
+	 * @since 2.6
+	 */
+	public void seekToEnd() {
+		getCallbacksAndTopics().forEach((cb, topics) -> cb.seekToEnd(topics));
+	}
+
+	/**
+	 * Seek all assigned partitions to the offset represented by the timestamp.
+	 * @param time the time to seek to.
+	 * @since 2.6
+	 */
+	public void seekToTimestamp(long time) {
+		getCallbacksAndTopics().forEach((cb, topcis) -> cb.seekToTimestamp(topcis, time));
+	}
+
 }

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

@@ -0,0 +1,139 @@
+/*
+ * Copyright 2020 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 static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+
+import java.util.Collections;
+import java.util.LinkedHashMap;
+import java.util.LinkedList;
+import java.util.Map;
+import java.util.concurrent.Callable;
+import java.util.concurrent.Executors;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+import org.apache.kafka.common.TopicPartition;
+import org.junit.jupiter.api.Test;
+
+import org.springframework.kafka.listener.ConsumerSeekAware.ConsumerSeekCallback;
+import org.springframework.kafka.test.utils.KafkaTestUtils;
+
+/**
+ * @author Gary Russell
+ * @since 2.6
+ *
+ */
+public class ConsumerSeekAwareTests {
+
+	@SuppressWarnings("unchecked")
+	@Test
+	void beginningEndAndBulkSeekToTimestamp() throws Exception {
+		class CSA extends AbstractConsumerSeekAware {
+		}
+		AbstractConsumerSeekAware csa = new CSA();
+		var exec1 = Executors.newSingleThreadExecutor();
+		var exec2 = Executors.newSingleThreadExecutor();
+		var cb1 = mock(ConsumerSeekCallback.class);
+		var cb2  = mock(ConsumerSeekCallback.class);
+		var first = new AtomicBoolean(true);
+		var map1 = new LinkedHashMap<>(Map.of(new TopicPartition("foo", 0), 0L, new TopicPartition("foo", 1), 0L));
+		var map2 = new LinkedHashMap<>(Map.of(new TopicPartition("foo", 2), 0L, new TopicPartition("foo", 3), 0L));
+		var register = new Callable<Void>() {
+
+			@Override
+			public Void call() {
+				if (first.getAndSet(false)) {
+					csa.registerSeekCallback(cb1);
+					csa.onPartitionsAssigned(map1, null);
+				}
+				else {
+					csa.registerSeekCallback(cb2);
+					csa.onPartitionsAssigned(map2, null);
+				}
+				return null;
+			}
+
+		};
+		exec1.submit(register).get();
+		exec2.submit(register).get();
+		csa.seekToBeginning();
+		verify(cb1).seekToBeginning(new LinkedList<>(map1.keySet()));
+		verify(cb2).seekToBeginning(new LinkedList<>(map2.keySet()));
+		csa.seekToEnd();
+		verify(cb1).seekToEnd(new LinkedList<>(map1.keySet()));
+		verify(cb2).seekToEnd(new LinkedList<>(map2.keySet()));
+		csa.seekToTimestamp(42L);
+		verify(cb1).seekToTimestamp(new LinkedList<>(map1.keySet()), 42L);
+		verify(cb2).seekToTimestamp(new LinkedList<>(map2.keySet()), 42L);
+		var revoke1 = new Callable<Void>() {
+
+			@Override
+			public Void call() {
+				if (!first.getAndSet(true)) {
+					csa.onPartitionsRevoked(Collections.singletonList(map1.keySet().iterator().next()));
+				}
+				else {
+					csa.onPartitionsRevoked(Collections.singletonList(map2.keySet().iterator().next()));
+				}
+				return null;
+			}
+
+		};
+		exec1.submit(revoke1).get();
+		exec2.submit(revoke1).get();
+		map1.remove(map1.keySet().iterator().next());
+		map2.remove(map2.keySet().iterator().next());
+		csa.seekToTimestamp(43L);
+		verify(cb1).seekToTimestamp(new LinkedList<>(map1.keySet()), 43L);
+		verify(cb2).seekToTimestamp(new LinkedList<>(map2.keySet()), 43L);
+		var revoke2 = new Callable<Void>() {
+
+			@Override
+			public Void call() {
+				if (first.getAndSet(false)) {
+					csa.onPartitionsRevoked(Collections.singletonList(map1.keySet().iterator().next()));
+				}
+				else {
+					csa.onPartitionsRevoked(Collections.singletonList(map2.keySet().iterator().next()));
+				}
+				return null;
+			}
+
+		};
+		exec1.submit(revoke2).get();
+		exec2.submit(revoke2).get();
+		assertThat(KafkaTestUtils.getPropertyValue(csa, "callbacks", Map.class)).isEmpty();
+		assertThat(KafkaTestUtils.getPropertyValue(csa, "callbacksToTopic", Map.class)).isEmpty();
+		var checkTL = new Callable<Void>() {
+
+			@Override
+			public Void call() throws Exception {
+				csa.unregisterSeekCallback();
+				assertThat(KafkaTestUtils.getPropertyValue(csa, "callbackForThread", ThreadLocal.class).get()).isNull();
+				return null;
+			}
+
+		};
+		exec1.submit(checkTL).get();
+		exec2.submit(checkTL).get();
+		exec1.shutdown();
+		exec2.shutdown();
+	}
+
+}

src/reference/asciidoc/kafka.adoc

@@ -2721,7 +2721,7 @@ NOTE: The `seekToBeginning` method that accepts a collection is useful, for exam
 ====
 [source, java]
 ----
-public class MyListener extends AbstractConsumerSeekAware {
+public class MyListener implements ConsumerSeekAware {
 
 ...
 
@@ -2851,6 +2851,41 @@ public class SeekToLastOnIdleListener extends AbstractConsumerSeekAware {
 ----
 ====
 
+Version 2.6 added convenience methods to the abstract class:
+
+* `seekToBeginning()` - seeks all assigned partitions to the beginning
+* `seekToEnd()` - seeks all assigned partitions to the end
+* `seekToTimestamp(long time)` - seeks all assigned partitions to the offset represented by that timestamp.
+
+Example:
+
+====
+[source, java]
+----
+public class MyListener extends AbstractConsumerSeekAware {
+
+    @KafkaListener(...)
+    void listn(...) {
+        ...
+    }
+}
+
+public class SomeOtherBean {
+
+    MyListener listener;
+
+    ...
+
+    void someMethod() {
+        this.listener.seekToTimestamp(System.currentTimeMillis - 60_000);
+    }
+
+}
+
+----
+====
+
+
 [[container-factory]]
 ==== Container factory
 

src/reference/asciidoc/whats-new.adoc

@@ -26,3 +26,6 @@ When using manual partition assignment, you can now specify a wildcard for deter
 In addition, if the listener implements `ConsumerSeekAware`, `onPartitionsAssigned()` is called after the manual assignment.
 (Also added in version 2.5.5).
 See <<manual-assignment>> for more information.
+
+Convenience methods have been added to `AbstractConsumerSeekAware` to make seeking easier.
+See <<seek>> for more information.