package cgl.iotcloud.transport.kafka.consumer; import cgl.iotcloud.core.msg.MessageContext; import cgl.iotcloud.core.transport.Manageable; import com.google.common.base.Joiner; import kafka.message.Message; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import java.util.*; import java.util.concurrent.BlockingQueue; public class KConsumer implements Manageable { private static Logger LOG = LoggerFactory.getLogger(KConsumer.class); public static class MessageAndRealOffset { public Message msg; public long offset; public MessageAndRealOffset(Message msg, long offset) { this.msg = msg; this.offset = offset; } } public static enum EmitState { EMITTED_MORE_LEFT, EMITTED_END, NO_EMITTED } private BlockingQueue<MessageContext> messageContexts; private boolean run = true; private String _uuid = UUID.randomUUID().toString(); ConsumerConfig _consumerConfig; PartitionCoordinator _coordinator; DynamicPartitionConnections _connections; ZkState _state; long _lastUpdateMs = 0; int _currPartitionIndex = 0; String _site; public KConsumer(String _site, BlockingQueue<MessageContext> messageContexts, ConsumerConfig consumerConfig) { this._site = _site; this._consumerConfig = consumerConfig; this.messageContexts = messageContexts; } private void close() { run = false; _state.close(); } private void nextTuple() { List<PartitionManager> managers = _coordinator.getMyManagedPartitions(); for (int i = 0; i < managers.size(); i++) { try { // in case the number of managers decreased _currPartitionIndex = _currPartitionIndex % managers.size(); EmitState state = managers.get(_currPartitionIndex).next(messageContexts); if (state != EmitState.EMITTED_MORE_LEFT) { _currPartitionIndex = (_currPartitionIndex + 1) % managers.size(); } if (state != EmitState.NO_EMITTED) { break; } } catch (FailedFetchException e) { LOG.warn("Fetch failed", e); _coordinator.refresh(); } } long now = System.currentTimeMillis(); if ((now - _lastUpdateMs) > _consumerConfig.stateUpdateIntervalMs) { commit(); } } private void deactivate() { commit(); } private void commit() { _lastUpdateMs = System.currentTimeMillis(); for (PartitionManager manager : _coordinator.getMyManagedPartitions()) { manager.commit(); } } public class Worker implements Runnable { @Override public void run() { while (run) { nextTuple(); } } } @Override public void start() { List<String> zkServers = _consumerConfig.zkServers; String servers = Joiner.on(",").join(zkServers); _state = new ZkState(servers, _consumerConfig.zkRoot); _connections = new DynamicPartitionConnections(_consumerConfig, KafkaUtils.makeBrokerReader(_consumerConfig)); // using TransactionalState like this is a hack _coordinator = new ZkCoordinator(_connections, _consumerConfig, _state, 0, 1, _uuid, _site); Thread t = new Thread(new Worker()); t.start(); } @Override public void stop() { deactivate(); run = false; close(); } }