fix: handle wal bloat (#1528)

Verify that replication connection is able to reconnect when faced with WAL bloat issues

Author: filipecabaco

lib/realtime/tenants/connect.ex

@@ -252,31 +252,10 @@ defmodule Realtime.Tenants.Connect do
   end
 
   def handle_continue(:start_replication, state) do
-    %{tenant: tenant} = state
-
-    with {:ok, replication_connection_pid} <- ReplicationConnection.start(tenant, self()) do
-      replication_connection_reference = Process.monitor(replication_connection_pid)
-
-      state = %{
-        state
-        | replication_connection_pid: replication_connection_pid,
-          replication_connection_reference: replication_connection_reference
-      }
-
-      {:noreply, state, {:continue, :setup_connected_user_events}}
-    else
-      {:error, :max_wal_senders_reached} ->
-        log_error("ReplicationMaxWalSendersReached", "Tenant database has reached the maximum number of WAL senders")
-        {:stop, :shutdown, state}
-
-      {:error, error} ->
-        log_error("StartReplicationFailed", error)
-        {:stop, :shutdown, state}
+    case start_replication_connection(state) do
+      {:ok, state} -> {:noreply, state, {:continue, :setup_connected_user_events}}
+      {:error, state} -> {:stop, :shutdown, state}
     end
-  rescue
-    error ->
-      log_error("StartReplicationFailed", error)
-      {:stop, :shutdown, state}
   end
 
   def handle_continue(:setup_connected_user_events, state) do
@@ -348,13 +327,30 @@ defmodule Realtime.Tenants.Connect do
     {:stop, :shutdown, state}
   end
 
+  @replication_recovery_backoff 1000
+
   # Handle replication connection termination
   def handle_info(
         {:DOWN, replication_connection_reference, _, _, _},
         %{replication_connection_reference: replication_connection_reference} = state
       ) do
-    Logger.warning("Replication connection has died")
-    {:stop, :shutdown, state}
+    log_warning("ReplicationConnectionDown", "Replication connection has been terminated")
+    Process.send_after(self(), :recover_replication_connection, @replication_recovery_backoff)
+    state = %{state | replication_connection_pid: nil, replication_connection_reference: nil}
+    {:noreply, state}
+  end
+
+  @replication_connection_query "SELECT 1 from pg_stat_activity where application_name='realtime_replication_connection'"
+  def handle_info(:recover_replication_connection, state) do
+    with %{num_rows: 0} <- Postgrex.query!(state.db_conn_pid, @replication_connection_query, []),
+         {:ok, state} <- start_replication_connection(state) do
+      {:noreply, state}
+    else
+      _ ->
+        log_error("ReplicationConnectionRecoveryFailed", "Replication connection recovery failed")
+        Process.send_after(self(), :recover_replication_connection, @replication_recovery_backoff)
+        {:noreply, state}
+    end
   end
 
   def handle_info(_, state), do: {:noreply, state}
@@ -414,4 +410,32 @@ defmodule Realtime.Tenants.Connect do
   defp tenant_suspended?(_), do: :ok
 
   defp rebalance_check_interval_in_ms(), do: Application.fetch_env!(:realtime, :rebalance_check_interval_in_ms)
+
+  defp start_replication_connection(state) do
+    %{tenant: tenant} = state
+
+    with {:ok, replication_connection_pid} <- ReplicationConnection.start(tenant, self()) do
+      replication_connection_reference = Process.monitor(replication_connection_pid)
+
+      state = %{
+        state
+        | replication_connection_pid: replication_connection_pid,
+          replication_connection_reference: replication_connection_reference
+      }
+
+      {:ok, state}
+    else
+      {:error, :max_wal_senders_reached} ->
+        log_error("ReplicationMaxWalSendersReached", "Tenant database has reached the maximum number of WAL senders")
+        {:error, state}
+
+      {:error, error} ->
+        log_error("StartReplicationFailed", error)
+        {:error, state}
+    end
+  rescue
+    error ->
+      log_error("StartReplicationFailed", error)
+      {:error, state}
+  end
 end

lib/realtime/tenants/replication_connection.ex

@@ -144,8 +144,8 @@ defmodule Realtime.Tenants.ReplicationConnection do
         port: connection_opts.port,
         socket_options: connection_opts.socket_options,
         ssl: connection_opts.ssl,
-        backoff_type: :stop,
         sync_connect: true,
+        auto_reconnect: false,
         parameters: [application_name: "realtime_replication_connection"]
       ]
 

mix.exs

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

test/integration/rt_channel_test.exs

@@ -25,6 +25,7 @@ defmodule Realtime.Integration.RtChannelTest do
   alias Realtime.Tenants
   alias Realtime.Tenants.Authorization
   alias Realtime.Tenants.Connect
+  alias Realtime.Tenants.ReplicationConnection
 
   alias RealtimeWeb.RealtimeChannel.Tracker
   alias RealtimeWeb.SocketDisconnect
@@ -2354,6 +2355,135 @@ defmodule Realtime.Integration.RtChannelTest do
     assert count == 2
   end
 
