feat: regional broadcasting (#1580)

When regional_broadcasting is true GenRpcPubSub uses a single node in each region to broadcast to the rest of the region. Nodes communicating inside the same region have a much lower latency.

Author: edgurgel

config/runtime.exs

@@ -71,6 +71,7 @@ broadcast_pool_size = Env.get_integer("BROADCAST_POOL_SIZE", 10)
 pubsub_adapter = System.get_env("PUBSUB_ADAPTER", "gen_rpc") |> String.to_atom()
 websocket_max_heap_size = div(Env.get_integer("WEBSOCKET_MAX_HEAP_SIZE", 50_000_000), :erlang.system_info(:wordsize))
 users_scope_shards = Env.get_integer("USERS_SCOPE_SHARDS", 5)
+regional_broadcasting = Env.get_boolean("REGIONAL_BROADCASTING", false)
 
 no_channel_timeout_in_ms =
   if config_env() == :test,
@@ -128,7 +129,8 @@ config :realtime,
   platform: platform,
   pubsub_adapter: pubsub_adapter,
   broadcast_pool_size: broadcast_pool_size,
-  users_scope_shards: users_scope_shards
+  users_scope_shards: users_scope_shards,
+  regional_broadcasting: regional_broadcasting
 
 if config_env() != :test && run_janitor? do
   config :realtime,

config/test.exs

@@ -31,6 +31,7 @@ config :realtime, RealtimeWeb.Endpoint,
   server: true
 
 config :realtime,
+  regional_broadcasting: true,
   region: "us-east-1",
   db_enc_key: "1234567890123456",
   jwt_claim_validators: System.get_env("JWT_CLAIM_VALIDATORS", "{}"),

lib/realtime/gen_rpc/pub_sub.ex

@@ -5,6 +5,8 @@ defmodule Realtime.GenRpcPubSub do
 
   @behaviour Phoenix.PubSub.Adapter
   alias Realtime.GenRpc
+  alias Realtime.GenRpcPubSub.Worker
+  alias Realtime.Nodes
   use Supervisor
 
   @impl true
@@ -45,36 +47,83 @@ defmodule Realtime.GenRpcPubSub do
   @impl true
   def broadcast(adapter_name, topic, message, dispatcher) do
     worker = worker_name(adapter_name, self())
-    GenRpc.abcast(Node.list(), worker, forward_to_local(topic, message, dispatcher), key: worker)
+
+    if Application.get_env(:realtime, :regional_broadcasting, false) do
+      my_region = Application.get_env(:realtime, :region)
+      # broadcast to all other nodes in the region
+
+      other_nodes = for node <- Realtime.Nodes.region_nodes(my_region), node != node(), do: node
+      GenRpc.abcast(other_nodes, worker, Worker.forward_to_local(topic, message, dispatcher), key: worker)
+
+      # send a message to a node in each region to forward to the rest of the region
+      other_region_nodes = nodes_from_other_regions(my_region, worker)
+
+      GenRpc.abcast(other_region_nodes, worker, Worker.forward_to_region(topic, message, dispatcher), key: worker)
+    else
+      GenRpc.abcast(Node.list(), worker, Worker.forward_to_local(topic, message, dispatcher), key: worker)
+    end
+
+    :ok
+  end
+
+  defp nodes_from_other_regions(my_region, key) do
+    Enum.flat_map(Nodes.all_node_regions(), fn
+      ^my_region ->
+        []
+
+      region ->
+        case Nodes.node_from_region(region, key) do
+          {:ok, node} -> [node]
+          _ -> []
+        end
+    end)
   end
 
   @impl true
   def direct_broadcast(adapter_name, node_name, topic, message, dispatcher) do
     worker = worker_name(adapter_name, self())
-    GenRpc.abcast([node_name], worker, forward_to_local(topic, message, dispatcher), key: worker)
+    GenRpc.abcast([node_name], worker, Worker.forward_to_local(topic, message, dispatcher), key: worker)
   end
-
-  defp forward_to_local(topic, message, dispatcher), do: {:ftl, topic, message, dispatcher}
 end
 
 defmodule Realtime.GenRpcPubSub.Worker do
   @moduledoc false
   use GenServer
 
+  def forward_to_local(topic, message, dispatcher), do: {:ftl, topic, message, dispatcher}
+  def forward_to_region(topic, message, dispatcher), do: {:ftr, topic, message, dispatcher}
+
   @doc false
-  def start_link({pubsub, worker}), do: GenServer.start_link(__MODULE__, pubsub, name: worker)
+  def start_link({pubsub, worker}), do: GenServer.start_link(__MODULE__, {pubsub, worker}, name: worker)
 
   @impl true
-  def init(pubsub) do
+  def init({pubsub, worker}) do
     Process.flag(:message_queue_data, :off_heap)
     Process.flag(:fullsweep_after, 20)
-    {:ok, pubsub}
+    {:ok, {pubsub, worker}}
   end
 
   @impl true
-  def handle_info({:ftl, topic, message, dispatcher}, pubsub) do
+  # Forward to local
+  def handle_info({:ftl, topic, message, dispatcher}, {pubsub, worker}) do
+    Phoenix.PubSub.local_broadcast(pubsub, topic, message, dispatcher)
+    {:noreply, {pubsub, worker}}
+  end
+
+  # Forward to the rest of the region
+  def handle_info({:ftr, topic, message, dispatcher}, {pubsub, worker}) do
+    # Forward to local first
     Phoenix.PubSub.local_broadcast(pubsub, topic, message, dispatcher)
-    {:noreply, pubsub}
+
+    # Then broadcast to the rest of my region
+    my_region = Application.get_env(:realtime, :region)
+    other_nodes = for node <- Realtime.Nodes.region_nodes(my_region), node != node(), do: node
+
+    if other_nodes != [] do
+      Realtime.GenRpc.abcast(other_nodes, worker, forward_to_local(topic, message, dispatcher), key: worker)
+    end
+
+    {:noreply, {pubsub, worker}}
   end
 
   @impl true

lib/realtime/nodes.ex

@@ -64,6 +64,27 @@ defmodule Realtime.Nodes do
 
   def region_nodes(nil), do: []
 
+  @doc """
+  Picks a node from a region based on the provided key
+  """
+  @spec node_from_region(String.t(), term()) :: {:ok, node} | {:error, :not_available}
+  def node_from_region(region, key) when is_binary(region) do
+    nodes = region_nodes(region)
+
+    case nodes do
+      [] ->
+        {:error, :not_available}
+
+      _ ->
+        member_count = Enum.count(nodes)
+        index = :erlang.phash2(key, member_count)
+
+        {:ok, Enum.fetch!(nodes, index)}
+    end
+  end
+
+  def node_from_region(_, _), do: {:error, :not_available}
+
   @doc """
   Picks the node to launch the Postgres connection on.
 
@@ -132,59 +153,9 @@ defmodule Realtime.Nodes do
     end
   end
 
-  @mapping_realtime_region_to_tenant_region_aws %{
-    "ap-southeast-1" => [
-      "ap-east-1",
-      "ap-northeast-1",
-      "ap-northeast-2",
-      "ap-south-1",
-      "ap-southeast-1"
-    ],
-    "ap-southeast-2" => ["ap-southeast-2"],
-    "eu-west-2" => [
-      "eu-central-1",
-      "eu-central-2",
-      "eu-north-1",
-      "eu-west-1",
-      "eu-west-2",
-      "eu-west-3"
-    ],
-    "us-east-1" => [
-      "ca-central-1",
-      "sa-east-1",
-      "us-east-1",
-      "us-east-2"
-    ],
-    "us-west-1" => ["us-west-1", "us-west-2"]
-  }
-  @mapping_realtime_region_to_tenant_region_fly %{
-    "iad" => ["ca-central-1", "sa-east-1", "us-east-1"],
-    "lhr" => ["eu-central-1", "eu-west-1", "eu-west-2", "eu-west-3"],
-    "sea" => ["us-west-1"],
-    "syd" => [
-      "ap-east-1",
-      "ap-northeast-1",
-      "ap-northeast-2",
-      "ap-south-1",
-      "ap-southeast-1",
-      "ap-southeast-2"
-    ]
-  }
+  @all_regions ~w(eu-west-2 us-east-1 us-west-1 ap-southeast-1 ap-southeast-2)
 
-  @doc """
-  Fetches the tenant regions for a given realtime reagion
-  """
-  @spec region_to_tenant_regions(String.t()) :: list() | nil
-  def region_to_tenant_regions(region) do
-    platform = Application.get_env(:realtime, :platform)
-
-    mappings =
-      case platform do
-        :aws -> @mapping_realtime_region_to_tenant_region_aws
-        :fly -> @mapping_realtime_region_to_tenant_region_fly
-        _ -> %{}
-      end
-
-    Map.get(mappings, region)
-  end
+  @spec all_node_regions() :: [String.t()]
+  @doc "List all the regions where nodes can be launched"
+  def all_node_regions(), do: @all_regions
 end

mix.exs

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

test/realtime/gen_rpc_pub_sub/worker_test.exs

@@ -0,0 +1,71 @@
+defmodule Realtime.GenRpcPubSub.WorkerTest do
+  use ExUnit.Case, async: true
+  alias Realtime.GenRpcPubSub.Worker
+  alias Realtime.GenRpc
+  alias Realtime.Nodes
+
+  use Mimic
+
+  @topic "test_topic"
+
+  setup do
+    worker = start_link_supervised!({Worker, {Realtime.PubSub, __MODULE__}})
+    %{worker: worker}
+  end
+
+  describe "forward to local" do
+    test "local broadcast", %{worker: worker} do
+      :ok = Phoenix.PubSub.subscribe(Realtime.PubSub, @topic)
+      send(worker, Worker.forward_to_local(@topic, "le message", Phoenix.PubSub))
+
+      assert_receive "le message"
+      refute_receive _any
+    end
+  end
+
+  describe "forward to region" do
+    setup %{worker: worker} do
+      GenRpc
+      |> stub()
+      |> allow(self(), worker)
+
+      Nodes
+      |> stub()
+      |> allow(self(), worker)
+
+      :ok
+    end
+
+    test "local broadcast + forward to other nodes", %{worker: worker} do
+      parent = self()
+      expect(Nodes, :region_nodes, fn "us-east-1" -> [node(), :node_us_2, :node_us_3] end)
+
+      expect(GenRpc, :abcast, fn [:node_us_2, :node_us_3],
+                                 Realtime.GenRpcPubSub.WorkerTest,
+                                 {:ftl, "test_topic", "le message", Phoenix.PubSub},
+                                 [key: Realtime.GenRpcPubSub.WorkerTest] ->
+        send(parent, :abcast_called)
+        :ok
+      end)
+
+      :ok = Phoenix.PubSub.subscribe(Realtime.PubSub, @topic)
+      send(worker, Worker.forward_to_region(@topic, "le message", Phoenix.PubSub))
+
+      assert_receive "le message"
+      assert_receive :abcast_called
+      refute_receive _any
+    end
+
+    test "local broadcast and no other nodes", %{worker: worker} do
+      expect(Nodes, :region_nodes, fn "us-east-1" -> [node()] end)
+
+      reject(GenRpc, :abcast, 4)
+
+      :ok = Phoenix.PubSub.subscribe(Realtime.PubSub, @topic)
+      send(worker, Worker.forward_to_region(@topic, "le message", Phoenix.PubSub))
+
+      assert_receive "le message"
+      refute_receive _any
+    end
+  end
+end