/**
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.activemq.broker;

import java.util.concurrent.Semaphore;
import java.util.concurrent.atomic.AtomicInteger;

import junit.framework.Test;
import org.apache.activemq.command.ActiveMQDestination;
import org.apache.activemq.command.ActiveMQQueue;
import org.apache.activemq.command.ActiveMQTopic;
import org.apache.activemq.command.ConnectionInfo;
import org.apache.activemq.command.ConsumerInfo;
import org.apache.activemq.command.Message;
import org.apache.activemq.command.MessageAck;
import org.apache.activemq.command.ProducerInfo;
import org.apache.activemq.command.SessionInfo;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

/**
 * BrokerBenchmark is used to get an idea of the raw performance of a broker.
 * Since the broker data structures using in message dispatching are under high
 * contention from client requests, it's performance should be monitored closely
 * since it typically is the biggest bottleneck in a high performance messaging
 * fabric. The benchmarks are run under all the following combinations options:
 * Queue vs. Topic, 1 vs. 10 producer threads, 1 vs. 10 consumer threads, and
 * Persistent vs. Non-Persistent messages. Message Acking uses client ack style
 * batch acking since that typically has the best ack performance.
 * 
 * @version $Revision: 1.9 $
 */
public class BrokerBenchmark extends BrokerTestSupport {
    private static final transient Log LOG = LogFactory.getLog(BrokerBenchmark.class);

    public int produceCount = Integer.parseInt(System.getProperty("PRODUCE_COUNT", "10000"));
    public ActiveMQDestination destination;
    public int prodcuerCount;
    public int consumerCount;
    public boolean deliveryMode;

    public void initCombosForTestPerformance() {
        addCombinationValues("destination", new Object[] {new ActiveMQQueue("TEST"), new ActiveMQTopic("TEST")});
        addCombinationValues("PRODUCER_COUNT", new Object[] {new Integer("1"), new Integer("10")});
        addCombinationValues("CONSUMER_COUNT", new Object[] {new Integer("1"), new Integer("10")});
        addCombinationValues("CONSUMER_COUNT", new Object[] {new Integer("1"), new Integer("10")});
        addCombinationValues("deliveryMode", new Object[] {Boolean.TRUE});
    }

    public void testPerformance() throws Exception {

        LOG.info("Running Benchmark for destination=" + destination + ", producers=" + prodcuerCount + ", consumers=" + consumerCount + ", deliveryMode=" + deliveryMode);
        final int consumeCount = destination.isTopic() ? consumerCount * produceCount : produceCount;

        final Semaphore consumersStarted = new Semaphore(1 - consumerCount);
        final Semaphore producersFinished = new Semaphore(1 - prodcuerCount);
        final Semaphore consumersFinished = new Semaphore(1 - consumerCount);
        final ProgressPrinter printer = new ProgressPrinter(produceCount + consumeCount, 10);

        // Start a producer and consumer

        profilerPause("Benchmark ready.  Start profiler ");

        long start = System.currentTimeMillis();

        final AtomicInteger receiveCounter = new AtomicInteger(0);
        for (int i = 0; i < consumerCount; i++) {
            new Thread() {
                public void run() {
                    try {

                        // Consume the messages
                        StubConnection connection = new StubConnection(broker);
                        ConnectionInfo connectionInfo = createConnectionInfo();
                        connection.send(connectionInfo);

                        SessionInfo sessionInfo = createSessionInfo(connectionInfo);
                        ConsumerInfo consumerInfo = createConsumerInfo(sessionInfo, destination);
                        consumerInfo.setPrefetchSize(1000);
                        connection.send(sessionInfo);
                        connection.send(consumerInfo);

                        consumersStarted.release();

                        while (receiveCounter.get() < consumeCount) {

                            int counter = 0;
                            // Get a least 1 message.
                            Message msg = receiveMessage(connection, 2000);
                            if (msg != null) {
                                printer.increment();
                                receiveCounter.incrementAndGet();

                                counter++;

                                // Try to piggy back a few extra message acks if
                                // they are ready.
                                Message extra = null;
                                while ((extra = receiveMessage(connection, 0)) != null) {
                                    msg = extra;
                                    printer.increment();
                                    receiveCounter.incrementAndGet();
                                    counter++;
                                }
                            }

                            if (msg != null) {
                                connection.send(createAck(consumerInfo, msg, counter, MessageAck.STANDARD_ACK_TYPE));
                            } else if (receiveCounter.get() < consumeCount) {
                                LOG.info("Consumer stall, waiting for message #" + receiveCounter.get() + 1);
                            }
                        }

                        connection.send(closeConsumerInfo(consumerInfo));
                    } catch (Throwable e) {
                        e.printStackTrace();
                    } finally {
                        consumersFinished.release();
                    }
                }

            }.start();
        }

        // Make sure that the consumers are started first to avoid sending
        // messages
        // before a topic is subscribed so that those messages are not missed.
        consumersStarted.acquire();

        // Send the messages in an async thread.
        for (int i = 0; i < prodcuerCount; i++) {
            new Thread() {
                public void run() {
                    try {
                        StubConnection connection = new StubConnection(broker);
                        ConnectionInfo connectionInfo = createConnectionInfo();
                        connection.send(connectionInfo);

                        SessionInfo sessionInfo = createSessionInfo(connectionInfo);
                        ProducerInfo producerInfo = createProducerInfo(sessionInfo);
                        connection.send(sessionInfo);
                        connection.send(producerInfo);

                        for (int i = 0; i < produceCount / prodcuerCount; i++) {
                            Message message = createMessage(producerInfo, destination);
                            message.setPersistent(deliveryMode);
                            message.setResponseRequired(false);
                            connection.send(message);
                            printer.increment();
                        }
                    } catch (Throwable e) {
                        e.printStackTrace();
                    } finally {
                        producersFinished.release();
                    }
                };
            }.start();
        }

        producersFinished.acquire();
        long end1 = System.currentTimeMillis();
        consumersFinished.acquire();
        long end2 = System.currentTimeMillis();

        LOG.info("Results for destination=" + destination + ", producers=" + prodcuerCount + ", consumers=" + consumerCount + ", deliveryMode=" + deliveryMode);
        LOG.info("Produced at messages/sec: " + (produceCount * 1000.0 / (end1 - start)));
        LOG.info("Consumed at messages/sec: " + (consumeCount * 1000.0 / (end2 - start)));
        profilerPause("Benchmark done.  Stop profiler ");
    }

    public static Test suite() {
        return suite(BrokerBenchmark.class);
    }

    public static void main(String[] args) {
        junit.textui.TestRunner.run(suite());
    }

}
