feat: speed up messages replication by avoiding re-encoding JSON (#1602)

* feat: speed up messages replication by avoiding re-encoding JSON

We can use the already JSON encoded payload and pass a Jason.Fragment
instead of decoding to then re-encode when the message gets sent to a
websocket

* fix: include payload id if not defined on realtime.send
* chore: increase timeout for socket monitor message

Author: edgurgel

lib/realtime/tenants/migrations.ex

@@ -75,7 +75,8 @@ defmodule Realtime.Tenants.Migrations do
     SubscriptionIndexBridgingDisabled,
     RunSubscriptionIndexBridgingDisabled,
     BroadcastSendErrorLogging,
-    CreateMessagesReplayIndex
+    CreateMessagesReplayIndex,
+    BroadcastSendIncludePayloadId
   }
 
   @migrations [
@@ -142,7 +143,8 @@ defmodule Realtime.Tenants.Migrations do
     {20_250_506_224_012, SubscriptionIndexBridgingDisabled},
     {20_250_523_164_012, RunSubscriptionIndexBridgingDisabled},
     {20_250_714_121_412, BroadcastSendErrorLogging},
-    {20_250_905_041_441, CreateMessagesReplayIndex}
+    {20_250_905_041_441, CreateMessagesReplayIndex},
+    {20_251_103_001_201, BroadcastSendIncludePayloadId}
   ]
 
   defstruct [:tenant_external_id, :settings]

lib/realtime/tenants/replication_connection.ex

@@ -159,6 +159,7 @@ defmodule Realtime.Tenants.ReplicationConnection do
 
   @impl true
   def init(%__MODULE__{tenant_id: tenant_id, monitored_pid: monitored_pid} = state) do
+    Process.flag(:fullsweep_after, 20)
     Logger.metadata(external_id: tenant_id, project: tenant_id)
     Process.monitor(monitored_pid)
 
@@ -315,7 +316,7 @@ defmodule Realtime.Tenants.ReplicationConnection do
            topic: topic,
            event: event,
            private: private,
-           payload: Map.put_new(payload, "id", id)
+           payload: Jason.Fragment.new(payload)
          },
          :ok <- BatchBroadcast.broadcast(nil, tenant, %{messages: [broadcast_message]}, true) do
       inserted_at = NaiveDateTime.from_iso8601!(inserted_at)
@@ -381,7 +382,6 @@ defmodule Realtime.Tenants.ReplicationConnection do
     |> Enum.zip(columns)
     |> Map.new(fn
       {nil, %{name: name}} -> {name, nil}
-      {value, %{name: name, type: "jsonb"}} -> {name, Jason.decode!(value)}
       {value, %{name: name, type: "bool"}} -> {name, value == "t"}
       {value, %{name: name}} -> {name, value}
     end)

lib/realtime/tenants/repo/migrations/20251103001201_broadcast_send_include_payload_id.ex

@@ -0,0 +1,41 @@
+defmodule Realtime.Tenants.Migrations.BroadcastSendIncludePayloadId do
+  @moduledoc false
+  use Ecto.Migration
+
+  # Include ID in the payload if not defined
+  def change do
+    execute("""
+    CREATE OR REPLACE FUNCTION realtime.send(payload jsonb, event text, topic text, private boolean DEFAULT true ) RETURNS void
+    AS $$
+    DECLARE
+      generated_id uuid;
+      final_payload jsonb;
+    BEGIN
+      BEGIN
+        -- Generate a new UUID for the id
+        generated_id := gen_random_uuid();
+
+        -- Check if payload has an 'id' key, if not, add the generated UUID
+        IF payload ? 'id' THEN
+          final_payload := payload;
+        ELSE
+          final_payload := jsonb_set(payload, '{id}', to_jsonb(generated_id));
+        END IF;
+
+        -- Set the topic configuration
+        EXECUTE format('SET LOCAL realtime.topic TO %L', topic);
+
+        -- Attempt to insert the message
+        INSERT INTO realtime.messages (id, payload, event, topic, private, extension)
+        VALUES (generated_id, final_payload, event, topic, private, 'broadcast');
+      EXCEPTION
+        WHEN OTHERS THEN
+          -- Capture and notify the error
+          RAISE WARNING 'ErrorSendingBroadcastMessage: %', SQLERRM;
+      END;
+    END;
+    $$
+    LANGUAGE plpgsql;
+    """)
+  end
+end

mix.exs

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

test/integration/rt_channel_test.exs

@@ -2589,7 +2589,7 @@ defmodule Realtime.Integration.RtChannelTest do
     Realtime.Tenants.Cache.invalidate_tenant_cache(external_id)
   end
 
-  defp assert_process_down(pid, timeout \\ 300) do
+  defp assert_process_down(pid, timeout \\ 1000) do
     ref = Process.monitor(pid)
     assert_receive {:DOWN, ^ref, :process, ^pid, _reason}, timeout
   end

test/realtime/tenants/replication_connection_test.exs

@@ -22,9 +22,8 @@ defmodule Realtime.Tenants.ReplicationConnectionTest do
     {:ok, db_conn} = Database.connect(tenant, "realtime_test", :stop)
     name = "supabase_realtime_messages_replication_slot_test"
     Postgrex.query(db_conn, "SELECT pg_drop_replication_slot($1)", [name])
-    Process.exit(db_conn, :normal)
 
-    %{tenant: tenant}
+    %{tenant: tenant, db_conn: db_conn}
   end
 
   describe "temporary process" do
