Add list_stream_publishers CLI command

Author: acogoluegnes

deps/rabbitmq_stream/include/rabbit_stream.hrl

@@ -80,4 +80,14 @@
   offset
   ]).
 
+-define(PUBLISHER_INFO_ITEMS, [
+  stream,
+  connection_pid,
+  publisher_id,
+  reference,
+  messages_published,
+  messages_confirmed,
+  messages_errored
+  ]).
+
 -define(STREAM_GUIDE_URL, <<"https://rabbitmq.com/stream.html">>).

deps/rabbitmq_stream/src/Elixir.RabbitMQ.CLI.Ctl.Commands.ListStreamConsumersCommand.erl

@@ -77,7 +77,7 @@ usage_additional() ->
     Prefix = <<" must be one of ">>,
     InfoItems =
         'Elixir.Enum':join(
-            lists:usort(?INFO_ITEMS), <<", ">>),
+            lists:usort(?CONSUMER_INFO_ITEMS), <<", ">>),
     [{<<"<column>">>, <<Prefix/binary, InfoItems/binary>>}].
 
 usage_doc_guides() ->

deps/rabbitmq_stream/src/Elixir.RabbitMQ.CLI.Ctl.Commands.ListStreamPublishersCommand.erl

@@ -0,0 +1,113 @@
+%% The contents of this file are subject to the Mozilla Public License
+%% Version 2.0 (the "License"); you may not use this file except in
+%% compliance with the License. You may obtain a copy of the License
+%% at https://www.mozilla.org/MPL/
+%%
+%% Software distributed under the License is distributed on an "AS IS"
+%% basis, WITHOUT WARRANTY OF ANY KIND, either express or implied. See
+%% the License for the specific language governing rights and
+%% limitations under the License.
+%%
+%% The Original Code is RabbitMQ.
+%%
+%% The Initial Developer of the Original Code is GoPivotal, Inc.
+%% Copyright (c) 2021 VMware, Inc. or its affiliates.  All rights reserved.
+
+-module('Elixir.RabbitMQ.CLI.Ctl.Commands.ListStreamPublishersCommand').
+
+-include("rabbit_stream.hrl").
+
+-behaviour('Elixir.RabbitMQ.CLI.CommandBehaviour').
+
+-export([formatter/0,
+         scopes/0,
+         switches/0,
+         aliases/0,
+         usage/0,
+         usage_additional/0,
+         usage_doc_guides/0,
+         banner/2,
+         validate/2,
+         merge_defaults/2,
+         run/2,
+         output/2,
+         description/0,
+         help_section/0]).
+
+formatter() ->
+    'Elixir.RabbitMQ.CLI.Formatters.Table'.
+
+scopes() ->
+    [ctl, diagnostics, streams].
+
+switches() ->
+    [{verbose, boolean}].
+
+aliases() ->
+    [{'V', verbose}].
+
+description() ->
+    <<"Lists all stream publishers for a vhost">>.
+
+help_section() ->
+    {plugin, stream}.
+
+validate(Args, _) ->
+    case 'Elixir.RabbitMQ.CLI.Ctl.InfoKeys':validate_info_keys(Args,
+                                                               ?PUBLISHER_INFO_ITEMS)
+    of
+        {ok, _} ->
+            ok;
+        Error ->
+            Error
+    end.
+
+merge_defaults([], Opts) ->
+    merge_defaults([rabbit_data_coercion:to_binary(Item)
+                    || Item <- ?PUBLISHER_INFO_ITEMS],
+                   Opts);
+merge_defaults(Args, Opts) ->
+    {Args, maps:merge(#{verbose => false, vhost => <<"/">>}, Opts)}.
+
+usage() ->
+    <<"list_stream_publishers [--vhost <vhost>] [<column> "
+      "...]">>.
+
+usage_additional() ->
+    Prefix = <<" must be one of ">>,
+    InfoItems =
+        'Elixir.Enum':join(
+            lists:usort(?CONSUMER_INFO_ITEMS), <<", ">>),
+    [{<<"<column>">>, <<Prefix/binary, InfoItems/binary>>}].
+
+usage_doc_guides() ->
+    [?STREAM_GUIDE_URL].
+
+run(Args,
+    #{node := NodeName,
+      vhost := VHost,
+      timeout := Timeout,
+      verbose := Verbose}) ->
+    InfoKeys =
+        case Verbose of
+            true ->
+                ?PUBLISHER_INFO_ITEMS;
+            false ->
+                'Elixir.RabbitMQ.CLI.Ctl.InfoKeys':prepare_info_keys(Args)
+        end,
+    Nodes = 'Elixir.RabbitMQ.CLI.Core.Helpers':nodes_in_cluster(NodeName),
+
+    'Elixir.RabbitMQ.CLI.Ctl.RpcStream':receive_list_items(NodeName,
+                                                           rabbit_stream,
+                                                           emit_publisher_info_all,
+                                                           [Nodes, VHost,
+                                                            InfoKeys],
+                                                           Timeout,
+                                                           InfoKeys,
+                                                           length(Nodes)).
+
+banner(_, _) ->
+    <<"Listing stream publishers ...">>.
+
+output(Result, _Opts) ->
+    'Elixir.RabbitMQ.CLI.DefaultOutput':output(Result).

deps/rabbitmq_stream/src/rabbit_stream.erl

@@ -27,6 +27,8 @@
          emit_connection_info_all/4,
          emit_consumer_info_all/5,
          emit_consumer_info_local/4,
+         emit_publisher_info_all/5,
+         emit_publisher_info_local/4,
          list/1]).
 
 -include_lib("rabbit_common/include/rabbit.hrl").
