Extract publishing IDs from batch publishing

In stream plugin, to e.g. send publish errors in case the stream
does not exist. Batches were not taken into account.

Author: acogoluegnes

deps/rabbitmq_stream/src/rabbit_stream_core.erl

@@ -12,7 +12,8 @@
 %% holds static or rarely changing fields
 -record(cfg, {}).
 -record(?MODULE,
-        {cfg :: #cfg{}, frames = [] :: [iodata()],
+        {cfg :: #cfg{},
+         frames = [] :: [iodata()],
          %% partial data
          data ::
              undefined |
@@ -145,52 +146,61 @@ all_commands(#?MODULE{commands = Commands0} = State) ->
 %% returns frames
 -spec incoming_data(binary(), state()) -> state().
 %% TODO: check max frame size
-incoming_data(<<>>, #?MODULE{frames = Frames, commands = Commands} = State) ->
+incoming_data(<<>>,
+              #?MODULE{frames = Frames, commands = Commands} = State) ->
     State#?MODULE{frames = [], commands = parse_frames(Frames, Commands)};
 incoming_data(<<Size:32, Frame:Size/binary, Rem/binary>>,
               #?MODULE{frames = Frames, data = undefined} = State) ->
     incoming_data(Rem,
                   State#?MODULE{frames = [Frame | Frames], data = undefined});
 incoming_data(<<Size:32, Rem/binary>>,
-              #?MODULE{frames = Frames, data = undefined,
-                       commands = Commands} = State) ->
+              #?MODULE{frames = Frames,
+                       data = undefined,
+                       commands = Commands} =
+                  State) ->
     %% not enough data to complete frame, stash and await more data
-    State#?MODULE{frames = [], data = {Size - byte_size(Rem), Rem},
+    State#?MODULE{frames = [],
+                  data = {Size - byte_size(Rem), Rem},
                   commands = parse_frames(Frames, Commands)};
 incoming_data(Data,
-              #?MODULE{frames = Frames, data = undefined,
-                       commands = Commands} = State)
-  when byte_size(Data) < 4 ->
+              #?MODULE{frames = Frames,
+                       data = undefined,
+                       commands = Commands} =
+                  State)
+    when byte_size(Data) < 4 ->
     %% not enough data to even know the size required
     %% just stash binary and hit last clause next
-    State#?MODULE{frames = [], data = Data,
+    State#?MODULE{frames = [],
+                  data = Data,
                   commands = parse_frames(Frames, Commands)};
 incoming_data(Data,
-              #?MODULE{frames = Frames, data = {Size, Partial},
-                       commands = Commands} = State) ->
+              #?MODULE{frames = Frames,
+                       data = {Size, Partial},
+                       commands = Commands} =
+                  State) ->
     case Data of
         <<Part:Size/binary, Rem/binary>> ->
             incoming_data(Rem,
                           State#?MODULE{frames =
-                                        [append_data(Partial, Part)
-                                         | Frames],
+                                            [append_data(Partial, Part)
+                                             | Frames],
                                         data = undefined});
         Rem ->
             State#?MODULE{frames = [],
                           data =
