WIP

Author: mkuratczyk

deps/rabbit/src/rabbit.erl

@@ -33,7 +33,8 @@
          pg_local_scope/1,
          pg_scope_amqp091_channel/0,
          pg_scope_amqp091_connection/0,
-         pg_scope_non_amqp_connection/0]).
+         pg_scope_non_amqp_connection/0,
+         pg_scope_direct_connection/0]).
 %% For CLI, testing and mgmt-agent.
 -export([set_log_level/1, log_locations/0, config_files/0]).
 -export([is_booted/1, is_booted/0, is_booting/1, is_booting/0]).
@@ -46,6 +47,7 @@
          pg_local_amqp091_channel/0,
          pg_local_amqp091_connection/0,
          pg_local_non_amqp_connection/0,
+         pg_local_direct_connection/0,
          prevent_startup_if_node_was_reset/0]).
 
 -rabbit_boot_step({pre_boot, [{description, "rabbit boot start"}]}).
@@ -316,6 +318,12 @@
                     {requires,    kernel_ready},
                     {enables,     core_initialized}]}).
 
+-rabbit_boot_step({pg_local_direct_connection,
+                   [{description, "local-only pg scope for direct connections"},
+                    {mfa,         {rabbit, pg_local_direct_connection, []}},
+                    {requires,    kernel_ready},
+                    {enables,     core_initialized}]}).
+
 %%---------------------------------------------------------------------------
 
 -include_lib("rabbit_common/include/rabbit.hrl").
@@ -1162,6 +1170,11 @@ pg_local_non_amqp_connection() ->
     persistent_term:put(pg_scope_non_amqp_connection, PgScope),
     rabbit_sup:start_child(pg_non_amqp_connection, pg, [PgScope]).
 
+pg_local_direct_connection() ->
+    PgScope = pg_local_scope(direct_connection),
+    persistent_term:put(pg_scope_direct_connection, PgScope),
+    rabbit_sup:start_child(pg_direct_connection, pg, [PgScope]).
+
 pg_local_scope(Prefix) ->
     list_to_atom(io_lib:format("~s_~s", [Prefix, node()])).
 
@@ -1174,6 +1187,9 @@ pg_scope_amqp091_connection() ->
 pg_scope_non_amqp_connection() ->
     persistent_term:get(pg_scope_non_amqp_connection).
 
+pg_scope_direct_connection() ->
+    persistent_term:get(pg_scope_direct_connection).
+
 -spec update_cluster_tags() -> 'ok'.
 
 update_cluster_tags() ->

deps/rabbit/src/rabbit_direct.erl

@@ -20,14 +20,14 @@
 -include_lib("rabbit_common/include/rabbit_misc.hrl").
 -include_lib("kernel/include/logger.hrl").
 
--define(TABLE, ?MODULE).
-
 %%----------------------------------------------------------------------------
 
+pg_scope() ->
+    rabbit:pg_scope_direct_connection().
+
 -spec boot() -> 'ok'.
 
 boot() ->
-    ?TABLE = ets:new(?TABLE, [set, public, named_table]),
     rabbit_sup:start_supervisor_child(
       rabbit_direct_client_sup, rabbit_client_sup,
       [{local, rabbit_direct_client_sup},
@@ -42,7 +42,7 @@ force_event_refresh(Ref) ->
 -spec list_local() -> [pid()].
 
 list_local() ->
-    [Pid || {Pid} <- ets:tab2list(?TABLE)].
+    pg:which_groups(pg_scope()).
 
 -spec list() -> [pid()].
 
@@ -190,7 +190,7 @@ connect1(User = #user{username = Username}, VHost, Pid, Infos) ->
             AuthzContext = proplists:get_value(variable_map, Infos, #{}),
             try rabbit_access_control:check_vhost_access(User, VHost,
                                                {ip, PeerHost}, AuthzContext) of
-                ok -> ets:insert(?TABLE, {Pid}),
+                ok -> ok = pg:join(pg_scope(), Pid, Pid),
                       rabbit_core_metrics:connection_created(Pid, Infos),
                       rabbit_event:notify(connection_created, Infos),
                       _ = rabbit_alarm:register(
@@ -250,7 +250,7 @@ start_channel(Number, ClientChannelPid, ConnPid, ConnName,
 -spec disconnect(pid(), rabbit_event:event_props()) -> 'ok'.
 
 disconnect(Pid, Infos) ->
-    ets:delete(?TABLE, Pid),
+    pg:leave(pg_scope(), Pid, Pid),
     rabbit_core_metrics:connection_closed(Pid),
     rabbit_event:notify(connection_closed, Infos).