Use pg_local to track AMQP 1.0 connections

Fixes #9371

Moves `pg_local` to `rabbit_common`

Author: lukebakken

deps/rabbit/BUILD.bazel

@@ -1004,11 +1004,6 @@ rabbitmq_suite(
     ],
 )
 
-rabbitmq_suite(
-    name = "unit_pg_local_SUITE",
-    size = "small",
-)
-
 rabbitmq_suite(
     name = "unit_plugin_directories_SUITE",
     size = "small",

deps/rabbit/app.bzl

@@ -41,7 +41,6 @@ def all_beam_files(name = "all_beam_files"):
             "src/mc_util.erl",
             "src/mirrored_supervisor.erl",
             "src/mirrored_supervisor_sups.erl",
-            "src/pg_local.erl",
             "src/pid_recomposition.erl",
             "src/rabbit.erl",
             "src/rabbit_access_control.erl",
@@ -292,7 +291,6 @@ def all_test_beam_files(name = "all_test_beam_files"):
             "src/mc_util.erl",
             "src/mirrored_supervisor.erl",
             "src/mirrored_supervisor_sups.erl",
-            "src/pg_local.erl",
             "src/pid_recomposition.erl",
             "src/rabbit.erl",
             "src/rabbit_access_control.erl",
@@ -555,7 +553,6 @@ def all_srcs(name = "all_srcs"):
             "src/mc_util.erl",
             "src/mirrored_supervisor.erl",
             "src/mirrored_supervisor_sups.erl",
-            "src/pg_local.erl",
             "src/pid_recomposition.erl",
             "src/rabbit.erl",
             "src/rabbit_access_control.erl",
@@ -1768,14 +1765,6 @@ def test_suite_beam_files(name = "test_suite_beam_files"):
         erlc_opts = "//:test_erlc_opts",
         deps = ["//deps/rabbit_common:erlang_app"],
     )
-    erlang_bytecode(
-        name = "unit_pg_local_SUITE_beam_files",
-        testonly = True,
-        srcs = ["test/unit_pg_local_SUITE.erl"],
-        outs = ["test/unit_pg_local_SUITE.beam"],
-        app_name = "rabbit",
-        erlc_opts = "//:test_erlc_opts",
-    )
     erlang_bytecode(
         name = "unit_plugin_directories_SUITE_beam_files",
         testonly = True,

deps/rabbit_common/BUILD.bazel

@@ -170,6 +170,11 @@ rabbitmq_suite(
     ],
 )
 
+rabbitmq_suite(
+    name = "unit_pg_local_SUITE",
+    size = "small",
+)
+
 rabbitmq_suite(
     name = "unit_priority_queue_SUITE",
     size = "small",

deps/rabbit_common/app.bzl

@@ -32,6 +32,7 @@ def all_beam_files(name = "all_beam_files"):
             "src/file_handle_cache_stats.erl",
             "src/mirrored_supervisor_locks.erl",
             "src/mnesia_sync.erl",
+            "src/pg_local.erl",
             "src/pmon.erl",
             "src/priority_queue.erl",
             "src/rabbit_amqp_connection.erl",
@@ -127,6 +128,7 @@ def all_test_beam_files(name = "all_test_beam_files"):
             "src/file_handle_cache_stats.erl",
             "src/mirrored_supervisor_locks.erl",
             "src/mnesia_sync.erl",
+            "src/pg_local.erl",
             "src/pmon.erl",
             "src/priority_queue.erl",
             "src/rabbit_amqp_connection.erl",
@@ -215,6 +217,7 @@ def all_srcs(name = "all_srcs"):
             "src/gen_server2.erl",
             "src/mirrored_supervisor_locks.erl",
             "src/mnesia_sync.erl",
+            "src/pg_local.erl",
             "src/pmon.erl",
             "src/priority_queue.erl",
             "src/rabbit_amqp_connection.erl",
@@ -348,6 +351,14 @@ def test_suite_beam_files(name = "test_suite_beam_files"):
         erlc_opts = "//:test_erlc_opts",
         deps = ["@proper//:erlang_app"],
     )
