change approach to make it reliant on the channel instead of waiting for pubsub

Author: filipecabaco

config/config.exs

@@ -12,7 +12,8 @@ config :realtime,
   ecto_repos: [Realtime.Repo],
   version: Mix.Project.config()[:version],
   replication_watchdog_interval: :timer.minutes(5),
-  replication_watchdog_timeout: :timer.minutes(1)
+  replication_watchdog_timeout: :timer.minutes(1),
+  replication_ready_timeout: :timer.minutes(1)
 
 # Configures the endpoint
 config :realtime, RealtimeWeb.Endpoint,

lib/realtime_web/channels/realtime_channel.ex

@@ -30,6 +30,7 @@ defmodule RealtimeWeb.RealtimeChannel do
   alias RealtimeWeb.RealtimeChannel.Tracker
 
   @confirm_token_ms_interval :timer.minutes(5)
+  @replication_ready_check_interval 10
   @fullsweep_after Application.compile_env!(:realtime, :websocket_fullsweep_after)
 
   @impl true
@@ -104,7 +105,6 @@ defmodule RealtimeWeb.RealtimeChannel do
 
       RealtimeWeb.Endpoint.subscribe(tenant_topic, metadata: metadata)
       RealtimeWeb.Endpoint.subscribe("realtime:operations:" <> tenant_id, metadata: metadata)
-      RealtimeWeb.Endpoint.subscribe(Connect.syn_topic(tenant_id))
       send(self(), :notify_replication_ready)
 
       is_new_api = new_api?(params)
@@ -136,7 +136,8 @@ defmodule RealtimeWeb.RealtimeChannel do
         tenant_topic: tenant_topic,
         channel_name: sub_topic,
         presence_enabled?: presence_enabled?,
-        replication_ready_notified?: false
+        replication_ready_notified?: false,
+        replication_ready_deadline: System.monotonic_time(:millisecond) + replication_ready_timeout()
       }
 
       socket =
@@ -293,19 +294,26 @@ defmodule RealtimeWeb.RealtimeChannel do
     {:noreply, assign(socket, %{pg_sub_ref: pg_sub_ref})}
   end
 
-  def handle_info(:notify_replication_ready, %{assigns: %{tenant: tenant_id}} = socket) do
-    case Connect.replication_status(tenant_id) do
-      {:ok, _replication_conn} -> {:noreply, push_replication_ready(socket)}
-      {:error, :not_connected} -> {:noreply, socket}
-    end
+  def handle_info(:notify_replication_ready, %{assigns: %{replication_ready_notified?: true}} = socket) do
+    {:noreply, socket}
   end
 
-  def handle_info(%{event: "ready", payload: %{replication_conn: pid}}, socket)
-      when is_pid(pid) do
-    {:noreply, push_replication_ready(socket)}
-  end
+  def handle_info(:notify_replication_ready, socket) do
+    %{assigns: %{tenant: tenant_id, channel_name: channel_name, replication_ready_deadline: deadline}} = socket
 
-  def handle_info(%{topic: "connect:" <> _}, socket), do: {:noreply, socket}
+    cond do
+      match?({:ok, _replication_conn}, Connect.replication_status(tenant_id)) ->
+        push_system_message("broadcast", socket, "ok", "Replication connection established", channel_name)
+        {:noreply, assign(socket, :replication_ready_notified?, true)}
+
+      System.monotonic_time(:millisecond) >= deadline ->
+        shutdown_response(socket, "Replication connection was not established in time")
+
+      true ->
+        Process.send_after(self(), :notify_replication_ready, @replication_ready_check_interval)
+        {:noreply, socket}
+    end
+  end
 
   def handle_info(_msg, %{assigns: %{policies: %Policies{broadcast: %BroadcastPolicies{read: false}}}} = socket) do
     Logger.warning("Broadcast message ignored")
@@ -740,12 +748,8 @@ defmodule RealtimeWeb.RealtimeChannel do
     {:stop, :normal, socket}
   end
 
