merge bug26396 into default

Author: michaelklishin

src/com/rabbitmq/client/impl/ChannelN.java

@@ -1186,6 +1186,23 @@ public AMQCommand rpc(Method method) throws IOException {
         return exnWrappingRpc(method);
     }
 
+    @Override
+    public void enqueueRpc(RpcContinuation k) {
+        synchronized (_channelMutex) {
+            super.enqueueRpc(k);
+            dispatcher.setUnlimited(true);
+        }
+    }
+
+    @Override
+    public RpcContinuation nextOutstandingRpc() {
+        synchronized (_channelMutex) {
+            RpcContinuation res = super.nextOutstandingRpc();
+            if (res != null) dispatcher.setUnlimited(false);
+            return res;
+        }
+    }
+
     private void handleAckNack(long seqNo, boolean multiple, boolean nack) {
         if (multiple) {
             unconfirmedSet.headSet(seqNo + 1).clear();

src/com/rabbitmq/client/impl/ConsumerDispatcher.java

@@ -62,6 +62,10 @@ public void quiesce() {
         this.shuttingDown = true;
     }
 
+    public void setUnlimited(boolean unlimited) {
+        this.workService.unlimit(channel, unlimited);
+    }
+
     public void handleConsumeOk(final Consumer delegate,
                                 final String consumerTag) {
         executeUnlessShuttingDown(

src/com/rabbitmq/client/impl/ConsumerWorkService.java

@@ -64,6 +64,10 @@ public void registerKey(Channel channel) {
         this.workPool.registerKey(channel);
     }
 
+    public void unlimit(Channel channel, boolean unlimited) {
+        this.workPool.unlimit(channel, unlimited);
+    }
+
     public void addWork(Channel channel, Runnable runnable) {
         if (this.workPool.addWorkItem(channel, runnable)) {
             this.executor.execute(new WorkPoolRunnable());

src/com/rabbitmq/client/impl/WorkPool.java

@@ -3,10 +3,11 @@
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.LinkedList;
 import java.util.Map;
 import java.util.Set;
-import java.util.concurrent.BlockingQueue;
-import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.Semaphore;
+import java.util.concurrent.atomic.AtomicInteger;
 
 /**
  * This is a generic implementation of the <q>Channels</q> specification
@@ -69,12 +70,97 @@
 public class WorkPool<K, W> {
     private static final int MAX_QUEUE_LENGTH = 1000;
 
+    // This is like a LinkedBlockingQueue of limited length except you can turn the limit
+    // on and off. And it only has the methods we need.
+    //
+    // This class is partly synchronised because:
+    //
+    // a) we cannot make put(T) synchronised as it may block indefinitely. Therefore we
+    //    only lock before modifying the list.
+    // b) we don't want to make setUnlimited() synchronised as it is called frequently by
+    //    the channel.
+    // c) anyway the issue with setUnlimited() is not that it be synchronised itself but
+    //    that calls to it should alternate between false and true. We assert this, but
+    //    it should not be able to go wrong because the RPC calls in AMQChannel and
+    //    ChannelN are all protected by the _channelMutex; we can't have more than one
+    //    outstanding RPC or finish the same RPC twice.
+
+    private class WorkQueue {
+        private LinkedList<W> list;
+        private boolean unlimited;
+        private int maxLengthWhenLimited;
+
+        private WorkQueue(int maxLengthWhenLimited) {
+            this.list = new LinkedList<W>();
+            this.unlimited = false; // Just for assertions
+            this.maxLengthWhenLimited = maxLengthWhenLimited;
+        }
+
+        public void put(W w) throws InterruptedException {
+            if (list.size() > maxLengthWhenLimited) {
+                acquireSemaphore();
+            }
+            synchronized (this) {
+                list.add(w);
+            }
+        }
+
+        public synchronized W poll() {
+            W res = list.poll();
+
+            if (list.size() <= maxLengthWhenLimited) {
+                releaseSemaphore();
+            }
+
+            return res;
+        }
+
+        public void setUnlimited(boolean unlimited) {
+            assert this.unlimited != unlimited;
+            this.unlimited = unlimited;
+            if (unlimited) {
+                increaseUnlimited();
+            }
+            else {
+                decreaseUnlimited();
+            }
+        }
+
+        public boolean isEmpty() {
+            return list.isEmpty();
+        }
+    }
+
     /** An injective queue of <i>ready</i> clients. */
     private final SetQueue<K> ready = new SetQueue<K>();
     /** The set of clients which have work <i>in progress</i>. */
     private final Set<K> inProgress = new HashSet<K>();
     /** The pool of registered clients, with their work queues. */
-    private final Map<K, BlockingQueue<W>> pool = new HashMap<K, BlockingQueue<W>>();
+    private final Map<K, WorkQueue> pool = new HashMap<K, WorkQueue>();
+
+    // The semaphore should only be used when unlimitedQueues == 0, otherwise we ignore it and
+    // thus don't block the connection.
+    private Semaphore semaphore = new Semaphore(1);
+    private AtomicInteger unlimitedQueues = new AtomicInteger(0);
+
+    private void acquireSemaphore() throws InterruptedException {
+        if (unlimitedQueues.get() == 0) {
+            semaphore.acquire();
+        }
+    }
+
+    private void releaseSemaphore() {
+        semaphore.release();
+    }
+
+    private void increaseUnlimited() {
+        unlimitedQueues.getAndIncrement();
+        semaphore.release();
+    }
+
+    private void decreaseUnlimited() {
+        unlimitedQueues.getAndDecrement();
+    }
 
     /**
      * Add client <code><b>key</b></code> to pool of item queues, with an empty queue.
@@ -86,7 +172,16 @@ public class WorkPool<K, W> {
     public void registerKey(K key) {
         synchronized (this) {
             if (!this.pool.containsKey(key)) {
-                this.pool.put(key, new LinkedBlockingQueue<W>(MAX_QUEUE_LENGTH));
+                this.pool.put(key, new WorkQueue(MAX_QUEUE_LENGTH));
+            }
+        }
+    }
+
+    public void unlimit(K key, boolean unlimited) {
+        synchronized (this) {
+            WorkQueue queue = this.pool.get(key);
+            if (queue != null) {
+                queue.setUnlimited(unlimited);
             }
         }
     }
@@ -128,7 +223,7 @@ public K nextWorkBlock(Collection<W> to, int size) {
         synchronized (this) {
             K nextKey = readyToInProgress();
             if (nextKey != null) {
-                BlockingQueue<W> queue = this.pool.get(nextKey);
+                WorkQueue queue = this.pool.get(nextKey);
                 drainTo(queue, to, size);
             }
             return nextKey;
@@ -137,13 +232,12 @@ public K nextWorkBlock(Collection<W> to, int size) {
 
     /**
      * Private implementation of <code><b>drainTo</b></code> (not implemented for <code><b>LinkedList&lt;W&gt;</b></code>s).
-     * @param <W> element type
      * @param deList to take (poll) elements from
      * @param c to add elements to
      * @param maxElements to take from deList
      * @return number of elements actually taken
      */
-    private static <W> int drainTo(BlockingQueue<W> deList, Collection<W> c, int maxElements) {
+    private int drainTo(WorkQueue deList, Collection<W> c, int maxElements) {
         int n = 0;
         while (n < maxElements) {
             W first = deList.poll();
@@ -165,7 +259,7 @@ private static <W> int drainTo(BlockingQueue<W> deList, Collection<W> c, int max
      * &mdash; <i>as a result of this work item</i>
      */
     public boolean addWorkItem(K key, W item) {
-        BlockingQueue<W> queue;
+        WorkQueue queue;
         synchronized (this) {
             queue = this.pool.get(key);
         }
@@ -213,7 +307,7 @@ public boolean finishWorkBlock(K key) {
     }
 
     private boolean moreWorkItems(K key) {
-        BlockingQueue<W> leList = this.pool.get(key);
+        WorkQueue leList = this.pool.get(key);
         return leList != null && !leList.isEmpty();
     }
 

test/src/com/rabbitmq/client/test/functional/FunctionalTests.java

@@ -40,9 +40,8 @@ public static void add(TestSuite suite) {
         suite.addTestSuite(Recover.class);
         suite.addTestSuite(Reject.class);
         suite.addTestSuite(Transactions.class);
-        // TODO uncomment after we fix bug 26396 one way or the other
-        //suite.addTestSuite(RequeueOnConnectionClose.class);
-        //suite.addTestSuite(RequeueOnChannelClose.class);
+        suite.addTestSuite(RequeueOnConnectionClose.class);
+        suite.addTestSuite(RequeueOnChannelClose.class);
         suite.addTestSuite(DurableOnTransient.class);
         suite.addTestSuite(NoRequeueOnCancel.class);
         suite.addTestSuite(Bug20004Test.class);