/**
 * 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();
   }
}