-  defp push_replication_ready(%{assigns: %{replication_ready_notified?: true}} = socket), do: socket
-
-  defp push_replication_ready(%{assigns: %{tenant: tenant_id, channel_name: channel_name}} = socket) do
-    RealtimeWeb.Endpoint.unsubscribe(Connect.syn_topic(tenant_id))
-    push_system_message("broadcast", socket, "ok", "Replication connection established", channel_name)
-    assign(socket, :replication_ready_notified?, true)
+  defp replication_ready_timeout do
+    Application.fetch_env!(:realtime, :replication_ready_timeout)
   end
 
   defp push_system_message(extension, socket, status, error, channel_name)

test/integration/rt_channel/billable_events_test.exs

@@ -76,7 +76,13 @@ defmodule Realtime.Integration.RtChannel.BillableEventsTest do
 
       # Join events
       assert_receive %Message{event: "phx_reply", payload: %{"status" => "ok"}, topic: ^topic}, 300
-      assert_receive %Message{topic: ^topic, event: "system"}, 5000
+
+      assert_receive %Message{
+                       topic: ^topic,
+                       event: "system",
+                       payload: %{"extension" => "postgres_changes", "status" => "ok"}
+                     },
+                     50000
 
       # Wait for RateCounter to run
       RateCounterHelper.tick_tenant_rate_counters!(tenant.external_id)
@@ -201,13 +207,25 @@ defmodule Realtime.Integration.RtChannel.BillableEventsTest do
 
       # Join events
       assert_receive %Message{event: "phx_reply", payload: %{"status" => "ok"}, topic: ^topic}, 300
-      assert_receive %Message{topic: ^topic, event: "system"}, 5000
+
+      assert_receive %Message{
+                       topic: ^topic,
+                       event: "system",
+                       payload: %{"extension" => "postgres_changes", "status" => "ok"}
+                     },
+                     5000
 
       # Add second user to test the "multiplication" of billable events
       {socket, _} = get_connection(tenant, serializer)
       WebsocketClient.join(socket, topic, %{config: config})
       assert_receive %Message{event: "phx_reply", payload: %{"status" => "ok"}, topic: ^topic}, 300
-      assert_receive %Message{topic: ^topic, event: "system"}, 5000
+
+      assert_receive %Message{
+                       topic: ^topic,
+                       event: "system",
+                       payload: %{"extension" => "postgres_changes", "status" => "ok"}
+                     },
+                     5000
 
       tenant = Tenants.get_tenant_by_external_id(tenant.external_id)
       {:ok, conn} = Database.connect(tenant, "realtime_test", :stop)
@@ -221,7 +239,7 @@ defmodule Realtime.Integration.RtChannel.BillableEventsTest do
                          event: "postgres_changes",
                          payload: %{"data" => %{"schema" => "public", "table" => "test", "type" => "INSERT"}}
                        },
-                       5000
+                       500
       end
 
       # Wait for RateCounter to run
@@ -249,7 +267,13 @@ defmodule Realtime.Integration.RtChannel.BillableEventsTest do
 
       # Join events
       assert_receive %Message{event: "phx_reply", payload: %{"status" => "ok"}, topic: ^topic}, 300
-      assert_receive %Message{topic: ^topic, event: "system"}, 5000
+
+      assert_receive %Message{
+                       topic: ^topic,
+                       event: "system",
+                       payload: %{"extension" => "postgres_changes", "status" => "error"}
+                     },
+                     5000
 
       # Wait for RateCounter to run
       RateCounterHelper.tick_tenant_rate_counters!(tenant.external_id)

test/realtime_web/channels/realtime_channel_replication_ready_test.exs

@@ -23,57 +23,69 @@ defmodule RealtimeWeb.RealtimeChannelReplicationReadyTest do
     assert_receive %Socket.Message{event: "system", payload: %{message: "Replication connection established"}}, 500
   end
 
