kafka spout

import java.util.HashMap;
import java.util.LinkedList;
import java.util.Map;
import java.util.Properties;
import java.util.Queue;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import backtype.storm.spout.SpoutOutputCollector;
import backtype.storm.task.TopologyContext;
import backtype.storm.topology.IRichSpout;
import backtype.storm.topology.OutputFieldsDeclarer;
import backtype.storm.tuple.Fields;
import backtype.storm.tuple.Values;

import com.newegg.mq.consumer.MessageAsynReceiver;
import com.newegg.mq.custom.bean.MessageData;
import com.newegg.storm.spout.fail.IFailHandler;
import com.newegg.storm.util.ConfigUtils;

/**
 * general KafkaSpout, compatible with kafka v0.8 and v0.7(jafka)
 * @author bw67
 *
 */
public class MqSpout implements IRichSpout {
	
	private static final Logger LOG = LoggerFactory.getLogger(MqSpout.class);
	private SpoutOutputCollector collector;
	private IFailHandler failHandler;
	private String topic; //one topic one thread
	private int bufferSize;
	private MessageAsynReceiver consumer;
	private Map<Integer, MessageData> inProgress;
	private Queue<Integer> queue;
	private int count;
	
	public MqSpout(){
		inProgress = new HashMap<Integer, MessageData>();
		queue = new LinkedList<Integer>();
	}
	
	@Override
	public void open(Map conf, TopologyContext context,
			SpoutOutputCollector collector) {
		
		this.collector = collector;	
		topic = ConfigUtils.getTopic((Map<String, Object>) conf);
        bufferSize = ConfigUtils.getMaxBufSize((Map<String, Object>) conf);
		createFailHandler((String) conf.get(ConfigUtils.CONFIG_FAIL_HANDLER));
		createConsumer(conf);
		failHandler.open(conf, context, collector);
		LOG.info("kafka spout opened, reading from topic {}, using failure policy {}", topic, failHandler.getIdentifier());
	}
	
    private void createFailHandler(final String failHandler) {
        if (failHandler == null) {
            this.failHandler = ConfigUtils.DEFAULT_FAIL_HANDLER;
        }
        else {
            this.failHandler = ConfigUtils.createFailHandlerFromString(failHandler);
        }
    }
    
    private void createConsumer(Map conf){
    	Properties consumerConfig = ConfigUtils.createKafkaConfig(conf);
    	consumer = new MessageAsynReceiver(consumerConfig){
			public void handleMessage(MessageData message) {
				consume(message);
			}		
		};
		consumer.connect();
		LOG.info("connecting kafka client to zookeeper at {} as client group {}",
	            consumerConfig.getProperty("zookeeper.connect"),
	            consumerConfig.getProperty("group.id"));
    }
	
	private void consume(MessageData message) {
		inProgress.put(count, message);
		queue.add(count);
	}
	
	private boolean fillBuffer() {
		if (!inProgress.isEmpty() || !queue.isEmpty()) {
            throw new IllegalStateException("cannot fill buffer when buffer or pending messages are non-empty");
        }
		
		count = 0;
        while (count < bufferSize && consumer.consumeNextMessageFrom(topic)) {
            count++;
        }
        if (inProgress.size() > 0) {
            // set queue to all currently pending kafka message ids
//            queue.addAll(inProgress.keySet());
            LOG.debug("buffer now has {} messages to be emitted", queue.size());
            // message(s) appended to buffer
            return true;
        }
        else {
            // no messages appended to buffer
            return false;
        }
		
	}

	@Override
	public void close() {
		// reset state by setting members to null
        consumer.close();
        failHandler.close();
		
	}

	@Override
	public void activate() {
		failHandler.activate();
		
	}

	@Override
	public void deactivate() {
		failHandler.deactivate();
		
	}

	@Override
	public void nextTuple() {
		// next tuple available when _queue contains ids or fillBuffer() is allowed and indicates more messages were available
        // see class documentation for implementation note on the rationale behind this condition
        if (!queue.isEmpty() || (inProgress.isEmpty() && fillBuffer())) {
            final Integer nextId = queue.poll();
            if (nextId != null) {
                final MessageData message = inProgress.get(nextId);
                // the next id from buffer should correspond to a message in the pending map
                if (message == null) {
                    throw new IllegalStateException("no pending message for next id " + nextId);
                }
                // message should be considered a single object from Values' point of view
                collector.emit(new Values((Object) message), nextId);
//                System.out.println(nextId + " " + message.getInDate() + " " + message.getTagType());
                LOG.debug("emitted kafka message id {}", nextId);
            }
        }
		
	}

	@Override
	public void ack(Object msgId) {
		if (msgId instanceof Integer) {
            final Integer id = (Integer) msgId;
            // message corresponding to o is no longer pending
            inProgress.remove(id);
            LOG.debug("kafka message {} acknowledged", id);
//            System.out.println("kafka message " + id + " acknowledged");
            if (inProgress.isEmpty()) {
                // commit offsets to zookeeper when pending is now empty
                // (buffer will be filled on next call to nextTuple())
                LOG.debug("all pending messages acknowledged, committing client offsets");
                consumer.commitOffsets();
//                System.out.println("offsets commited!");
            }
            // notify fail handler of tuple success
            failHandler.ack(id);
        }
		
	}

	@Override
	public void fail(Object msgId) {
		if (msgId instanceof Integer) {
            final Integer id = (Integer) msgId;
            // delegate decision of replaying the message to failure policy
            if (failHandler.shouldReplay(id)) {
                LOG.debug("kafka message id {} failed in topology, adding to buffer again", id);
                queue.add(id);
            }
            else {
                LOG.debug("kafka message id {} failed in topology, delegating failure to policy", id);
                // remove message from pending; _failHandler will take action if needed
                failHandler.fail(id, inProgress.remove(id));
            }
        }
		
	}

	@Override
	public void declareOutputFields(OutputFieldsDeclarer declarer) {
		declarer.declare(new Fields("message"));
		
	}

	@Override
	public Map<String, Object> getComponentConfiguration() {
		return null;
	}

}

reference: https://github.com/HolmesNL/kafka-spout