+    erlang_bytecode(
+        name = "unit_pg_local_SUITE_beam_files",
+        testonly = True,
+        srcs = ["test/unit_pg_local_SUITE.erl"],
+        outs = ["test/unit_pg_local_SUITE.beam"],
+        app_name = "rabbit",
+        erlc_opts = "//:test_erlc_opts",
+    )
     erlang_bytecode(
         name = "unit_priority_queue_SUITE_beam_files",
         testonly = True,

deps/rabbit/src/pg_local.erl renamed to deps/rabbit_common/src/pg_local.erl

@@ -6,10 +6,11 @@
 %%
 -module(rabbit_amqp1_0).
 
--export([connection_info_local/1,
-         emit_connection_info_local/3,
+-export([emit_connection_info_local/3,
          emit_connection_info_all/4,
-         list/0]).
+         list/0,
+         register_connection/1,
+         unregister_connection/1]).
 
 emit_connection_info_all(Nodes, Items, Ref, AggregatorPid) ->
     Pids = [spawn_link(Node, rabbit_amqp1_0, emit_connection_info_local,
@@ -26,17 +27,22 @@ emit_connection_info_local(Items, Ref, AggregatorPid) ->
       end,
       list()).
 
-connection_info_local(Items) ->
-    Connections = list(),
-    [rabbit_amqp1_0_reader:info(Pid, Items) || Pid <- Connections].
-
 list() ->
-    [ReaderPid
-     || {_, TcpPid, _, [tcp_listener_sup]} <- supervisor:which_children(rabbit_sup),
-        {_, RanchEPid, _, [ranch_embedded_sup]} <- supervisor:which_children(TcpPid),
-        {_, RanchLPid, _, [ranch_listener_sup]} <- supervisor:which_children(RanchEPid),
-        {_, RanchCSPid, _, [ranch_conns_sup_sup]} <- supervisor:which_children(RanchLPid),
-        {_, RanchCPid, _, [ranch_conns_sup]} <- supervisor:which_children(RanchCSPid),
-        {rabbit_connection_sup, ConnPid, _, _} <- supervisor:which_children(RanchCPid),
-        {reader, ReaderPid, _, _} <- supervisor:which_children(ConnPid)
-    ].
+    pg_local:get_members(rabbit_amqp10_connections).
+    %% [ReaderPid
+    %%  || {_, TcpPid, _, [tcp_listener_sup]} <- supervisor:which_children(rabbit_sup),
+    %%     {_, RanchEPid, _, [ranch_embedded_sup]} <- supervisor:which_children(TcpPid),
+    %%     {_, RanchLPid, _, [ranch_listener_sup]} <- supervisor:which_children(RanchEPid),
+    %%     {_, RanchCSPid, _, [ranch_conns_sup_sup]} <- supervisor:which_children(RanchLPid),
+    %%     {_, RanchCPid, _, [ranch_conns_sup]} <- supervisor:which_children(RanchCSPid),
+    %%     {rabbit_connection_sup, ConnPid, _, _} <- supervisor:which_children(RanchCPid),
+    %%     {reader, ReaderPid, _, _} <- supervisor:which_children(ConnPid)
+    %% ].
+
+-spec register_connection(pid()) -> ok.
+register_connection(Pid) ->
+    pg_local:join(rabbit_amqp10_connections, Pid).
+
+-spec unregister_connection(pid()) -> ok.
+unregister_connection(Pid) ->
+    pg_local:leave(rabbit_amqp10_connections, Pid).

deps/rabbit/test/unit_pg_local_SUITE.erl renamed to deps/rabbit_common/test/unit_pg_local_SUITE.erl

@@ -238,10 +238,12 @@ update_last_blocked_by(Throttle) ->
 
 close_connection(State = #v1{connection = #v1_connection{
                                              timeout_sec = TimeoutSec}}) ->
+    Pid = self(),
     erlang:send_after((if TimeoutSec > 0 andalso
                           TimeoutSec < ?CLOSING_TIMEOUT -> TimeoutSec;
                           true                          -> ?CLOSING_TIMEOUT
-                       end) * 1000, self(), terminate_connection),
+                       end) * 1000, Pid, terminate_connection),
+    rabbit_amqp1_0:unregister_connection(Pid),
     State#v1{connection_state = closed}.
 
 handle_dependent_exit(ChPid, Reason, State) ->
@@ -434,6 +436,7 @@ handle_1_0_connection_frame(#'v1_0.open'{ max_frame_size = ClientFrameMax,
                         container_id   = {utf8, rabbit_nodes:cluster_name()},
                         properties     = server_properties()}),
     Conserve = rabbit_alarm:register(self(), {?MODULE, conserve_resources, []}),
+    rabbit_amqp1_0:register_connection(self()),
     control_throttle(
       State1#v1{throttle = Throttle#throttle{alarmed_by = Conserve}});