fix: wait for Connect.lookup_or_start_connection to be ready (#1466)

This helps with cold starts as multiple WebSocket connections might
happen at the same time

Author: edgurgel

lib/realtime/syn_handler.ex

@@ -4,6 +4,17 @@ defmodule Realtime.SynHandler do
   """
   require Logger
   alias RealtimeWeb.Endpoint
+  alias Realtime.Tenants.Connect
+
+  @behaviour :syn_event_handler
+
+  @impl true
+  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})
+  end
+
+  def on_registry_process_updated(_scope, _name, _pid, _meta, _reason), do: :ok
 
   @doc """
   When processes registered with :syn are unregistered, either manually or by stopping, this
@@ -14,6 +25,7 @@ defmodule Realtime.SynHandler do
   We want to log conflict resolutions to know when more than one process on the cluster
   was started, and subsequently stopped because :syn handled the conflict.
   """
+  @impl true
   def on_process_unregistered(mod, name, _pid, _meta, reason) do
     case reason do
       :syn_conflict_resolution ->
@@ -27,6 +39,7 @@ defmodule Realtime.SynHandler do
     :ok
   end
 
+  @impl true
   def resolve_registry_conflict(mod, name, {pid1, %{region: region}, time1}, {pid2, _, time2}) do
     platform_region = Realtime.Nodes.platform_region_translator(region)
 

lib/realtime/tenants/connect.ex

@@ -66,7 +66,6 @@ defmodule Realtime.Tenants.Connect do
         call_external_node(tenant_id, opts)
 
       {:error, :tenant_database_connection_initializing} ->
-        Process.sleep(100)
         call_external_node(tenant_id, opts)
 
       {:error, :initializing} ->
@@ -84,8 +83,8 @@ defmodule Realtime.Tenants.Connect do
           | {:error, :tenant_database_connection_initializing}
   def get_status(tenant_id) do
     case :syn.lookup(__MODULE__, tenant_id) do
-      {_, %{conn: nil}} ->
-        {:error, :initializing}
+      {_pid, %{conn: nil}} ->
+        wait_for_connection(tenant_id)
 
       {_, %{conn: conn}} ->
         {:ok, conn}
@@ -100,6 +99,28 @@ defmodule Realtime.Tenants.Connect do
     end
   end
 
+  def syn_topic(tenant_id), do: "connect:#{tenant_id}"
+
+  defp wait_for_connection(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
+    case :syn.lookup(__MODULE__, tenant_id) do
+      {_pid, %{conn: conn}} when is_pid(conn) ->
+        {:ok, conn}
+
+      _ ->
+        # Wait for up to 5 seconds for the ready event
+        receive do
+          %{event: "ready", payload: %{conn: conn}} -> {:ok, conn}
+        after
+          5_000 -> {:error, :initializing}
+        end
+    end
+  after
+    RealtimeWeb.Endpoint.unsubscribe(syn_topic(tenant_id))
+  end
+
   @doc """
   Connects to a tenant's database and stores the DBConnection in the process :syn metadata
   """

mix.exs

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

test/integration/rt_channel_test.exs

@@ -660,7 +660,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)
-          refute_receive %Message{event: "broadcast", payload: ^payload, topic: ^topic}, 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: "broadcast", payload: ^payload, topic: ^topic}, 6000
         end)
 
       assert log =~ "UnableToHandleBroadcast"
@@ -837,7 +838,8 @@ defmodule Realtime.Integration.RtChannelTest do
           WebsocketClient.send_event(socket, topic, "presence", payload)
 
           refute_receive %Message{event: "presence_diff"}, 500
-          refute_receive %Message{event: "phx_leave", topic: ^topic}
+          # 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
         end)
 
       assert log =~ "UnableToHandlePresence"
@@ -2251,11 +2253,10 @@ defmodule Realtime.Integration.RtChannelTest do
     {:ok, db_conn} = Database.connect(tenant, "realtime_test", :stop)
     clean_table(db_conn, "realtime", "messages")
     topic = Map.get(context, :topic, random_string())
-    message = message_fixture(tenant, %{topic: topic})
 
-    if policies = context[:policies], do: create_rls_policies(db_conn, policies, message)
+    if policies = context[:policies], do: create_rls_policies(db_conn, policies, %{topic: topic})
 
-    %{topic: message.topic}
+    %{topic: topic}
   end
 
   defp setup_trigger(%{tenant: tenant, topic: topic}) do

test/realtime/gen_rpc_test.exs

@@ -9,12 +9,15 @@ defmodule Realtime.GenRpcTest do
   setup context do
     {:ok, node} = Clustered.start(nil, extra_config: context[:extra_config] || [])
 
-    :telemetry.attach(__MODULE__, [:realtime, :rpc], &__MODULE__.handle_telemetry/4, pid: self())
-    on_exit(fn -> :telemetry.detach(__MODULE__) end)
     %{node: node}
   end
 
   describe "call/5" do
+    setup do
+      :telemetry.attach(__MODULE__, [:realtime, :rpc], &__MODULE__.handle_telemetry/4, pid: self())
+      on_exit(fn -> :telemetry.detach(__MODULE__) end)
+    end
+
     test "returns the result calling local node" do
       current_node = node()
 
@@ -172,6 +175,11 @@ defmodule Realtime.GenRpcTest do
   end
 
   describe "multicall/4" do
+    setup do
+      :telemetry.attach(__MODULE__, [:realtime, :rpc], &__MODULE__.handle_telemetry/4, pid: self())
+      on_exit(fn -> :telemetry.detach(__MODULE__) end)
+    end
+
     test "returns the result of the function call per node", %{node: node} do
       current_node = node()
 

test/realtime/tenants/connect_test.exs

@@ -35,6 +35,62 @@ defmodule Realtime.Tenants.ConnectTest do
     refute_receive {:DOWN, ^ref, :process, ^pid, _reason}, timeout
   end
 
+  describe "handle cold start" do
+    test "multiple proccesses succeed together", %{tenant: tenant} do
+      parent = self()
+
+      # Let's slow down Connect starting
+      expect(Database, :check_tenant_connection, fn t ->
+        :timer.sleep(1000)
+        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)
+
+      # This one should block and wait for the first Connect
+      {:ok, pid} = Connect.lookup_or_start_connection(tenant.external_id)
+
+      assert_receive {:ok, ^pid}
+      assert_receive {:ok, ^pid}
+      assert_receive {:ok, ^pid}
+    end
+
+    test "more than 5 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)
+        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)
+
+      # 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
+    end
+  end
+
   describe "region rebalancing" do
     test "rebalancing needed process stops", %{tenant: tenant} do
       external_id = tenant.external_id
@@ -346,14 +402,6 @@ defmodule Realtime.Tenants.ConnectTest do
       assert log =~ "ReplicationMaxWalSendersReached"
     end
 
-    test "syn with no connection", %{tenant: tenant} do
-      external_id = tenant.external_id
-      expect(:syn, :lookup, 2, fn Connect, ^external_id -> {nil, %{conn: nil}} end)
-
-      assert {:error, :tenant_database_unavailable} = Connect.lookup_or_start_connection(external_id)
-      assert {:error, :initializing} = Connect.get_status(external_id)
-    end
-
     test "handle rpc errors gracefully" do
       expect(Realtime.Nodes, :get_node_for_tenant, fn _ -> {:ok, :potato@nohost, "us-east-1"} end)
 

test/test_helper.exs

@@ -56,10 +56,11 @@ Mimic.copy(Realtime.RateCounter)
 Mimic.copy(Realtime.Tenants.Authorization)
 Mimic.copy(Realtime.Tenants.Cache)
 Mimic.copy(Realtime.Tenants.Connect)
+Mimic.copy(Realtime.Database)
 Mimic.copy(Realtime.Tenants.Migrations)
 Mimic.copy(Realtime.Tenants.Rebalancer)
+Mimic.copy(Realtime.Tenants.ReplicationConnection)
 Mimic.copy(RealtimeWeb.ChannelsAuthorization)
 Mimic.copy(RealtimeWeb.Endpoint)
 Mimic.copy(RealtimeWeb.JwtVerification)
-Mimic.copy(Realtime.Tenants.ReplicationConnection)
 Mimic.copy(RealtimeWeb.TenantBroadcaster)