fix: on ReplicationPoller send only subscription IDs relevant to the node (#1583)

Author: edgurgel

lib/extensions/postgres_cdc_rls/replication_poller.ex

@@ -206,12 +206,13 @@ defmodule Extensions.PostgresCdcRls.ReplicationPoller do
 
           case collect_subscription_nodes(subscribers_nodes_table, change.subscription_ids) do
             {:ok, nodes} ->
-              for node <- nodes do
+              for {node, subscription_ids} <- nodes do
                 TenantBroadcaster.pubsub_direct_broadcast(
                   node,
                   tenant_id,
                   topic,
-                  change,
+                  # Send only the subscription IDs relevant to this node
+                  %{change | subscription_ids: MapSet.new(subscription_ids)},
                   MessageDispatcher,
                   :postgres_changes
                 )
@@ -236,10 +237,16 @@ defmodule Extensions.PostgresCdcRls.ReplicationPoller do
   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} ->
+    Enum.reduce_while(subscription_ids, {:ok, %{}}, 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}}
+        [{_, node}] ->
+          updated_acc =
+            Map.update(acc, node, [subscription_id], fn existing_ids -> [subscription_id | existing_ids] end)
+
+          {:cont, {:ok, updated_acc}}
+
+        _ ->
+          {:halt, {:error, :node_not_found}}
       end
     end)
   rescue

mix.exs

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

test/realtime/extensions/cdc_rls/replication_poller_test.exs

@@ -296,7 +296,8 @@ defmodule Realtime.Extensions.PostgresCdcRls.ReplicationPollerTest do
                false,
                [
                  sub1 = <<71, 36, 83, 212, 168, 9, 17, 240, 165, 186, 118, 202, 193, 157, 232, 187>>,
-                 sub2 = <<251, 188, 190, 118, 168, 119, 17, 240, 188, 87, 118, 202, 193, 157, 232, 187>>
+                 sub2 = <<251, 188, 190, 118, 168, 119, 17, 240, 188, 87, 118, 202, 193, 157, 232, 187>>,
+                 sub3 = <<49, 59, 209, 112, 173, 77, 17, 240, 191, 41, 118, 202, 193, 157, 232, 187>>
                ],
                []
              ]
@@ -306,9 +307,10 @@ defmodule Realtime.Extensions.PostgresCdcRls.ReplicationPollerTest do
            messages: []
          }}
 
-      # Both subscriptions have node information
+      # All subscriptions have node information
       :ets.insert(args["subscribers_nodes_table"], {sub1, node()})
       :ets.insert(args["subscribers_nodes_table"], {sub2, :"[email protected]"})
+      :ets.insert(args["subscribers_nodes_table"], {sub3, node()})
 
       expect(Replications, :list_changes, fn _, _, _, _, _ -> results end)
       reject(&TenantBroadcaster.pubsub_broadcast/5)
@@ -345,9 +347,10 @@ defmodule Realtime.Extensions.PostgresCdcRls.ReplicationPollerTest do
 
       assert Enum.count(calls) == 2
 
-      Enum.each(calls, fn [node, _, _, _, _, _] ->
-        assert node in [node(), :"[email protected]"]
-      end)
+      node_subs = Enum.map(calls, fn [node, _, _, change, _, _] -> {node, change.subscription_ids} end)
+
+      assert {node(), MapSet.new([sub1, sub3])} in node_subs
+      assert {:"[email protected]", MapSet.new([sub2])} in node_subs
     end
   end