/**
 * 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.bugs;

import java.io.Serializable;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference;

import javax.jms.Connection;
import javax.jms.JMSException;
import javax.jms.MessageConsumer;
import javax.jms.MessageListener;
import javax.jms.MessageProducer;
import javax.jms.ObjectMessage;
import javax.jms.Queue;
import javax.jms.Session;

import junit.framework.TestCase;
import org.apache.activemq.ActiveMQConnectionFactory;
import org.apache.activemq.ActiveMQPrefetchPolicy;
import org.apache.activemq.broker.BrokerService;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Test case demonstrating situation where messages are not delivered to
 * consumers.
 */
public class QueueWorkerPrefetchTest extends TestCase implements
        MessageListener {
    private static final Logger LOG = LoggerFactory
            .getLogger(QueueWorkerPrefetchTest.class);
    private static final int BATCH_SIZE = 10;
    private static final long WAIT_TIMEOUT = 1000 * 10;

    /** The connection URL. */
    private static final String BROKER_BIND_ADDRESS = "tcp://localhost:0";

    /**
     * The queue prefetch size to use. A value greater than 1 seems to make
     * things work.
     */
    private static final int QUEUE_PREFETCH_SIZE = 1;

    /**
     * The number of workers to use. A single worker with a prefetch of 1 works.
     */
    private static final int NUM_WORKERS = 2;

    /** Embedded JMS broker. */
    private BrokerService broker;

    /** The master's producer object for creating work items. */
    private MessageProducer workItemProducer;

    /** The master's consumer object for consuming ack messages from workers. */
    private MessageConsumer masterItemConsumer;

    /** The number of acks received by the master. */
    private final AtomicLong acksReceived = new AtomicLong(0);

    private final AtomicReference<CountDownLatch> latch = new AtomicReference<CountDownLatch>();

    private String connectionUri;

    /** Messages sent to the work-item queue. */
    private static class WorkMessage implements Serializable {
        private static final long serialVersionUID = 1L;
        private final int id;

        public WorkMessage(int id) {
            this.id = id;
        }

        @Override
        public String toString() {
            return "Work: " + id;
        }
    }

    /**
     * The worker process. Consume messages from the work-item queue, possibly
     * creating more messages to submit to the work-item queue. For each work
     * item, send an ack to the master.
     */
    private static class Worker implements MessageListener {
        /**
         * Counter shared between workers to decided when new work-item messages
         * are created.
         */
        private static AtomicInteger counter = new AtomicInteger(0);

        /** Session to use. */
        private Session session;

        /** Producer for sending ack messages to the master. */
        private MessageProducer masterItemProducer;

        /** Producer for sending new work items to the work-items queue. */
        private MessageProducer workItemProducer;

        public Worker(Session session) throws JMSException {
            this.session = session;
            masterItemProducer = session.createProducer(session
                    .createQueue("master-item"));
            Queue workItemQueue = session.createQueue("work-item");
            workItemProducer = session.createProducer(workItemQueue);
            MessageConsumer workItemConsumer = session
                    .createConsumer(workItemQueue);
            workItemConsumer.setMessageListener(this);
        }

        public void onMessage(javax.jms.Message message) {
            try {
                WorkMessage work = (WorkMessage) ((ObjectMessage) message)
                        .getObject();

                long c = counter.incrementAndGet();

                // Don't create a new work item for every BATCH_SIZE message. */
                if (c % BATCH_SIZE != 0) {
                    // Send new work item to work-item queue.
                    workItemProducer.send(session
                            .createObjectMessage(new WorkMessage(work.id + 1)));
                }

                // Send ack to master.
                masterItemProducer.send(session.createObjectMessage(work));
            } catch (JMSException e) {
                throw new IllegalStateException("Something has gone wrong", e);
            }
        }

        /** Close of JMS resources used by worker. */
        public void close() throws JMSException {
            masterItemProducer.close();
            workItemProducer.close();
            session.close();
        }
    }

    /** Master message handler. Process ack messages. */
    public void onMessage(javax.jms.Message message) {
        long acks = acksReceived.incrementAndGet();
        latch.get().countDown();
        if (acks % 1 == 0) {
            LOG.info("Master now has ack count of: " + acksReceived);
        }
    }

    protected void setUp() throws Exception {
        // Create the message broker.
        super.setUp();
        broker = new BrokerService();
        broker.setPersistent(false);
        broker.setUseJmx(true);
        broker.addConnector(BROKER_BIND_ADDRESS);
        broker.start();
        broker.waitUntilStarted();

        connectionUri = broker.getTransportConnectors().get(0).getPublishableConnectString();
    }

    protected void tearDown() throws Exception {
        // Shut down the message broker.
        broker.deleteAllMessages();
        broker.stop();
        super.tearDown();
    }

    public void testActiveMQ() throws Exception {
        // Create the connection to the broker.
        ActiveMQConnectionFactory connectionFactory = new ActiveMQConnectionFactory(connectionUri);
        ActiveMQPrefetchPolicy prefetchPolicy = new ActiveMQPrefetchPolicy();
        prefetchPolicy.setQueuePrefetch(QUEUE_PREFETCH_SIZE);
        connectionFactory.setPrefetchPolicy(prefetchPolicy);
        Connection connection = connectionFactory.createConnection();
        connection.start();

        Session masterSession = connection.createSession(false,
                Session.AUTO_ACKNOWLEDGE);
        workItemProducer = masterSession.createProducer(masterSession
                .createQueue("work-item"));
        masterItemConsumer = masterSession.createConsumer(masterSession
                .createQueue("master-item"));
        masterItemConsumer.setMessageListener(this);

        // Create the workers.
        Worker[] workers = new Worker[NUM_WORKERS];
        for (int i = 0; i < NUM_WORKERS; i++) {
            workers[i] = new Worker(connection.createSession(false,
                    Session.AUTO_ACKNOWLEDGE));
        }

        // Send a message to the work queue, and wait for the BATCH_SIZE acks
        // from the workers.
        acksReceived.set(0);
        latch.set(new CountDownLatch(BATCH_SIZE));
        workItemProducer.send(masterSession
                .createObjectMessage(new WorkMessage(1)));

        if (!latch.get().await(WAIT_TIMEOUT, TimeUnit.MILLISECONDS)) {
            fail("First batch only received " + acksReceived + " messages");
        }

        LOG.info("First batch received");

        // Send another message to the work queue, and wait for the next 1000 acks. It is
        // at this point where the workers never get notified of this message, as they
        // have a large pending queue. Creating a new worker at this point however will
        // receive this new message.
        acksReceived.set(0);
        latch.set(new CountDownLatch(BATCH_SIZE));
        workItemProducer.send(masterSession
                .createObjectMessage(new WorkMessage(1)));

        if (!latch.get().await(WAIT_TIMEOUT, TimeUnit.MILLISECONDS)) {
            fail("Second batch only received " + acksReceived + " messages");
        }

        LOG.info("Second batch received");

        // Cleanup all JMS resources.
        for (int i = 0; i < NUM_WORKERS; i++) {
            workers[i].close();
        }
        masterSession.close();
        connection.close();
    }
}