@@ -70,7 +69,7 @@ defmodule Realtime.Tenants.ReplicationConnectionTest do
       assert {:error, _} = ReplicationConnection.start(tenant, self())
     end
 
-    test "starts a handler for the tenant and broadcasts", %{tenant: tenant} do
+    test "starts a handler for the tenant and broadcasts", %{tenant: tenant, db_conn: db_conn} do
       start_link_supervised!(
         {ReplicationConnection, %ReplicationConnection{tenant_id: tenant.external_id, monitored_pid: self()}},
         restart: :transient
@@ -100,7 +99,6 @@ defmodule Realtime.Tenants.ReplicationConnectionTest do
           "event" => "INSERT",
           "meta" => %{"id" => row.id},
           "payload" => %{
-            "id" => row.id,
             "value" => value
           },
           "type" => "broadcast"
@@ -122,7 +120,6 @@ defmodule Realtime.Tenants.ReplicationConnectionTest do
           })
         end
 
-      {:ok, db_conn} = Database.connect(tenant, "realtime_test", :stop)
       {:ok, _} = Realtime.Repo.insert_all_entries(db_conn, messages, Message)
 
       messages_received =
@@ -140,9 +137,8 @@ defmodule Realtime.Tenants.ReplicationConnectionTest do
                      "event" => "broadcast",
                      "payload" => %{
                        "event" => "INSERT",
-                       "meta" => %{"id" => id},
+                       "meta" => %{"id" => _id},
                        "payload" => %{
-                         "id" => id,
                          "value" => ^value
                        }
                      },
@@ -231,7 +227,7 @@ defmodule Realtime.Tenants.ReplicationConnectionTest do
       assert logs =~ "UnableToBroadcastChanges: %{messages: [%{payload: [\"Payload size exceeds tenant limit\"]}]}"
     end
 
-    test "payload without id", %{tenant: tenant} do
+    test "payload without id", %{tenant: tenant, db_conn: db_conn} do
       start_link_supervised!(
         {ReplicationConnection, %ReplicationConnection{tenant_id: tenant.external_id, monitored_pid: self()}},
         restart: :transient
@@ -241,15 +237,16 @@ defmodule Realtime.Tenants.ReplicationConnectionTest do
       tenant_topic = Tenants.tenant_topic(tenant.external_id, topic, false)
       subscribe(tenant_topic, topic)
 
-      fixture =
-        message_fixture(tenant, %{
-          "topic" => topic,
-          "private" => true,
-          "event" => "INSERT",
-          "payload" => %{"value" => "something"}
-        })
+      value = "something"
+      event = "INSERT"
+
+      Postgrex.query!(
+        db_conn,
+        "SELECT realtime.send (json_build_object ('value', $1 :: text)::jsonb, $2 :: text, $3 :: text, TRUE::bool);",
+        [value, event, topic]
+      )
 
-      fixture_id = fixture.id
+      {:ok, [%{id: id}]} = Repo.all(db_conn, from(m in Message), Message)
 
       assert_receive {:socket_push, :text, data}, 500
       message = data |> IO.iodata_to_binary() |> Jason.decode!()
@@ -258,7 +255,7 @@ defmodule Realtime.Tenants.ReplicationConnectionTest do
                "event" => "broadcast",
                "payload" => %{
                  "event" => "INSERT",
-                 "meta" => %{"id" => ^fixture_id},
+                 "meta" => %{"id" => ^id},
                  "payload" => payload,
                  "type" => "broadcast"
                },
@@ -268,11 +265,11 @@ defmodule Realtime.Tenants.ReplicationConnectionTest do
 
       assert payload == %{
                "value" => "something",
-               "id" => fixture_id
+               "id" => id
              }
     end
 
-    test "payload including id", %{tenant: tenant} do
+    test "payload including id", %{tenant: tenant, db_conn: db_conn} do
       start_link_supervised!(
         {ReplicationConnection, %ReplicationConnection{tenant_id: tenant.external_id, monitored_pid: self()}},
         restart: :transient
@@ -282,25 +279,27 @@ defmodule Realtime.Tenants.ReplicationConnectionTest do
       tenant_topic = Tenants.tenant_topic(tenant.external_id, topic, false)
       subscribe(tenant_topic, topic)
 
-      payload = %{"value" => "something", "id" => "123456"}
+      id = "123456"
+      value = "something"
+      event = "INSERT"
 
-      %{id: fixture_id} =
-        message_fixture(tenant, %{
-          "topic" => topic,
-          "private" => true,
-          "event" => "INSERT",
-          "payload" => payload
-        })
+      Postgrex.query!(
+        db_conn,
+        "SELECT realtime.send (json_build_object ('value', $1 :: text, 'id', $2 :: text)::jsonb, $3 :: text, $4 :: text, TRUE::bool);",
+        [value, id, event, topic]
+      )
+
+      {:ok, [%{id: message_id}]} = Repo.all(db_conn, from(m in Message), Message)
 
       assert_receive {:socket_push, :text, data}, 500
       message = data |> IO.iodata_to_binary() |> Jason.decode!()
 
       assert %{
                "event" => "broadcast",
                "payload" => %{
-                 "meta" => %{"id" => ^fixture_id},
+                 "meta" => %{"id" => ^message_id},
                  "event" => "INSERT",
-                 "payload" => ^payload,
+                 "payload" => %{"value" => "something", "id" => ^id},
                  "type" => "broadcast"
                },
                "ref" => nil,