QQ: add new consumer cancel option: 'remove'

This option immediately removes and returns all messages for a
consumer instead of the softer 'cancel' option which keeps the
consumer around until all pending messages have been either
settled or returned.

This involves a change to the rabbit_queue_type:cancel/5 API
to rabbit_queue_type:cancel/3.

Author: kjnilsson

deps/rabbit/src/rabbit_amqp_session.erl

@@ -1136,8 +1136,9 @@ handle_control(Detach = #'v1_0.detach'{handle = ?UINT(HandleInt)},
               QName = rabbit_misc:r(Vhost, queue, QNameBin),
               case rabbit_amqqueue:lookup(QName) of
                   {ok, Q} ->
-                      %%TODO Consider adding a new rabbit_queue_type:remove_consumer API that - from the point of view of
-                      %% the queue process - behaves as if our session process terminated: All messages checked out
+                      %% TODO: Consider adding a new rabbit_queue_type:remove_consumer
+                      %% API that - from the point of view of the queue process -
+                      %% behaves as if our session process terminated: All messages checked out
                       %% to this consumer should be re-queued automatically instead of us requeueing them here after cancelling
                       %% consumption.
                       %% For AMQP legacy (and STOMP / MQTT) consumer cancellation not requeueing messages is a good approach as
@@ -1149,7 +1150,9 @@ handle_control(Detach = #'v1_0.detach'{handle = ?UINT(HandleInt)},
                       %% first detaching and then re-attaching to the same session with the same link handle (the handle
                       %% becomes available for re-use once a link is closed): This will result in the same consumer tag,
                       %% and we ideally disallow "updating" an AMQP consumer.
-                      case rabbit_queue_type:cancel(Q, Ctag, undefined, Username, QStates0) of
+                      Spec = #{consumer_tag => Ctag,
+                               user => Username},
+                      case rabbit_queue_type:cancel(Q, Spec, QStates0) of
                           {ok, QStates1} ->
                               {Unsettled1, MsgIds} = remove_link_from_outgoing_unsettled_map(Ctag, Unsettled0),
                               case MsgIds of

deps/rabbit/src/rabbit_amqqueue.erl

@@ -1772,8 +1772,10 @@ basic_consume(Q, NoAck, ChPid, LimiterPid,
                    rabbit_queue_type:state()) ->
     {ok, rabbit_queue_type:state()} | {error, term()}.
 basic_cancel(Q, ConsumerTag, OkMsg, ActingUser, QStates) ->
-    rabbit_queue_type:cancel(Q, ConsumerTag,
-                             OkMsg, ActingUser, QStates).
+    %% TODO: is this function used anywhere?
+    rabbit_queue_type:cancel(Q, #{consumer_tag => ConsumerTag,
+                                  ok_msg => OkMsg,
+                                  user => ActingUser}, QStates).
 
 -spec notify_decorators(amqqueue:amqqueue()) -> 'ok'.
 

deps/rabbit/src/rabbit_channel.erl

@@ -1454,8 +1454,9 @@ handle_method(#'basic.cancel'{consumer_tag = ConsumerTag, nowait = NoWait},
                    fun () -> {error, not_found} end,
                    fun () ->
                            rabbit_queue_type:cancel(
-                             Q, ConsumerTag, ok_msg(NoWait, OkMsg),
-                             Username, QueueStates0)
+                             Q, #{consumer_tag => ConsumerTag,
+                                  ok_msg => ok_msg(NoWait, OkMsg),
+                                  user => Username}, QueueStates0)
                    end) of
                 {ok, QueueStates} ->
                     rabbit_global_counters:consumer_deleted(amqp091),

deps/rabbit/src/rabbit_classic_queue.erl

@@ -37,7 +37,7 @@
          close/1,
          update/2,
          consume/3,
-         cancel/5,
+         cancel/3,
          handle_event/3,
          deliver/3,
          settle/5,
@@ -276,7 +276,9 @@ consume_backwards_compat({credited, credit_api_v1}, Args) ->
      [{<<"x-credit">>, table, [{<<"credit">>, long, 0},
                                {<<"drain">>,  bool, false}]} | Args]}.
 
