WIP Support returns

Author: ansd

deps/rabbit/src/rabbit_fifo.erl

@@ -200,7 +200,8 @@ init(#{name := Name,
                                                    resource = Resource}}),
     case State#?STATE.cfg#cfg.filter_enabled of
         true ->
-            State#?STATE{messages = rabbit_fifo_filter_q:new()};
+            State#?STATE{messages = rabbit_fifo_filter_q:new(),
+                         returns = gb_trees:empty()};
         false ->
             State
     end.
@@ -506,7 +507,7 @@ apply(#{index := Idx} = Meta,
 apply(#{index := Index}, #purge{},
       #?STATE{cfg = #cfg{filter_enabled = FilterEnabled},
               messages_total = Total,
-              returns = Returns,
+              returns = Returns0,
               ra_indexes = Indexes0
              } = State0) ->
     NumReady = messages_ready(State0),
@@ -517,7 +518,7 @@ apply(#{index := Index}, #purge{},
                       %% No message is awaiting acknowledgement.
                       %% Optimization: empty all 'ra_indexes'.
                       rabbit_fifo_index:empty();
-                  _ ->
+                  _ when not FilterEnabled ->
                       %% Some messages are checked out to consumers
                       %% awaiting acknowledgement.
                       %% Therefore we cannot empty all 'ra_indexes'.
@@ -526,16 +527,20 @@ apply(#{index := Index}, #purge{},
                       %% not part of the 'ra_indexes'.
                       lqueue:fold(fun(?MSG(I, _), Acc) ->
                                           rabbit_fifo_index:delete(I, Acc)
-                                  end, Indexes0, Returns)
+                                  end, Indexes0, Returns0);
+                  _ ->
+                      Indexes0
               end,
-    Messages = case FilterEnabled of
-                   true -> rabbit_fifo_filter_q:new();
-                   false -> rabbit_fifo_q:new()
-               end,
+    {Messages, Returns} = case FilterEnabled of
+                              true ->
+                                  {rabbit_fifo_filter_q:new(), gb_trees:empty()};
+                              false ->
+                                  {rabbit_fifo_q:new(), lqueue:new()}
+                          end,
     State1 = State0#?STATE{ra_indexes = Indexes,
                            messages = Messages,
                            messages_total = Total - NumReady,
-                           returns = lqueue:new(),
+                           returns = Returns,
                            msg_bytes_enqueue = 0
                           },
     Effects0 = [{aux, force_checkpoint}, garbage_collection],
@@ -890,7 +895,7 @@ overview(#?STATE{consumers = Cons,
                  enqueue_count = EnqCount,
                  msg_bytes_enqueue = EnqueueBytes,
                  msg_bytes_checkout = CheckoutBytes,
-                 cfg = Cfg,
+                 cfg = #cfg{filter_enabled = FilterEnabled} = Cfg,
                  dlx = DlxState,
                  messages = Messages,
                  returns = Returns,
@@ -917,12 +922,16 @@ overview(#?STATE{consumers = Cons,
                       _ ->
                           #{}
                   end,
-    MsgsRet = lqueue:len(Returns),
-    #{num_hi := MsgsHi,
-      num_no := MsgsNo} = case Cfg#cfg.filter_enabled of
-                              true -> rabbit_fifo_filter_q:overview(Messages);
-                              false -> rabbit_fifo_q:overview(Messages)
-                          end,
+    {#{num_hi := MsgsHi,
+       num_no := MsgsNo},
+     MsgsRet} = case FilterEnabled of
+                    true ->
+                        {rabbit_fifo_filter_q:overview(Messages),
+                         gb_trees:size(Returns)};
+                    false ->
+                        {rabbit_fifo_q:overview(Messages),
+                         lqueue:len(Returns)}
+                end,
     Overview = #{type => ?STATE,
                  config => Conf,
                  num_consumers => map_size(Cons),
@@ -1401,8 +1410,8 @@ is_v4() ->
 
 messages_ready(#?STATE{cfg = #cfg{filter_enabled = true},
                        messages = M,
-                       returns = _Todo}) ->
-    rabbit_fifo_filter_q:len(M);
+                       returns = R}) ->
+    rabbit_fifo_filter_q:size(M) + gb_trees:size(R);
 messages_ready(#?STATE{messages = M,
                        returns = R}) ->
     rabbit_fifo_q:len(M) + lqueue:len(R).
