feat: selective broadcast for postgres changes (#1573)

Selective broadcast for Postgres Changes. We use the fact that we know where each subscription subscribed from and do a direct broadcast for each node instead of spamming the whole cluster.

If the subscription node is not readily available (for whatever reason) we fallback to sending to the whole cluster.

Author: edgurgel

lib/extensions/postgres_cdc_rls/replication_poller.ex

@@ -21,15 +21,15 @@ defmodule Extensions.PostgresCdcRls.ReplicationPoller do
   alias Realtime.RateCounter
   alias Realtime.Tenants
 
+  alias RealtimeWeb.TenantBroadcaster
+
   def start_link(opts), do: GenServer.start_link(__MODULE__, opts)
 
   @impl true
   def init(args) do
     tenant_id = args["id"]
     Logger.metadata(external_id: tenant_id, project: tenant_id)
 
-    %Realtime.Api.Tenant{} = Tenants.Cache.get_tenant_by_external_id(tenant_id)
-
     rate_counter_args = Tenants.db_events_per_second_rate(tenant_id, 4000)
 
     RateCounter.new(rate_counter_args)
@@ -50,7 +50,8 @@ defmodule Extensions.PostgresCdcRls.ReplicationPoller do
       retry_count: 0,
       slot_name: args["slot_name"] <> slot_name_suffix(),
       tenant_id: tenant_id,
-      rate_counter_args: rate_counter_args
+      rate_counter_args: rate_counter_args,
+      subscribers_nodes_table: args["subscribers_nodes_table"]
     }
 
     {:ok, _} = Registry.register(__MODULE__.Registry, tenant_id, %{})
@@ -84,6 +85,7 @@ defmodule Extensions.PostgresCdcRls.ReplicationPoller do
           max_changes: max_changes,
           conn: conn,
           tenant_id: tenant_id,
+          subscribers_nodes_table: subscribers_nodes_table,
           rate_counter_args: rate_counter_args
         } = state
       ) do
@@ -94,7 +96,7 @@ defmodule Extensions.PostgresCdcRls.ReplicationPoller do
     {time, list_changes} = :timer.tc(Replications, :list_changes, args)
     record_list_changes_telemetry(time, tenant_id)
 
-    case handle_list_changes_result(list_changes, tenant_id, rate_counter_args) do
+    case handle_list_changes_result(list_changes, subscribers_nodes_table, tenant_id, rate_counter_args) do
       {:ok, row_count} ->
         Backoff.reset(backoff)
 
@@ -187,6 +189,7 @@ defmodule Extensions.PostgresCdcRls.ReplicationPoller do
             rows: [_ | _] = rows,
             num_rows: rows_count
           }},
+         subscribers_nodes_table,
          tenant_id,
          rate_counter_args
        ) do
@@ -201,15 +204,47 @@ defmodule Extensions.PostgresCdcRls.ReplicationPoller do
             change <- columns |> Enum.zip(row) |> generate_record() |> List.wrap() do
           topic = "realtime:postgres:" <> tenant_id
 
-          RealtimeWeb.TenantBroadcaster.pubsub_broadcast(tenant_id, topic, change, MessageDispatcher, :postgres_changes)
+          case collect_subscription_nodes(subscribers_nodes_table, change.subscription_ids) do
+            {:ok, nodes} ->
+              for node <- nodes do
+                TenantBroadcaster.pubsub_direct_broadcast(
+                  node,
+                  tenant_id,
+                  topic,
+                  change,
+                  MessageDispatcher,
+                  :postgres_changes
+                )
+              end
+
+            {:error, :node_not_found} ->
+              TenantBroadcaster.pubsub_broadcast(
+                tenant_id,
+                topic,
+                change,
+                MessageDispatcher,
+                :postgres_changes
+              )
+          end
         end
     end
 
     {:ok, rows_count}
   end
 
