com.pinterest.doctorkafka.util.ReplicaStatsUtil Maven / Gradle / Ivy

Go to download
package com.pinterest.doctorkafka.util;

import org.apache.kafka.clients.consumer.KafkaConsumer;
import org.apache.kafka.clients.consumer.OffsetAndTimestamp;
import org.apache.kafka.common.TopicPartition;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

public class ReplicaStatsUtil {

  public static Map getProcessingStartOffsets(
      KafkaConsumer kafkaConsumer,
      String brokerStatsTopic,
      long startTimestampInMillis) {
    List tpList = kafkaConsumer.partitionsFor(brokerStatsTopic).stream()
                  .map(p->new TopicPartition(p.topic(), p.partition())).collect(Collectors.toList());
    Map partitionMap = new HashMap<>();
    for (TopicPartition topicPartition : tpList) {
      partitionMap.put(topicPartition, startTimestampInMillis);
    }

    Map offsetsForTimes = kafkaConsumer
        .offsetsForTimes(partitionMap);
    for (Map.Entry entry : offsetsForTimes.entrySet()) {
      partitionMap.put(entry.getKey(), entry.getValue().offset());
    }
    return partitionMap;
  }
}