fix: region rebalancing for Postgres Changes (#1464)

edgurgel · web-flow · commit 1efdf781904e · 2025-07-22T21:43:37.000+12:00
diff --git a/lib/extensions/postgres_cdc_rls/subscription_manager.ex b/lib/extensions/postgres_cdc_rls/subscription_manager.ex
@@ -5,6 +5,7 @@ defmodule Extensions.PostgresCdcRls.SubscriptionManager do
   use GenServer
   use Realtime.Logs
 
+  alias Realtime.Tenants.Rebalancer
   alias Extensions.PostgresCdcRls, as: Rls
 
   alias Realtime.Database
@@ -29,7 +30,8 @@ defmodule Extensions.PostgresCdcRls.SubscriptionManager do
       :no_users_ref,
       no_users_ts: nil,
       oids: %{},
-      check_oid_ref: nil
+      check_oid_ref: nil,
+      check_region_interval: nil
     ]
 
     @type t :: %__MODULE__{
@@ -44,7 +46,8 @@ defmodule Extensions.PostgresCdcRls.SubscriptionManager do
               queue: :queue.queue()
             },
             no_users_ref: reference(),
-            no_users_ts: non_neg_integer() | nil
+            no_users_ts: non_neg_integer() | nil,
+            check_region_interval: non_neg_integer
           }
   end
 
@@ -79,6 +82,9 @@ defmodule Extensions.PostgresCdcRls.SubscriptionManager do
 
     oids = Subscriptions.fetch_publication_tables(conn, publication)
 
+    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,
@@ -89,7 +95,8 @@ defmodule Extensions.PostgresCdcRls.SubscriptionManager do
         ref: check_delete_queue(),
         queue: :queue.new()
       },
-      no_users_ref: check_no_users()
+      no_users_ref: check_no_users(),
+      check_region_interval: check_region_interval
     }
 
     send(self(), :check_oids)
@@ -200,6 +207,22 @@ defmodule Extensions.PostgresCdcRls.SubscriptionManager do
     {:noreply, %{state | no_users_ts: ts_new, no_users_ref: check_no_users()}}
   end
 
+  def handle_info({:check_region, previous_nodes_set}, state) do
+    current_nodes_set = MapSet.new(Node.list())
+
+    case Rebalancer.check(previous_nodes_set, current_nodes_set, state.id) do
+      :ok ->
+        # Let's check again in the future
+        send_region_check_message(state.check_region_interval)
+        {:noreply, state}
+
+      {:error, :wrong_region} ->
+        Logger.warning("Rebalancing Postgres Changes replication for a closer region")
+        Rls.handle_stop(state.id, 15_000)
+        {:noreply, state}
+    end
+  end
+
   def handle_info(msg, state) do
     log_error("UnhandledProcessMessage", msg)
 
@@ -216,4 +239,10 @@ defmodule Extensions.PostgresCdcRls.SubscriptionManager do
 
   defp check_delete_queue(timeout \\ @timeout),
     do: Process.send_after(self(), :check_delete_queue, timeout)
+
+  defp send_region_check_message(check_region_interval) do
+    Process.send_after(self(), {:check_region, MapSet.new(Node.list())}, check_region_interval)
+  end
+
+  defp rebalance_check_interval_in_ms(), do: Application.fetch_env!(:realtime, :rebalance_check_interval_in_ms)
 end
