From ac88c9ddb7636179b7634c1e9b1335024080afde Mon Sep 17 00:00:00 2001
From: Stas <sts@abc3.dev>
Date: Tue, 10 Dec 2024 13:47:38 +0100
Subject: [PATCH] feat: add metrics for active proxy clients (#508)

---
 config/config.exs                  |  3 ++-
 config/dev.exs                     |  3 ++-
 lib/supavisor/application.ex       |  4 +++
 lib/supavisor/client_handler.ex    | 17 ++++++++-----
 lib/supavisor/monitoring/tenant.ex | 40 ++++++++++++++++++++++++++++++
 5 files changed, 59 insertions(+), 8 deletions(-)

diff --git a/config/config.exs b/config/config.exs
index 46291625..8e22c329 100644
--- a/config/config.exs
+++ b/config/config.exs
@@ -39,7 +39,8 @@ config :logger, :console,
     :type,
     :app_name,
     :peer_ip,
-    :local
+    :local,
+    :proxy
   ]
 
 # Use Jason for JSON parsing in Phoenix
diff --git a/config/dev.exs b/config/dev.exs
index eefe173f..d5d19bdd 100644
--- a/config/dev.exs
+++ b/config/dev.exs
@@ -74,7 +74,8 @@ config :logger, :console,
     :type,
     :app_name,
     :peer_ip,
-    :local
+    :local,
+    :proxy
   ]
 
 # Set a higher stacktrace during development. Avoid configuring such
diff --git a/lib/supavisor/application.ex b/lib/supavisor/application.ex
index 2b7321b1..4bf08935 100644
--- a/lib/supavisor/application.ex
+++ b/lib/supavisor/application.ex
@@ -81,6 +81,10 @@ defmodule Supavisor.Application do
        keys: :duplicate,
        name: Supavisor.Registry.TenantClients,
        partitions: System.schedulers_online()},
+      {Registry,
+       keys: :duplicate,
+       name: Supavisor.Registry.TenantProxyClients,
+       partitions: System.schedulers_online()},
       {Cluster.Supervisor, [topologies, [name: Supavisor.ClusterSupervisor]]},
       Supavisor.Repo,
       # Start the Telemetry supervisor
diff --git a/lib/supavisor/client_handler.ex b/lib/supavisor/client_handler.ex
index 40069974..fa08d32e 100644
--- a/lib/supavisor/client_handler.ex
+++ b/lib/supavisor/client_handler.ex
@@ -14,6 +14,8 @@ defmodule Supavisor.ClientHandler do
   @switch_active_count Application.compile_env(:supavisor, :switch_active_count)
   @subscribe_retries Application.compile_env(:supavisor, :subscribe_retries)
   @timeout_subscribe 500
+  @clients_registry Supavisor.Registry.TenantClients
+  @proxy_clients_registry Supavisor.Registry.TenantProxyClients
 
   alias Supavisor.{
     DbHandler,
@@ -248,7 +250,7 @@ defmodule Supavisor.ClientHandler do
           local: data.local
         )
 
-        Registry.register(Supavisor.Registry.TenantClients, id, [])
+        Registry.register(@clients_registry, id, [])
 
         {:ok, addr} = HandlerHelpers.addr_from_sock(sock)
 
@@ -424,9 +426,6 @@ defmodule Supavisor.ClientHandler do
       :proxy ->
         case Supavisor.get_pool_ranch(data.id) do
           {:ok, %{port: port, host: host}} ->
-            # Not sure if manager is alive yet
-            # manager_ref = Process.monitor(opts.workers.manager)
-
             auth =
               Map.merge(data.auth, %{
                 port: port,
@@ -437,8 +436,8 @@ defmodule Supavisor.ClientHandler do
                 upstream_verify: nil
               })
 
-            # {:keep_state, %{data | auth: auth, manager: manager_ref},
-            # {:next_event, :internal, :connect_db}}
+            Logger.metadata(proxy: true)
+            set_proxy_registry(data.id)
 
             {:keep_state, %{data | auth: auth}, {:next_event, :internal, :connect_db}}
 
@@ -1156,4 +1155,10 @@ defmodule Supavisor.ClientHandler do
     HandlerHelpers.activate(data.sock)
     0
   end
+
+  @spec set_proxy_registry(Supavisor.id()) :: {:ok, pid()} | {:error, term()}
+  defp set_proxy_registry(id) do
+    Registry.unregister(@clients_registry, id)
+    Registry.register(@proxy_clients_registry, id, [])
+  end
 end
diff --git a/lib/supavisor/monitoring/tenant.ex b/lib/supavisor/monitoring/tenant.ex
index e7f5e586..fe3c5a17 100644
--- a/lib/supavisor/monitoring/tenant.ex
+++ b/lib/supavisor/monitoring/tenant.ex
@@ -14,6 +14,7 @@ defmodule Supavisor.PromEx.Plugins.Tenant do
 
     [
       concurrent_connections(poll_rate),
+      concurrent_proxy_connections(poll_rate),
       concurrent_tenants(poll_rate)
     ]
   end
@@ -219,6 +220,45 @@ defmodule Supavisor.PromEx.Plugins.Tenant do
     )
   end
 
+  defp concurrent_proxy_connections(poll_rate) do
+    Polling.build(
+      :supavisor_concurrent_proxy_connections,
+      poll_rate,
+      {__MODULE__, :execute_tenant_proxy_metrics, []},
+      [
+        last_value(
+          [:supavisor, :proxy, :connections, :active],
+          event_name: [:supavisor, :proxy, :connections],
+          description: "The total count of active proxy clients for a tenant.",
+          measurement: :active,
+          tags: @tags
+        )
+      ]
+    )
+  end
+
+  def execute_tenant_proxy_metrics do
+    Registry.select(Supavisor.Registry.TenantProxyClients, [{{:"$1", :_, :_}, [], [:"$1"]}])
+    |> Enum.frequencies()
+    |> Enum.each(&emit_proxy_telemetry_for_tenant/1)
+  end
+
+  @spec emit_proxy_telemetry_for_tenant({S.id(), non_neg_integer()}) :: :ok
+  def emit_proxy_telemetry_for_tenant({{{type, tenant}, user, mode, db_name, search_path}, count}) do
+    :telemetry.execute(
+      [:supavisor, :proxy, :connections],
+      %{active: count},
+      %{
+        tenant: tenant,
+        user: user,
+        mode: mode,
+        type: type,
+        db_name: db_name,
+        search_path: search_path
+      }
+    )
+  end
+
   defp concurrent_tenants(poll_rate) do
     Polling.build(
       :supavisor_concurrent_tenants,