+  describe "WAL bloat handling" do
+    setup %{tenant: tenant} do
+      topic = random_string()
+      {:ok, db_conn} = Database.connect(tenant, "realtime_test", :stop)
+
+      %{rows: [[max_wal_size]]} = Postgrex.query!(db_conn, "SHOW max_wal_size", [])
+      %{rows: [[wal_keep_size]]} = Postgrex.query!(db_conn, "SHOW wal_keep_size", [])
+      %{rows: [[max_slot_wal_keep_size]]} = Postgrex.query!(db_conn, "SHOW max_slot_wal_keep_size", [])
+
+      assert max_wal_size == "32MB"
+      assert wal_keep_size == "32MB"
+      assert max_slot_wal_keep_size == "32MB"
+
+      Postgrex.query!(db_conn, "CREATE TABLE IF NOT EXISTS wal_test (id INT, data TEXT)", [])
+
+      Postgrex.query!(
+        db_conn,
+        """
+          CREATE OR REPLACE FUNCTION wal_test_trigger_func() RETURNS TRIGGER AS $$
+          BEGIN
+            PERFORM realtime.send(json_build_object ('value', 'test' :: text)::jsonb, 'test', '#{topic}', false);
+            RETURN NULL;
+          END;
+          $$ LANGUAGE plpgsql;
+        """,
+        []
+      )
+
+      Postgrex.query!(db_conn, "DROP TRIGGER IF EXISTS wal_test_trigger ON wal_test", [])
+
+      Postgrex.query!(
+        db_conn,
+        """
+          CREATE TRIGGER wal_test_trigger
+          AFTER INSERT OR UPDATE OR DELETE ON wal_test
+          FOR EACH ROW
+          EXECUTE FUNCTION wal_test_trigger_func()
+        """,
+        []
+      )
+
+      GenServer.stop(db_conn)
+
+      on_exit(fn ->
+        {:ok, db_conn} = Database.connect(tenant, "realtime_test", :stop)
+
+        Postgrex.query!(db_conn, "DROP TABLE IF EXISTS wal_test CASCADE", [])
+      end)
+
+      %{topic: topic}
+    end
+
+    test "track PID changes during WAL bloat creation", %{tenant: tenant, topic: topic} do
+      {socket, _} = get_connection(tenant, "authenticated")
+      config = %{broadcast: %{self: true}, private: false}
+      full_topic = "realtime:#{topic}"
+
+      active_slot_query =
+        "SELECT active_pid FROM pg_replication_slots where active_pid is not null and slot_name = 'supabase_realtime_messages_replication_slot_'"
+
+      WebsocketClient.join(socket, full_topic, %{config: config})
+
+      assert_receive %Message{event: "phx_reply", payload: %{"status" => "ok"}}, 500
+      assert_receive %Message{event: "presence_state"}, 500
+
+      assert Connect.ready?(tenant.external_id)
+
+      {:ok, db_conn} = Connect.lookup_or_start_connection(tenant.external_id)
+
+      original_connect_pid = Connect.whereis(tenant.external_id)
+      original_replication_pid = ReplicationConnection.whereis(tenant.external_id)
+      %{rows: [[original_db_pid]]} = Postgrex.query!(db_conn, active_slot_query, [])
+
+      tasks =
+        for _ <- 1..5 do
+          Task.async(fn ->
+            {:ok, bloat_conn} = Database.connect(tenant, "realtime_bloat", :stop)
+
+            Postgrex.transaction(bloat_conn, fn conn ->
+              Postgrex.query(conn, "INSERT INTO wal_test SELECT generate_series(1, 100000), repeat('x', 2000)", [])
+              {:error, "test"}
+            end)
+
+            Process.exit(bloat_conn, :normal)
+          end)
+        end
+
+      Task.await_many(tasks, 20000)
+
+      # Kill all pending transactions still running
+      Postgrex.query!(
+        db_conn,
+        "SELECT pg_terminate_backend(pid) from pg_stat_activity where application_name='realtime_bloat'",
+        []
+      )
+
+      # Does it recover?
+      assert Connect.ready?(tenant.external_id)
+      {:ok, db_conn} = Connect.lookup_or_start_connection(tenant.external_id)
+      Process.sleep(1000)
+      %{rows: [[new_db_pid]]} = Postgrex.query!(db_conn, active_slot_query, [])
+
+      assert new_db_pid != original_db_pid
+      assert ^original_connect_pid = Connect.whereis(tenant.external_id)
+      assert original_replication_pid != ReplicationConnection.whereis(tenant.external_id)
+
+      # Check if socket is still connected
+      payload = %{"event" => "TEST", "payload" => %{"msg" => 1}, "type" => "broadcast"}
+      WebsocketClient.send_event(socket, full_topic, "broadcast", payload)
+      assert_receive %Message{event: "broadcast", payload: ^payload, topic: ^full_topic}, 500
+
+      # Check if we are receiving the message from replication connection
+      Postgrex.query!(db_conn, "INSERT INTO wal_test VALUES (1, 'test')", [])
+
+      assert_receive %Phoenix.Socket.Message{
+                       event: "broadcast",
+                       payload: %{
+                         "event" => "test",
+                         "payload" => %{"value" => "test"},
+                         "type" => "broadcast"
+                       },
+                       join_ref: nil,
+                       ref: nil,
+                       topic: ^full_topic
+                     },
+                     5000
+    end
+  end
+
   defp mode(%{mode: :distributed}) do
     tenant = Api.get_tenant_by_external_id("dev_tenant")
 