-                          {Size - byte_size(Rem),
-                           append_data(Partial, Rem)},
+                              {Size - byte_size(Rem),
+                               append_data(Partial, Rem)},
                           commands = parse_frames(Frames, Commands)}
     end;
 incoming_data(Data, #?MODULE{data = Partial} = State)
-  when is_binary(Partial) ->
+    when is_binary(Partial) ->
     incoming_data(<<Partial/binary, Data/binary>>,
                   State#?MODULE{data = undefined}).
 
 parse_frames(Frames, Queue) ->
-    lists:foldr(fun(Frame, Acc) ->
-                        queue:in(parse_command(Frame), Acc)
-                end, Queue, Frames).
+    lists:foldr(fun(Frame, Acc) -> queue:in(parse_command(Frame), Acc)
+                end,
+                Queue, Frames).
 
 -spec frame(command()) -> iodata().
 frame({publish_confirm, PublisherId, PublishingIds}) ->

deps/rabbitmq_stream/src/rabbit_stream_reader.erl

@@ -836,18 +836,27 @@ handle_inbound_data(Transport,
     CoreState1 = rabbit_stream_core:incoming_data(Data, CoreState0),
     {Commands, CoreState} = rabbit_stream_core:all_commands(CoreState1),
     lists:foldl(fun(Command, {C, S}) ->
-                        HandleFrameFun(Transport, C, S, Command)
+                   HandleFrameFun(Transport, C, S, Command)
                 end,
-                {Connection,
-                 State#stream_connection_state{data = CoreState}},
+                {Connection, State#stream_connection_state{data = CoreState}},
                 Commands).
 
 publishing_ids_from_messages(<<>>) ->
     [];
 publishing_ids_from_messages(<<PublishingId:64,
-                               MessageSize:32,
+                               0:1,
+                               MessageSize:31,
                                _Message:MessageSize/binary,
                                Rest/binary>>) ->
+    [PublishingId | publishing_ids_from_messages(Rest)];
+publishing_ids_from_messages(<<PublishingId:64,
+                               1:1,
+                               _CompressionType:3,
+                               _Unused:4,
+                               _MessageCount:16,
+                               BatchSize:32,
+                               _Batch:BatchSize/binary,
+                               Rest/binary>>) ->
     [PublishingId | publishing_ids_from_messages(Rest)].
 
 handle_frame_pre_auth(Transport,
@@ -1077,7 +1086,7 @@ handle_frame_pre_auth(Transport,
                 #{<<"advertised_host">> => AdvertisedHost,
                   <<"advertised_port">> => AdvertisedPort},
 
-                rabbit_log:info("sending open response ok ~s", [VirtualHost]),
+            rabbit_log:info("sending open response ok ~s", [VirtualHost]),
             Frame =
                 rabbit_stream_core:frame({response, CorrelationId,
                                           {open, ?RESPONSE_CODE_OK,

deps/rabbitmq_stream/test/commands_SUITE.erl

@@ -18,7 +18,6 @@
 -include("rabbit_stream.hrl").
 
 -define(WAIT, 5000).
-
 -define(COMMAND_LIST_CONNECTIONS,
         'Elixir.RabbitMQ.CLI.Ctl.Commands.ListStreamConnectionsCommand').
 -define(COMMAND_LIST_CONSUMERS,

deps/rabbitmq_stream/test/rabbit_stream_SUITE.erl

@@ -265,7 +265,8 @@ test_authenticate(Transport, S, C0) ->
     OpenFrame =
         rabbit_stream_core:frame({request, 3, {open, VirtualHost}}),
     ok = Transport:send(S, OpenFrame),
-    {{response, 3, {open, ?RESPONSE_CODE_OK, _ConnectionProperties}}, C4} =
+    {{response, 3, {open, ?RESPONSE_CODE_OK, _ConnectionProperties}},
+     C4} =
         receive_commands(Transport, S, C3),
     C4.
 
@@ -322,7 +323,6 @@ test_subscribe(Transport, S, SubscriptionId, Stream, C0) ->
     C.
 
 test_deliver(Transport, S, SubscriptionId, COffset, Body, C0) ->
-
     ct:pal("test_deliver ", []),
     {{deliver, SubscriptionId, Chunk}, C} =
         receive_commands(Transport, S, C0),
@@ -372,9 +372,9 @@ receive_commands(Transport, S, C0) ->
                     C1 = rabbit_stream_core:incoming_data(Data, C0),
                     case rabbit_stream_core:next_command(C1) of
                         empty ->
-                            {ok, Data2} =  Transport:recv(S, 0, 5000),
+                            {ok, Data2} = Transport:recv(S, 0, 5000),
                             rabbit_stream_core:next_command(
-                              rabbit_stream_core:incoming_data(Data2, C1));
+                                rabbit_stream_core:incoming_data(Data2, C1));
                         Res ->
                             Res
                     end;

deps/rabbitmq_stream/test/rabbit_stream_core_SUITE.erl

@@ -163,6 +163,7 @@ roundtrip_metadata_no_leader(_Config) ->
 test_roundtrip(Cmd) ->
     Init = rabbit_stream_core:init(undefined),
     Frame = iolist_to_binary(rabbit_stream_core:frame(Cmd)),
-    {[Cmd], _} = rabbit_stream_core:all_commands(
-                   rabbit_stream_core:incoming_data(Frame, Init)),
+    {[Cmd], _} =
+        rabbit_stream_core:all_commands(
+            rabbit_stream_core:incoming_data(Frame, Init)),
     ok.