fix: move DB setup to happen after Connect.init

This change reduces the impact of slow DB setup impacting other tenants
trying to connect at the same time that landed on the same partition

Author: edgurgel

lib/realtime/syn_handler.ex

@@ -10,9 +10,9 @@ defmodule Realtime.SynHandler do
   @behaviour :syn_event_handler
 
   @impl true
-  def on_registry_process_updated(Connect, tenant_id, _pid, %{conn: conn}, :normal) when is_pid(conn) do
+  def on_registry_process_updated(Connect, tenant_id, pid, %{conn: conn}, :normal) when is_pid(conn) do
     # Update that a database connection is ready
-    Endpoint.local_broadcast(Connect.syn_topic(tenant_id), "ready", %{conn: conn})
+    Endpoint.local_broadcast(Connect.syn_topic(tenant_id), "ready", %{pid: pid, conn: conn})
   end
 
   def on_registry_process_updated(PostgresCdcRls, tenant_id, _pid, meta, _reason) do
@@ -38,7 +38,7 @@ defmodule Realtime.SynHandler do
     end
 
     topic = topic(mod)
-    Endpoint.local_broadcast(topic <> ":" <> name, topic <> "_down", nil)
+    Endpoint.local_broadcast(topic <> ":" <> name, topic <> "_down", %{pid: pid, reason: reason})
 
     :ok
   end

lib/realtime/tenants/connect.ex

@@ -19,7 +19,6 @@ defmodule Realtime.Tenants.Connect do
   alias Realtime.Tenants.Connect.GetTenant
   alias Realtime.Tenants.Connect.Piper
   alias Realtime.Tenants.Connect.RegisterProcess
-  alias Realtime.Tenants.Connect.StartCounters
   alias Realtime.Tenants.Migrations
   alias Realtime.Tenants.ReplicationConnection
   alias Realtime.UsersCounter
@@ -83,14 +82,13 @@ defmodule Realtime.Tenants.Connect do
           | {:error, :tenant_database_connection_initializing}
   def get_status(tenant_id) do
     case :syn.lookup(__MODULE__, tenant_id) do
-      {_pid, %{conn: nil}} ->
-        wait_for_connection(tenant_id)
+      {pid, %{conn: nil}} ->
+        wait_for_connection(pid, tenant_id)
 
       {_, %{conn: conn}} ->
         {:ok, conn}
 
       :undefined ->
-        Logger.warning("Connection process starting up")
         {:error, :tenant_database_connection_initializing}
 
       error ->
@@ -101,7 +99,7 @@ defmodule Realtime.Tenants.Connect do
 
   def syn_topic(tenant_id), do: "connect:#{tenant_id}"
 
-  defp wait_for_connection(tenant_id) do
+  defp wait_for_connection(pid, tenant_id) do
     RealtimeWeb.Endpoint.subscribe(syn_topic(tenant_id))
 
     # We do a lookup after subscribing because we could've missed a message while subscribing
@@ -112,9 +110,18 @@ defmodule Realtime.Tenants.Connect do
       _ ->
         # Wait for up to 5 seconds for the ready event
         receive do
-          %{event: "ready", payload: %{conn: conn}} -> {:ok, conn}
+          %{event: "ready", payload: %{pid: ^pid, conn: conn}} ->
+            {:ok, conn}
+
+          %{event: "connect_down", payload: %{pid: ^pid, reason: {:shutdown, :tenant_db_too_many_connections}}} ->
+            {:error, :tenant_db_too_many_connections}
+
+          %{event: "connect_down", payload: %{pid: ^pid, reason: _reason}} ->
+            metadata = [external_id: tenant_id, project: tenant_id]
+            log_error("UnableToConnectToTenantDatabase", "Unable to connect to tenant database", metadata)
+            {:error, :tenant_database_unavailable}
         after
-          5_000 -> {:error, :initializing}
+          15_000 -> {:error, :initializing}
         end
     end
   after
@@ -139,16 +146,6 @@ defmodule Realtime.Tenants.Connect do
       {:error, {:already_started, _}} ->
         get_status(tenant_id)
 
