package org.apache.hedwig.client.loadtest;

import com.google.protobuf.ByteString;
import com.google.protobuf.InvalidProtocolBufferException;
import org.apache.hedwig.client.HedwigClient;
import org.apache.hedwig.client.api.MessageHandler;
import org.apache.hedwig.client.api.Subscriber;
import org.apache.hedwig.client.conf.ClientConfiguration;
import org.apache.hedwig.exceptions.PubSubException;
import org.apache.hedwig.protocol.LoadTest;
import org.apache.hedwig.protocol.LoadTest.LoadTestMessage;
import org.apache.hedwig.protocol.PubSubProtocol;
import org.apache.hedwig.protocol.PubSubProtocol.Message;
import org.apache.hedwig.protocol.PubSubProtocol.SubscribeRequest.CreateOrAttach;
import org.apache.hedwig.util.Callback;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.LinkedList;
import java.util.concurrent.TimeUnit;

public class LoadTestReceiver extends LoadTestBase {
    private static Logger logger = LoggerFactory.getLogger(LoadTestReceiver.class);
    private int numSubscribers;
    private int subscriberStartIndex;
    private LinkedList<SingleSubscriber> subList = new LinkedList<SingleSubscriber>();
    public LoadTestReceiver(TopicProvider topicProvider, LoadTestUtils ltUtil,
                            ClientConfiguration conf,
                            int numSubscribers, int subscriberStartIndex) {
        super(topicProvider, ltUtil, conf, "receive");
        this.numSubscribers = numSubscribers;
        this.subscriberStartIndex = subscriberStartIndex;
    }

    private class SingleSubscriber implements Runnable {
        final HedwigClient client;
        final ByteString subId;
        public SingleSubscriber(ByteString subId) {
            this.subId = subId;
            this.client = new HedwigClient(conf);
        }

        private void startDelivery(ByteString topic) {
            try {
                client.getSubscriber().startDelivery(topic, subId, new MessageHandler() {
                    @Override
                    public void deliver(ByteString topic, ByteString subscriberId, Message msg,
                                        Callback<Void> callback, Object context) {
                        if (logger.isDebugEnabled()) {
                            logger.debug("Received message:" + msg + " for topic:" + topic.toString() + " for" +
                                    " subscriber:" + subscriberId.toStringUtf8());
                        }
                        try {
                            LoadTestMessage ltm = LoadTestMessage.parseFrom(msg.getBody().toByteArray());
                            long latencyMillis = System.currentTimeMillis() - ltm.getTimestamp();
                            stat.requestComplete(TimeUnit.MILLISECONDS.toMicros(latencyMillis));
                        } catch (InvalidProtocolBufferException e) {
                            logger.error("Error while parsing message:" + msg, e);
                        }
                        callback.operationFinished(context, null);
                    }
                });
            } catch (Exception e) {
                logger.error("Exception while starting delivery for topic:" + topic.toStringUtf8() +
                        " for subscriber:" + subId.toStringUtf8(), e);
            }
        }
        public void stop() {
            for (ByteString topic : topicProvider.getTopicList()) {
                try {
                    client.getSubscriber().closeSubscription(topic, subId);
                } catch (Exception e) {
                    logger.error("Exception while unsubscribing from topic:" + topic.toStringUtf8() +
                            " for subscriber:" + subId.toStringUtf8());
                }
            }
        }

        public void run() {
            Subscriber subscriber = client.getSubscriber();
            for (final ByteString topic : topicProvider.getTopicList()) {
                subscriber.asyncSubscribe(topic, subId, CreateOrAttach.CREATE_OR_ATTACH, new Callback<Void>() {
                    @Override
                    public void operationFinished(Object ctx, Void resultOfOperation) {
                        if (logger.isDebugEnabled()) {
                            logger.debug("Subscription for topic:" + topic.toStringUtf8() + " for " +
                                    "subscriber:" + subId.toStringUtf8() + " succeeded");
                        }
                        startDelivery(topic);
                    }
                    @Override
                    public void operationFailed(Object ctx, PubSubException exception) {
                        logger.error("Subscription to topic:" + topic.toStringUtf8() + " for " +
                                "subscriber:" + subId.toStringUtf8() + " failed", exception);
                    }
                }, null);
            }
        }
    }

    public void start() {
        for (int i = subscriberStartIndex; i < subscriberStartIndex+numSubscribers; i++) {
            ByteString subId = ltUtil.getSubscriberFromNumber(i);
            SingleSubscriber singleSub = new SingleSubscriber(subId);
            subList.add(singleSub);
            singleSub.run();
        }
    }

    public void stop() {
        for (SingleSubscriber sub : subList) {
            sub.stop();
        }
    }
}

@@ -127,32 +127,71 @@
 ufferException;%0A
+import com.twitter.common.stats.Stats;%0A
 import org.apach
@@ -220,16 +220,16 @@
 Client;%0A
-
 import o
@@ -889,16 +889,115 @@
 imeUnit;
+%0Aimport java.util.concurrent.atomic.AtomicLong;%0Aimport java.util.concurrent.atomic.AtomicLongArray;
 %0A%0Apublic
@@ -1979,24 +1979,149 @@
 ng topic) %7B%0A
+            final AtomicLong lastId = new AtomicLong(-1);%0A            Stats.export(%22lastId%22 + subId.toStringUtf8(), lastId);%0A
             
@@ -2668,32 +2668,32 @@
 oStringUtf8());%0A
-
                 
@@ -2694,32 +2694,473 @@
               %7D%0A
+                        long newId = msg.getMsgId().getLocalComponent();%0A                        long oldId = lastId.get();%0A                        if (oldId != -1 && oldId+1 != newId) %7B%0A                            logger.warn(%22Last id is %7B%7D and new id is %7B%7D. Skipped %7B%7D messages.%22, new Object%5B%5D %7BoldId, newId, newId-oldId%7D);%0A                            stat.incErrors();%0A                        %7D%0A                        lastId.set(newId);%0A
                 

Add logging for skipped messages in the load test Log if we skip any messages. Sequence ids should be strictly increasing by 1.