-  test "pushes the system message when the syn ready broadcast arrives after join", %{tenant: tenant} do
-    stub(Connect, :lookup_or_start_connection, fn _ -> {:ok, self()} end)
-    stub(Connect, :replication_status, fn _ -> {:error, :not_connected} end)
+  test "pushes the system message once replication becomes ready while polling", %{tenant: tenant} do
+    {:ok, counter} = Agent.start_link(fn -> 0 end)
 
-    assert {:ok, _, _} = join(tenant)
+    stub(Connect, :lookup_or_start_connection, fn _ -> {:ok, self()} end)
 
-    refute_receive %Socket.Message{event: "system", payload: %{message: "Replication connection established"}}, 200
+    stub(Connect, :replication_status, fn _ ->
+      case Agent.get_and_update(counter, fn n -> {n, n + 1} end) do
+        n when n < 3 -> {:error, :not_connected}
+        _ -> {:ok, self()}
+      end
+    end)
 
-    signal_ready(tenant, self())
+    assert {:ok, _, _} = join(tenant)
 
     assert_receive %Socket.Message{event: "system", payload: %{message: "Replication connection established"}}, 500
   end
 
-  test "ignores syn ready broadcasts without a replication connection", %{tenant: tenant} do
+  test "does not push while replication is unavailable", %{tenant: tenant} do
     stub(Connect, :lookup_or_start_connection, fn _ -> {:ok, self()} end)
     stub(Connect, :replication_status, fn _ -> {:error, :not_connected} end)
 
     assert {:ok, _, _} = join(tenant)
 
-    signal_ready(tenant, nil)
-
     refute_receive %Socket.Message{event: "system", payload: %{message: "Replication connection established"}}, 300
   end
 
-  test "notifies at most once and stops listening after the first signal", %{tenant: tenant} do
+  test "notifies at most once", %{tenant: tenant} do
     stub(Connect, :lookup_or_start_connection, fn _ -> {:ok, self()} end)
     stub(Connect, :replication_status, fn _ -> {:ok, self()} end)
 
     assert {:ok, _, _} = join(tenant)
 
     assert_receive %Socket.Message{event: "system", payload: %{message: "Replication connection established"}}, 500
 
-    signal_ready(tenant, self())
-
     refute_receive %Socket.Message{event: "system", payload: %{message: "Replication connection established"}}, 300
   end
 
+  test "stops the channel when replication is not established before the timeout", %{tenant: tenant} do
+    previous = Application.get_env(:realtime, :replication_ready_timeout)
+    Application.put_env(:realtime, :replication_ready_timeout, 50)
+    on_exit(fn -> Application.put_env(:realtime, :replication_ready_timeout, previous) end)
+
+    stub(Connect, :lookup_or_start_connection, fn _ -> {:ok, self()} end)
+    stub(Connect, :replication_status, fn _ -> {:error, :not_connected} end)
+
+    assert {:ok, _, socket} = join(tenant)
+    ref = Process.monitor(socket.channel_pid)
+
+    assert_receive %Socket.Message{
+                     event: "system",
+                     payload: %{status: "error", message: "Replication connection was not established in time"}
+                   },
+                   500
+
+    assert_receive {:DOWN, ^ref, :process, _, _}, 500
+  end
+
   defp join(tenant) do
     jwt = generate_jwt_token(tenant)
     {:ok, socket} = connect(UserSocket, %{}, conn_opts(tenant, jwt))
     subscribe_and_join(socket, "realtime:test", %{"config" => %{}})
   end
 
-  defp signal_ready(tenant, replication_conn) do
-    RealtimeWeb.Endpoint.local_broadcast(
-      Connect.syn_topic(tenant.external_id),
-      "ready",
-      %{pid: self(), conn: self(), replication_conn: replication_conn}
-    )
-  end
-
   defp conn_opts(tenant, token) do
     [
       connect_info: %{