fix: change metrics cleaner to also clean up Connect shutdown (#1674)

Once Connect for the tenant shuts down clean up after threshold has
passed

Author: edgurgel

lib/realtime/metrics_cleaner.ex

@@ -6,18 +6,26 @@ defmodule Realtime.MetricsCleaner do
 
   defstruct [:check_ref, :interval]
 
-  def handle_event([:beacon, :users, :group, :vacant], _measurements, %{group: tenant_id}, cleaner_table) do
-    :ets.insert(cleaner_table, {tenant_id, DateTime.to_unix(DateTime.utc_now(), :second)})
+  def handle_beacon_event([:beacon, :users, :group, :vacant], _, %{group: tenant_id}, vacant_websockets) do
+    :ets.insert(vacant_websockets, {tenant_id, DateTime.to_unix(DateTime.utc_now(), :second)})
   end
 
-  def handle_event([:beacon, :users, :group, :occupied], _measurements, %{group: tenant_id}, cleaner_table) do
-    :ets.delete(cleaner_table, tenant_id)
+  def handle_beacon_event([:beacon, :users, :group, :occupied], _, %{group: tenant_id}, vacant_websockets) do
+    :ets.delete(vacant_websockets, tenant_id)
+  end
+
+  def handle_syn_event([:syn, Realtime.Tenants.Connect, :unregistered], _, %{name: tenant_id}, disconnected_tenants) do
+    :ets.insert(disconnected_tenants, {tenant_id, DateTime.to_unix(DateTime.utc_now(), :second)})
+  end
+
+  def handle_syn_event([:syn, Realtime.Tenants.Connect, :registered], _, %{name: tenant_id}, disconnected_tenants) do
+    :ets.delete(disconnected_tenants, tenant_id)
   end
 
   def start_link(opts), do: GenServer.start_link(__MODULE__, opts)
 
-  # 5 minutes
-  @default_vacant_metric_threshold_in_seconds 300
+  # 10 minutes
+  @default_vacant_metric_threshold_in_seconds 600
 
   @impl true
   def init(opts) do
@@ -30,32 +38,54 @@ defmodule Realtime.MetricsCleaner do
 
     Logger.info("Starting MetricsCleaner")
 
-    tid = :ets.new(:metrics_cleaner, [:set, :public, read_concurrency: false, write_concurrency: :auto])
+    vacant_websockets = :ets.new(:vacant_websockets, [:set, :public, read_concurrency: false, write_concurrency: :auto])
+
+    disconnected_tenants =
+      :ets.new(:disconnected_tenants, [:set, :public, read_concurrency: false, write_concurrency: :auto])
 
     :ok =
       :telemetry.attach_many(
-        self(),
+        [self(), :vacant_websockets],
         [[:beacon, :users, :group, :occupied], [:beacon, :users, :group, :vacant]],
-        &__MODULE__.handle_event/4,
-        tid
+        &__MODULE__.handle_beacon_event/4,
+        vacant_websockets
+      )
+
+    :ok =
+      :telemetry.attach_many(
+        [self(), :disconnected_tenants],
+        [[:syn, Realtime.Tenants.Connect, :registered], [:syn, Realtime.Tenants.Connect, :unregistered]],
+        &__MODULE__.handle_syn_event/4,
+        disconnected_tenants
       )
 
     {:ok,
      %{
        check_ref: check(interval),
        interval: interval,
        vacant_metric_threshold_in_seconds: vacant_metric_threshold_in_seconds,
-       tid: tid
+       vacant_websockets: vacant_websockets,
+       disconnected_tenants: disconnected_tenants
      }}
   end
 
+  @impl true
+  def terminate(_reason, _state) do
+    :telemetry.detach([self(), :vacant_websockets])
+    :telemetry.detach([self(), :disconnected_tenants])
+    :ok
+  end
+
   @impl true
   def handle_info(:check, %{interval: interval} = state) do
     Process.cancel_timer(state.check_ref)
 
     {exec_time, _} =
       :timer.tc(
-        fn -> loop_and_cleanup_metrics_table(state.tid, state.vacant_metric_threshold_in_seconds) end,
+        fn ->
+          loop_and_cleanup_metrics_table(state.vacant_websockets, state.vacant_metric_threshold_in_seconds)
+          loop_and_cleanup_metrics_table(state.disconnected_tenants, state.vacant_metric_threshold_in_seconds)
+        end,
         :millisecond
       )
 

lib/realtime/syn_handler.ex

@@ -29,6 +29,11 @@ defmodule Realtime.SynHandler do
     end
   end
 
+  @impl true
+  def on_process_registered(scope, name, _pid, _meta, _reason) do
+    :telemetry.execute([:syn, scope, :registered], %{}, %{name: name})
+  end
+
   @doc """
   When processes registered with :syn are unregistered, either manually or by stopping, this
   callback is invoked.
@@ -40,6 +45,8 @@ defmodule Realtime.SynHandler do
   """
   @impl true
   def on_process_unregistered(scope, name, pid, _meta, reason) do
+    :telemetry.execute([:syn, scope, :unregistered], %{}, %{name: name})
+
     case Atom.to_string(scope) do
       @postgres_cdc_scope_prefix <> _ = scope ->
         Endpoint.local_broadcast(PostgresCdc.syn_topic(name), scope <> "_down", %{pid: pid, reason: reason})

mix.exs

@@ -4,7 +4,7 @@ defmodule Realtime.MixProject do
   def project do
     [
       app: :realtime,
-      version: "2.71.3",
+      version: "2.71.4",
       elixir: "~> 1.18",
       elixirc_paths: elixirc_paths(Mix.env()),
       start_permanent: Mix.env() == :prod,

test/realtime/metrics_cleaner_test.exs

@@ -2,8 +2,9 @@ defmodule Realtime.MetricsCleanerTest do
   use Realtime.DataCase, async: true
 
   alias Realtime.MetricsCleaner
+  alias Realtime.Tenants.Connect
 
-  describe "metrics cleanup" do
+  describe "metrics cleanup - vacant websockets" do
     test "cleans up metrics for users that have been disconnected" do
       :telemetry.execute(
         [:realtime, :connections],
@@ -48,15 +49,15 @@ defmodule Realtime.MetricsCleanerTest do
       # Wait for clean up to run
       Process.sleep(200)
 
-      # Nothing changes
+      # Nothing changes yet (threshold not reached)
       metrics = Realtime.PromEx.get_metrics() |> IO.iodata_to_binary()
 
       assert String.contains?(metrics, "tenant=\"occupied-tenant\"")
       assert String.contains?(metrics, "tenant=\"vacant-tenant1\"")
       assert String.contains?(metrics, "tenant=\"vacant-tenant2\"")
 
-      # Wait for clean up to run again
-      Process.sleep(2100)
+      # Wait for threshold to pass and cleanup to run
+      Process.sleep(2200)
 
       # vacant tenant metrics are now gone
       metrics = Realtime.PromEx.get_metrics() |> IO.iodata_to_binary()
@@ -65,5 +66,127 @@ defmodule Realtime.MetricsCleanerTest do
       refute String.contains?(metrics, "tenant=\"vacant-tenant1\"")
       refute String.contains?(metrics, "tenant=\"vacant-tenant2\"")
     end
+
+    test "does not clean up metrics if websockets reconnect before threshold" do
+      :telemetry.execute(
+        [:realtime, :connections],
+        %{connected: 1, connected_cluster: 10, limit: 100},
+        %{tenant: "reconnect-tenant"}
+      )
+
+      pid = spawn_link(fn -> Process.sleep(:infinity) end)
+
+      Beacon.join(:users, "reconnect-tenant", pid)
+
+      metrics = Realtime.PromEx.get_metrics() |> IO.iodata_to_binary()
+      assert String.contains?(metrics, "tenant=\"reconnect-tenant\"")
+
+      start_supervised!(
+        {MetricsCleaner, [metrics_cleaner_schedule_timer_in_ms: 100, vacant_metric_threshold_in_seconds: 1]}
+      )
+
+      # Disconnect
+      Beacon.leave(:users, "reconnect-tenant", pid)
+      Process.sleep(500)
+
+      # Reconnect before threshold
+      pid2 = spawn_link(fn -> Process.sleep(:infinity) end)
+      Beacon.join(:users, "reconnect-tenant", pid2)
+
+      # Wait for cleanup to run
+      Process.sleep(2200)
+
+      # Metrics should still be present
+      metrics = Realtime.PromEx.get_metrics() |> IO.iodata_to_binary()
+      assert String.contains?(metrics, "tenant=\"reconnect-tenant\"")
+    end
+  end
+
+  describe "metrics cleanup - disconnected tenants" do
+    test "cleans up metrics for tenants that have been unregistered" do
+      :telemetry.execute(
+        [:realtime, :connections],
+        %{connected: 1, connected_cluster: 10, limit: 100},
+        %{tenant: "connected-tenant"}
+      )
+
+      :telemetry.execute(
+        [:realtime, :connections],
+        %{connected: 0, connected_cluster: 20, limit: 100},
+        %{tenant: "disconnected-tenant1"}
+      )
+
+      :telemetry.execute(
+        [:realtime, :connections],
+        %{connected: 0, connected_cluster: 20, limit: 100},
+        %{tenant: "disconnected-tenant2"}
+      )
+
+      metrics = Realtime.PromEx.get_metrics() |> IO.iodata_to_binary()
+
+      assert String.contains?(metrics, "tenant=\"connected-tenant\"")
+      assert String.contains?(metrics, "tenant=\"disconnected-tenant1\"")
+      assert String.contains?(metrics, "tenant=\"disconnected-tenant2\"")
+
+      start_supervised!(
+        {MetricsCleaner, [metrics_cleaner_schedule_timer_in_ms: 100, vacant_metric_threshold_in_seconds: 1]}
+      )
+
+      # Simulate tenant registration (connected)
+      :telemetry.execute([:syn, Connect, :registered], %{}, %{name: "connected-tenant"})
+
+      # Simulate tenant unregistration (disconnected)
+      :telemetry.execute([:syn, Connect, :unregistered], %{}, %{name: "disconnected-tenant1"})
+      :telemetry.execute([:syn, Connect, :unregistered], %{}, %{name: "disconnected-tenant2"})
+
+      # Wait for clean up to run
+      Process.sleep(200)
+
+      # Nothing changes yet (threshold not reached)
+      metrics = Realtime.PromEx.get_metrics() |> IO.iodata_to_binary()
+
+      assert String.contains?(metrics, "tenant=\"connected-tenant\"")
+      assert String.contains?(metrics, "tenant=\"disconnected-tenant1\"")
+      assert String.contains?(metrics, "tenant=\"disconnected-tenant2\"")
+
+      # Wait for threshold to pass and cleanup to run
+      Process.sleep(2200)
+
+      # disconnected tenant metrics are now gone
+      metrics = Realtime.PromEx.get_metrics() |> IO.iodata_to_binary()
+
+      assert String.contains?(metrics, "tenant=\"connected-tenant\"")
+      refute String.contains?(metrics, "tenant=\"disconnected-tenant1\"")
+      refute String.contains?(metrics, "tenant=\"disconnected-tenant2\"")
+    end
+
+    test "does not clean up metrics if tenant reconnects before threshold" do
+      :telemetry.execute(
+        [:realtime, :connections],
+        %{connected: 1, connected_cluster: 10, limit: 100},
+        %{tenant: "reconnect-tenant"}
+      )
+
+      metrics = Realtime.PromEx.get_metrics() |> IO.iodata_to_binary()
+      assert String.contains?(metrics, "tenant=\"reconnect-tenant\"")
+
+      start_supervised!(
+        {MetricsCleaner, [metrics_cleaner_schedule_timer_in_ms: 100, vacant_metric_threshold_in_seconds: 1]}
+      )
+
+      # Simulate tenant unregistration
+      :telemetry.execute([:syn, Connect, :unregistered], %{}, %{name: "reconnect-tenant"})
+      Process.sleep(500)
+
+      # Re-register before threshold
+      :telemetry.execute([:syn, Connect, :registered], %{}, %{name: "reconnect-tenant"})
+
+      # Wait for cleanup to run
+      Process.sleep(2200)
+
+      # Metrics should still be present
+      metrics = Realtime.PromEx.get_metrics() |> IO.iodata_to_binary()
+      assert String.contains?(metrics, "tenant=\"reconnect-tenant\"")
+    end
   end
 end

test/realtime/syn_handler_test.exs

@@ -227,11 +227,68 @@ defmodule Realtime.SynHandlerTest do
     end
   end
 
+  describe "on_process_registered/5" do
+    test "emits telemetry event for process registration" do
+      pid = self()
+      meta = %{some: :meta}
+      reason = :normal
+
+      # Attach a test handler to capture the telemetry event
+      test_pid = self()
+      handler_id = [:test, :syn_handler, :registered]
+
+      :telemetry.attach(
+        handler_id,
+        [:syn, @mod, :registered],
+        fn event, measurements, metadata, _config ->
+          send(test_pid, {:telemetry_event, event, measurements, metadata})
+        end,
+        nil
+      )
+
+      on_exit(fn -> :telemetry.detach(handler_id) end)
+
+      assert SynHandler.on_process_registered(@mod, @name, pid, meta, reason) == :ok
+
+      assert_receive {:telemetry_event, [:syn, @mod, :registered], %{}, %{name: @name}}
+    end
+  end
+
   describe "on_process_unregistered/5" do
     setup do
       RealtimeWeb.Endpoint.subscribe("#{@topic}:#{@name}")
     end
 
+    test "emits telemetry event for process unregistration" do
+      reason = :normal
+      pid = self()
+
+      # Attach a test handler to capture the telemetry event
+      test_pid = self()
+      handler_id = [:test, :syn_handler, :unregistered]
+
+      :telemetry.attach(
+        handler_id,
+        [:syn, @mod, :unregistered],
+        fn event, measurements, metadata, _config ->
+          send(test_pid, {:telemetry_event, event, measurements, metadata})
+        end,
+        nil
+      )
+
+      on_exit(fn -> :telemetry.detach(handler_id) end)
+
+      capture_log(fn ->
+        assert SynHandler.on_process_unregistered(@mod, @name, pid, %{}, reason) == :ok
+      end)
+
+      assert_receive {:telemetry_event, [:syn, @mod, :unregistered], %{}, %{name: @name}}
+
+      topic = "#{@topic}:#{@name}"
+      event = "#{@topic}_down"
+      assert_receive %Phoenix.Socket.Broadcast{topic: ^topic, event: ^event, payload: %{reason: ^reason, pid: ^pid}}
+    end
+
     test "it handles :syn_conflict_resolution reason" do
       reason = :syn_conflict_resolution
       pid = self()