Move the Semaphore up to be scoped to the WorkPool. If any WorkQueue aka channel is unlimited then they all are.

Author: Simon MacMullen

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

@@ -1190,15 +1190,15 @@ public AMQCommand rpc(Method method) throws IOException {
     public void enqueueRpc(RpcContinuation k) {
         synchronized (_channelMutex) {
             super.enqueueRpc(k);
-            dispatcher.setLimited(false);
+            dispatcher.setUnlimited(true);
         }
     }
 
     @Override
     public RpcContinuation nextOutstandingRpc() {
         synchronized (_channelMutex) {
             RpcContinuation res = super.nextOutstandingRpc();
-            if (res != null) dispatcher.setLimited(true);
+            if (res != null) dispatcher.setUnlimited(false);
             return res;
         }
     }

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

@@ -62,8 +62,8 @@ public void quiesce() {
         this.shuttingDown = true;
     }
 
-    public void setLimited(boolean limited) {
-        this.workService.limit(channel, limited);
+    public void setUnlimited(boolean unlimited) {
+        this.workService.unlimit(channel, unlimited);
     }
 
     public void handleConsumeOk(final Consumer delegate,

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

@@ -58,8 +58,8 @@ public void registerKey(Channel channel) {
         this.workPool.registerKey(channel);
     }
 
-    public void limit(Channel channel, boolean limited) {
-        this.workPool.limit(channel, limited);
+    public void unlimit(Channel channel, boolean unlimited) {
+        this.workPool.unlimit(channel, unlimited);
     }
 
     public void addWork(Channel channel, Runnable runnable) {

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

@@ -7,6 +7,7 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.Semaphore;
+import java.util.concurrent.atomic.AtomicInteger;
 
 /**
  * This is a generic implementation of the <q>Channels</q> specification
@@ -76,52 +77,52 @@ public class WorkPool<K, W> {
     //
     // 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 setLimited() synchronised as it is called frequently by
+    // b) we don't want to make setUnlimited() synchronised as it is called frequently by
     //    the channel.
-    // c) anyway the issue with setLimited() is not that it be synchronised itself but
+    // 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 static class WorkQueue<T> {
-        private Semaphore semaphore;
-        private LinkedList<T> list;
-        private boolean limited;
+    private class WorkQueue {
+        private LinkedList<W> list;
+        private boolean unlimited;
         private int maxLengthWhenLimited;
 
         private WorkQueue(int maxLengthWhenLimited) {
-            this.semaphore = new Semaphore(1);
-            this.list = new LinkedList<T>();
-            this.limited = true;
+            this.list = new LinkedList<W>();
+            this.unlimited = false; // Just for assertions
             this.maxLengthWhenLimited = maxLengthWhenLimited;
         }
 
-        public void put(T t) throws InterruptedException {
-            if (limited && list.size() > maxLengthWhenLimited) {
-                assert !semaphore.hasQueuedThreads();
-                semaphore.acquire();
+        public void put(W w) throws InterruptedException {
+            if (list.size() > maxLengthWhenLimited) {
+                acquireSemaphore();
             }
             synchronized (this) {
-                list.add(t);
+                list.add(w);
             }
         }
 
-        public synchronized T poll() {
-            T res = list.poll();
+        public synchronized W poll() {
+            W res = list.poll();
 
-            if (list.size() <= maxLengthWhenLimited && semaphore.hasQueuedThreads()) {
-                semaphore.release();
+            if (list.size() <= maxLengthWhenLimited) {
+                releaseSemaphore();
             }
 
             return res;
         }
 
-        public void setLimited(boolean limited) {
-            assert this.limited != limited;
-            this.limited = limited;
-            if (!limited && semaphore.hasQueuedThreads()) {
-                semaphore.release();
+        public void setUnlimited(boolean unlimited) {
+            assert this.unlimited != unlimited;
+            this.unlimited = unlimited;
+            if (unlimited) {
+                increaseUnlimited();
+            }
+            else {
+                decreaseUnlimited();
             }
         }
 
@@ -135,7 +136,31 @@ public boolean isEmpty() {
     /** 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, WorkQueue<W>> pool = new HashMap<K, WorkQueue<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.
@@ -147,16 +172,16 @@ public boolean isEmpty() {
     public void registerKey(K key) {
         synchronized (this) {
             if (!this.pool.containsKey(key)) {
-                this.pool.put(key, new WorkQueue<W>(MAX_QUEUE_LENGTH));
+                this.pool.put(key, new WorkQueue(MAX_QUEUE_LENGTH));
             }
         }
     }
 
-    public void limit(K key, boolean limited) {
+    public void unlimit(K key, boolean unlimited) {
         synchronized (this) {
-            WorkQueue<W> queue = this.pool.get(key);
+            WorkQueue queue = this.pool.get(key);
             if (queue != null) {
-                queue.setLimited(limited);
+                queue.setUnlimited(unlimited);
             }
         }
     }
@@ -198,7 +223,7 @@ public K nextWorkBlock(Collection<W> to, int size) {
         synchronized (this) {
             K nextKey = readyToInProgress();
             if (nextKey != null) {
-                WorkQueue<W> queue = this.pool.get(nextKey);
+                WorkQueue queue = this.pool.get(nextKey);
                 drainTo(queue, to, size);
             }
             return nextKey;
@@ -207,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(WorkQueue<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();
@@ -235,7 +259,7 @@ private static <W> int drainTo(WorkQueue<W> deList, Collection<W> c, int maxElem
      * &mdash; <i>as a result of this work item</i>
      */
     public boolean addWorkItem(K key, W item) {
-        WorkQueue<W> queue;
+        WorkQueue queue;
         synchronized (this) {
             queue = this.pool.get(key);
         }
@@ -283,7 +307,7 @@ public boolean finishWorkBlock(K key) {
     }
 
     private boolean moreWorkItems(K key) {
-        WorkQueue<W> leList = this.pool.get(key);
+        WorkQueue leList = this.pool.get(key);
         return leList != null && !leList.isEmpty();
     }