Merge pull request #3423 from rabbitmq/streams-sort-partitions

Sort stream partitions using binding parameter

Author: acogoluegnes

deps/rabbitmq_stream/BUILD.bazel

@@ -80,6 +80,10 @@ suites = [
         PACKAGE,
         name = "config_schema_SUITE",
     ),
+    rabbitmq_integration_suite(
+        PACKAGE,
+        name = "rabbit_stream_utils_SUITE",
+    ),
     rabbitmq_integration_suite(
         PACKAGE,
         name = "rabbit_stream_SUITE",

deps/rabbitmq_stream/src/rabbit_stream_manager.erl

@@ -369,8 +369,10 @@ handle_call({route, RoutingKey, VirtualHost, SuperStream}, _From,
                   [] ->
                       {ok, no_route};
                   Routes ->
-                      %% FIXME filter non-stream resources
-                      {ok, [Stream || #resource{name = Stream} <- Routes]}
+                      {ok,
+                       [Stream
+                        || #resource{name = Stream} = R <- Routes,
+                           is_resource_stream_queue(R)]}
               end
           catch
               exit:Error ->
@@ -383,9 +385,13 @@ handle_call({partitions, VirtualHost, SuperStream}, _From, State) ->
     ExchangeName = rabbit_misc:r(VirtualHost, exchange, SuperStream),
     Res = try
               rabbit_exchange:lookup_or_die(ExchangeName),
-              %% FIXME make sure queue is a stream
-              %% TODO bindings could be sorted by partition number, by using a binding argument
-              %% this would make the spreading of messages stable
+              UnorderedBindings =
+                  [Binding
+                   || Binding = #binding{destination = D}
+                          <- rabbit_binding:list_for_source(ExchangeName),
+                      is_resource_stream_queue(D)],
+              OrderedBindings =
+                  rabbit_stream_utils:sort_partitions(UnorderedBindings),
               {ok,
                lists:foldl(fun (#binding{destination =
                                              #resource{kind = queue, name = Q}},
@@ -394,7 +400,7 @@ handle_call({partitions, VirtualHost, SuperStream}, _From, State) ->
                                (_Binding, Acc) ->
                                    Acc
                            end,
-                           [], rabbit_binding:list_for_source(ExchangeName))}
+                           [], OrderedBindings)}
           catch
               exit:Error ->
                   rabbit_log:error("Error while looking up exchange ~p, ~p",
@@ -445,3 +451,13 @@ is_stream_queue(Q) ->
         _ ->
             false
     end.
+
+is_resource_stream_queue(#resource{kind = queue} = Resource) ->
+    case rabbit_amqqueue:lookup(Resource) of
+        {ok, Q} ->
+            is_stream_queue(Q);
+        _ ->
+            false
+    end;
+is_resource_stream_queue(_) ->
+    false.

deps/rabbitmq_stream/src/rabbit_stream_reader.erl

@@ -174,7 +174,8 @@ callback_mode() ->
 terminate(Reason, State,
           #statem_data{transport = Transport,
                        connection = #stream_connection{socket = Socket},
-                       connection_state = ConnectionState} = StatemData) ->
+                       connection_state = ConnectionState} =
+              StatemData) ->
     close(Transport, Socket, ConnectionState),
     rabbit_networking:unregister_non_amqp_connection(self()),
     notify_connection_closed(StatemData),
@@ -1053,7 +1054,8 @@ close_sent(enter, _OldState,
                                                StateTimeout}}) ->
     {keep_state_and_data, {state_timeout, StateTimeout, close}};
 close_sent(state_timeout, close, #statem_data{}) ->
-    rabbit_log_connection:warning("Closing connection because of timeout in state '~s' likely due to lack of client action.",
+    rabbit_log_connection:warning("Closing connection because of timeout in state "
+                                  "'~s' likely due to lack of client action.",
                                   [?FUNCTION_NAME]),
     stop;
 close_sent(info, {tcp, S, Data},
@@ -1081,7 +1083,8 @@ close_sent(info, {tcp_closed, S}, _StatemData) ->
                                 [S, self()]),
     stop;
 close_sent(info, {tcp_error, S, Reason}, #statem_data{}) ->
-    rabbit_log_connection:error("Stream protocol connection socket error: ~p [~w] [~w]",
+    rabbit_log_connection:error("Stream protocol connection socket error: ~p [~w] "
+                                "[~w]",
                                 [Reason, S, self()]),
     stop;
 close_sent(info, {resource_alarm, IsThereAlarm},

deps/rabbitmq_stream/src/rabbit_stream_utils.erl

@@ -25,10 +25,13 @@
          check_configure_permitted/3,
          check_write_permitted/3,
          check_read_permitted/3,
-         extract_stream_list/2]).
+         extract_stream_list/2,
+         sort_partitions/1]).
 
 -define(MAX_PERMISSION_CACHE_SIZE, 12).
 
+-include_lib("rabbit_common/include/rabbit.hrl").
+
 enforce_correct_stream_name(Name) ->
     % from rabbit_channel
     StrippedName =
@@ -213,3 +216,23 @@ extract_stream_list(<<>>, Streams) ->
 extract_stream_list(<<Length:16, Stream:Length/binary, Rest/binary>>,
                     Streams) ->
     extract_stream_list(Rest, [Stream | Streams]).
+
+-spec sort_partitions([#binding{}]) -> [#binding{}].
+sort_partitions(Partitions) ->
+    lists:sort(fun(#binding{args = Args1}, #binding{args = Args2}) ->
+                  Arg1 =
+                      rabbit_misc:table_lookup(Args1,
+                                               <<"x-stream-partition-order">>),
+                  Arg2 =
+                      rabbit_misc:table_lookup(Args2,
+                                               <<"x-stream-partition-order">>),
+                  case {Arg1, Arg2} of
+                      {{_, Order1}, {_, Order2}} ->
+                          rabbit_data_coercion:to_integer(Order1)
+                          =< rabbit_data_coercion:to_integer(Order2);
+                      {undefined, {_, _Order2}} -> false;
+                      {{_, _Order1}, undefined} -> true;
+                      _ -> true
+                  end
+               end,
+               Partitions).

deps/rabbitmq_stream/test/rabbit_stream_utils_SUITE.erl

@@ -0,0 +1,73 @@
+-module(rabbit_stream_utils_SUITE).
+
+-compile(nowarn_export_all).
+-compile(export_all).
+
+-include_lib("eunit/include/eunit.hrl").
+-include_lib("rabbit_common/include/rabbit.hrl").
+
+%%%===================================================================
+%%% Common Test callbacks
+%%%===================================================================
+
+all() ->
+    [{group, tests}].
+
+suite() ->
+    [{timetrap, {seconds, 30}}].
+
+groups() ->
+    [{tests, [], [sort_partitions]}].
+
+init_per_suite(Config) ->
+    Config.
+
+end_per_suite(_Config) ->
+    ok.
+
+group(_GroupName) ->
+    [].
+
+init_per_group(_GroupName, Config) ->
+    Config.
+
+end_per_group(_GroupName, _Config) ->
+    ok.
+
+init_per_testcase(_TestCase, Config) ->
+    Config.
+
+end_per_testcase(_TestCase, _Config) ->
+    ok.
+
+%%%===================================================================
+%%% Test cases
+%%%===================================================================
+
+sort_partitions(_Config) ->
+    [] = rabbit_stream_utils:sort_partitions([]),
+    ?assertEqual([<<"a">>, <<"b">>, <<"c">>],
+                 [S
+                  || #binding{destination = #resource{name = S}}
+                         <- rabbit_stream_utils:sort_partitions([binding(<<"c">>,
+                                                                         2),
+                                                                 binding(<<"b">>,
+                                                                         1),
+                                                                 binding(<<"a">>,
+                                                                         0)])]),
+    ?assertEqual([<<"a">>, <<"c">>, <<"no-order-field">>],
+                 [S
+                  || #binding{destination = #resource{name = S}}
+                         <- rabbit_stream_utils:sort_partitions([binding(<<"c">>,
+                                                                         10),
+                                                                 binding(<<"no-order-field">>),
+                                                                 binding(<<"a">>,
+                                                                         0)])]),
+    ok.
+
+binding(Destination, Order) ->
+    #binding{destination = #resource{name = Destination},
+             args = [{<<"x-stream-partition-order">>, signedint, Order}]}.
+
+binding(Destination) ->
+    #binding{destination = #resource{name = Destination}, args = []}.