/*******************************************************************************
* Copyright © 2012-2015 eBay Software Foundation
* This program is dual licensed under the MIT and Apache 2.0 licenses.
* Please see LICENSE for more information.
*******************************************************************************/
package com.ebay.jetstream.event.channel.kafka;
import static kafka.api.OffsetRequest.CurrentVersion;
import static kafka.api.OffsetRequest.EarliestTime;
import static kafka.api.OffsetRequest.LargestTimeString;
import static kafka.api.OffsetRequest.LatestTime;
import static kafka.api.OffsetRequest.SmallestTimeString;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicBoolean;
import kafka.api.FetchRequest;
import kafka.api.FetchRequestBuilder;
import kafka.api.PartitionOffsetRequestInfo;
import kafka.cluster.Broker;
import kafka.common.ErrorMapping;
import kafka.common.OffsetOutOfRangeException;
import kafka.common.TopicAndPartition;
import kafka.javaapi.FetchResponse;
import kafka.javaapi.OffsetRequest;
import kafka.javaapi.OffsetResponse;
import kafka.javaapi.consumer.SimpleConsumer;
import kafka.javaapi.message.ByteBufferMessageSet;
import kafka.message.Message;
import kafka.message.MessageAndOffset;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import com.ebay.jetstream.event.JetstreamEvent;
import com.ebay.jetstream.event.JetstreamReservedKeys;
/**
* @author xiaojuwu1
*
*/
public class PartitionReader implements Comparable<PartitionReader> {
private static final Logger LOGGER = LoggerFactory.getLogger(PartitionReader.class
.getName());
private final static String kmKey = JetstreamReservedKeys.EventKafkaMeta
.toString();
private final KafkaController.ZkConnector m_zkConnector;
private final String m_groupId;
private final String m_topic;
private final int m_partition;
private final String m_clientId;
private final KafkaConsumerConfig m_config;
private final KafkaMessageSerializer m_serializer;
private final int m_retries = 3;
private final String m_pstatePath;
private final String m_offsetPath;
private Broker m_leader;
private SimpleConsumer m_consumer;
private long m_readOffset = -1;
private int m_nextBatchSizeBytes = -1;
/**
* whether the partition is taken or to be released
*/
private final AtomicBoolean m_taken = new AtomicBoolean(true);
/**
* is lost in zookeeper, determines onStreamTermination/onException to be
* called
*/
private final AtomicBoolean m_lost = new AtomicBoolean(false);
/**
* to control the wait time when nothing fetched
*/
private long m_nextFetchInMs = 0;
/**
* onIdle control
*/
private long m_idleTimestamp = 0;
private boolean m_onIdled = false;
public PartitionReader(String topic, int partition,
KafkaConsumerConfig config,
KafkaController.ZkConnector zkConnector,
KafkaMessageSerializer serializer) {
this.m_groupId = config.getGroupId();
this.m_topic = topic;
this.m_partition = partition;
this.m_pstatePath = partitionStatePath(topic, partition);
this.m_offsetPath = offsetPath(m_groupId, topic, partition);
this.m_clientId = this.m_topic + "_" + this.m_partition;
this.m_config = config;
this.m_serializer = serializer;
this.m_zkConnector = zkConnector;
init();
}
public void initOffset(boolean isNew) {
if (LOGGER.isInfoEnabled())
LOGGER.info( "Init offset for " + m_clientId);
if (isNew) {
m_readOffset = fetchResetOffset(SmallestTimeString());
// commit the init offset first, useful to revert
commitOffset();
} else {
if (offsetExists())
m_readOffset = readZkOffset();
if (m_readOffset < 0) {
m_readOffset = fetchResetOffset(m_config.getAutoOffsetReset());
// commit the init offset first, useful to revert
commitOffset();
}
}
if (m_readOffset < 0)
throw new RuntimeException(
"Fatal errors in getting init offset for " + m_clientId);
if (LOGGER.isInfoEnabled())
LOGGER.info( m_clientId + ": init offset is set to "
+ m_readOffset);
}
public void resetOffset() {
if (LOGGER.isInfoEnabled())
LOGGER.info( "Reset offset for " + m_clientId);
m_readOffset = fetchResetOffset(m_config.getAutoOffsetReset());
if (m_readOffset < 0)
throw new RuntimeException("Fatal errors in resetting offset for "
+ m_clientId);
commitOffset();
if (LOGGER.isInfoEnabled())
LOGGER.info( m_clientId + ": offset is reset to "
+ m_readOffset);
}
public void revertOffset() {
long oldOffset = m_readOffset;
m_readOffset = readZkOffset();
if (m_readOffset < 0) {
LOGGER.warn( m_clientId
+ "No offset node in zookeeper.");
resetOffset();
}
if (LOGGER.isInfoEnabled())
LOGGER.info( m_clientId + ": offset is reverted from "
+ oldOffset + " to " + m_readOffset);
}
public long getReadOffset() {
return m_readOffset;
}
public void setReadOffset(long offset) {
if (LOGGER.isInfoEnabled())
LOGGER.info( m_clientId + "'s readOffset is set to "
+ offset + " from " + m_readOffset);
m_readOffset = offset;
}
public void setNextBatchSizeBytes(int nextBatchSizeBytes) {
if (nextBatchSizeBytes < 0)
return;
this.m_nextBatchSizeBytes = nextBatchSizeBytes;
}
public boolean isTaken() {
return m_taken.get();
}
/**
* set to release and ConsumeTask will do it in next round
*/
public void setToRelease() {
m_taken.set(false);
}
public void setLost() {
m_lost.set(true);
}
public boolean isLost() {
return m_lost.get();
}
public String getTopic() {
return m_topic;
}
public int getPartition() {
return m_partition;
}
public String getId() {
return m_clientId;
}
public long calcWaitTime() {
return m_nextFetchInMs - System.currentTimeMillis();
}
public long getNextFetchInMs() {
return m_nextFetchInMs;
}
public void setNextFetchInMs(long nextFetchInMs) {
this.m_nextFetchInMs = nextFetchInMs;
}
public void fistIdle() {
this.m_idleTimestamp = System.currentTimeMillis();
}
public long getIdleTime() {
return System.currentTimeMillis() - m_idleTimestamp;
}
public boolean isIdled() {
return this.m_idleTimestamp > 0;
}
public void setOnIdled() {
this.m_idleTimestamp = 0;
this.m_onIdled = true;
}
public boolean isOnIdled() {
return this.m_onIdled;
}
public void clearIdleStats() {
this.m_idleTimestamp = 0;
this.m_onIdled = false;
}
/**
* consumer events from kafka by one FetchRequest
*/
public List<JetstreamEvent> readEvents() throws OffsetOutOfRangeException {
List<JetstreamEvent> events = new ArrayList<JetstreamEvent>();
if (m_nextBatchSizeBytes < 0)
m_nextBatchSizeBytes = m_config.getBatchSizeBytes();
if (m_nextBatchSizeBytes == 0) {
// nextBatchSize only affects one fetch
m_nextBatchSizeBytes = m_config.getBatchSizeBytes();
return events;
}
for (int i = 0; i < 2; i++) {
for (int j = 0; j < m_retries; j++) {
FetchRequest req = new FetchRequestBuilder()
.clientId(m_clientId)
.addFetch(m_topic, m_partition, m_readOffset,
m_nextBatchSizeBytes).build();
FetchResponse fetchResponse = null;
try {
fetchResponse = m_consumer.fetch(req);
} catch (Exception e) {
LOGGER.error( "Error occurs when fetching.", e);
continue;
}
if (fetchResponse == null)
continue;
if (fetchResponse.hasError()) {
short code = fetchResponse.errorCode(m_topic, m_partition);
if (code == ErrorMapping.OffsetOutOfRangeCode()) {
long smallest = -1L, largest = -1L;
try {
smallest = fetchResetOffset(SmallestTimeString());
largest = fetchResetOffset(LargestTimeString());
} catch (Exception e) {
}
throw new OffsetOutOfRangeException(m_clientId
+ ": readOffset=" + m_readOffset + " smallest="
+ smallest + " largest=" + largest);
} else
continue;
} else {
ByteBufferMessageSet messageSet = fetchResponse.messageSet(
m_topic, m_partition);
boolean hasMessage = messageSet.iterator().hasNext();
if (!hasMessage && !readToTheEnd()) {
m_nextBatchSizeBytes = Math.min(
m_nextBatchSizeBytes * 2,
m_config.getMaxBatchSizeBytes());
continue;
}
for (MessageAndOffset messageAndOffset : messageSet) {
long currentOffset = messageAndOffset.offset();
if (currentOffset < m_readOffset) {
continue;
}
m_readOffset = messageAndOffset.nextOffset();
Message message = messageAndOffset.message();
ByteBuffer k = message.key();
ByteBuffer p = message.payload();
byte[] key = null;
if (k != null) {
key = new byte[k.limit()];
k.get(key);
}
byte[] payload = new byte[p.limit()];
p.get(payload);
JetstreamEvent event = m_serializer
.decode(key, payload);
EventKafkaMetadata meta = new EventKafkaMetadata(
m_topic, m_partition, currentOffset);
event.put(kmKey, meta.encode());
events.add(event);
}
// nextBatchSize only affects one fetch
m_nextBatchSizeBytes = m_config.getBatchSizeBytes();
return events;
}
}
// cannot get events after retries, reinit and try again
reinit();
}
throw new RuntimeException("Fail to read events from " + m_clientId
+ " with offset " + m_readOffset);
}
public void commitOffset() {
if (!m_zkConnector.exists(m_offsetPath))
m_zkConnector.create(m_offsetPath, false);
m_zkConnector.writeString(m_offsetPath, String.valueOf(m_readOffset));
if (LOGGER.isInfoEnabled())
LOGGER.info( m_clientId + ": offset is committed with "
+ m_readOffset);
}
public void close() {
m_consumer.close();
}
private boolean offsetExists() {
return m_zkConnector.exists(m_offsetPath);
}
public long readZkOffset() {
try {
if (!offsetExists())
return -1;
String offsetStr = m_zkConnector.readString(m_offsetPath);
if (offsetStr == null)
return -1;
return Long.parseLong(offsetStr);
} catch (RuntimeException e) {
throw e;
} catch (Exception e) {
throw new RuntimeException(e);
}
}
private boolean readToTheEnd() {
return m_readOffset >= getLargestOffset();
}
private long getLargestOffset() {
return fetchResetOffset(LargestTimeString());
}
private long fetchResetOffset(String reset) {
long time = LatestTime();
if (reset != null && reset.equals(SmallestTimeString()))
time = EarliestTime();
Map<TopicAndPartition, PartitionOffsetRequestInfo> requestInfo = new HashMap<TopicAndPartition, PartitionOffsetRequestInfo>();
TopicAndPartition tp = new TopicAndPartition(m_topic, m_partition);
PartitionOffsetRequestInfo info = new PartitionOffsetRequestInfo(time,
1);
requestInfo.put(tp, info);
OffsetRequest request = new OffsetRequest(requestInfo,
CurrentVersion(), m_clientId);
for (int i = 0; i < 2; i++) {
for (int j = 0; j < m_retries; j++) {
OffsetResponse response = m_consumer.getOffsetsBefore(request);
if (response.hasError()) {
short errorCode = response.errorCode(m_topic, m_partition);
LOGGER.warn(
"Error when fetch offset from kafka, errorCode="
+ errorCode);
continue;
}
long[] offsets = response.offsets(m_topic, m_partition);
if (offsets.length <= 0)
continue;
return offsets[0];
}
// cannot get offset after retries, reinit and try again
reinit();
}
throw new RuntimeException("Fail to get resetOffset " + reset
+ " after retries for " + m_clientId);
}
private void reinit() {
if (m_consumer != null)
m_consumer.close();
init();
}
private void init() {
m_taken.set(true);
this.m_leader = getLeader();
this.m_consumer = new SimpleConsumer(m_leader.host(), m_leader.port(),
m_config.getSocketTimeoutMs(),
m_config.getSocketReceiveBufferBytes(), m_clientId);
this.m_nextBatchSizeBytes = m_config.getBatchSizeBytes();
}
private Broker getLeader() {
try {
Map<String, Object> state = m_zkConnector.readJSON(m_pstatePath);
Integer brokerId = ((Number) state.get("leader")).intValue();
Map<String, Object> brokerData = m_zkConnector
.readJSON(brokerPath(brokerId));
if (brokerData == null)
throw new RuntimeException("Broker info not found for "
+ m_clientId);
String host = (String) brokerData.get("host");
Integer port = ((Number) brokerData.get("port")).intValue();
return new Broker(brokerId, host, port);
} catch (RuntimeException e) {
throw e;
} catch (Exception e) {
throw new RuntimeException(e);
}
}
private String partitionStatePath(String topic, int partition) {
StringBuffer sb = new StringBuffer("/brokers/topics/");
sb.append(topic).append("/partitions/");
sb.append(partition).append("/state");
return sb.toString();
}
private String brokerPath(int brokerId) {
StringBuffer sb = new StringBuffer("/brokers/ids/");
sb.append(brokerId);
return sb.toString();
}
private String offsetPath(String groupId, String topic, int partition) {
StringBuffer sb = new StringBuffer("/consumers/");
sb.append(groupId).append("/offsets/");
sb.append(topic).append("/").append(partition);
return sb.toString();
}
@Override
public int compareTo(PartitionReader p) {
if (m_nextFetchInMs < p.getNextFetchInMs())
return -1;
else if (m_nextFetchInMs > p.getNextFetchInMs())
return 1;
else
return 0;
}
}