OffsetResetter implements ConsumerRebalanceListener { @Override public void onPartitionsAssigned(Collection<TopicPartition> partitions) { Map<Integer, Collection<Long>> offsetsPerPartition = new HashMap<>(); for (PartitionedWriter<MESSAGE_KIND> writer : writers) { final List<Integer> partitionsId = partitions.stream() .map(TopicPartition::partition) .collect(Collectors.toList()); try { final Map<Integer, Long> startingOffsets = writer.getStartingOffsets(partitionsId); startingOffsets.forEach((part, offset) -> offsetsPerPartition.computeIfAbsent(part, ignored -> new HashSet<>()).add(offset)); } catch (IOException e) { LOGGER.warn("Couldn't get offset for partitions {}, will resume from earliest", partitionsId.stream().map(String::valueOf).collect(Collectors.joining(", ")), e); partitions.forEach(part -> offsetsPerPartition.computeIfAbsent(part.partition(), ignored -> new HashSet<>()) .add(UNKNOWN_OFFSET)); } } Map<Integer, Long> minOffsetPerPartition = offsetsPerPartition .entrySet() .stream() .collect(Collectors.toMap(Map.Entry::getKey, offsets -> offsets .getValue() .stream() .filter(offset -> offset != UNKNOWN_OFFSET) .mapToLong(l -> l) .min() .orElse(UNKNOWN_OFFSET)) ); for (TopicPartition topicPartition: partitions) { int partition = topicPartition.partition(); long startingOffset = minOffsetPerPartition.getOrDefault(partition, UNKNOWN_OFFSET); synchronized (consumer) { if (startingOffset == UNKNOWN_OFFSET) { consumer.seekToBeginning(Collections.singleton(topicPartition)); LOGGER.warn("Resuming consumption of partition {} from the beginning. " + "This should not happen unless this is the first time this app runs", partition); } else { consumer.seek(topicPartition, startingOffset); LOGGER.info("Resuming consumption of partition {} from offset {}", partition, startingOffset); } } } } OffsetResetter(Consumer<K, V> consumer, java.util.function.Consumer<Integer> partitionRevokedConsumer,
                          Collection<PartitionedWriter<MESSAGE_KIND>> writers); @Override void onPartitionsRevoked(Collection<TopicPartition> partitions); @Override void onPartitionsAssigned(Collection<TopicPartition> partitions);  }
@Test public void partitionAssignedValidOffset() throws IOException { final Consumer<Long, String> consumer = mock(Consumer.class); final PartitionedWriter firstWriter = mock(PartitionedWriter.class); final PartitionedWriter secondWriter = mock(PartitionedWriter.class); final PartitionedWriter newEventWriter = mock(PartitionedWriter.class); final OffsetResetter offsetResetter = new OffsetResetter<>(consumer, mock(java.util.function.Consumer.class), Arrays.asList(firstWriter, secondWriter, newEventWriter)); final TopicPartition firstPartition = new TopicPartition(TOPIC, 1); final TopicPartition secondPartition = new TopicPartition(TOPIC, 2); final List<TopicPartition> partitions = Arrays.asList(firstPartition, secondPartition); final Map<Integer, Long> firstOffsets = new HashMap<>(); firstOffsets.put(firstPartition.partition(), 10L); firstOffsets.put(secondPartition.partition(), OffsetComputer.NO_OFFSET); when(firstWriter.getStartingOffsets(any())).thenReturn(firstOffsets); final Map<Integer, Long> secondOffsets = new HashMap<>(); secondOffsets.put(firstPartition.partition(), 15L); secondOffsets.put(secondPartition.partition(), OffsetComputer.NO_OFFSET); when(secondWriter.getStartingOffsets(any())).thenReturn(secondOffsets); final Map<Integer, Long> thirdOffsets = new HashMap<>(); thirdOffsets.put(firstPartition.partition(), OffsetComputer.NO_OFFSET); thirdOffsets.put(secondPartition.partition(), OffsetComputer.NO_OFFSET); when(newEventWriter.getStartingOffsets(any())).thenReturn(thirdOffsets); offsetResetter.onPartitionsAssigned(partitions); verify(consumer, times(1)).seek(eq(firstPartition), eq(10L)); verify(consumer, times(1)).seekToBeginning(Collections.singleton(secondPartition)); verifyNoMoreInteractions(consumer); }
@Test public void partitionsAssignedCannotFetchOffset() throws IOException { final Consumer<Long, String> consumer = mock(Consumer.class); final PartitionedWriter successfulWriter = mock(PartitionedWriter.class); final PartitionedWriter exceptionalWriter = mock(PartitionedWriter.class); final OffsetResetter offsetResetter = new OffsetResetter<>(consumer, mock(java.util.function.Consumer.class), Arrays.asList(successfulWriter, exceptionalWriter)); final TopicPartition partition = new TopicPartition(TOPIC, 1); final List<TopicPartition> partitions = Collections.singletonList(partition); when(successfulWriter.getStartingOffsets(any())).thenReturn(new HashMap<>()); when(exceptionalWriter.getStartingOffsets(any())).thenThrow(new IOException("Ayo")); offsetResetter.onPartitionsAssigned(partitions); verify(consumer, times(1)).seekToBeginning(Collections.singleton(partition)); verifyNoMoreInteractions(consumer); }
@Test public void partitionsAssignedNoWriter() { final Consumer<Long, String> consumer = mock(Consumer.class); final OffsetResetter offsetResetter = new OffsetResetter<>(consumer, mock(java.util.function.Consumer.class), Collections.emptyList()); final TopicPartition partition = new TopicPartition(TOPIC, 1); final List<TopicPartition> partitions = Collections.singletonList(partition); offsetResetter.onPartitionsAssigned(partitions); verify(consumer, times(1)).seekToBeginning(Collections.singleton(partition)); verifyNoMoreInteractions(consumer); }
@Test public void partitionsAssignedNoTopic() { final Consumer<Long, String> consumer = mock(Consumer.class); final OffsetResetter offsetResetter = new OffsetResetter<>(consumer, mock(java.util.function.Consumer.class), Collections.singleton(mock(PartitionedWriter.class))); offsetResetter.onPartitionsAssigned(Collections.emptyList()); verifyZeroInteractions(consumer); }