@@ -1523,7 +1532,8 @@ activate_next_consumer(#?STATE{cfg = #cfg{consumer_strategy = competing}} = Stat
 activate_next_consumer(#?STATE{cfg = #cfg{filter_enabled = FilterEnabled},
                                consumers = Cons0,
                                waiting_consumers = Waiting0,
-                               messages = Messages
+                               messages = Messages,
+                               return_count = ReturnCount
                               } = State0,
                        Effects0) ->
     %% invariant, the waiting list always need to be sorted by consumers that are
@@ -1555,7 +1565,8 @@ activate_next_consumer(#?STATE{cfg = #cfg{filter_enabled = FilterEnabled},
                                                  Consumer,
                                                  ServiceQueue,
                                                  FilterEnabled,
-                                                 Messages),
+                                                 Messages,
+                                                 ReturnCount),
             State = State0#?STATE{consumers = Cons0#{NextCKey => Consumer},
                                   service_queue = ServiceQueue1,
                                   consumers_q = queue_consumer(NextCKey,
@@ -1580,7 +1591,8 @@ activate_next_consumer(#?STATE{cfg = #cfg{filter_enabled = FilterEnabled},
                                                  Consumer,
                                                  ServiceQueue,
                                                  FilterEnabled,
-                                                 Messages),
+                                                 Messages,
+                                                 ReturnCount),
             Cons1 = Cons0#{NextCKey => Consumer},
             Cons = maps:remove(ActiveCKey, Cons1),
             ConQ1 = queue_consumer(NextCKey, ?CONSUMER_PRIORITY(Consumer), ConQ0),
@@ -1677,6 +1689,8 @@ queue_filtering_consumers(#?STATE{cfg = #cfg{filter_enabled = true},
     State#?STATE{service_queue = ConsQ,
                  %% Rotating should help dispatching messages across multiple consumers.
                  %% Note however that this doesn't guarantee round robin.
+                 %% TODO We might need some randomness whether we rotate to avoid
+                 %% certain patterns where always the same consumer gets messages assigned.
                  consumers_q = priority_queue:rotate(ConsQ)};
 queue_filtering_consumers(State) ->
     State.
@@ -1830,26 +1844,27 @@ enqueue(Idx, Ts, RawMsg, {_MetaSize, BodySize}, _MsgMeta,
 
 return(#{} = Meta, ConsumerKey, MsgIds, IncrDelCount, Anns,
        Checked, Effects0, State0)
- when is_map(Anns) ->
+  when is_map(Anns) ->
     %% We requeue in the same order as messages got returned by the client.
     {State1, Effects1} =
-        lists:foldl(
-          fun(MsgId, Acc = {S0, E0}) ->
-                  case Checked of
-                      #{MsgId := Msg} ->
-                          return_one(Meta, MsgId, Msg, IncrDelCount, Anns,
-                                     S0, E0, ConsumerKey);
-                      #{} ->
-                          Acc
-                  end
-          end, {State0, Effects0}, MsgIds),
+    lists:foldl(
+      fun(MsgId, Acc = {S0, E0}) ->
+              case Checked of
+                  #{MsgId := Msg} ->
+                      return_one(Meta, MsgId, Msg, IncrDelCount, Anns,
+                                 S0, E0, ConsumerKey);
+                  #{} ->
+                      Acc
+              end
+      end, {State0, Effects0}, MsgIds),
     State2 = case State1#?STATE.consumers of
                  #{ConsumerKey := Con} ->
                      update_or_remove_con(Meta, ConsumerKey, Con, State1);
                  _ ->
                      State1
              end,
-    checkout(Meta, State0, State2, Effects1).
+    State3 = queue_filtering_consumers(State2),
+    checkout(Meta, State0, State3, Effects1).
 
 % used to process messages that are finished
 complete(Meta, ConsumerKey, [MsgId],
@@ -1987,10 +2002,12 @@ annotate_msg(Header, Msg0) ->
     end.
 
 return_one(Meta, MsgId, ?MSG(_, _) = Msg0, DelivFailed, Anns,
-           #?STATE{returns = Returns,
+           #?STATE{returns = Returns0,
+                   return_count = ReturnCount0,
                    consumers = Consumers,
                    dlx = DlxState0,
-                   cfg = #cfg{delivery_limit = DeliveryLimit,
+                   cfg = #cfg{filter_enabled = FilterEnabled,
+                              delivery_limit = DeliveryLimit,
                               dead_letter_handler = DLH}} = State0,
            Effects0, ConsumerKey) ->
     #consumer{checked_out = Checked0} = Con0 = maps:get(ConsumerKey, Consumers),