@@ -131,6 +133,25 @@ emit_consumer_info_local(VHost, Items, Ref, AggregatorPid) ->
                                                        end,
                                                        list(VHost)).
 
+emit_publisher_info_all(Nodes, VHost, Items, Ref, AggregatorPid) ->
+    Pids =
+        [spawn_link(Node,
+                    rabbit_stream,
+                    emit_publisher_info_local,
+                    [VHost, Items, Ref, AggregatorPid])
+         || Node <- Nodes],
+    rabbit_control_misc:await_emitters_termination(Pids),
+    ok.
+
+emit_publisher_info_local(VHost, Items, Ref, AggregatorPid) ->
+    rabbit_control_misc:emitting_map_with_exit_handler(AggregatorPid,
+                                                       Ref,
+                                                       fun(Pid) ->
+                                                          rabbit_stream_reader:publishers_info(Pid,
+                                                                                               Items)
+                                                       end,
+                                                       list(VHost)).
+
 list(VHost) ->
     [Client
      || {_, ListSupPid, _, _}

deps/rabbitmq_stream/src/rabbit_stream_reader.erl

@@ -141,6 +141,7 @@
          init/1,
          info/2,
          consumers_info/2,
+         publishers_info/2,
          in_vhost/2]).
 
 start_link(KeepaliveSup, Transport, Ref, Opts) ->
