package com.liveramp.kafka_service.consumer;

import java.util.Properties;

import org.apache.kafka.clients.consumer.KafkaConsumer;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.liveramp.kafka_service.consumer.persist_helpers.KafkaOffsetHelper;
import com.liveramp.kafka_service.consumer.rebalance_callbacks.BaseRebalanceCallback;

public abstract class BaseConsumer<K, V> {
  private static final Logger LOG = LoggerFactory.getLogger(BaseConsumer.class);

  private final KafkaOffsetHelper offsetHelper;
  private final KafkaConsumer<byte[], byte[]> consumer;

  protected BaseConsumer(Properties config, KafkaOffsetHelper offsetHelper) {
    config.setProperty("enable.auto.commit", "false");
    BaseRebalanceCallback baseRebalanceCallback = BaseRebalanceCallback.create(offsetHelper);
    this.offsetHelper = offsetHelper;
    this.consumer = new KafkaConsumer<>(config, baseRebalanceCallback);
  }

  public void startSubscription(String topic) {
    consumer.subscribe(topic);
  }

  public void stopSubscription(String topic) {
    consumer.unsubscribe(topic);
  }

  public void close() {
    consumer.close();
  }

}

@@ -60,20 +60,261 @@
 til.
-Properties;%0A
+Map;%0Aimport java.util.Properties;%0Aimport java.util.Queue;%0A%0Aimport com.google.common.collect.Maps;%0Aimport com.google.common.collect.Queues;%0Aimport org.apache.kafka.clients.consumer.ConsumerRecord;%0Aimport org.apache.kafka.clients.consumer.ConsumerRecords;
 %0Aimp
@@ -712,16 +712,68 @@
 .class);
+%0A  private static final long DEFAULT_TIMEOUT = 3000;
 %0A%0A  priv
@@ -848,32 +848,86 @@
 mer%3C
-byte%5B%5D, byte%5B%5D%3E consumer
+K, V%3E consumer;%0A  private final Map%3CString, Queue%3CConsumerRecord%3CK, V%3E%3E%3E cache
 ;%0A%0A 
@@ -1262,16 +1262,52 @@
 lback);%0A
+    this.cache = Maps.newHashMap();%0A
   %7D%0A%0A  p
@@ -1322,24 +1322,16 @@
 id s
-tartS
 ubscri
-ption
+be
 (Str
@@ -1366,32 +1366,100 @@
 bscribe(topic);%0A
+    cache.put(topic, Queues.%3CConsumerRecord%3CK, V%3E%3EnewArrayDeque());%0A
   %7D%0A%0A  public vo
@@ -1465,13 +1465,11 @@
 oid 
-stopS
+uns
 ubsc
@@ -1525,16 +1525,41 @@
 topic);%0A
+    cache.remove(topic);%0A
   %7D%0A%0A  p
@@ -1606,10 +1606,795 @@
 );%0A  %7D%0A%0A
+  public boolean hasNext(String topic) %7B%0A    if (!cache.containsKey(topic)) %7B%0A      return false;%0A    %7D%0A%0A    if (cache.get(topic).size() %3E 0) %7B%0A      return true;%0A    %7D%0A%0A    Map%3CString, ConsumerRecords%3CK, V%3E%3E topicToRecordsMap = consumer.poll(DEFAULT_TIMEOUT);%0A    for (Map.Entry%3CString, ConsumerRecords%3CK, V%3E%3E topicToRecords : topicToRecordsMap.entrySet()) %7B%0A      cache.get(topicToRecords.getKey()).addAll(topicToRecords.getValue().records());%0A    %7D%0A%0A    return cache.get(topic).size() %3E 0;%0A  %7D%0A%0A  public V next(String topic) %7B%0A    ConsumerRecord%3CK, V%3E record = cache.get(topic).remove();%0A%0A    try %7B%0A      offsetHelper.persistOffset(record.topicAndPartition(), record.offset());%0A      return record.value();%0A    %7D catch (Exception e) %7B%0A      throw new RuntimeException(e);%0A    %7D%0A  %7D%0A
 %7D%0A

Complete BaseConsumer