-  defp handle_list_changes_result({:ok, _}, _, _), do: {:ok, 0}
-  defp handle_list_changes_result({:error, reason}, _, _), do: {:error, reason}
+  defp handle_list_changes_result({:ok, _}, _, _, _), do: {:ok, 0}
+  defp handle_list_changes_result({:error, reason}, _, _, _), do: {:error, reason}
+
+  defp collect_subscription_nodes(subscribers_nodes_table, subscription_ids) do
+    Enum.reduce_while(subscription_ids, {:ok, MapSet.new()}, fn subscription_id, {:ok, acc} ->
+      case :ets.lookup(subscribers_nodes_table, subscription_id) do
+        [{_, node}] -> {:cont, {:ok, MapSet.put(acc, node)}}
+        _ -> {:halt, {:error, :node_not_found}}
+      end
+    end)
+  rescue
+    _ -> {:error, :node_not_found}
+  end
 
   def generate_record([
         {"wal",

lib/extensions/postgres_cdc_rls/subscription_manager.ex

@@ -24,7 +24,8 @@ defmodule Extensions.PostgresCdcRls.SubscriptionManager do
     defstruct [
       :id,
       :publication,
-      :subscribers_tid,
+      :subscribers_pids_table,
+      :subscribers_nodes_table,
       :conn,
       :delete_queue,
       :no_users_ref,
@@ -37,7 +38,8 @@ defmodule Extensions.PostgresCdcRls.SubscriptionManager do
     @type t :: %__MODULE__{
             id: String.t(),
             publication: String.t(),
-            subscribers_tid: :ets.tid(),
+            subscribers_pids_table: :ets.tid(),
+            subscribers_nodes_table: :ets.tid(),
             conn: Postgrex.conn(),
             oids: map(),
             check_oid_ref: reference() | nil,
@@ -67,7 +69,12 @@ defmodule Extensions.PostgresCdcRls.SubscriptionManager do
 
   @impl true
   def handle_continue({:connect, args}, _) do
-    %{"id" => id, "publication" => publication, "subscribers_tid" => subscribers_tid} = args
+    %{
+      "id" => id,
+      "publication" => publication,
+      "subscribers_pids_table" => subscribers_pids_table,
+      "subscribers_nodes_table" => subscribers_nodes_table
+    } = args
 
     subscription_manager_settings = Database.from_settings(args, "realtime_subscription_manager")
 
@@ -85,31 +92,35 @@ defmodule Extensions.PostgresCdcRls.SubscriptionManager do
     check_region_interval = Map.get(args, :check_region_interval, rebalance_check_interval_in_ms())
     send_region_check_message(check_region_interval)
 
-    state = %State{
-      id: id,
-      conn: conn,
-      publication: publication,
-      subscribers_tid: subscribers_tid,
-      oids: oids,
-      delete_queue: %{
-        ref: check_delete_queue(),
-        queue: :queue.new()
-      },
-      no_users_ref: check_no_users(),
-      check_region_interval: check_region_interval
-    }
+    state =
+      %State{
+        id: id,
+        conn: conn,
+        publication: publication,
+        subscribers_pids_table: subscribers_pids_table,
+        subscribers_nodes_table: subscribers_nodes_table,
+        oids: oids,
+        delete_queue: %{
+          ref: check_delete_queue(),
+          queue: :queue.new()
+        },
+        no_users_ref: check_no_users(),
+        check_region_interval: check_region_interval
+      }
 
     send(self(), :check_oids)
     {:noreply, state}
   end
 
   @impl true
   def handle_info({:subscribed, {pid, id}}, state) do
-    case :ets.match(state.subscribers_tid, {pid, id, :"$1", :_}) do
-      [] -> :ets.insert(state.subscribers_tid, {pid, id, Process.monitor(pid), node(pid)})
+    case :ets.match(state.subscribers_pids_table, {pid, id, :"$1", :_}) do
+      [] -> :ets.insert(state.subscribers_pids_table, {pid, id, Process.monitor(pid), node(pid)})
       _ -> :ok
     end
 
+    :ets.insert(state.subscribers_nodes_table, {UUID.string_to_binary!(id), node(pid)})
+
     {:noreply, %{state | no_users_ts: nil}}
   end
 
@@ -132,7 +143,7 @@ defmodule Extensions.PostgresCdcRls.SubscriptionManager do
             Process.demonitor(ref, [:flush])
             send(pid, :postgres_subscribe)
           end
-          |> :ets.foldl([], state.subscribers_tid)
+          |> :ets.foldl([], state.subscribers_pids_table)
 
           new_oids
       end
@@ -142,19 +153,25 @@ defmodule Extensions.PostgresCdcRls.SubscriptionManager do
 
   def handle_info(
         {:DOWN, _ref, :process, pid, _reason},
-        %State{subscribers_tid: tid, delete_queue: %{queue: q}} = state
+        %State{
+          subscribers_pids_table: subscribers_pids_table,
+          subscribers_nodes_table: subscribers_nodes_table,
+          delete_queue: %{queue: q}
+        } = state
       ) do
     q1 =
-      case :ets.take(tid, pid) do
+      case :ets.take(subscribers_pids_table, pid) do
         [] ->
           q
 
         values ->
           for {_pid, id, _ref, _node} <- values, reduce: q do
             acc ->
-              id
-              |> UUID.string_to_binary!()
-              |> :queue.in(acc)
+              bin_id = UUID.string_to_binary!(id)
+
+              :ets.delete(subscribers_nodes_table, bin_id)
+
+              :queue.in(bin_id, acc)
           end
       end
 
@@ -187,7 +204,7 @@ defmodule Extensions.PostgresCdcRls.SubscriptionManager do
     {:noreply, %{state | delete_queue: %{ref: ref, queue: q1}}}
   end
 
-  def handle_info(:check_no_users, %{subscribers_tid: tid, no_users_ts: ts} = state) do
+  def handle_info(:check_no_users, %{subscribers_pids_table: tid, no_users_ts: ts} = state) do
     Helpers.cancel_timer(state.no_users_ref)
 
     ts_new =

lib/extensions/postgres_cdc_rls/subscriptions_checker.ex

@@ -17,13 +17,14 @@ defmodule Extensions.PostgresCdcRls.SubscriptionsChecker do
 
   defmodule State do
     @moduledoc false
-    defstruct [:id, :conn, :check_active_pids, :subscribers_tid, :delete_queue]
+    defstruct [:id, :conn, :check_active_pids, :subscribers_pids_table, :subscribers_nodes_table, :delete_queue]
 
     @type t :: %__MODULE__{
             id: String.t(),
             conn: Postgrex.conn(),
             check_active_pids: reference(),
-            subscribers_tid: :ets.tid(),
+            subscribers_pids_table: :ets.tid(),
+            subscribers_nodes_table: :ets.tid(),
             delete_queue: %{
               ref: reference(),
               queue: :queue.queue()
@@ -47,7 +48,11 @@ defmodule Extensions.PostgresCdcRls.SubscriptionsChecker do
 
   @impl true
   def handle_continue({:connect, args}, _) do
-    %{"id" => id, "subscribers_tid" => subscribers_tid} = args
+    %{
+      "id" => id,
+      "subscribers_pids_table" => subscribers_pids_table,
+      "subscribers_nodes_table" => subscribers_nodes_table
+    } = args
 
     realtime_subscription_checker_settings =
       Database.from_settings(args, "realtime_subscription_checker")
@@ -58,7 +63,8 @@ defmodule Extensions.PostgresCdcRls.SubscriptionsChecker do
       id: id,
       conn: conn,
       check_active_pids: check_active_pids(),
-      subscribers_tid: subscribers_tid,
+      subscribers_pids_table: subscribers_pids_table,
+      subscribers_nodes_table: subscribers_nodes_table,
       delete_queue: %{
         ref: nil,
         queue: :queue.new()
@@ -69,18 +75,14 @@ defmodule Extensions.PostgresCdcRls.SubscriptionsChecker do
   end
 
   @impl true
-  def handle_info(
-        :check_active_pids,
-        %State{check_active_pids: ref, subscribers_tid: tid, delete_queue: delete_queue, id: id} =
-          state
-      ) do
+  def handle_info(:check_active_pids, %State{check_active_pids: ref, delete_queue: delete_queue, id: id} = state) do
     Helpers.cancel_timer(ref)
 
     ids =
-      tid
+      state.subscribers_pids_table
       |> subscribers_by_node()
       |> not_alive_pids_dist()
-      |> pop_not_alive_pids(tid, id)
+      |> pop_not_alive_pids(state.subscribers_pids_table, state.subscribers_nodes_table, id)
 
     new_delete_queue =
       if length(ids) > 0 do
@@ -128,10 +130,10 @@ defmodule Extensions.PostgresCdcRls.SubscriptionsChecker do
 
   ## Internal functions
 
-  @spec pop_not_alive_pids([pid()], :ets.tid(), binary()) :: [Ecto.UUID.t()]
-  def pop_not_alive_pids(pids, tid, tenant_id) do
+  @spec pop_not_alive_pids([pid()], :ets.tid(), :ets.tid(), binary()) :: [Ecto.UUID.t()]
+  def pop_not_alive_pids(pids, subscribers_pids_table, subscribers_nodes_table, tenant_id) do
     Enum.reduce(pids, [], fn pid, acc ->
-      case :ets.lookup(tid, pid) do
+      case :ets.lookup(subscribers_pids_table, pid) do
         [] ->
           Telemetry.execute(
             [:realtime, :subscriptions_checker, :pid_not_found],
@@ -149,8 +151,11 @@ defmodule Extensions.PostgresCdcRls.SubscriptionsChecker do
               %{tenant_id: tenant_id}
             )
 
-            :ets.delete(tid, pid)
-            UUID.string_to_binary!(postgres_id)
+            :ets.delete(subscribers_pids_table, pid)
+            bin_id = UUID.string_to_binary!(postgres_id)
+
+            :ets.delete(subscribers_nodes_table, bin_id)
+            bin_id
           end ++ acc
       end
     end)

lib/extensions/postgres_cdc_rls/worker_supervisor.ex

@@ -19,12 +19,19 @@ defmodule Extensions.PostgresCdcRls.WorkerSupervisor do
     Logger.metadata(external_id: tenant, project: tenant)
     unless Api.get_tenant_by_external_id(tenant, :primary), do: raise(Exception)
 
-    tid_args = Map.merge(args, %{"subscribers_tid" => :ets.new(__MODULE__, [:public, :bag])})
+    subscribers_pids_table = :ets.new(__MODULE__, [:public, :bag])
+    subscribers_nodes_table = :ets.new(__MODULE__, [:public, :set])
+
+    tid_args =
+      Map.merge(args, %{
+        "subscribers_pids_table" => subscribers_pids_table,
+        "subscribers_nodes_table" => subscribers_nodes_table
+      })
 
     children = [
       %{
         id: ReplicationPoller,
-        start: {ReplicationPoller, :start_link, [args]},
+        start: {ReplicationPoller, :start_link, [tid_args]},
         restart: :transient
       },
       %{

lib/realtime/gen_rpc.ex

@@ -26,6 +26,36 @@ defmodule Realtime.GenRpc do
     :ok
   end
 
+  @doc """
+  Fire and forget apply(mod, func, args) on one node
+
+  Options:
+
+  - `:key` - Optional key to consistently select the same gen_rpc client to guarantee some message order between nodes
+  """
+  @spec cast(node, module, atom, list(any), keyword()) :: :ok
+  def cast(node, mod, func, args, opts \\ [])
+
+  # Local
+  def cast(node, mod, func, args, _opts) when node == node() do
+    :erpc.cast(node, mod, func, args)
+    :ok
+  end
+
+  def cast(node, mod, func, args, opts)
+      when is_atom(node) and is_atom(mod) and is_atom(func) and is_list(args) and is_list(opts) do
+    key = Keyword.get(opts, :key, nil)
+
+    # Ensure this node is part of the connected nodes
+    if node in Node.list() do
+      node_key = rpc_node(node, key)
+
+      :gen_rpc.cast(node_key, mod, func, args)
+    end
+
+    :ok
+  end
+
   @doc """
   Fire and forget apply(mod, func, args) on all nodes