merge bug19749 into bug20097

and replace a few new instances of gen_server calls/cast in
rabbit_amqqueue with gen_server2.

5 files changed, 385 insertions, 89 deletions

diff --git a/ebin/rabbit.app b/ebin/rabbit.app
index 5ecd247b..64afb5ff 100644
--- a/ebin/rabbit.app
+++ b/ebin/rabbit.app
@@ -19,6 +19,7 @@
              rabbit_framing_channel,
              rabbit_framing,
              rabbit_heartbeat,
+             rabbit_limiter,
              rabbit_load,
              rabbit_log,
              rabbit_memsup_linux,
diff --git a/src/rabbit_amqqueue.erl b/src/rabbit_amqqueue.erl
index cf4c324d..5d612fbb 100644
--- a/src/rabbit_amqqueue.erl
+++ b/src/rabbit_amqqueue.erl
@@ -37,9 +37,9 @@
          stat/1, stat_all/0, deliver/5, redeliver/2, requeue/3, ack/4]).
 -export([list/1, info/1, info/2, info_all/1, info_all/2]).
 -export([claim_queue/2]).
--export([basic_get/3, basic_consume/7, basic_cancel/4]).
--export([notify_sent/2]).
--export([commit_all/2, rollback_all/2, notify_down_all/2]).
+-export([basic_get/3, basic_consume/8, basic_cancel/4]).
+-export([notify_sent/2, unblock/2]).
+-export([commit_all/2, rollback_all/2, notify_down_all/2, limit_all/3]).
 -export([on_node_down/1]).
 
 -import(mnesia).
@@ -91,15 +91,17 @@
 -spec(commit_all/2 :: ([pid()], txn()) -> ok_or_errors()).
 -spec(rollback_all/2 :: ([pid()], txn()) -> ok_or_errors()).
 -spec(notify_down_all/2 :: ([pid()], pid()) -> ok_or_errors()).
+-spec(limit_all/3 :: ([pid()], pid(), pid()) -> ok_or_errors()).
 -spec(claim_queue/2 :: (amqqueue(), pid()) -> 'ok' | 'locked').
 -spec(basic_get/3 :: (amqqueue(), pid(), bool()) ->
              {'ok', non_neg_integer(), msg()} | 'empty').
--spec(basic_consume/7 ::
-      (amqqueue(), bool(), pid(), pid(), ctag(), bool(), any()) ->
+-spec(basic_consume/8 ::
+      (amqqueue(), bool(), pid(), pid(), pid(), ctag(), bool(), any()) ->
              'ok' | {'error', 'queue_owned_by_another_connection' |
                      'exclusive_consume_unavailable'}).
 -spec(basic_cancel/4 :: (amqqueue(), pid(), ctag(), any()) -> 'ok').
 -spec(notify_sent/2 :: (pid(), pid()) -> 'ok').
+-spec(unblock/2 :: (pid(), pid()) -> 'ok').
 -spec(internal_delete/1 :: (queue_name()) -> 'ok' | not_found()).
 -spec(on_node_down/1 :: (erlang_node()) -> 'ok').
 -spec(pseudo_queue/2 :: (binary(), pid()) -> amqqueue()).
@@ -259,16 +261,22 @@ notify_down_all(QPids, ChPid) ->
       fun (QPid) -> gen_server2:call(QPid, {notify_down, ChPid}, Timeout) end,
       QPids).
 