diff --git a/mix.exs b/mix.exs
@@ -4,7 +4,7 @@ defmodule Realtime.MixProject do
   def project do
     [
       app: :realtime,
-      version: "2.41.4",
+      version: "2.41.5",
       elixir: "~> 1.17.3",
       elixirc_paths: elixirc_paths(Mix.env()),
       start_permanent: Mix.env() == :prod,
diff --git a/test/realtime/extensions/cdc_rls/cdc_rls_test.exs b/test/realtime/extensions/cdc_rls/cdc_rls_test.exs
@@ -1,15 +1,21 @@
 defmodule Realtime.Extensions.CdcRlsTest do
   # async: false due to usage of dev_tenant
+  # Also global mimic mock
   use RealtimeWeb.ChannelCase, async: false
   use Mimic
 
-  alias Realtime.Api.Tenant
-  alias Realtime.Database
+  import ExUnit.CaptureLog
+
+  setup :set_mimic_global
+
   alias Extensions.PostgresCdcRls
   alias PostgresCdcRls.SubscriptionManager
   alias Postgrex
   alias Realtime.Api
+  alias Realtime.Api.Tenant
+  alias Realtime.Database
   alias Realtime.PostgresCdc
+  alias Realtime.Tenants.Rebalancer
   alias RealtimeWeb.ChannelsAuthorization
 
   @cdc "postgres_cdc_rls"
@@ -125,6 +131,45 @@ defmodule Realtime.Extensions.CdcRlsTest do
     end
   end
 
+  describe "Region rebalancing" do
+    setup do
+      tenant = Containers.checkout_tenant(run_migrations: true)
+      %Tenant{extensions: extensions, external_id: external_id} = tenant
+      postgres_extension = PostgresCdc.filter_settings("postgres_cdc_rls", extensions)
+
+      args =
+        postgres_extension
+        |> Map.put("id", external_id)
+        |> Map.put(:check_region_interval, 100)
+
+      %{tenant_id: tenant.id, args: args}
+    end
+
+    test "rebalancing needed process stops", %{tenant_id: tenant_id, args: args} do
+      log =
+        capture_log(fn ->
+          expect(Rebalancer, :check, 1, fn _, _, tenant_id -> {:error, :wrong_region} end)
+          reject(&Rebalancer.check/3)
+
+          {:ok, pid} = PostgresCdcRls.start(args)
+          ref = Process.monitor(pid)
+
+          assert_receive {:DOWN, ^ref, :process, ^pid, _reason}, 3000
+        end)
+
+      assert log =~ "Rebalancing Postgres Changes replication for a closer region"
+    end
+
+    test "rebalancing not needed process stays up", %{tenant_id: tenant_id, args: args} do
+      stub(Rebalancer, :check, fn _, _, ^tenant_id -> :ok end)
+
+      {:ok, pid} = PostgresCdcRls.start(args)
+      ref = Process.monitor(pid)
+
+      refute_receive {:DOWN, ^ref, :process, ^pid, _reason}, 1000
+    end
+  end
+
   for adapter <- [:phoenix, :gen_rpc] do
     describe "integration #{adapter}" do
       @describetag adapter: adapter
@@ -192,10 +237,6 @@ defmodule Realtime.Extensions.CdcRlsTest do
         metadata = [metadata: subscription_metadata]
         :ok = PostgresCdc.subscribe(PostgresCdcRls, pg_change_params, external_id, metadata)
 
-        # Mimicking how Phoenix.Channel subscribes when a channel joins using fastlane
-        fastlane = {:fastlane, self(), serializer, []}
-        :ok = Phoenix.PubSub.subscribe(Realtime.PubSub, topic, metadata: fastlane)
-
         # First time it will return nil
         PostgresCdcRls.handle_connect(args)
         # Wait for it to start
@@ -287,10 +328,6 @@ defmodule Realtime.Extensions.CdcRlsTest do
         metadata = [metadata: subscription_metadata]
         :ok = PostgresCdc.subscribe(PostgresCdcRls, pg_change_params, external_id, metadata)
 
-        # Mimicking how Phoenix.Channel subscribes when a channel joins using fastlane
-        fastlane = {:fastlane, self(), serializer, []}
-        :ok = Phoenix.PubSub.subscribe(Realtime.PubSub, topic, metadata: fastlane)
-
         # Now subscribe to the Postgres Changes
         {:ok, _} = PostgresCdcRls.handle_after_connect(response, postgres_extension, pg_change_params)
         assert %Postgrex.Result{rows: [[1]]} = Postgrex.query!(conn, "select count(*) from realtime.subscription", [])

Original file line number	Diff line number	Diff line change
`@@ -4,7 +4,7 @@ defmodule Realtime.MixProject do`
`4`	`4`	`def project do`
`5`	`5`	`[`
`6`	`6`	`app: :realtime,`
`7`		`- version: "2.41.4",`
	`7`	`+ version: "2.41.5",`
`8`	`8`	`elixir: "~> 1.17.3",`
`9`	`9`	`elixirc_paths: elixirc_paths(Mix.env()),`
`10`	`10`	`start_permanent: Mix.env() == :prod,`