@@ -2007,20 +2024,30 @@ return_one(Meta, MsgId, ?MSG(_, _) = Msg0, DelivFailed, Anns,
             Checked = maps:remove(MsgId, Checked0),
             Con = Con0#consumer{checked_out = Checked,
                                 credit = increase_credit(Con0, 1)},
-            {add_bytes_return(
-               Header,
-               State0#?STATE{consumers = Consumers#{ConsumerKey := Con},
-                             returns = lqueue:in(Msg, Returns)}),
-             Effects0}
+            ReturnCount = ReturnCount0 + 1,
+            Returns = case FilterEnabled of
+                          true ->
+                              gb_trees:insert(ReturnCount, Msg, Returns0);
+                          false ->
+                              lqueue:in(Msg, Returns0)
+                      end,
+            State1 = State0#?STATE{consumers = Consumers#{ConsumerKey := Con},
+                                   returns = Returns,
+                                   return_count = ReturnCount},
+            State = add_bytes_return(Header, State1),
+            {State, Effects0}
     end.
 
 return_all(Meta, #?STATE{consumers = Cons} = State0, Effects0, ConsumerKey,
            #consumer{checked_out = Checked} = Con, DelivFailed) ->
-    State = State0#?STATE{consumers = Cons#{ConsumerKey => Con}},
-    lists:foldl(fun ({MsgId, Msg}, {S, E}) ->
-                        return_one(Meta, MsgId, Msg, DelivFailed, #{},
-                                   S, E, ConsumerKey)
-                end, {State, Effects0}, lists:sort(maps:to_list(Checked))).
+    State1 = State0#?STATE{consumers = Cons#{ConsumerKey => Con}},
+    {State2, Effects} = lists:foldl(
+                          fun ({MsgId, Msg}, {S, E}) ->
+                                  return_one(Meta, MsgId, Msg, DelivFailed, #{},
+                                             S, E, ConsumerKey)
+                          end, {State1, Effects0}, lists:sort(maps:to_list(Checked))),
+    State = queue_filtering_consumers(State2),
+    {State, Effects}.
 
 checkout(Meta, OldState, State0, Effects0) ->
     checkout(Meta, OldState, State0, Effects0, ok).
@@ -2148,21 +2175,58 @@ add_delivery_effects(Effects0, AccMap, State) ->
                         end, Efs, chunk_disk_msgs(DiskMsgs, 0, [[]]))
               end, Effects0, AccMap).
 
-%%TODO check returns first
 take_next_consumer_msg(RaTs,
                        Con0 = #consumer{cfg = #consumer_cfg{filter = Filter},
-                                        max_filtered_idxs = Idxs0},
+                                        scanned_returns = ScannedReturns},
+                       State0 = #?STATE{returns = Returns0,
+                                        return_count = ReturnCount})
+  when ScannedReturns < ReturnCount ->
+    case filter_returns(RaTs, Filter, ScannedReturns, Returns0) of
+        none ->
+            Con = Con0#consumer{scanned_returns = ReturnCount},
+            take_next_consumer_msg(RaTs, Con, State0);
+        {Key, Msg} ->
+            Returns = gb_trees:delete(Key, Returns0),
+            Con = Con0#consumer{scanned_returns = Key},
+            State = State0#?MODULE{returns = Returns},
+            {Msg, Con, State}
+    end;
+take_next_consumer_msg(RaTs,
+                       Con0 = #consumer{cfg = #consumer_cfg{filter = Filter},
+                                        scanned_idxs = Idxs0},
                        State0 = #?STATE{messages = Messages0}) ->
     case rabbit_fifo_filter_q:take(RaTs, Filter, Idxs0, Messages0) of
         {empty, Idxs} ->
-            Con = Con0#consumer{max_filtered_idxs = Idxs},
+            Con = Con0#consumer{scanned_idxs = Idxs},
             {empty, Con};
         {Msg, Idxs, Messages} ->
-            Con = Con0#consumer{max_filtered_idxs = Idxs},
+            Con = Con0#consumer{scanned_idxs = Idxs},
             State = State0#?STATE{messages = Messages},
             {Msg, Con, State}
     end.
 