+limit_all(QPids, ChPid, LimiterPid) ->
+    safe_pmap_ok(
+      fun (_) -> ok end,
+      fun (QPid) -> gen_server2:cast(QPid, {limit, ChPid, LimiterPid}) end,
+      QPids).
+    
 claim_queue(#amqqueue{pid = QPid}, ReaderPid) ->
     gen_server2:call(QPid, {claim_queue, ReaderPid}).
 
 basic_get(#amqqueue{pid = QPid}, ChPid, NoAck) ->
     gen_server2:call(QPid, {basic_get, ChPid, NoAck}).
 
-basic_consume(#amqqueue{pid = QPid}, NoAck, ReaderPid, ChPid,
+basic_consume(#amqqueue{pid = QPid}, NoAck, ReaderPid, ChPid, LimiterPid,
               ConsumerTag, ExclusiveConsume, OkMsg) ->
-    gen_server2:call(QPid, {basic_consume, NoAck, ReaderPid, ChPid,
-                           ConsumerTag, ExclusiveConsume, OkMsg}).
+    gen_server2:call(QPid, {basic_consume, NoAck, ReaderPid, ChPid, 
+                           LimiterPid, ConsumerTag, ExclusiveConsume, OkMsg}).
 
 basic_cancel(#amqqueue{pid = QPid}, ChPid, ConsumerTag, OkMsg) ->
     ok = gen_server2:call(QPid, {basic_cancel, ChPid, ConsumerTag, OkMsg}).
@@ -276,6 +284,9 @@ basic_cancel(#amqqueue{pid = QPid}, ChPid, ConsumerTag, OkMsg) ->
 notify_sent(QPid, ChPid) ->
     gen_server2:cast(QPid, {notify_sent, ChPid}).
 
+unblock(QPid, ChPid) ->
+    gen_server2:cast(QPid, {unblock, ChPid}).
+
 internal_delete(QueueName) ->
     rabbit_misc:execute_mnesia_transaction(
       fun () ->
diff --git a/src/rabbit_amqqueue_process.erl b/src/rabbit_amqqueue_process.erl
index 709e355e..5199fb87 100644
--- a/src/rabbit_amqqueue_process.erl
+++ b/src/rabbit_amqqueue_process.erl
@@ -62,9 +62,10 @@
 %% These are held in our process dictionary
 -record(cr, {consumers,
              ch_pid,
+             limiter_pid,
              monitor_ref,
              unacked_messages,
-             is_overload_protection_active,
+             is_limit_active,
              unsent_message_count}).
 
 -define(INFO_KEYS,
@@ -131,7 +132,7 @@ ch_record(ChPid) ->
                     ch_pid = ChPid,
                     monitor_ref = MonitorRef,
                     unacked_messages = dict:new(),
-                    is_overload_protection_active = false,
+                    is_limit_active = false,
                     unsent_message_count = 0},
             put(Key, C),
             C;
@@ -144,20 +145,16 @@ store_ch_record(C = #cr{ch_pid = ChPid}) ->
 all_ch_record() ->
     [C || {{ch, _}, C} <- get()].
 
-update_store_and_maybe_block_ch(
-  C = #cr{is_overload_protection_active = Active,
-          unsent_message_count = Count}) ->
-    {Result, NewActive} =
-        if
-            not(Active) and (Count > ?UNSENT_MESSAGE_LIMIT) ->
-                {block_ch, true};
-            Active and (Count == 0) ->
-                {unblock_ch, false};
-            true ->
-                {ok, Active}
-        end,
-    store_ch_record(C#cr{is_overload_protection_active = NewActive}),
-    Result.
+is_ch_blocked(#cr{unsent_message_count = Count, is_limit_active = Limited}) ->
+    Limited orelse Count > ?UNSENT_MESSAGE_LIMIT.
+
+ch_record_state_transition(OldCR, NewCR) ->
+    BlockedOld = is_ch_blocked(OldCR),
+    BlockedNew = is_ch_blocked(NewCR),
+    if BlockedOld andalso not(BlockedNew) -> unblock;
+       BlockedNew andalso not(BlockedOld) -> block;
+       true                               -> ok
+    end.
 
 deliver_immediately(Message, Delivered,
                     State = #q{q = #amqqueue{name = QName},
@@ -168,26 +165,37 @@ deliver_immediately(Message, Delivered,
         {{value, QEntry = {ChPid, #consumer{tag = ConsumerTag,
                                             ack_required = AckRequired}}},
          RoundRobinTail} ->
-            rabbit_channel:deliver(
-              ChPid, ConsumerTag, AckRequired,
-              {QName, self(), NextId, Delivered, Message}),
-            C = #cr{unsent_message_count = Count,
+            C = #cr{limiter_pid = LimiterPid,
+                    unsent_message_count = Count,
                     unacked_messages = UAM} = ch_record(ChPid),
-            NewUAM = case AckRequired of
-                         true  -> dict:store(NextId, Message, UAM);
-                         false -> UAM
-                     end,
-            NewConsumers =
-                case update_store_and_maybe_block_ch(
-                       C#cr{unsent_message_count = Count + 1,
-                            unacked_messages = NewUAM}) of
-                    ok       -> queue:in(QEntry, RoundRobinTail);
-                    block_ch -> block_consumers(ChPid, RoundRobinTail)
-                end,
-            {offered, AckRequired, State#q{round_robin = NewConsumers,
-                                           next_msg_id = NextId +1}};
+            case not(AckRequired) orelse rabbit_limiter:can_send(
+                                           LimiterPid, self()) of
+                true ->
+                    rabbit_channel:deliver(
+                      ChPid, ConsumerTag, AckRequired,
+                      {QName, self(), NextId, Delivered, Message}),
+                    NewUAM = case AckRequired of
+                                 true  -> dict:store(NextId, Message, UAM);
+                                 false -> UAM
+                             end,
+                    NewC = C#cr{unsent_message_count = Count + 1,
+                                unacked_messages = NewUAM},
+                    store_ch_record(NewC),
+                    NewConsumers =
+                        case ch_record_state_transition(C, NewC) of
+                            ok    -> queue:in(QEntry, RoundRobinTail);
+                            block -> block_consumers(ChPid, RoundRobinTail)
+                        end,
+                    {offered, AckRequired, State#q{round_robin = NewConsumers,
+                                                   next_msg_id = NextId + 1}};
+                false ->
+                    store_ch_record(C#cr{is_limit_active = true}),
+                    NewConsumers = block_consumers(ChPid, RoundRobinTail),
+                    deliver_immediately(Message, Delivered,
+                                        State#q{round_robin = NewConsumers})
+            end;
         {empty, _} ->
-            not_offered
+            {not_offered, State}
     end.
 
 attempt_delivery(none, Message, State) ->
@@ -198,8 +206,8 @@ attempt_delivery(none, Message, State) ->
             persist_message(none, qname(State), Message),
             persist_delivery(qname(State), Message, false),
             {true, State1};
-        not_offered ->
-            {false, State}
+        {not_offered, State1} ->
+            {false, State1}
     end;
 attempt_delivery(Txn, Message, State) ->
     persist_message(Txn, qname(State), Message),
@@ -237,16 +245,22 @@ block_consumer(ChPid, ConsumerTag, RoundRobin) ->
                               (CP /= ChPid) or (CT /= ConsumerTag)
                       end, queue:to_list(RoundRobin))).
 
-possibly_unblock(C = #cr{consumers = Consumers, ch_pid = ChPid},
-                 State = #q{round_robin = RoundRobin}) ->
-    case update_store_and_maybe_block_ch(C) of
-        ok ->
+possibly_unblock(State, ChPid, Update) ->
+    case lookup_ch(ChPid) of
+        not_found ->
             State;
-        unblock_ch ->
-            run_poke_burst(State#q{round_robin =
-                                   unblock_consumers(ChPid, Consumers, RoundRobin)})
+        C ->
+            NewC = Update(C),
+            store_ch_record(NewC),
+            case ch_record_state_transition(C, NewC) of
+                ok      -> State;
+                unblock -> NewRR = unblock_consumers(ChPid,
+                                                     NewC#cr.consumers,
+                                                     State#q.round_robin),
+                           run_poke_burst(State#q{round_robin = NewRR})
+            end
     end.
-
+    
 check_auto_delete(State = #q{q = #amqqueue{auto_delete = false}}) ->
     {continue, State};
 check_auto_delete(State = #q{has_had_consumers = false}) ->
@@ -301,7 +315,7 @@ handle_ch_down(DownPid, State = #q{exclusive_consumer = Holder,
                     {stop, normal, NewState}
             end
     end.
-                
+
 cancel_holder(ChPid, ConsumerTag, {ChPid, ConsumerTag}) ->
     none;
 cancel_holder(_ChPid, _ConsumerTag, Holder) ->
@@ -334,8 +348,8 @@ run_poke_burst(MessageBuffer, State) ->
                 {offered, false, NewState} ->
                     persist_auto_ack(qname(State), Message),
                     run_poke_burst(BufferTail, NewState);
-                not_offered ->
-                    State#q{message_buffer = MessageBuffer}
+                {not_offered, NewState} ->
+                    NewState#q{message_buffer = MessageBuffer}
             end;
         {empty, _} ->
             State#q{message_buffer = MessageBuffer}
@@ -585,8 +599,8 @@ handle_call({basic_get, ChPid, NoAck}, _From,
             reply(empty, State)
     end;
 
-handle_call({basic_consume, NoAck, ReaderPid, ChPid, ConsumerTag,
-             ExclusiveConsume, OkMsg},
+handle_call({basic_consume, NoAck, ReaderPid, ChPid, LimiterPid,
+             ConsumerTag, ExclusiveConsume, OkMsg},
             _From, State = #q{owner = Owner,
                               exclusive_consumer = ExistingHolder,
                               round_robin = RoundRobin}) ->
@@ -600,8 +614,13 @@ handle_call({basic_consume, NoAck, ReaderPid, ChPid, ConsumerTag,
                 ok ->
                     C = #cr{consumers = Consumers} = ch_record(ChPid),
                     Consumer = #consumer{tag = ConsumerTag, ack_required = not(NoAck)},
-                    C1 = C#cr{consumers = [Consumer | Consumers]},
-                    store_ch_record(C1),
+                    store_ch_record(C#cr{consumers = [Consumer | Consumers],
+                                         limiter_pid = LimiterPid}),
+                    if Consumers == [] ->
+                            ok = rabbit_limiter:register(LimiterPid, self());
+                       true ->
+                            ok
+                    end,
                     State1 = State#q{has_had_consumers = true,
                                      exclusive_consumer =
                                        if
@@ -621,12 +640,16 @@ handle_call({basic_cancel, ChPid, ConsumerTag, OkMsg}, _From,
         not_found ->
             ok = maybe_send_reply(ChPid, OkMsg),
             reply(ok, State);
-        C = #cr{consumers = Consumers} ->
+        C = #cr{consumers = Consumers, limiter_pid = LimiterPid} ->
             NewConsumers = lists:filter
                              (fun (#consumer{tag = CT}) -> CT /= ConsumerTag end,
                               Consumers),
-            C1 = C#cr{consumers = NewConsumers},
-            store_ch_record(C1),
+            store_ch_record(C#cr{consumers = NewConsumers}),
+            if NewConsumers == [] ->
+                    ok = rabbit_limiter:unregister(LimiterPid, self());
+               true ->
+                    ok
+            end,
             ok = maybe_send_reply(ChPid, OkMsg),
             case check_auto_delete(
                    State#q{exclusive_consumer = cancel_holder(ChPid,
@@ -729,14 +752,33 @@ handle_cast({requeue, MsgIds, ChPid}, State) ->
                       [{Message, true} || Message <- Messages], State))
     end;
 
+handle_cast({unblock, ChPid}, State) ->
+    noreply(
+      possibly_unblock(State, ChPid,
+                       fun (C) -> C#cr{is_limit_active = false} end));
+
 handle_cast({notify_sent, ChPid}, State) ->
-    case lookup_ch(ChPid) of
-        not_found -> noreply(State);
-        T = #cr{unsent_message_count =Count} ->
-            noreply(possibly_unblock(
-                      T#cr{unsent_message_count = Count - 1},
-                      State))
-    end.
+    noreply(
+      possibly_unblock(State, ChPid,
+                       fun (C = #cr{unsent_message_count = Count}) ->
+                               C#cr{unsent_message_count = Count - 1}
+                       end));
+
+handle_cast({limit, ChPid, LimiterPid}, State) ->
+    noreply(
+      possibly_unblock(
+        State, ChPid,
+        fun (C = #cr{consumers = Consumers,
+                     limiter_pid = OldLimiterPid,
+                     is_limit_active = Limited}) ->
+                if Consumers =/= [] andalso OldLimiterPid == undefined ->
+                        ok = rabbit_limiter:register(LimiterPid, self());
+                   true ->
+                        ok
+                end,
+                NewLimited = Limited andalso LimiterPid =/= undefined,
+                C#cr{limiter_pid = LimiterPid, is_limit_active = NewLimited}
+        end)).
 
 handle_info({'DOWN', MonitorRef, process, DownPid, _Reason},
             State = #q{owner = {DownPid, MonitorRef}}) ->
diff --git a/src/rabbit_channel.erl b/src/rabbit_channel.erl
index ca2782c7..513d3050 100644
--- a/src/rabbit_channel.erl
+++ b/src/rabbit_channel.erl
@@ -39,7 +39,7 @@
 %% callbacks
 -export([init/2, handle_message/2]).
 
--record(ch, {state, proxy_pid, reader_pid, writer_pid,
+-record(ch, {state, proxy_pid, reader_pid, writer_pid, limiter_pid,
              transaction_id, tx_participants, next_tag,
              uncommitted_ack_q, unacked_message_q,
              username, virtual_host,
@@ -100,6 +100,7 @@ init(ProxyPid, [ReaderPid, WriterPid, Username, VHost]) ->
         proxy_pid               = ProxyPid,
         reader_pid              = ReaderPid,
         writer_pid              = WriterPid,
+        limiter_pid             = undefined,
         transaction_id          = none,
         tx_participants         = sets:new(),
         next_tag                = 1,
@@ -151,6 +152,12 @@ handle_message({conserve_memory, Conserve}, State) ->
            State#ch.writer_pid, #'channel.flow'{active = not(Conserve)}),
     State;
 
+handle_message({'EXIT', Pid, Reason}, State = #ch{proxy_pid = Pid}) ->
+    terminate(Reason, State);
+
+handle_message({'EXIT', _Pid, normal}, State) ->
+    State;
+
 handle_message({'EXIT', _Pid, Reason}, State) ->
     terminate(Reason, State);
 
@@ -159,13 +166,15 @@ handle_message(Other, State) ->
 
 %%---------------------------------------------------------------------------
 
-terminate(Reason, State = #ch{writer_pid = WriterPid}) ->
+terminate(Reason, State = #ch{writer_pid = WriterPid,
+                              limiter_pid = LimiterPid}) ->
     Res = notify_queues(internal_rollback(State)),
     case Reason of
         normal -> ok = Res;
         _      -> ok
     end,
     rabbit_writer:shutdown(WriterPid),
+    rabbit_limiter:shutdown(LimiterPid),
     exit(Reason).
 
 return_ok(State, true, _Msg)  -> {noreply, State};
@@ -288,7 +297,8 @@ handle_method(#'basic.ack'{delivery_tag = DeliveryTag,
     {Acked, Remaining} = collect_acks(UAMQ, DeliveryTag, Multiple),
     Participants = ack(State#ch.proxy_pid, TxnKey, Acked),
     {noreply, case TxnKey of
-                  none -> State#ch{unacked_message_q = Remaining};
+                  none -> ok = notify_limiter(State#ch.limiter_pid, Acked),
+                          State#ch{unacked_message_q = Remaining};
                   _    -> NewUAQ = queue:join(State#ch.uncommitted_ack_q,
                                               Acked),
                           add_tx_participants(
@@ -332,6 +342,7 @@ handle_method(#'basic.consume'{queue = QueueNameBin,
                                nowait = NoWait},
               _, State = #ch{ proxy_pid = ProxyPid,
                               reader_pid = ReaderPid,
+                              limiter_pid = LimiterPid,
                               consumer_mapping = ConsumerMapping }) ->
     case dict:find(ConsumerTag, ConsumerMapping) of
         error ->
@@ -349,7 +360,7 @@ handle_method(#'basic.consume'{queue = QueueNameBin,
                    QueueName,
                    fun (Q) ->
                            rabbit_amqqueue:basic_consume(
-                             Q, NoAck, ReaderPid, ProxyPid,
+                             Q, NoAck, ReaderPid, ProxyPid, LimiterPid,
                              ActualConsumerTag, ExclusiveConsume,
                              ok_msg(NoWait, #'basic.consume_ok'{
                                       consumer_tag = ActualConsumerTag}))
@@ -414,9 +425,32 @@ handle_method(#'basic.cancel'{consumer_tag = ConsumerTag,
             end
     end;
 
-handle_method(#'basic.qos'{}, _, State) ->
-    %% FIXME: Need to implement QOS
-    {reply, #'basic.qos_ok'{}, State};
+handle_method(#'basic.qos'{global = true}, _, _State) ->
+    rabbit_misc:protocol_error(not_implemented, "global=true", []);
+
+handle_method(#'basic.qos'{prefetch_size = Size}, _, _State) when Size /= 0 ->
+    rabbit_misc:protocol_error(not_implemented, 
+                               "prefetch_size!=0 (~w)", [Size]);
+
+handle_method(#'basic.qos'{prefetch_count = PrefetchCount},
+              _, State = #ch{ limiter_pid = LimiterPid,
+                              proxy_pid = ProxyPid }) ->
+    NewLimiterPid = case {LimiterPid, PrefetchCount} of
+                        {undefined, 0} ->
+                            undefined;
+                        {undefined, _} ->
+                            LPid = rabbit_limiter:start_link(ProxyPid),
+                            ok = limit_queues(LPid, State),
+                            LPid;
+                        {_, 0} ->
+                            ok = rabbit_limiter:shutdown(LimiterPid),
+                            ok = limit_queues(undefined, State),
+                            undefined;
+                        {_, _} ->
+                            LimiterPid
+                    end,
+    ok = rabbit_limiter:limit(NewLimiterPid, PrefetchCount),
+    {reply, #'basic.qos_ok'{}, State#ch{limiter_pid = NewLimiterPid}};
 
 handle_method(#'basic.recover'{requeue = true},
               _, State = #ch{ transaction_id = none,
@@ -762,7 +796,9 @@ internal_commit(State = #ch{transaction_id = TxnKey,
                             tx_participants = Participants}) ->
     case rabbit_amqqueue:commit_all(sets:to_list(Participants),
                                     TxnKey) of
-        ok              -> new_tx(State);
+        ok              -> ok = notify_limiter(State#ch.limiter_pid,
+                                               State#ch.uncommitted_ack_q),
+                           new_tx(State);
         {error, Errors} -> rabbit_misc:protocol_error(
                              internal_error, "commit failed: ~w", [Errors])
     end.
@@ -800,18 +836,35 @@ fold_per_queue(F, Acc0, UAQ) ->
               Acc0, D).
 
 notify_queues(#ch{proxy_pid = ProxyPid, consumer_mapping = Consumers}) ->
-    rabbit_amqqueue:notify_down_all(
-      [QPid || QueueName <- 
-                   sets:to_list(
-                     dict:fold(fun (_ConsumerTag, QueueName, S) ->
-                                       sets:add_element(QueueName, S)
-                               end, sets:new(), Consumers)),
-               case rabbit_amqqueue:lookup(QueueName) of
-                   {ok, Q} -> QPid = Q#amqqueue.pid, true;
-                   %% queue has been deleted in the meantime
-                   {error, not_found} -> QPid = none, false
-               end],
-      ProxyPid).
+    rabbit_amqqueue:notify_down_all(consumer_queues(Consumers), ProxyPid).
+
+limit_queues(LPid, #ch{proxy_pid = ProxyPid, consumer_mapping = Consumers}) ->
+    rabbit_amqqueue:limit_all(consumer_queues(Consumers), ProxyPid, LPid).
+
+consumer_queues(Consumers) ->
+    [QPid || QueueName <- 
+                 sets:to_list(
+                   dict:fold(fun (_ConsumerTag, QueueName, S) ->
+                                     sets:add_element(QueueName, S)
+                             end, sets:new(), Consumers)),
+             case rabbit_amqqueue:lookup(QueueName) of
+                 {ok, Q} -> QPid = Q#amqqueue.pid, true;
+                 %% queue has been deleted in the meantime
+                 {error, not_found} -> QPid = none, false
+             end].
+
+%% tell the limiter about the number of acks that have been received
+%% for messages delivered to subscribed consumers, but not acks for
+%% messages sent in a response to a basic.get.
+notify_limiter(undefined, _Acked) ->
+    ok;
+notify_limiter(LimiterPid, Acked) ->
+    case lists:foldl(fun ({_, none, _}, Acc) -> Acc;
+                         ({_, _, _}, Acc)    -> Acc + 1
+                     end, 0, queue:to_list(Acked)) of
+        0     -> ok;
+        Count -> rabbit_limiter:ack(LimiterPid, Count)
+    end.
 
 is_message_persistent(#content{properties = #'P_basic'{
                                  delivery_mode = Mode}}) ->
diff --git a/src/rabbit_limiter.erl b/src/rabbit_limiter.erl
new file mode 100644
index 00000000..3776edd0
--- /dev/null
+++ b/src/rabbit_limiter.erl
@@ -0,0 +1,189 @@
+%%   The contents of this file are subject to the Mozilla Public License
+%%   Version 1.1 (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.mozilla.org/MPL/
+%%
+%%   Software distributed under the License is distributed on an "AS IS"
+%%   basis, WITHOUT WARRANTY OF ANY KIND, either express or implied. See the
+%%   License for the specific language governing rights and limitations
+%%   under the License.
+%%
+%%   The Original Code is RabbitMQ.
+%%
+%%   The Initial Developers of the Original Code are LShift Ltd,
+%%   Cohesive Financial Technologies LLC, and Rabbit Technologies Ltd.
+%%
+%%   Portions created before 22-Nov-2008 00:00:00 GMT by LShift Ltd,
+%%   Cohesive Financial Technologies LLC, or Rabbit Technologies Ltd
+%%   are Copyright (C) 2007-2008 LShift Ltd, Cohesive Financial
+%%   Technologies LLC, and Rabbit Technologies Ltd.
+%%
+%%   Portions created by LShift Ltd are Copyright (C) 2007-2009 LShift
+%%   Ltd. Portions created by Cohesive Financial Technologies LLC are
+%%   Copyright (C) 2007-2009 Cohesive Financial Technologies
+%%   LLC. Portions created by Rabbit Technologies Ltd are Copyright
+%%   (C) 2007-2009 Rabbit Technologies Ltd.
+%%
+%%   All Rights Reserved.
+%%
+%%   Contributor(s): ______________________________________.
+%%
+
+-module(rabbit_limiter).
+
+-behaviour(gen_server).
+
+-export([init/1, terminate/2, code_change/3, handle_call/3, handle_cast/2,
+         handle_info/2]).
+-export([start_link/1, shutdown/1]).
+-export([limit/2, can_send/2, ack/2, register/2, unregister/2]).
+
+%%----------------------------------------------------------------------------
+
+-ifdef(use_specs).
+
+-spec(start_link/1 :: (pid()) -> pid()).
+-spec(shutdown/1 :: (pid()) -> 'ok').
+-spec(limit/2 :: (pid(), non_neg_integer()) -> 'ok').
+-spec(can_send/2 :: (pid(), pid()) -> bool()).
+-spec(ack/2 :: (pid(), non_neg_integer()) -> 'ok').
+-spec(register/2 :: (pid(), pid()) -> 'ok').
+-spec(unregister/2 :: (pid(), pid()) -> 'ok').
+
+-endif.
+
+%%----------------------------------------------------------------------------
+
+-record(lim, {prefetch_count = 0,
+              ch_pid,
+              queues = dict:new(),
+              volume = 0}).
+
+%%----------------------------------------------------------------------------
+%% API
+%%----------------------------------------------------------------------------
+
+start_link(ChPid) ->
+    {ok, Pid} = gen_server:start_link(?MODULE, [ChPid], []),
+    Pid.
+
+shutdown(undefined) ->
+    ok;
+shutdown(LimiterPid) ->
+    gen_server:cast(LimiterPid, shutdown).
+
+limit(undefined, 0) ->
+    ok;
+limit(LimiterPid, PrefetchCount) ->
+    gen_server:cast(LimiterPid, {limit, PrefetchCount}).
+
+%% Ask the limiter whether the queue can deliver a message without
+%% breaching a limit
+can_send(undefined, _QPid) ->
+    true;
+can_send(LimiterPid, QPid) ->
+    rabbit_misc:with_exit_handler(
+      fun () -> true end,
+      fun () -> gen_server:call(LimiterPid, {can_send, QPid}) end).
+
+%% Let the limiter know that the channel has received some acks from a
+%% consumer
+ack(undefined, _Count) -> ok;
+ack(LimiterPid, Count) -> gen_server:cast(LimiterPid, {ack, Count}).
+
+register(undefined, _QPid) -> ok;
+register(LimiterPid, QPid) -> gen_server:cast(LimiterPid, {register, QPid}).
+
+unregister(undefined, _QPid) -> ok;
+unregister(LimiterPid, QPid) -> gen_server:cast(LimiterPid, {unregister, QPid}).
+
+%%----------------------------------------------------------------------------
+%% gen_server callbacks
+%%----------------------------------------------------------------------------
+
+init([ChPid]) ->
+    {ok, #lim{ch_pid = ChPid} }.
+
+handle_call({can_send, QPid}, _From, State = #lim{volume = Volume}) ->
+    case limit_reached(State) of
+        true  -> {reply, false, limit_queue(QPid, State)};
+        false -> {reply, true, State#lim{volume = Volume + 1}}
+    end.
+
+handle_cast(shutdown, State) ->
+    {stop, normal, State};
+
+handle_cast({limit, PrefetchCount}, State) ->
+    {noreply, maybe_notify(State, State#lim{prefetch_count = PrefetchCount})};
+
+handle_cast({ack, Count}, State = #lim{volume = Volume}) ->
+    NewVolume = if Volume == 0 -> 0;
+                   true        -> Volume - Count
+                end,
+    {noreply, maybe_notify(State, State#lim{volume = NewVolume})};
+
+handle_cast({register, QPid}, State) ->
+    {noreply, remember_queue(QPid, State)};
+
+handle_cast({unregister, QPid}, State) ->
+    {noreply, forget_queue(QPid, State)}.
+
+handle_info({'DOWN', _MonitorRef, _Type, QPid, _Info}, State) ->
+    {noreply, forget_queue(QPid, State)}.
+
+terminate(_, _) ->
+    ok.
+
+code_change(_, State, _) ->
+    State.
+
+%%----------------------------------------------------------------------------
+%% Internal plumbing
+%%----------------------------------------------------------------------------
+
+maybe_notify(OldState, NewState) ->
+    case limit_reached(OldState) andalso not(limit_reached(NewState)) of
+        true  -> notify_queues(NewState);
+        false -> NewState
+    end.
+
+limit_reached(#lim{prefetch_count = Limit, volume = Volume}) ->
+    Limit =/= 0 andalso Volume >= Limit.
+
+remember_queue(QPid, State = #lim{queues = Queues}) ->
+    case dict:is_key(QPid, Queues) of
+        false -> MRef = erlang:monitor(process, QPid),
+                 State#lim{queues = dict:store(QPid, {MRef, false}, Queues)};
+        true  -> State
+    end.
+
+forget_queue(QPid, State = #lim{ch_pid = ChPid, queues = Queues}) ->
+    case dict:find(QPid, Queues) of
+        {ok, {MRef, _}} ->
+            true = erlang:demonitor(MRef),
+            ok = rabbit_amqqueue:unblock(QPid, ChPid),
+            State#lim{queues = dict:erase(QPid, Queues)};
+        error -> State
+    end.
+
+limit_queue(QPid, State = #lim{queues = Queues}) ->
+    UpdateFun = fun ({MRef, _}) -> {MRef, true} end,
+    State#lim{queues = dict:update(QPid, UpdateFun, Queues)}.
+
+notify_queues(State = #lim{ch_pid = ChPid, queues = Queues}) ->
+    {QList, NewQueues} =
+        dict:fold(fun (_QPid, {_, false}, Acc) -> Acc;
+                      (QPid, {MRef, true}, {L, D}) ->
+                          {[QPid | L], dict:store(QPid, {MRef, false}, D)}
+                  end, {[], Queues}, Queues),
+    case length(QList) of
+        0 -> ok;
+        L ->
+            %% We randomly vary the position of queues in the list,
+            %% thus ensuring that each queue has an equal chance of
+            %% being notified first.
+            {L1, L2} = lists:split(random:uniform(L), QList),
+            [ok = rabbit_amqqueue:unblock(Q, ChPid) || Q <- L2 ++ L1],
+            ok
+    end,
+    State#lim{queues = NewQueues}.