Support single active consumer in quorum queue

acogoluegnes · acogoluegnes · commit b4d354546d44 · 2018-12-14T18:09:30.000+01:00
Uses a buffer list for non-active consumers. The active consumer is stored in the usual consumers structure, so the logic around servicing consumers is kept the same. [#162582065] Fixes #1799
diff --git a/src/rabbit_fifo.erl b/src/rabbit_fifo.erl
@@ -226,7 +226,12 @@
          %% This is done so that consumers are still served in a deterministic
          %% order on recovery.
          prefix_msg_counts = {0, 0} :: {Return :: non_neg_integer(),
-                                        PrefixMsgs :: non_neg_integer()}
+                                        PrefixMsgs :: non_neg_integer()},
+         %% whether single active consumer is on or not for this queue
+         single_active_consumer_on = false :: boolean(),
+         %% waiting consumers, one is picked active consumer is cancelled or dies
+         %% used only when single active consumer is on
+         waiting_consumers = [] :: list()
         }).
 
 -opaque state() :: #state{}.
@@ -262,9 +267,11 @@ update_state(Conf, State) ->
     DLH = maps:get(dead_letter_handler, Conf, undefined),
     BLH = maps:get(become_leader_handler, Conf, undefined),
     SHI = maps:get(shadow_copy_interval, Conf, ?SHADOW_COPY_INTERVAL),
+    SingleActiveConsumerOn = maps:get(single_active_consumer_on, Conf, false),
     State#state{dead_letter_handler = DLH,
                 become_leader_handler = BLH,
-                shadow_copy_interval = SHI}.
+                shadow_copy_interval = SHI,
+                single_active_consumer_on = SingleActiveConsumerOn}.
 
 % msg_ids are scoped per consumer
 % ra_indexes holds all raft indexes for enqueues currently on queue
