Merge pull request #367 from rabbitmq/rabbitmq-server-366

fixes batch_publish_delivered for priority queues

2 files changed, 51 insertions, 9 deletions

diff --git a/src/rabbit_mirror_queue_sync.erl b/src/rabbit_mirror_queue_sync.erl
index 534ef1afad..62fc718f79 100644
--- a/src/rabbit_mirror_queue_sync.erl
+++ b/src/rabbit_mirror_queue_sync.erl
@@ -352,13 +352,46 @@ batch_publish(Batch, MA, BQ, BQS) ->
     BQS1 = BQ:batch_publish(Batch, none, noflow, BQS),
     {MA, BQS1}.
 
+%% TODO
+%%
+%% The case clause in this function assumes that we are either dealing
+%% with a backing_queue that returns acktags as integers, or a
+%% priority queue.
+%% A possible fix to this would be to add a function
+%% to the BQ API where we pass a list of messages and acktags and the
+%% BQ implementation knows how to zip them together.
 batch_publish_delivered(Batch, MA, BQ, BQS) ->
     {AckTags, BQS1} = BQ:batch_publish_delivered(Batch, none, noflow, BQS),
     MA1 =
-        lists:foldl(fun ({{Msg, _}, AckTag}, MAs) ->
-                            [{Msg#basic_message.id, AckTag} | MAs]
-                    end, MA, lists:zip(Batch, AckTags)),
+        case hd(AckTags) of
+            HeadTag when is_integer(HeadTag) ->
+                lists:foldl(fun ({{Msg, _}, AckTag}, MAs) ->
+                                    [{msg_id(Msg), AckTag} | MAs]
+                            end, MA, lists:zip(Batch, AckTags));
+            _AckTags  ->
+                %% priority queue processing of acktags
+                BatchByPriority = batch_by_priority(Batch),
+                lists:foldl(fun (Acks, MAs) ->
+                                    {P, _AckTag} = hd(Acks),
+                                    Pubs = orddict:fetch(P, BatchByPriority),
+                                    MAs0 = zip_msgs_and_tags(Pubs, Acks),
+                                    MAs ++ MAs0
+                            end, MA, AckTags)
+        end,
     {MA1, BQS1}.
 
+batch_by_priority(Batch) ->
+    rabbit_priority_queue:partition_publish_delivered_batch(Batch).
+
+zip_msgs_and_tags(Pubs, AckTags) ->
+    lists:zipwith(
+      fun (Pub, AckTag) ->
+              {Msg, _Props} = Pub,
+              {msg_id(Msg), AckTag}
+      end, Pubs, AckTags).
+
 props(Props) ->
     Props#message_properties{needs_confirming = false}.
+
+msg_id(#basic_message{ id = Id }) ->
+    Id.
diff --git a/src/rabbit_priority_queue.erl b/src/rabbit_priority_queue.erl
index 28cee163a1..4d638b334a 100644
--- a/src/rabbit_priority_queue.erl
+++ b/src/rabbit_priority_queue.erl
@@ -42,6 +42,9 @@
          handle_pre_hibernate/1, resume/1, msg_rates/1,
          info/2, invoke/3, is_duplicate/2]).
 
+%% for rabbit_mirror_queue_sync.
+-export([partition_publish_delivered_batch/1]).
+
 -record(state, {bq, bqss}).
 -record(passthrough, {bq, bqs}).
 
@@ -205,8 +208,7 @@ publish(Msg, MsgProps, IsDelivered, ChPid, Flow,
     ?passthrough1(publish(Msg, MsgProps, IsDelivered, ChPid, Flow, BQS)).
 
 batch_publish(Publishes, ChPid, Flow, State = #state{bq = BQ}) ->
-    PubDict = publishes_by_priority(
-                Publishes, fun ({Msg, _, _}) -> Msg end),
+    PubDict = partition_publish_batch(Publishes),
     lists:foldl(
       fun ({Priority, Pubs}, St) ->
               pick1(fun (_P, BQSN) ->
@@ -228,8 +230,7 @@ publish_delivered(Msg, MsgProps, ChPid, Flow,
     ?passthrough2(publish_delivered(Msg, MsgProps, ChPid, Flow, BQS)).
 
 batch_publish_delivered(Publishes, ChPid, Flow, State = #state{bq = BQ}) ->
-    PubDict = publishes_by_priority(
-                Publishes, fun ({Msg, _}) -> Msg end),
+    PubDict = partition_publish_delivered_batch(Publishes),
     {PrioritiesAndAcks, State1} =
         lists:foldl(
           fun ({Priority, Pubs}, {PriosAndAcks, St}) ->
@@ -238,7 +239,7 @@ batch_publish_delivered(Publishes, ChPid, Flow, State = #state{bq = BQ}) ->
                                     {AckTags, BQSN1} =
                                         BQ:batch_publish_delivered(
                                           Pubs, ChPid, Flow, BQSN),
-                                    {{P, AckTags}, BQSN1}
+                                    {priority_on_acktags(P, AckTags), BQSN1}
                             end, Priority, St),
                   {[PriosAndAcks1 | PriosAndAcks], St1}
           end, {[], State}, orddict:to_list(PubDict)),
@@ -566,7 +567,15 @@ a(State = #state{bqss = BQSs}) ->
     end.
 
 %%----------------------------------------------------------------------------
-publishes_by_priority(Publishes, ExtractMsg) ->
+partition_publish_batch(Publishes) ->
+    partition_publishes(
+      Publishes, fun ({Msg, _, _}) -> Msg end).
+
+partition_publish_delivered_batch(Publishes) ->
+    partition_publishes(
+      Publishes, fun ({Msg, _}) -> Msg end).
+
+partition_publishes(Publishes, ExtractMsg) ->
     lists:foldl(fun (Pub, Dict) ->
                         Msg = ExtractMsg(Pub),
                         rabbit_misc:orddict_cons(priority2(Msg), Pub, Dict)