com.networknt.eventuate.kafka.consumer.OffsetTracker Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of eventuate-server-kafka Show documentation

A module that is responsible for Kafka communication

There is a newer version: 2.0.11

Show newest version

package com.networknt.eventuate.kafka.consumer;

import org.apache.kafka.clients.consumer.OffsetAndMetadata;
import org.apache.kafka.common.TopicPartition;

import java.util.HashMap;
import java.util.Map;
import java.util.Set;

/**
 * Keeps track of message offsets that are (a) being processed and (b) have been processed and can be committed
 */
public class OffsetTracker {

  private Map state = new HashMap<>();

  @Override
  public String toString() {
    return "OffsetTracker{" +
            "state=" + state +
            '}';
  }

  TopicPartitionOffsets fetch(TopicPartition topicPartition) {
    TopicPartitionOffsets tpo = state.get(topicPartition);
    if (tpo == null) {
      tpo = new TopicPartitionOffsets();
      state.put(topicPartition, tpo);
    }
    return tpo;
  }
  void noteUnprocessed(TopicPartition topicPartition, long offset) {
    fetch(topicPartition).noteUnprocessed(offset);
  }

  void noteProcessed(TopicPartition topicPartition, long offset) {
    fetch(topicPartition).noteProcessed(offset);
  }

  public Map offsetsToCommit() {
    Map result = new HashMap<>();
    state.forEach((tp, tpo) -> {
      tpo.offsetToCommit().ifPresent(offset -> {
        result.put(tp, new OffsetAndMetadata(offset + 1, ""));
      });
    });
    return result;
  }

  public void noteOffsetsCommitted(Map offsetsToCommit) {
    offsetsToCommit.forEach((tp, om) -> {
      fetch(tp).noteOffsetCommitted(om.offset());
    });
  }

  public Map> getPending() {
    Map> result = new HashMap<>();
    state.forEach((tp, tpo) -> {
      Set pending = tpo.getPending();
      if (!pending.isEmpty())
        result.put(tp, pending);
    });
    return result;
  }
}