@@ -667,6 +674,54 @@ num_checked_out(#state{consumers = Cons}) ->
                 end, 0, maps:values(Cons)).
 
 cancel_consumer(ConsumerId,
+                {Effects0, #state{single_active_consumer_on = false} = S0}) ->
+    %% general case, single active consumer off
+    cancel_consumer0(ConsumerId, {Effects0, S0});
+cancel_consumer(ConsumerId,
+                {Effects0, #state{single_active_consumer_on = true,
+                                  waiting_consumers = WaitingConsumers  } = S0}) when length(WaitingConsumers) == 0 ->
+    %% single active consumer on, no consumers are waiting
+    cancel_consumer0(ConsumerId, {Effects0, S0});
+cancel_consumer(ConsumerId,
+                {Effects0, #state{consumers = Cons0,
+                                  single_active_consumer_on = true,
+                                  waiting_consumers = WaitingConsumers0  } = State0}) ->
+    %% single active consumer on, consumers are waiting
+    case maps:take(ConsumerId, Cons0) of
+        {_CurrentActiveConsumer, _} ->
+            % The active consumer is to be removed
+            % Cancel it
+            {Effects1, State1} = case maps:take(ConsumerId, Cons0) of
+                {#consumer{checked_out = Checked0}, _} ->
+                    S = return_all(State0, Checked0),
+                    Effects = cancel_consumer_effects(ConsumerId, S, Effects0),
+                    {Effects, State0};
+                error ->
+                    {Effects0, State0}
+            end,
+            % Take another one from the waiting consumers and put it in consumers
+            {NewActiveConsumerId, NewActiveConsumer} = lists:nth(1, WaitingConsumers0),
+            WaitingConsumers1 = lists:delete({NewActiveConsumerId, NewActiveConsumer}, WaitingConsumers0),
+            #state{service_queue = ServiceQueue} = State0,
+            ServiceQueue1 = maybe_queue_consumer(NewActiveConsumerId, NewActiveConsumer, ServiceQueue),
+            State2 = State1#state{consumers = #{NewActiveConsumerId => NewActiveConsumer},
+                                  service_queue = ServiceQueue1,
+                                  waiting_consumers = WaitingConsumers1},
+            {Effects1, State2};
+        error ->
+            % The cancelled consumer is not the active one
+            % Just remove it from idle_consumers
+            case lists:keytake(ConsumerId, 1, WaitingConsumers0) of
+                {value, {ConsumerId, #consumer{checked_out = Checked0}}, WaitingConsumers1} ->
+                    S = return_all(State0, Checked0),
+                    Effects = cancel_consumer_effects(ConsumerId, S, Effects0),
+                    {Effects, #state{waiting_consumers = WaitingConsumers1}};
+                false ->
+                    {Effects0, State0}
+            end
+    end.
+
+cancel_consumer0(ConsumerId,
                 {Effects0, #state{consumers = C0} = S0}) ->
     case maps:take(ConsumerId, C0) of
         {#consumer{checked_out = Checked0}, Cons} ->
@@ -677,7 +732,7 @@ cancel_consumer(ConsumerId,
                     {[{aux, inactive} | Effects], S#state{consumers = Cons}};
                 _ ->
                     {Effects, S#state{consumers = Cons}}
-                end;
+            end;
         error ->
             % already removed - do nothing
             {Effects0, S0}
@@ -1053,23 +1108,42 @@ uniq_queue_in(Key, Queue) ->
     end.
 
 
-update_consumer(ConsumerId, Meta, {Life, Credit, Mode},
+update_consumer(ConsumerId, Meta, Spec,
+                #state{single_active_consumer_on = false} = State0) ->
+    %% general case, single active consumer off
+    update_consumer0(ConsumerId, Meta, Spec, State0);
+update_consumer(ConsumerId, Meta, Spec,
                 #state{consumers = Cons0,
-                       service_queue = ServiceQueue0} = State0) ->
+                       single_active_consumer_on = true} = State0) when map_size(Cons0) == 0 ->
+    %% single active consumer on, no one is consuming yet
+    update_consumer0(ConsumerId, Meta, Spec, State0);
+update_consumer(ConsumerId, Meta, {Life, Credit, Mode},
+                #state{single_active_consumer_on = true,
+                       waiting_consumers = WaitingConsumers0} = State0) ->
+    %% single active consumer on and one active consumer already
+    %% adding the new consumer to the waiting list
+    Consumer = #consumer{lifetime = Life, meta = Meta,
+                         credit = Credit, credit_mode = Mode},
+    WaitingConsumers1 = WaitingConsumers0 ++ [{ConsumerId, Consumer}],
+    State0#state{waiting_consumers = WaitingConsumers1}.
+
+update_consumer0(ConsumerId, Meta, {Life, Credit, Mode},
+                 #state{consumers = Cons0,
+                        service_queue = ServiceQueue0} = State0) ->
     %% TODO: this logic may not be correct for updating a pre-existing consumer
     Init = #consumer{lifetime = Life, meta = Meta,
                      credit = Credit, credit_mode = Mode},
     Cons = maps:update_with(ConsumerId,
-                             fun(S) ->
-                                     %% remove any in-flight messages from
-                                     %% the credit update
-                                     N = maps:size(S#consumer.checked_out),
-                                     C = max(0, Credit - N),
-                                     S#consumer{lifetime = Life,
-                                                credit = C}
-                             end, Init, Cons0),
+                            fun(S) ->
+                                %% remove any in-flight messages from
+                                %% the credit update
+                                N = maps:size(S#consumer.checked_out),
+                                C = max(0, Credit - N),
+                                S#consumer{lifetime = Life,
+                                    credit = C}
+                            end, Init, Cons0),
     ServiceQueue = maybe_queue_consumer(ConsumerId, maps:get(ConsumerId, Cons),
-                                        ServiceQueue0),
+        ServiceQueue0),
 
     State0#state{consumers = Cons, service_queue = ServiceQueue}.
 
diff --git a/src/rabbit_quorum_queue.erl b/src/rabbit_quorum_queue.erl
@@ -37,7 +37,8 @@
 -export([policy_changed/2]).
 -export([cleanup_data_dir/0]).
 
--include_lib("rabbit_common/include/rabbit.hrl").
+%%-include_lib("rabbit_common/include/rabbit.hrl").
+-include_lib("rabbit.hrl").
 -include_lib("stdlib/include/qlc.hrl").
 
 -type ra_server_id() :: {Name :: atom(), Node :: node()}.
@@ -151,7 +152,14 @@ ra_machine_config(Q = #amqqueue{name = QName}) ->
     #{dead_letter_handler => dlx_mfa(Q),
       queue_resource => QName,
       become_leader_handler => {?MODULE, become_leader, [QName]},
-      metrics_handler => {?MODULE, update_metrics, [QName]}}.
+      metrics_handler => {?MODULE, update_metrics, [QName]},
+      single_active_consumer_on => single_active_consumer_on(Q)}.
+
+single_active_consumer_on(#amqqueue{arguments = QArguments}) ->
+    case rabbit_misc:table_lookup(QArguments, <<"x-single-active-consumer">>) of
+        {bool, true} -> true;
+        _            -> false
+    end.
 
 cancel_consumer_handler(QName, {ConsumerTag, ChPid}) ->
     Node = node(ChPid),
diff --git a/test/single_active_consumer_SUITE.erl b/test/single_active_consumer_SUITE.erl
@@ -24,16 +24,22 @@
 
 all() ->
     [
-        {group, default}
+        {group, classic_queue}, {group, quorum_queue}
     ].
 
 groups() ->
     [
-        {default, [], [
+        {classic_queue, [], [
             all_messages_go_to_one_consumer,
             fallback_to_another_consumer_when_first_one_is_cancelled,
             fallback_to_another_consumer_when_exclusive_consumer_channel_is_cancelled,
             amqp_exclusive_consume_fails_on_exclusive_consumer_queue
+        ]},
+        {quorum_queue, [], [
+            all_messages_go_to_one_consumer,
+            fallback_to_another_consumer_when_first_one_is_cancelled,
+            fallback_to_another_consumer_when_exclusive_consumer_channel_is_cancelled
+            %% amqp_exclusive_consume_fails_on_exclusive_consumer_queue % Exclusive consume not implemented in QQ
         ]}
     ].
 
@@ -48,18 +54,37 @@ init_per_suite(Config) ->
 
 end_per_suite(Config) ->
     rabbit_ct_helpers:run_teardown_steps(Config,
-        rabbit_ct_client_helpers:teardown_steps() ++
-        rabbit_ct_broker_helpers:teardown_steps()).
+    rabbit_ct_client_helpers:teardown_steps() ++
+    rabbit_ct_broker_helpers:teardown_steps()).
+
+init_per_group(classic_queue, Config) ->
+    [{single_active_consumer_queue_declare,
+        #'queue.declare'{arguments = [
+            {<<"x-single-active-consumer">>, bool, true},
+            {<<"x-queue-type">>, longstr, <<"classic">>}
+        ],
+            auto_delete = true}
+        } | Config];
+init_per_group(quorum_queue, Config) ->
+    [{single_active_consumer_queue_declare,
+        #'queue.declare'{arguments = [
+            {<<"x-single-active-consumer">>, bool, true},
+            {<<"x-queue-type">>, longstr, <<"quorum">>}
+        ],
+            durable = true, exclusive = false, auto_delete = false}
+    } | Config].
+
+end_per_group(_, Config) ->
+    Config.
 
 init_per_testcase(Testcase, Config) ->
     rabbit_ct_helpers:testcase_started(Config, Testcase).