-      {:error, {:shutdown, :tenant_db_too_many_connections}} ->
-        {:error, :tenant_db_too_many_connections}
-
-      {:error, {:shutdown, :tenant_not_found}} ->
-        {:error, :tenant_not_found}
-
-      {:error, :shutdown} ->
-        log_error("UnableToConnectToTenantDatabase", "Unable to connect to tenant database", metadata)
-        {:error, :tenant_database_unavailable}
-
       {:error, error} ->
         log_error("UnableToConnectToTenantDatabase", error, metadata)
         {:error, :tenant_database_unavailable}
@@ -209,30 +206,33 @@ defmodule Realtime.Tenants.Connect do
   def init(%{tenant_id: tenant_id} = state) do
     Logger.metadata(external_id: tenant_id, project: tenant_id)
 
+    {:ok, state, {:continue, :db_connect}}
+  end
+
+  @impl true
+  def handle_continue(:db_connect, state) do
     pipes = [
       GetTenant,
       CheckConnection,
-      StartCounters,
       RegisterProcess
     ]
 
     case Piper.run(pipes, state) do
       {:ok, acc} ->
-        {:ok, acc, {:continue, :run_migrations}}
+        {:noreply, acc, {:continue, :run_migrations}}
 
       {:error, :tenant_not_found} ->
-        {:stop, {:shutdown, :tenant_not_found}}
+        {:stop, {:shutdown, :tenant_not_found}, state}
 
       {:error, :tenant_db_too_many_connections} ->
-        {:stop, {:shutdown, :tenant_db_too_many_connections}}
+        {:stop, {:shutdown, :tenant_db_too_many_connections}, state}
 
       {:error, error} ->
         log_error("UnableToConnectToTenantDatabase", error)
-        {:stop, :shutdown}
+        {:stop, :shutdown, state}
     end
   end
 
-  @impl true
   def handle_continue(:run_migrations, state) do
     %{tenant: tenant, db_conn_pid: db_conn_pid} = state
     Logger.warning("Tenant #{tenant.external_id} is initializing: #{inspect(node())}")
@@ -375,6 +375,7 @@ defmodule Realtime.Tenants.Connect do
 
   ## Private functions
   defp call_external_node(tenant_id, opts) do
+    Logger.warning("Connection process starting up")
     rpc_timeout = Keyword.get(opts, :rpc_timeout, @rpc_timeout_default)
 
     with tenant <- Tenants.Cache.get_tenant_by_external_id(tenant_id),

lib/realtime/tenants/connect/check_connection.ex

@@ -2,16 +2,14 @@ defmodule Realtime.Tenants.Connect.CheckConnection do
   @moduledoc """
   Check tenant database connection.
   """
-  alias Realtime.Database
 
   @behaviour Realtime.Tenants.Connect.Piper
   @impl true
   def run(acc) do
     %{tenant: tenant} = acc
 
-    case Database.check_tenant_connection(tenant) do
+    case Realtime.Database.check_tenant_connection(tenant) do
       {:ok, conn} ->
-        Process.link(conn)
         db_conn_reference = Process.monitor(conn)
         {:ok, %{acc | db_conn_pid: conn, db_conn_reference: db_conn_reference}}
 

lib/realtime/tenants/connect/start_counters.ex

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

mix.exs

@@ -653,8 +653,8 @@ defmodule Realtime.Integration.RtChannelTest do
           :syn.update_registry(Connect, tenant.external_id, fn _pid, meta -> %{meta | conn: nil} end)
           payload = %{"event" => "TEST", "payload" => %{"msg" => 1}, "type" => "broadcast"}
           WebsocketClient.send_event(service_role_socket, topic, "broadcast", payload)
-          # Waiting more than 5 seconds as this is the amount of time we will wait for the Connection to be ready
-          refute_receive %Message{event: "broadcast", payload: ^payload, topic: ^topic}, 6000
+          # Waiting more than 15 seconds as this is the amount of time we will wait for the Connection to be ready
+          refute_receive %Message{event: "broadcast", payload: ^payload, topic: ^topic}, 16000
         end)
 
       assert log =~ "UnableToHandleBroadcast"