@@ -663,6 +664,9 @@ listen_loop_post_auth(Transport,
         {'$gen_call', From, {consumers_info, Items}} ->
             gen_server:reply(From, consumers_infos(Items, State)),
             listen_loop_post_auth(Transport, Connection, State, Configuration);
+        {'$gen_call', From, {publishers_info, Items}} ->
+            gen_server:reply(From, publishers_infos(Items, Connection)),
+            listen_loop_post_auth(Transport, Connection, State, Configuration);
         emit_stats ->
             Connection1 = emit_stats(Connection, State),
             listen_loop_post_auth(Transport, Connection1, State, Configuration);
@@ -2438,6 +2442,39 @@ consumer_i(connection_pid, _) ->
 consumer_i(stream, #consumer{stream = S}) ->
     S.
 
+publishers_info(Pid, InfoItems) ->
+    case InfoItems -- ?PUBLISHER_INFO_ITEMS of
+        [] ->
+            gen_server2:call(Pid, {publishers_info, InfoItems});
+        UnknownItems ->
+            throw({bad_argument, UnknownItems})
+    end.
+
+publishers_infos(Items,
+                 #stream_connection{publishers = Publishers}) ->
+    [[{Item, publisher_i(Item, Publisher)} || Item <- Items]
+     || Publisher <- maps:values(Publishers)].
+
+publisher_i(stream, #publisher{stream = S}) ->
+    S;
+publisher_i(connection_pid, _) ->
+    self();
+publisher_i(publisher_id, #publisher{publisher_id = Id}) ->
+    Id;
+publisher_i(reference, #publisher{reference = undefined}) ->
+    <<"">>;
+publisher_i(reference, #publisher{reference = Ref}) ->
+    Ref;
+publisher_i(messages_published,
+            #publisher{message_counters = Counters}) ->
+    messages_published(Counters);
+publisher_i(messages_confirmed,
+            #publisher{message_counters = Counters}) ->
+    messages_confirmed(Counters);
+publisher_i(messages_errored,
+            #publisher{message_counters = Counters}) ->
+    messages_errored(Counters).
+
 info(Pid, InfoItems) ->
     case InfoItems -- ?INFO_ITEMS of
         [] ->

deps/rabbitmq_stream/test/commands_SUITE.erl

@@ -19,15 +19,20 @@
         'Elixir.RabbitMQ.CLI.Ctl.Commands.ListStreamConnectionsCommand').
 -define(COMMAND_LIST_CONSUMERS,
         'Elixir.RabbitMQ.CLI.Ctl.Commands.ListStreamConsumersCommand').
+-define(COMMAND_LIST_PUBLISHERS,
+        'Elixir.RabbitMQ.CLI.Ctl.Commands.ListStreamPublishersCommand').
 
 all() ->
-    [{group, list_connections}, {group, list_consumers}].
+    [{group, list_connections}, {group, list_consumers},
+     {group, list_publishers}].
 
 groups() ->
     [{list_connections, [],
       [list_connections_merge_defaults, list_connections_run]},
      {list_consumers, [],
-      [list_consumers_merge_defaults, list_consumers_run]}].
+      [list_consumers_merge_defaults, list_consumers_run]},
+     {list_publishers, [],
+      [list_publishers_merge_defaults, list_publishers_run]}].
 
 init_per_suite(Config) ->
     Config1 =
@@ -195,12 +200,84 @@ list_consumers_run(Config) ->
     wait_until_consumer_count(Config, 0),
     ok.
 
+list_publishers_merge_defaults(_Config) ->
+    DefaultItems =
+        [rabbit_data_coercion:to_binary(Item)
+         || Item <- ?PUBLISHER_INFO_ITEMS],
+    {DefaultItems, #{verbose := false}} =
+        ?COMMAND_LIST_PUBLISHERS:merge_defaults([], #{}),
+
+    {[<<"other_key">>], #{verbose := true}} =
+        ?COMMAND_LIST_PUBLISHERS:merge_defaults([<<"other_key">>],
+                                                #{verbose => true}),
+
+    {[<<"other_key">>], #{verbose := false}} =
+        ?COMMAND_LIST_PUBLISHERS:merge_defaults([<<"other_key">>],
+                                                #{verbose => false}).
+
+list_publishers_run(Config) ->
+    Node = rabbit_ct_broker_helpers:get_node_config(Config, 0, nodename),
+    Opts =
+        #{node => Node,
+          timeout => 10000,
+          verbose => false,
+          vhost => <<"/">>},
+
+    %% No connections, no publishers
+    [] = to_list(?COMMAND_LIST_PUBLISHERS:run([], Opts)),
+
+    StreamPort = rabbit_stream_SUITE:get_stream_port(Config),
+    S1 = start_stream_connection(StreamPort),
+    wait_until_connection_count(Config, 1),
+
+    Stream = <<"list_publishers_run">>,
+    create_stream(S1, Stream),
+    PubId = 42,
+    declare_publisher(S1, PubId, Stream),
+
+    wait_until_publisher_count(Config, 1),
+
+    S2 = start_stream_connection(StreamPort),
+    wait_until_connection_count(Config, 2),
+    declare_publisher(S2, PubId, Stream),
+
+    wait_until_publisher_count(Config, 2),
+
+    %% Verbose returns all keys
+    InfoItems = ?PUBLISHER_INFO_ITEMS,
+    Infos = lists:map(fun(El) -> atom_to_binary(El, utf8) end, InfoItems),
+    AllKeys = to_list(?COMMAND_LIST_PUBLISHERS:run(Infos, Opts)),
+    Verbose =
+        to_list(?COMMAND_LIST_PUBLISHERS:run([], Opts#{verbose => true})),
+    ?assertEqual(AllKeys, Verbose),
+    %% There are two consumers
+    [[First], [_Second]] = AllKeys,
+
+    %% Keys are info items
+    ?assertEqual(length(InfoItems), length(First)),
+
+    {Keys, _} = lists:unzip(First),
+
+    ?assertEqual([], Keys -- InfoItems),
+    ?assertEqual([], InfoItems -- Keys),
+
+    delete_stream(S1, Stream),
+    metadata_update_stream_deleted(S1, Stream),
+    metadata_update_stream_deleted(S2, Stream),
+    close(S1),
+    close(S2),
+    wait_until_publisher_count(Config, 0),
+    ok.
+
 create_stream(S, Stream) ->
     rabbit_stream_SUITE:test_create_stream(S, Stream).
 
 subscribe(S, SubId, Stream) ->
     rabbit_stream_SUITE:test_subscribe(S, SubId, Stream).
 
+declare_publisher(S, PubId, Stream) ->
+    rabbit_stream_SUITE:test_declare_publisher(S, PubId, Stream).
+
 delete_stream(S, Stream) ->
     rabbit_stream_SUITE:test_delete_stream(S, Stream).
 
@@ -243,6 +320,15 @@ wait_until_consumer_count(Config, Expected) ->
         test_utils:wait_until(fun() -> consumer_count(Config) == Expected
                               end).
 
+publisher_count(Config) ->
+    command_result_count(?COMMAND_LIST_PUBLISHERS:run([<<"stream">>],
+                                                      options(Config))).
+
+wait_until_publisher_count(Config, Expected) ->
+    ok =
+        test_utils:wait_until(fun() -> publisher_count(Config) == Expected
+                              end).
+
 start_stream_connection(Port) ->
     {ok, S} =
         gen_tcp:connect("localhost", Port, [{active, false}, {mode, binary}]),