Fix lock semantics for the exclusive connection lock when running against a pooler (#2791)

We have recently observed some source in Electric Cloud struggle due to
lock and replication slot contention issues which in theory shouldn't
have happened at all. The cause of the issues was an imperfect
validation logic applied to the creation of new sources. As a
consequence, some sources get created with a pooled connection URL.

More concretely, our validation can successfully weed out pooled
connection strings from Supabase, but Neon's pool allows establishing
replication connections by proxying them directly to the DB. As a
result, Electric ends up acquiring an advisory lock on a pooled
connection where the pooler runs in transaction mode, leading to a bunch
of undesirable effects that prevent Electric's normal operation.

---

This PR addresses the cause of the issues by making the `LockConnection`
process open a replication connection to the DB or fail if that cannot
be achieved. In this way, either the advisory lock will be taken on a
long-living connection as expected, or the connection won't be able to
open and the source will immediately transition into an error state,
instead of assigning an advisory lock to a random pooled connection as
it currently does.

Author: alco

.changeset/eight-pigs-swim.md

@@ -0,0 +1,5 @@
+---
+"@core/sync-service": patch
+---
+
+Open a replication connection to acquire the exclusive connection lock, fixing the lock semantics for cases where Electric runs with a pooled connection string.

packages/sync-service/lib/electric/connection/manager.ex

@@ -111,6 +111,8 @@ defmodule Electric.Connection.Manager do
       :pg_system_identifier,
       # PostgreSQL timeline ID
       :pg_timeline_id,
+      # PostgreSQL's current WAL flush position at lock connection startup
+      :pg_wal_flush_lsn,
       # ID used for process labeling and sibling discovery
       :stack_id,
       # Registry used for stack events
@@ -212,8 +214,12 @@ defmodule Electric.Connection.Manager do
     GenServer.cast(manager, :replication_client_streamed_first_message)
   end
 
-  def pg_info_looked_up(manager, pg_info) do
-    GenServer.cast(manager, {:pg_info_looked_up, pg_info})
+  def pg_system_info_obtained(manager, system_info) do
+    GenServer.cast(manager, {:pg_system_info_obtained, system_info})
+  end
+
+  def pg_info_obtained(manager, pg_info) do
+    GenServer.cast(manager, {:pg_info_obtained, pg_info})
   end
 
   # Used for testing the responsiveness of the manager process
@@ -832,24 +838,35 @@ defmodule Electric.Connection.Manager do
     {:noreply, state}
   end
 
-  def handle_cast({:pg_info_looked_up, {server_version, system_identifier, timeline_id}}, state) do
+  def handle_cast({:pg_system_info_obtained, info}, state) do
+    {:noreply,
+     %State{
+       state
+       | pg_system_identifier: info.system_identifier,
+         pg_timeline_id: info.timeline_id,
+         pg_wal_flush_lsn: info.current_wal_flush_lsn
+     }}
+  end
+
+  def handle_cast({:pg_info_obtained, %{server_version_num: server_version}}, state) do
+    Logger.info(
+      "Postgres server version = #{server_version}, " <>
+        "system identifier = #{state.pg_system_identifier}, " <>
+        "timeline_id = #{state.pg_timeline_id}"
+    )
+
     :telemetry.execute(
       [:electric, :postgres, :info_looked_up],
       %{
         pg_version: server_version,
-        pg_system_identifier: system_identifier,
-        pg_timeline_id: timeline_id
+        pg_system_identifier: state.pg_system_identifier,
+        pg_timeline_id: state.pg_timeline_id,
+        pg_wal_flush_lsn: state.pg_wal_flush_lsn
       },
       %{stack_id: state.stack_id}
     )
 
-    {:noreply,
-     %State{
-       state
-       | pg_version: server_version,
-         pg_system_identifier: system_identifier,
-         pg_timeline_id: timeline_id
-     }}
+    {:noreply, %State{state | pg_version: server_version}}
   end
 
   @impl true

packages/sync-service/lib/electric/db_connection_error.ex

@@ -132,6 +132,15 @@ defmodule Electric.DbConnectionError do
     maybe_database_does_not_exist(error) || unknown_error(error)
   end
 
+  def from_error(%Postgrex.Error{postgres: %{code: :syntax_error, pg_code: "42601"}} = error) do
+    %DbConnectionError{
+      message: error.postgres.message,
+      type: :syntax_error,
+      original_error: error,
+      retry_may_fix?: false
+    }
+  end
+
   def from_error(error), do: unknown_error(error)
 
   def format_original_error(%DbConnectionError{original_error: error}) do

packages/sync-service/lib/electric/postgres/lock_connection.ex

@@ -36,11 +36,23 @@ defmodule Electric.Postgres.LockConnection do
   def start_link(opts) do
     {connection_opts, init_opts} = Keyword.pop(opts, :connection_opts)
 
+    # Start the lock connection in logical replication mode to side-step any connection pooler
+    # that may be sitting between us and the Postgres server.
+    #
+    # We cannot get desired semantics of session-level advisory locks when connecting to
+    # Postgres through a pooler that runs in transaction mode (such as PGBouncer running in
+    # front of Neon). Starting a connection in replication mode ensures that it will be a
+    # direct connection to the database, so it can take a session-level advisory lock whose
+    # lifetime will be tied to the connection's lifetime.
+    connection_opts =
+      connection_opts
+      |> Electric.Utils.deobfuscate_password()
+      |> connection_opts_with_logical_replication()
+
     Postgrex.SimpleConnection.start_link(
       __MODULE__,
       init_opts,
-      [timeout: :infinity, auto_reconnect: false, sync_connect: false] ++
-        Electric.Utils.deobfuscate_password(connection_opts)
+      [timeout: :infinity, auto_reconnect: false, sync_connect: false] ++ connection_opts
     )
   end
 
@@ -58,6 +70,8 @@ defmodule Electric.Postgres.LockConnection do
     Logger.metadata(metadata)
     Electric.Telemetry.Sentry.set_tags_context(metadata)
 
+    Logger.debug("Opening lock connection")
+
     state = %State{
       connection_manager: opts.connection_manager,
       lock_name: opts.lock_name,
@@ -71,11 +85,26 @@ defmodule Electric.Postgres.LockConnection do
   @impl true
   def handle_connect(state) do
     notify_connection_opened(state)
-    send(self(), :acquire_lock)
+
+    # Verify that the connection has been opened in replication mode.
+    #
+    # If there's a pooler running in front of the Postgres server, it may have simply ignored
+    # the replication=database connection parameter, defeating the purpose of us requesting the
+    # replication mode in the first place which is to get the desired session-level locking
+    # semantics.
+    #
+    # Issuing a statement that would cause a syntax error on a regular connection is a surefire
+    # way to ensure the connection is running in the correct mode.
+    send(self(), :identify_system)
+
     {:noreply, state}
   end
 
   @impl true
+  def handle_info(:identify_system, state) do
+    {:query, "IDENTIFY_SYSTEM", state}
+  end
+
   def handle_info(:acquire_lock, state) do
     if state.lock_acquired do
       notify_lock_acquired(state)
@@ -91,6 +120,32 @@ defmodule Electric.Postgres.LockConnection do
   end
 
   @impl true
+  def handle_result([%Postgrex.Result{command: :identify} = result], state) do
+    # [db] postgres:postgres=> IDENTIFY_SYSTEM;
+    #       systemid       │ timeline │  xlogpos  │  dbname
+    # ─────────────────────┼──────────┼───────────┼──────────
+    #  7506979529870965272 │        1 │ 0/220AE10 │ postgres
+    # (1 row)
+    [[systemid, timeline, xlogpos, _dbname]] = result.rows
+
+    notify_system_identified(state, %{
+      system_identifier: systemid,
+      timeline_id: timeline,
+      current_wal_flush_lsn: xlogpos
+    })
+
+    # Now proceed to the actual lock acquisition.
+    send(self(), :acquire_lock)
+
+    {:noreply, state}
+  end
+
+  def handle_result(%Postgrex.Error{postgres: %{code: :syntax_error}} = error, _state) do
+    # Postgrex.SimpleConnection does not support {:stop, ...} or {:shutdown, ...} return values
+    # from callback functions, so we raise here and let the connection manager handle the error.
+    raise error
+  end
+
   def handle_result([%Postgrex.Result{columns: ["pg_advisory_lock"]}], state) do
     Logger.info("Lock acquired from postgres with name #{state.lock_name}")
     notify_lock_acquired(state)
@@ -116,6 +171,10 @@ defmodule Electric.Postgres.LockConnection do
     Connection.Manager.lock_connection_started(manager)
   end
 
+  defp notify_system_identified(%State{connection_manager: manager}, info) do
+    Connection.Manager.pg_system_info_obtained(manager, info)
+  end
+
   defp notify_lock_acquisition_error(error, %State{connection_manager: manager}) do
     Connection.Manager.exclusive_connection_lock_acquisition_failed(manager, error)
   end
@@ -143,4 +202,12 @@ defmodule Electric.Postgres.LockConnection do
        do: true
 
   defp is_expected_error?(_), do: false
+
+  defp connection_opts_with_logical_replication(connection_opts) do
+    update_in(
+      connection_opts,
+      [:parameters],
+      fn params -> params |> List.wrap() |> Keyword.put(:replication, "database") end
+    )
+  end
 end

packages/sync-service/lib/electric/postgres/replication_client.ex

@@ -184,10 +184,14 @@ defmodule Electric.Postgres.ReplicationClient do
     {current_step, next_step, extra_info, return_val} =
       ConnectionSetup.process_query_result(result_list_or_error, state)
 
+    if current_step == :query_pg_info,
+      do: notify_pg_info_obtained(state, extra_info)
+
     if current_step == :create_slot and extra_info == :created_new_slot,
       do: notify_created_new_slot(state)
 
-    if next_step == :ready_to_stream, do: notify_ready_to_stream(state)
+    if next_step == :ready_to_stream,
+      do: notify_ready_to_stream(state)
 
     return_val
   end
@@ -375,25 +379,28 @@ defmodule Electric.Postgres.ReplicationClient do
 
   defp update_applied_wal(state, wal) when is_number(wal), do: state
 
-  defp notify_connection_opened(%State{connection_manager: connection_manager} = state) do
-    :ok = Electric.Connection.Manager.replication_client_started(connection_manager)
+  defp notify_connection_opened(%State{connection_manager: manager} = state) do
+    :ok = Electric.Connection.Manager.replication_client_started(manager)
     state
   end
 
-  defp notify_created_new_slot(%State{connection_manager: connection_manager} = state) do
-    :ok = Electric.Connection.Manager.replication_client_created_new_slot(connection_manager)
+  defp notify_pg_info_obtained(%State{connection_manager: manager} = state, pg_info) do
+    :ok = Electric.Connection.Manager.pg_info_obtained(manager, pg_info)
     state
   end
 
-  defp notify_ready_to_stream(%State{connection_manager: connection_manager} = state) do
-    :ok = Electric.Connection.Manager.replication_client_ready_to_stream(connection_manager)
+  defp notify_created_new_slot(%State{connection_manager: manager} = state) do
+    :ok = Electric.Connection.Manager.replication_client_created_new_slot(manager)
     state
   end
 
-  defp notify_seen_first_message(%State{connection_manager: connection_manager} = state) do
-    :ok =
-      Electric.Connection.Manager.replication_client_streamed_first_message(connection_manager)
+  defp notify_ready_to_stream(%State{connection_manager: manager} = state) do
+    :ok = Electric.Connection.Manager.replication_client_ready_to_stream(manager)
+    state
+  end
 
+  defp notify_seen_first_message(%State{connection_manager: manager} = state) do
+    :ok = Electric.Connection.Manager.replication_client_streamed_first_message(manager)
     state
   end
 end

packages/sync-service/lib/electric/postgres/replication_client/connection_setup.ex

@@ -55,32 +55,13 @@ defmodule Electric.Postgres.ReplicationClient.ConnectionSetup do
 
   defp pg_info_query(state) do
     Logger.debug("ReplicationClient step: pg_info_query")
-
-    query = """
-    SELECT
-      current_setting('server_version_num') server_version_num,
-      (pg_control_system()).system_identifier,
-      (pg_control_checkpoint()).timeline_id
-    """
-
+    query = "SELECT current_setting('server_version_num') server_version_num"
     {:query, query, state}
   end
 
   defp pg_info_result([%Postgrex.Result{} = result], state) do
-    %{rows: [[version_str, system_identifier, timeline_id]]} = result
-
-    Logger.info(
-      "Postgres server version = #{version_str}, " <>
-        "system identifier = #{system_identifier}, " <>
-        "timeline_id = #{timeline_id}"
-    )
-
-    Electric.Connection.Manager.pg_info_looked_up(
-      state.connection_manager,
-      {String.to_integer(version_str), system_identifier, timeline_id}
-    )
-
-    state
+    %{rows: [[version_str]]} = result
+    {%{server_version_num: String.to_integer(version_str)}, state}
   end
 
   defp create_publication_query(state) do

packages/sync-service/test/electric/postgres/replication_client_test.exs

@@ -37,7 +37,7 @@ defmodule Electric.Postgres.ReplicationClientTest do
 
     defp process_message({:"$gen_cast", :replication_client_started}), do: nil
     defp process_message({:"$gen_cast", :replication_client_created_new_slot}), do: nil
-    defp process_message({:"$gen_cast", {:pg_info_looked_up, _}}), do: nil
+    defp process_message({:"$gen_cast", {:pg_info_obtained, _}}), do: nil
 
     defp process_message({:"$gen_cast", :replication_client_streamed_first_message}),
       do: {self(), :streaming_started}