-
 end_per_testcase(Testcase, Config) ->
     rabbit_ct_helpers:testcase_finished(Config, Testcase).
 
 all_messages_go_to_one_consumer(Config) ->
     {C, Ch} = connection_and_channel(Config),
-    Q = queue_declare(Ch),
+    Q = queue_declare(Ch, Config),
     NbMessages = 5,
     ConsumerPid = spawn(?MODULE, consume, [{self(), {maps:new(), 0}, NbMessages}]),
     #'basic.consume_ok'{consumer_tag = CTag1} =
@@ -85,36 +110,52 @@ all_messages_go_to_one_consumer(Config) ->
 
 fallback_to_another_consumer_when_first_one_is_cancelled(Config) ->
     {C, Ch} = connection_and_channel(Config),
-    Q = queue_declare(Ch),
+    Q = queue_declare(Ch, Config),
     NbMessages = 10,
     ConsumerPid = spawn(?MODULE, consume, [{self(), {maps:new(), 0}, NbMessages}]),
     #'basic.consume_ok'{consumer_tag = CTag1} =
         amqp_channel:subscribe(Ch, #'basic.consume'{queue = Q, no_ack = true}, ConsumerPid),
     #'basic.consume_ok'{consumer_tag = CTag2} =
         amqp_channel:subscribe(Ch, #'basic.consume'{queue = Q, no_ack = true}, ConsumerPid),
-    #'basic.consume_ok'{consumer_tag = _CTag3} =
+    #'basic.consume_ok'{consumer_tag = CTag3} =
         amqp_channel:subscribe(Ch, #'basic.consume'{queue = Q, no_ack = true}, ConsumerPid),
 
     Publish = #'basic.publish'{exchange = <<>>, routing_key = Q},
     [amqp_channel:cast(Ch, Publish, #amqp_msg{payload = <<"foobar">>}) || _X <- lists:seq(1, NbMessages div 2)],
 
-    #'basic.cancel_ok'{} = amqp_channel:call(Ch, #'basic.cancel'{consumer_tag = CTag1}),
+    {MessagesPerConsumer1, _} = wait_for_messages(NbMessages div 2),
+    FirstActiveConsumerInList = maps:keys(maps:filter(fun(_CTag, MessageCount) -> MessageCount > 0 end, MessagesPerConsumer1)),
+    ?assertEqual(1, length(FirstActiveConsumerInList)),
+
+    FirstActiveConsumer = lists:nth(1, FirstActiveConsumerInList),
+    #'basic.cancel_ok'{} = amqp_channel:call(Ch, #'basic.cancel'{consumer_tag = FirstActiveConsumer}),
+
+    {cancel_ok, FirstActiveConsumer} = wait_for_cancel_ok(),
 
     [amqp_channel:cast(Ch, Publish, #amqp_msg{payload = <<"foobar">>}) || _X <- lists:seq(NbMessages div 2 + 1, NbMessages - 1)],
 
-    #'basic.cancel_ok'{} = amqp_channel:call(Ch, #'basic.cancel'{consumer_tag = CTag2}),
+    {MessagesPerConsumer2, _} = wait_for_messages(NbMessages div 2 - 1),
+    SecondActiveConsumerInList = maps:keys(maps:filter(
+        fun(CTag, MessageCount) -> MessageCount > 0 andalso CTag /= FirstActiveConsumer end,
+        MessagesPerConsumer2)
+    ),
+    ?assertEqual(1, length(SecondActiveConsumerInList)),
+    SecondActiveConsumer = lists:nth(1, SecondActiveConsumerInList),
+
+    #'basic.cancel_ok'{} = amqp_channel:call(Ch, #'basic.cancel'{consumer_tag = SecondActiveConsumer}),
 
     amqp_channel:cast(Ch, Publish, #amqp_msg{payload = <<"foobar">>}),
+    wait_for_messages(1),
+
+    LastActiveConsumer = lists:nth(1, lists:delete(FirstActiveConsumer, lists:delete(SecondActiveConsumer, [CTag1, CTag2, CTag3]))),
 
     receive
         {consumer_done, {MessagesPerConsumer, MessageCount}} ->
             ?assertEqual(NbMessages, MessageCount),
             ?assertEqual(3, maps:size(MessagesPerConsumer)),
-            ?assertEqual(NbMessages div 2, maps:get(CTag1, MessagesPerConsumer)),
-            Counts = maps:values(MessagesPerConsumer),
-            ?assert(lists:member(NbMessages div 2, Counts)),
-            ?assert(lists:member(NbMessages div 2 - 1, Counts)),
-            ?assert(lists:member(1, Counts))
+            ?assertEqual(NbMessages div 2, maps:get(FirstActiveConsumer, MessagesPerConsumer)),
+            ?assertEqual(NbMessages div 2 - 1, maps:get(SecondActiveConsumer, MessagesPerConsumer)),
+            ?assertEqual(1, maps:get(LastActiveConsumer, MessagesPerConsumer))
     after 1000 ->
         throw(failed)
     end,
@@ -127,7 +168,7 @@ fallback_to_another_consumer_when_exclusive_consumer_channel_is_cancelled(Config
     {C1, Ch1} = connection_and_channel(Config),
     {C2, Ch2} = connection_and_channel(Config),
     {C3, Ch3} = connection_and_channel(Config),
-    Q = queue_declare(Ch),
+    Q = queue_declare(Ch, Config),
     NbMessages = 10,
     Consumer1Pid = spawn(?MODULE, consume, [{self(), {maps:new(), 0}, NbMessages div 2}]),
     Consumer2Pid = spawn(?MODULE, consume, [{self(), {maps:new(), 0}, NbMessages div 2 - 1}]),
@@ -168,7 +209,7 @@ fallback_to_another_consumer_when_exclusive_consumer_channel_is_cancelled(Config
 
 amqp_exclusive_consume_fails_on_exclusive_consumer_queue(Config) ->
     {C, Ch} = connection_and_channel(Config),
-    Q = queue_declare(Ch),
+    Q = queue_declare(Ch, Config),
     ?assertExit(
         {{shutdown, {server_initiated_close, 403, _}}, _},
         amqp_channel:call(Ch, #'basic.consume'{queue = Q, exclusive = true})
@@ -181,9 +222,8 @@ connection_and_channel(Config) ->
     {ok, Ch} = amqp_connection:open_channel(C),
     {C, Ch}.
 
-queue_declare(Channel) ->
-    Declare = #'queue.declare'{arguments = [{"x-single-active-consumer", bool, true}],
-        auto_delete = true},
+queue_declare(Channel, Config) ->
+    Declare = ?config(single_active_consumer_queue_declare, Config),
     #'queue.declare_ok'{queue = Q} = amqp_channel:call(Channel, Declare),
     Q.
 
@@ -198,11 +238,12 @@ consume({Parent, {MessagesPerConsumer, MessageCount}, CountDown}) ->
                 {maps:update_with(CTag, fun(V) -> V + 1 end, MessagesPerConsumer),
                     MessageCount + 1}};
         {#'basic.deliver'{consumer_tag = CTag}, _Content} ->
-            consume({Parent,
-                {maps:update_with(CTag, fun(V) -> V + 1 end, MessagesPerConsumer),
-                    MessageCount + 1},
-                CountDown - 1});
-        #'basic.cancel_ok'{} ->
+            NewState = {maps:update_with(CTag, fun(V) -> V + 1 end, MessagesPerConsumer),
+                MessageCount + 1},
+            Parent ! {message, NewState},
+            consume({Parent, NewState, CountDown - 1});
+        #'basic.cancel_ok'{consumer_tag = CTag} ->
+            Parent ! {cancel_ok, CTag},
             consume({Parent, {MessagesPerConsumer, MessageCount}, CountDown});
         _ ->
             consume({Parent, {MessagesPerConsumer, MessageCount}, CountDown})
@@ -216,7 +257,30 @@ consume_results() ->
         {consumer_done, {MessagesPerConsumer, MessageCount}} ->
             {MessagesPerConsumer, MessageCount};
         {consumer_timeout, {MessagesPerConsumer, MessageCount}} ->
-            {MessagesPerConsumer, MessageCount}
+            {MessagesPerConsumer, MessageCount};
+        _ ->
+            consume_results()
     after 1000 ->
         throw(failed)
+    end.
+
+wait_for_messages(ExpectedCount) ->
+    wait_for_messages(ExpectedCount, {}).
+
+wait_for_messages(0, State) ->
+    State;
+wait_for_messages(ExpectedCount, _) ->
+    receive
+        {message, {MessagesPerConsumer, MessageCount}} ->
+            wait_for_messages(ExpectedCount - 1, {MessagesPerConsumer, MessageCount})
+    after 5000 ->
+        throw(message_waiting_timeout)
+    end.
+
+wait_for_cancel_ok() ->
+    receive
+        {cancel_ok, CTag} ->
+            {cancel_ok, CTag}
+    after 5000 ->
+        throw(consumer_cancel_ok_timeout)
     end.