@@ -831,7 +831,7 @@ defmodule Realtime.Integration.RtChannelTest do
 
           refute_receive %Message{event: "presence_diff"}, 500
           # Waiting more than 5 seconds as this is the amount of time we will wait for the Connection to be ready
-          refute_receive %Message{event: "phx_leave", topic: ^topic}, 6000
+          refute_receive %Message{event: "phx_leave", topic: ^topic}, 16000
         end)
 
       assert log =~ "UnableToHandlePresence"

test/integration/rt_channel_test.exs

@@ -168,32 +168,40 @@ defmodule Realtime.SynHandlerTest do
 
     test "it handles :syn_conflict_resolution reason" do
       reason = :syn_conflict_resolution
+      pid = self()
 
       log =
         capture_log(fn ->
-          assert SynHandler.on_process_unregistered(@mod, @name, self(), %{}, reason) == :ok
+          assert SynHandler.on_process_unregistered(@mod, @name, pid, %{}, reason) == :ok
         end)
 
       topic = "#{@topic}:#{@name}"
       event = "#{@topic}_down"
 
       assert log =~ "#{@mod} terminated due to syn conflict resolution: #{inspect(@name)} #{inspect(self())}"
-      assert_receive %Phoenix.Socket.Broadcast{topic: ^topic, event: ^event, payload: nil}
+      assert_receive %Phoenix.Socket.Broadcast{topic: ^topic, event: ^event, payload: %{reason: ^reason, pid: ^pid}}
     end
 
     test "it handles other reasons" do
       reason = :other_reason
+      pid = self()
 
       log =
         capture_log(fn ->
-          assert SynHandler.on_process_unregistered(@mod, @name, self(), %{}, reason) == :ok
+          assert SynHandler.on_process_unregistered(@mod, @name, pid, %{}, reason) == :ok
         end)
 
       topic = "#{@topic}:#{@name}"
       event = "#{@topic}_down"
 
       refute log =~ "#{@mod} terminated: #{inspect(@name)} #{node()}"
-      assert_receive %Phoenix.Socket.Broadcast{topic: ^topic, event: ^event, payload: nil}, 500
+
+      assert_receive %Phoenix.Socket.Broadcast{
+                       topic: ^topic,
+                       event: ^event,
+                       payload: %{reason: ^reason, pid: ^pid}
+                     },
+                     500
     end
   end
 end

test/realtime/syn_handler_test.exs

@@ -78,31 +78,27 @@ defmodule Realtime.Tenants.ConnectTest do
       assert_receive {:ok, ^pid}
     end
 
-    test "more than 5 seconds passed error out", %{tenant: tenant} do
+    test "more than 15 seconds passed error out", %{tenant: tenant} do
       parent = self()
 
       # Let's slow down Connect starting
       expect(Database, :check_tenant_connection, fn t ->
-        :timer.sleep(5500)
+        Process.sleep(15500)
         call_original(Database, :check_tenant_connection, [t])
       end)
 
       connect = fn -> send(parent, Connect.lookup_or_start_connection(tenant.external_id)) end
 
-      # Start an early connect
-      spawn(connect)
-      :timer.sleep(100)
-
-      # Start others
       spawn(connect)
       spawn(connect)
 
-      {:error, :tenant_database_unavailable} = Connect.lookup_or_start_connection(tenant.external_id)
+      {:error, :initializing} = Connect.lookup_or_start_connection(tenant.external_id)
+      # The above call waited 15 seconds
+      assert_receive {:error, :initializing}
+      assert_receive {:error, :initializing}
 
-      # Only one will succeed the others timed out waiting
-      assert_receive {:error, :tenant_database_unavailable}
-      assert_receive {:error, :tenant_database_unavailable}
-      assert_receive {:ok, _pid}, 7000
+      # This one will succeed
+      {:ok, _pid} = Connect.lookup_or_start_connection(tenant.external_id)
     end
   end