GH-599: Fix initial seek

Fixes https://github.com/spring-projects/spring-kafka/issues/599

Previously, initial seeks using `TopicPartitionInitialOffset`s only worked with
a provided `offset`. The `SeekPosition` field was ignored, and only used for
subsequent seek operations.

`initPartitionsIfNeeded()` now processes both styles of initial offset.
This commit is contained in:
Gary Russell
2018-03-05 19:50:44 -05:00
committed by Artem Bilan
parent 12c1ad22f9
commit c9f24b0af5
2 changed files with 76 additions and 5 deletions

View File

@@ -27,6 +27,7 @@ import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
@@ -420,7 +421,8 @@ public class KafkaMessageListenerContainer<K, V> extends AbstractMessageListener
this.definedPartitions = new HashMap<>(topicPartitions.size());
for (TopicPartitionInitialOffset topicPartition : topicPartitions) {
this.definedPartitions.put(topicPartition.topicPartition(),
new OffsetMetadata(topicPartition.initialOffset(), topicPartition.isRelativeToCurrent()));
new OffsetMetadata(topicPartition.initialOffset(), topicPartition.isRelativeToCurrent(),
topicPartition.getPosition()));
}
consumer.assign(new ArrayList<>(this.definedPartitions.keySet()));
}
@@ -647,7 +649,12 @@ public class KafkaMessageListenerContainer<K, V> extends AbstractMessageListener
this.count = 0;
this.last = System.currentTimeMillis();
if (isRunning() && this.definedPartitions != null) {
initPartitionsIfNeeded();
try {
initPartitionsIfNeeded();
}
catch (Exception e) {
this.logger.error("Failed to set initial offsets", e);
}
}
long lastReceive = System.currentTimeMillis();
long lastAlertAt = lastReceive;
@@ -1186,9 +1193,27 @@ public class KafkaMessageListenerContainer<K, V> extends AbstractMessageListener
/*
* Note: initial position setting is only supported with explicit topic assignment.
* When using auto assignment (subscribe), the ConsumerRebalanceListener is not
* called until we poll() the consumer.
* called until we poll() the consumer. Users can use a ConsumerAwareRebalanceListener
* or a ConsumerSeekAware listener in that case.
*/
for (Entry<TopicPartition, OffsetMetadata> entry : this.definedPartitions.entrySet()) {
Map<TopicPartition, OffsetMetadata> partitions = new HashMap<>(this.definedPartitions);
Set<TopicPartition> beginnings = partitions.entrySet().stream()
.filter(e -> SeekPosition.BEGINNING.equals(e.getValue().seekPosition))
.map(e -> e.getKey())
.collect(Collectors.toSet());
beginnings.forEach(k -> partitions.remove(k));
Set<TopicPartition> ends = partitions.entrySet().stream()
.filter(e -> SeekPosition.END.equals(e.getValue().seekPosition))
.map(e -> e.getKey())
.collect(Collectors.toSet());
ends.forEach(k -> partitions.remove(k));
if (beginnings.size() > 0) {
this.consumer.seekToBeginning(beginnings);
}
if (ends.size() > 0) {
this.consumer.seekToEnd(ends);
}
for (Entry<TopicPartition, OffsetMetadata> entry : partitions.entrySet()) {
TopicPartition topicPartition = entry.getKey();
OffsetMetadata metadata = entry.getValue();
Long offset = metadata.offset;
@@ -1378,9 +1403,12 @@ public class KafkaMessageListenerContainer<K, V> extends AbstractMessageListener
private final boolean relativeToCurrent;
OffsetMetadata(Long offset, boolean relativeToCurrent) {
private final SeekPosition seekPosition;
OffsetMetadata(Long offset, boolean relativeToCurrent, SeekPosition seekPosition) {
this.offset = offset;
this.relativeToCurrent = relativeToCurrent;
this.seekPosition = seekPosition;
}
}

View File

@@ -35,6 +35,7 @@ import java.util.BitSet;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
@@ -78,6 +79,7 @@ import org.springframework.kafka.listener.adapter.FilteringMessageListenerAdapte
import org.springframework.kafka.listener.config.ContainerProperties;
import org.springframework.kafka.support.Acknowledgment;
import org.springframework.kafka.support.TopicPartitionInitialOffset;
import org.springframework.kafka.support.TopicPartitionInitialOffset.SeekPosition;
import org.springframework.kafka.support.serializer.JsonDeserializer;
import org.springframework.kafka.support.serializer.JsonSerializer;
import org.springframework.kafka.test.rule.KafkaEmbedded;
@@ -1675,6 +1677,47 @@ public class KafkaMessageListenerContainerTests {
container.stop();
}
@SuppressWarnings({ "unchecked", "rawtypes" })
@Test
public void testInitialSeek() throws Exception {
ConsumerFactory<Integer, String> cf = mock(ConsumerFactory.class);
Consumer<Integer, String> consumer = mock(Consumer.class);
given(cf.createConsumer(isNull(), eq("clientId"), isNull())).willReturn(consumer);
ConsumerRecords<Integer, String> emptyRecords = new ConsumerRecords<>(Collections.emptyMap());
final CountDownLatch latch = new CountDownLatch(1);
given(consumer.poll(anyLong())).willAnswer(i -> {
latch.countDown();
Thread.sleep(50);
return emptyRecords;
});
TopicPartitionInitialOffset[] topicPartition = new TopicPartitionInitialOffset[] {
new TopicPartitionInitialOffset("foo", 0, SeekPosition.BEGINNING),
new TopicPartitionInitialOffset("foo", 1, SeekPosition.END),
new TopicPartitionInitialOffset("foo", 2, 0L),
new TopicPartitionInitialOffset("foo", 3, Long.MAX_VALUE),
new TopicPartitionInitialOffset("foo", 4, SeekPosition.BEGINNING),
new TopicPartitionInitialOffset("foo", 5, SeekPosition.END),
};
ContainerProperties containerProps = new ContainerProperties(topicPartition);
containerProps.setAckMode(AckMode.RECORD);
containerProps.setClientId("clientId");
containerProps.setMessageListener((MessageListener) r -> { });
KafkaMessageListenerContainer<Integer, String> container =
new KafkaMessageListenerContainer<>(cf, containerProps);
container.start();
assertThat(latch.await(10, TimeUnit.SECONDS)).isTrue();
ArgumentCaptor<Collection<TopicPartition>> captor = ArgumentCaptor.forClass(List.class);
verify(consumer).seekToBeginning(captor.capture());
assertThat(captor.getValue()
.equals(new HashSet<>(Arrays.asList(new TopicPartition("foo", 0), new TopicPartition("foo", 4)))));
verify(consumer).seekToEnd(captor.capture());
assertThat(captor.getValue()
.equals(new HashSet<>(Arrays.asList(new TopicPartition("foo", 1), new TopicPartition("foo", 5)))));
verify(consumer).seek(new TopicPartition("foo", 2), 0L);
verify(consumer).seek(new TopicPartition("foo", 3), Long.MAX_VALUE);
container.stop();
}
private Consumer<?, ?> spyOnConsumer(KafkaMessageListenerContainer<Integer, String> container) {
Consumer<?, ?> consumer = spy(
KafkaTestUtils.getPropertyValue(container, "listenerConsumer.consumer", Consumer.class));