-cancel(Q, ConsumerTag, OkMsg, ActingUser, State) ->
+cancel(Q, #{consumer_tag := ConsumerTag,
+            user := ActingUser} = Spec, State) ->
+    OkMsg = maps:get(ok_msg, Spec, undefined),
     QPid = amqqueue:get_pid(Q),
     case delegate:invoke(QPid, {gen_server2, call,
                                 [{basic_cancel, self(), ConsumerTag,

deps/rabbit/src/rabbit_fifo.erl

@@ -477,13 +477,15 @@ apply(#{index := Index,
             end
     end;
 apply(#{index := Idx} = Meta,
-      #checkout{spec = cancel,
-                consumer_id = ConsumerId}, State0) ->
+      #checkout{spec = Spec,
+                consumer_id = ConsumerId}, State0)
+  when Spec == cancel orelse
+       Spec == remove ->
     case consumer_key_from_id(ConsumerId, State0) of
         {ok, ConsumerKey} ->
             {State1, Effects1} = activate_next_consumer(
                                    cancel_consumer(Meta, ConsumerKey, State0, [],
-                                                   consumer_cancel)),
+                                                   Spec)),
             Reply = {ok, consumer_cancel_info(ConsumerKey, State1)},
             {State, _, Effects} = checkout(Meta, State0, State1, Effects1),
             update_smallest_raft_index(Idx, Reply, State, Effects);
@@ -1536,14 +1538,14 @@ maybe_return_all(#{system_time := Ts} = Meta, ConsumerKey,
                  #consumer{cfg = CCfg} = Consumer, S0,
                  Effects0, Reason) ->
     case Reason of
-        consumer_cancel ->
+        cancel ->
             {update_or_remove_con(
                Meta, ConsumerKey,
                Consumer#consumer{cfg = CCfg#consumer_cfg{lifetime = once},
                                  credit = 0,
                                  status = cancelled},
                S0), Effects0};
-        down ->
+        _ ->
             {S1, Effects1} = return_all(Meta, S0, Effects0, ConsumerKey, Consumer),
             {S1#?STATE{consumers = maps:remove(ConsumerKey, S1#?STATE.consumers),
                        last_active = Ts},
@@ -2436,9 +2438,16 @@ make_enqueue(Pid, Seq, Msg) ->
 make_register_enqueuer(Pid) ->
     #register_enqueuer{pid = Pid}.
 
--spec make_checkout(consumer_id(),
-                    checkout_spec(), consumer_meta()) -> protocol().
-make_checkout({_, _} = ConsumerId, Spec, Meta) ->
+-spec make_checkout(consumer_id(), checkout_spec(), consumer_meta()) ->
+    protocol().
+make_checkout({_, _} = ConsumerId, Spec0, Meta) ->
+    Spec = case is_v4() of
+               false when Spec0 == remove ->
+                   %% if v4 is not active, fall back to cancel spec
+                   make_checkout(ConsumerId, cancel, Meta);
+               _ ->
+                   Spec0
+           end,
     #checkout{consumer_id = ConsumerId,
               spec = Spec, meta = Meta}.
 

deps/rabbit/src/rabbit_fifo.hrl

@@ -80,7 +80,7 @@
                           credited | simple_prefetch} |
 
                          {dequeue, settled | unsettled} |
-                         cancel |
+                         cancel | remove |
                          %% new v4 format
                          {once | auto, credit_mode()}.
 

deps/rabbit/src/rabbit_fifo_client.erl

@@ -16,6 +16,7 @@
          init/2,
          checkout/4,
          cancel_checkout/2,
+         cancel_checkout/3,
          enqueue/3,
          enqueue/4,
          dequeue/4,