test/realtime/tenants/connect_test.exs

@@ -348,11 +348,13 @@ defmodule Realtime.Tenants.ConnectTest do
       assert replication_connection_before == replication_connection_after
     end
 
-    test "on replication connection postgres pid being stopped, also kills the Connect module", %{tenant: tenant} do
+    test "on replication connection postgres pid being stopped, Connect module recovers it", %{tenant: tenant} do
       assert {:ok, db_conn} = Connect.lookup_or_start_connection(tenant.external_id)
       assert Connect.ready?(tenant.external_id)
 
       replication_connection_pid = ReplicationConnection.whereis(tenant.external_id)
+      Process.monitor(replication_connection_pid)
+
       assert Process.alive?(replication_connection_pid)
       pid = Connect.whereis(tenant.external_id)
 
@@ -362,21 +364,33 @@ defmodule Realtime.Tenants.ConnectTest do
         []
       )
 
-      assert_process_down(replication_connection_pid)
-      assert_process_down(pid)
+      assert_receive {:DOWN, _, :process, ^replication_connection_pid, _}
+
+      Process.sleep(1500)
+      new_replication_connection_pid = ReplicationConnection.whereis(tenant.external_id)
+
+      assert replication_connection_pid != new_replication_connection_pid
+      assert Process.alive?(new_replication_connection_pid)
+      assert Process.alive?(pid)
     end
 
-    test "on replication connection exit, also kills the Connect module", %{tenant: tenant} do
+    test "on replication connection exit, Connect module recovers it", %{tenant: tenant} do
       assert {:ok, _db_conn} = Connect.lookup_or_start_connection(tenant.external_id)
       assert Connect.ready?(tenant.external_id)
 
       replication_connection_pid = ReplicationConnection.whereis(tenant.external_id)
+      Process.monitor(replication_connection_pid)
       assert Process.alive?(replication_connection_pid)
       pid = Connect.whereis(tenant.external_id)
       Process.exit(replication_connection_pid, :kill)
+      assert_receive {:DOWN, _, :process, ^replication_connection_pid, _}
 
-      assert_process_down(replication_connection_pid)
-      assert_process_down(pid)
+      Process.sleep(1500)
+      new_replication_connection_pid = ReplicationConnection.whereis(tenant.external_id)
+
+      assert replication_connection_pid != new_replication_connection_pid
+      assert Process.alive?(new_replication_connection_pid)
+      assert Process.alive?(pid)
     end
 
     test "handles max_wal_senders by logging the correct operational code", %{tenant: tenant} do

test/realtime/tenants/replication_connection_test.exs

@@ -331,6 +331,26 @@ defmodule Realtime.Tenants.ReplicationConnectionTest do
 
       assert {:error, :max_wal_senders_reached} = ReplicationConnection.start(tenant, self())
     end
+
+    test "handles WAL pressure gracefully", %{tenant: tenant} do
+      {:ok, replication_pid} = ReplicationConnection.start(tenant, self())
+
+      {:ok, conn} = Database.connect(tenant, "realtime_test", :stop)
+      on_exit(fn -> Process.exit(conn, :normal) end)
+
+      large_payload = String.duplicate("x", 10 * 1024 * 1024)
+
+      for i <- 1..5 do
+        message_fixture_with_conn(tenant, conn, %{
+          "topic" => "stress_#{i}",
+          "private" => true,
+          "event" => "INSERT",
+          "payload" => %{"data" => large_payload}
+        })
+      end
+
+      assert Process.alive?(replication_pid)
+    end
   end
 
   describe "whereis/1" do
@@ -409,4 +429,20 @@ defmodule Realtime.Tenants.ReplicationConnectionTest do
     ref = Process.monitor(pid)
     assert_receive {:DOWN, ^ref, :process, ^pid, _reason}, timeout
   end
+
+  defp message_fixture_with_conn(_tenant, conn, override) do
+    create_attrs = %{
+      "topic" => random_string(),
+      "extension" => "broadcast"
+    }
+
+    override = override |> Enum.map(fn {k, v} -> {"#{k}", v} end) |> Map.new()
+
+    {:ok, message} =
+      create_attrs
+      |> Map.merge(override)
+      |> TenantConnection.create_message(conn)
+
+    message
+  end
 end

test/support/containers.ex

@@ -267,7 +267,13 @@ defmodule Containers do
         @image,
         "postgres",
         "-c",
-        "config_file=/etc/postgresql/postgresql.conf"
+        "config_file=/etc/postgresql/postgresql.conf",
+        "-c",
+        "wal_keep_size=32MB",
+        "-c",
+        "max_wal_size=32MB",
+        "-c",
+        "max_slot_wal_keep_size=32MB"
       ])
   end
 end