+filter_returns(RaTs, Filter, Scanned, Tree) ->
+    From = Scanned + 1,
+    Iter = gb_trees:iterator_from(From, Tree),
+    filter_returns0(RaTs, Filter, gb_trees:next(Iter)).
+
+filter_returns0(_RaTs, _Filter, none) ->
+    none;
+filter_returns0(RaTs, Filter,
+                {ReturnCount, ?MSG(_Idx, #{meta := Meta} = Hdr) = Msg, Iter}) ->
+    case rabbit_fifo:get_header(expiry, Hdr) of
+        ExpiryTs when is_integer(ExpiryTs) andalso RaTs >= ExpiryTs ->
+            %% Message expired.
+            filter_returns0(RaTs, Filter, gb_trees:next(Iter));
+        _ ->
+            case rabbit_fifo_filter:filter(Filter, Meta) of
+                true ->
+                    {ReturnCount, Msg};
+                false ->
+                    filter_returns0(RaTs, Filter, gb_trees:next(Iter))
+            end
+    end.
+
 take_next_msg(#?STATE{returns = Returns0,
                       messages = Messages0,
                       ra_indexes = Indexes0
@@ -2241,7 +2305,7 @@ checkout_one(Meta = #{system_time := Ts},
              Effects0) ->
     case priority_queue:out(SQ0) of
         {empty, _} ->
-            Activity = case rabbit_fifo_filter_q:len(Messages0) of
+            Activity = case rabbit_fifo_filter_q:size(Messages0) of
                            0 -> nochange;
                            _ -> inactive
                        end,
@@ -2515,7 +2579,8 @@ update_or_remove_con(_Meta, ConsumerKey,
                      #?STATE{cfg = #cfg{filter_enabled = FilterEnabled},
                              consumers = Cons,
                              service_queue = ServiceQueue,
-                             messages = Messages
+                             messages = Messages,
+                             return_count = ReturnCount
                             } = State) ->
     %%TODO On the hot path, we want maps:update/3.
     State#?STATE{consumers = maps:put(ConsumerKey, Con, Cons),
@@ -2525,15 +2590,17 @@ update_or_remove_con(_Meta, ConsumerKey,
                  %% priority_queue:member/2 check.
                  service_queue = maybe_queue_consumer(
                                    ConsumerKey, Con, ServiceQueue,
-                                   FilterEnabled, Messages)}.
+                                   FilterEnabled, Messages, ReturnCount)}.
 
 maybe_queue_consumer(Key, #consumer{credit = Credit,
-                                    max_filtered_idxs = Idxs,
+                                    scanned_idxs = Idxs,
+                                    scanned_returns = ScannedReturns,
                                     status = up,
                                     cfg = #consumer_cfg{priority = P}},
-                     ServiceQueue, true, Messages)
+                     ServiceQueue, true, Messages, ReturnCount)
   when Credit > 0 ->
-    case rabbit_fifo_filter_q:is_fully_scanned(Idxs, Messages) of
+    case ScannedReturns >= ReturnCount andalso
+         rabbit_fifo_filter_q:is_fully_scanned(Idxs, Messages) of
         true ->
             ServiceQueue;
         false ->
@@ -2542,10 +2609,11 @@ maybe_queue_consumer(Key, #consumer{credit = Credit,
 maybe_queue_consumer(Key, #consumer{credit = Credit,
                                     status = up,
                                     cfg = #consumer_cfg{priority = P}},
-                     ServiceQueue, false, _Messages)
+                     ServiceQueue, false, _Messages, _ReturnCount)
   when Credit > 0 ->
     queue_consumer(Key, P, ServiceQueue);
-maybe_queue_consumer(_Key, _Consumer, ServiceQueue, _FilterEnabled, _Messages) ->
+maybe_queue_consumer(_Key, _Consumer,
+                     ServiceQueue, _FilterEnabled, _Messages, _ReturnCount) ->
     ServiceQueue.
 
 queue_consumer(Key, Prio, ServiceQueue) ->
@@ -2684,13 +2752,14 @@ credit_active_consumer(
   #?STATE{cfg = #cfg{filter_enabled = FilterEnabled},
           consumers = Cons0,
           service_queue = ServiceQueue0,
-          messages = Messages} = State0) ->
+          messages = Messages,
+          return_count = ReturnCount} = State0) ->
     LinkCreditSnd = link_credit_snd(DeliveryCountRcv, LinkCreditRcv,
                                     DeliveryCountSnd, Cfg),
     %% grant the credit
     Con1 = Con0#consumer{credit = LinkCreditSnd},
-    ServiceQueue = maybe_queue_consumer(
-                     ConsumerKey, Con1, ServiceQueue0, FilterEnabled, Messages),
+    ServiceQueue = maybe_queue_consumer(ConsumerKey, Con1, ServiceQueue0,
+                                        FilterEnabled, Messages, ReturnCount),
     State1 = State0#?STATE{service_queue = ServiceQueue,
                            consumers = maps:update(ConsumerKey, Con1, Cons0)},
     {State2, ok, Effects} = checkout(Meta, State0, State1, []),