@@ -459,6 +460,9 @@ credit(ConsumerTag, DeliveryCount, Credit, Drain, Echo,
     State = State0#state{consumers = CDels},
     {send_command(ServerId, undefined, Cmd, normal, State), []}.
 
+cancel_checkout(ConsumerTag, State) ->
+    cancel_checkout(ConsumerTag, cancel, State).
+
 %% @doc Cancels a checkout with the rabbit_fifo queue  for the consumer tag
 %%
 %% This is a synchronous call. I.e. the call will block until the command
@@ -468,18 +472,35 @@ credit(ConsumerTag, DeliveryCount, Credit, Drain, Echo,
 %% @param State The {@module} state.
 %%
 %% @returns `{ok, State}' or `{error | timeout, term()}'
--spec cancel_checkout(rabbit_types:ctag(), state()) ->
+-spec cancel_checkout(rabbit_types:ctag(), Reason :: cancel | remove, state()) ->
     {ok, state()} | {error | timeout, term()}.
-cancel_checkout(ConsumerTag, #state{consumers = Consumers} = State0) ->
+cancel_checkout(ConsumerTag, Reason,
+                #state{consumers = Consumers,
+                       unsent_commands = Unsent} = State0)
+  when is_atom(Reason) ->
     case Consumers of
-        #{ConsumerTag := #consumer{}} ->
+        #{ConsumerTag := #consumer{key = Cid}} ->
             Servers = sorted_servers(State0),
             ConsumerId = {ConsumerTag, self()},
-            %% TODO: send any pending commands for consumer
-            %% checkout always uses the ConsumerId, rather than the key
-            Cmd = rabbit_fifo:make_checkout(ConsumerId, cancel, #{}),
-            State = State0#state{consumers = maps:remove(ConsumerTag, Consumers)},
-            case try_process_command(Servers, Cmd, State0) of
+            %% send any pending commands for consumer before cancelling
+            Commands = case Unsent of
+                           #{Cid := {Settled, Returns, Discards}} ->
+                                 add_command(Cid, settle, Settled,
+                                             add_command(Cid, return, Returns,
+                                                         add_command(Cid, discard,
+                                                                     Discards, [])));
+                           _ ->
+                               []
+                       end,
+            ServerId = pick_server(State0),
+            %% send all the settlements, discards and returns
+            State1 = lists:foldl(fun (C, S0) ->
+                                        send_command(ServerId, undefined, C,
+                                                     normal, S0)
+                                end, State0, Commands),
+            Cmd = rabbit_fifo:make_checkout(ConsumerId, Reason, #{}),
+            State = State1#state{consumers = maps:remove(ConsumerTag, Consumers)},
+            case try_process_command(Servers, Cmd, State) of
                 {ok, _, Leader} ->
                     {ok, State#state{leader = Leader}};
                 Err ->

deps/rabbit/src/rabbit_queue_type.erl

@@ -39,7 +39,7 @@
          %% stateful client API
          new/2,
          consume/3,
-         cancel/5,
+         cancel/3,
          handle_down/4,
          handle_event/3,
          module/2,
@@ -124,7 +124,11 @@
                           exclusive_consume => boolean(),
                           args => rabbit_framing:amqp_table(),
                           ok_msg := term(),
-                          acting_user :=  rabbit_types:username()}.
+                          acting_user := rabbit_types:username()}.
+-type cancel_spec() :: #{consumer_tag := rabbit_types:ctag(),
+                         reason => cancel | remove,
+                         ok_msg => term(),
+                         user := rabbit_types:username()}.
 
 -type delivery_options() :: #{correlation => correlation(),
                               atom() => term()}.
@@ -134,6 +138,7 @@
 -export_type([state/0,
               consume_mode/0,
               consume_spec/0,
+              cancel_spec/0,
               delivery_options/0,
               action/0,
               actions/0,
@@ -194,9 +199,7 @@
     {protocol_error, Type :: atom(), Reason :: string(), Args :: term()}.
 
 -callback cancel(amqqueue:amqqueue(),
-                 rabbit_types:ctag(),
-                 term(),
-                 rabbit_types:username(),
+                 cancel_spec(),
                  queue_state()) ->
     {ok, queue_state()} | {error, term()}.
 
@@ -453,17 +456,15 @@ consume(Q, Spec, State) ->
             Err
     end.
 
-%% TODO switch to cancel spec api
+%% TODO: switch to cancel spec api
 -spec cancel(amqqueue:amqqueue(),
-             rabbit_types:ctag(),
-             term(),
-             rabbit_types:username(),
+             cancel_spec(),
              state()) ->
     {ok, state()} | {error, term()}.
-cancel(Q, Tag, OkMsg, ActiveUser, Ctxs) ->
+cancel(Q, Spec, Ctxs) ->
     #ctx{state = State0} = Ctx = get_ctx(Q, Ctxs),
     Mod = amqqueue:get_type(Q),
-    case Mod:cancel(Q, Tag, OkMsg, ActiveUser, State0) of
+    case Mod:cancel(Q, Spec, State0) of
         {ok, State} ->
             {ok, set_ctx(Q, Ctx#ctx{state = State}, Ctxs)};
         Err ->

deps/rabbit/src/rabbit_quorum_queue.erl

@@ -24,7 +24,7 @@
          delete/4,
          delete_immediately/1]).
 -export([state_info/1, info/2, stat/1, infos/1, infos/2]).
--export([settle/5, dequeue/5, consume/3, cancel/5]).
+-export([settle/5, dequeue/5, consume/3, cancel/3]).
 -export([credit_v1/5, credit/7]).
 -export([purge/1]).
 -export([stateless_deliver/2, deliver/3]).
@@ -130,7 +130,7 @@
 -define(DELETE_TIMEOUT, 5000).
 -define(MEMBER_CHANGE_TIMEOUT, 20_000).
 -define(SNAPSHOT_INTERVAL, 8192). %% the ra default is 4096
--define(UNLIMITED_PREFETCH_COUNT, 2000). %% something large for ra
+% -define(UNLIMITED_PREFETCH_COUNT, 2000). %% something large for ra
 
 %%----------- QQ policies ---------------------------------------------------
 
@@ -906,8 +906,8 @@ consume(Q, Spec, QState0) when ?amqqueue_is_quorum(Q) ->
             {ok, QState}
     end.
 
-cancel(_Q, ConsumerTag, OkMsg, _ActingUser, State) ->
-    maybe_send_reply(self(), OkMsg),
+cancel(_Q, #{consumer_tag := ConsumerTag} = Spec, State) ->
+    maybe_send_reply(self(), maps:get(ok_msg, Spec, undefined)),
     rabbit_fifo_client:cancel_checkout(quorum_ctag(ConsumerTag), State).
 
 emit_consumer_created(ChPid, CTag, Exclusive, AckRequired, QName, PrefetchCount, Args, Ref, ActingUser) ->

deps/rabbit/src/rabbit_stream_queue.erl

@@ -19,7 +19,7 @@
          recover/2,
          is_recoverable/1,
          consume/3,
-         cancel/5,
+         cancel/3,
          handle_event/3,
          deliver/3,
          settle/5,
@@ -446,8 +446,10 @@ begin_stream(#stream_client{name = QName,
                    reader_options = Options},
     {ok, State#stream_client{readers = Readers0#{Tag => Str0}}}.
 
-cancel(_Q, ConsumerTag, OkMsg, ActingUser, #stream_client{readers = Readers0,
-                                                          name = QName} = State) ->
+cancel(_Q, #{consumer_tag := ConsumerTag,
+             user := ActingUser} = Spec,
+       #stream_client{readers = Readers0,
+                      name = QName} = State) ->
     case maps:take(ConsumerTag, Readers0) of
         {#stream{log = Log}, Readers} ->
             ok = close_log(Log),
@@ -457,7 +459,7 @@ cancel(_Q, ConsumerTag, OkMsg, ActingUser, #stream_client{readers = Readers0,
                                  {channel, self()},
                                  {queue, QName},
                                  {user_who_performed_action, ActingUser}]),
-            maybe_send_reply(self(), OkMsg),
+            maybe_send_reply(self(), maps:get(ok_msg, Spec, undefined)),
             {ok, State#stream_client{readers = Readers}};
         error ->
             {ok, State}