Verify durable slot handover and deactivation.

Author: cpursley

README.md

@@ -29,7 +29,7 @@ by adding `walex` to your list of dependencies in `mix.exs`:
 ```elixir
 def deps do
   [
-    {:walex, "~> 4.2.0"}
+    {:walex, "~> 4.4.0"}
   ]
 end
 ```

lib/walex/replication/query_builder.ex

@@ -1,21 +1,21 @@
 defmodule WalEx.Replication.QueryBuilder do
-  def publication_exists(state) do
-    "SELECT 1 FROM pg_publication WHERE pubname = '#{state.publication}' LIMIT 1;"
+  def publication_exists(%{publication: publication}) do
+    "SELECT 1 FROM pg_publication WHERE pubname = '#{publication}' LIMIT 1;"
   end
 
-  def slot_exists(state) do
-    "SELECT active FROM pg_replication_slots WHERE slot_name = '#{state.slot_name}' LIMIT 1;"
+  def slot_exists(%{slot_name: slot_name}) do
+    "SELECT active FROM pg_replication_slots WHERE slot_name = '#{slot_name}' LIMIT 1;"
   end
 
-  def create_temporary_slot(state) do
-    "CREATE_REPLICATION_SLOT #{state.slot_name} TEMPORARY LOGICAL pgoutput NOEXPORT_SNAPSHOT;"
+  def create_temporary_slot(%{slot_name: slot_name}) do
+    "CREATE_REPLICATION_SLOT #{slot_name} TEMPORARY LOGICAL pgoutput NOEXPORT_SNAPSHOT;"
   end
 
-  def create_durable_slot(state) do
-    "CREATE_REPLICATION_SLOT #{state.slot_name} LOGICAL pgoutput NOEXPORT_SNAPSHOT;"
+  def create_durable_slot(%{slot_name: slot_name}) do
+    "CREATE_REPLICATION_SLOT #{slot_name} LOGICAL pgoutput NOEXPORT_SNAPSHOT;"
   end
 
-  def start_replication_slot(state) do
-    "START_REPLICATION SLOT #{state.slot_name} LOGICAL 0/0 (proto_version '1', publication_names '#{state.publication}')"
+  def start_replication_slot(%{slot_name: slot_name, publication: publication}) do
+    "START_REPLICATION SLOT #{slot_name} LOGICAL 0/0 (proto_version '1', publication_names '#{publication}')"
   end
 end

lib/walex/replication/server.ex

@@ -11,35 +11,18 @@ defmodule WalEx.Replication.Server do
   alias WalEx.Decoder
   alias WalEx.Replication.QueryBuilder
 
+  require Logger
+
+  @max_retries 10
+  @initial_backoff 1000
+
   def start_link(opts) do
     app_name = Keyword.get(opts, :app_name)
     opts = set_pgx_replication_conn_opts(app_name)
 
     Postgrex.ReplicationConnection.start_link(__MODULE__, [app_name: app_name], opts)
   end
 
-  defp set_pgx_replication_conn_opts(app_name) do
-    database_configs_keys = [
-      :hostname,
-      :username,
-      :password,
-      :port,
-      :database,
-      :ssl,
-      :ssl_opts,
-      :socket_options
-    ]
-
-    extra_opts = [auto_reconnect: true]
-    database_configs = WalEx.Config.get_configs(app_name, database_configs_keys)
-
-    replications_name = [
-      name: WalExRegistry.set_name(:set_gen_server, __MODULE__, app_name)
-    ]
-
-    extra_opts ++ database_configs ++ replications_name
-  end
-
   @impl true
   def init(opts) do
     app_name = Keyword.get(opts, :app_name)
@@ -90,7 +73,7 @@ defmodule WalEx.Replication.Server do
   def handle_result(results, %{step: :publication_exists} = state) do
     case results do
       [%Postgrex.Result{num_rows: 0}] ->
-        raise "Publication doesn't exists. publication: #{inspect(state.publication)}"
+        raise "Publication doesn't exist. publication: #{inspect(state.publication)}"
 
       _ ->
         raise "Unexpected result when checking if publication exists. #{inspect(results)}"
@@ -110,11 +93,17 @@ defmodule WalEx.Replication.Server do
       ) do
     case active do
       "f" ->
-        query = QueryBuilder.start_replication_slot(state)
-        {:stream, query, [], %{state | step: :streaming}}
+        Logger.info("Activating inactive replication slot: #{state.slot_name}")
+        start_replication_with_retry(state, 0, @initial_backoff)
 
       "t" ->
-        raise "Durable slot already active"
+        Logger.info(
+          "Replication slot #{state.slot_name} is active. Waiting for it to become inactive."
+        )
+
+        schedule_slot_check()
+
+        {:noreply, state}
     end
   end
 
@@ -125,8 +114,7 @@ defmodule WalEx.Replication.Server do
 
   @impl true
   def handle_result([%Postgrex.Result{} | _results], state = %{step: :create_slot}) do
-    query = QueryBuilder.start_replication_slot(state)
-    {:stream, query, [], %{state | step: :streaming}}
+    start_replication_with_retry(state, 0, @initial_backoff)
   end
 
   @impl true
@@ -136,7 +124,22 @@ defmodule WalEx.Replication.Server do
   end
 
   @impl true
-  # https://www.postgresql.org/docs/14/protocol-replication.html
+  def handle_result(
+        %Postgrex.Error{postgres: %{code: :object_in_use}},
+        state = %{step: {:start_replication, retry_count, backoff}}
+      ) do
+    Logger.warning("Replication slot in use, retrying... (attempt #{retry_count + 1})")
+    Process.sleep(backoff)
+    start_replication_with_retry(state, retry_count + 1, backoff * 2)
+  end
+
+  @impl true
+  def handle_result(_, state = %{step: {:start_replication, _retry_count, _backoff}}) do
+    Logger.info("Successfully started replication slot: #{state.slot_name}")
+    {:noreply, %{state | step: :streaming}}
+  end
+
+  @impl true
   def handle_data(<<?w, _wal_start::64, _wal_end::64, _clock::64, rest::binary>>, state) do
     rest
     |> Decoder.decode_message()
@@ -145,6 +148,7 @@ defmodule WalEx.Replication.Server do
     {:noreply, state}
   end
 
+  @impl true
   def handle_data(<<?k, wal_end::64, _clock::64, reply>>, state) do
     messages =
       case reply do
@@ -155,6 +159,55 @@ defmodule WalEx.Replication.Server do
     {:noreply, messages, state}
   end
 
+  @impl true
+  def handle_info(:check_slot_status, state) do
+    query = QueryBuilder.slot_exists(state)
+    {:query, query, %{state | step: :slot_exists}}
+  end
+
+  defp set_pgx_replication_conn_opts(app_name) do
+    database_configs_keys = [
+      :hostname,
+      :username,
+      :password,
+      :port,
+      :database,
+      :ssl,
+      :ssl_opts,
+      :socket_options
+    ]
+
+    extra_opts = [auto_reconnect: true]
+    database_configs = WalEx.Config.get_configs(app_name, database_configs_keys)
+
+    replications_name = [
+      name: WalExRegistry.set_name(:set_gen_server, __MODULE__, app_name)
+    ]
+
+    extra_opts ++ database_configs ++ replications_name
+  end
+
+  defp start_replication_with_retry(state, retry_count, backoff)
+       when retry_count < @max_retries do
+    query = QueryBuilder.start_replication_slot(state)
+    {:stream, query, [], %{state | step: {:start_replication, retry_count, backoff}}}
+  end
+
+  defp start_replication_with_retry(state, _retry_count, _backoff) do
+    Logger.warning(
+      "Failed to start replication slot after maximum retries. Scheduling another check."
+    )
+
+    schedule_slot_check()
+
+    {:noreply, state}
+  end
+
+  defp schedule_slot_check() do
+    # Check again after 5 seconds
+    Process.send_after(self(), :check_slot_status, 5000)
+  end
+
   @epoch DateTime.to_unix(~U[2000-01-01 00:00:00Z], :microsecond)
   defp current_time, do: System.os_time(:microsecond) - @epoch
 end

mix.exs

@@ -4,7 +4,7 @@ defmodule WalEx.MixProject do
   def project do
     [
       app: :walex,
-      version: "4.3.0",
+      version: "4.4.0",
       elixir: "~> 1.15",
       build_embedded: Mix.env() == :prod,
       start_permanent: Mix.env() == :prod,

test/walex/database_test.exs

@@ -163,7 +163,7 @@ defmodule WalEx.DatabaseTest do
         "temporary" => false
       }
 
-      stopped_slot = Map.replace(durable_slot, "active", false)
+      inactive_slot = %{durable_slot | "active" => false}
 
       start_supervised!({WalExSupervisor, durable_opts},
         restart: :temporary,
@@ -174,27 +174,39 @@ defmodule WalEx.DatabaseTest do
 
       other_app_opts = Keyword.replace!(durable_opts, :name, :other_app)
 
-      assert {:error, {{:shutdown, error}, _}} =
-               start_supervised({WalExSupervisor, other_app_opts},
-                 restart: :temporary,
-                 id: :other_app_supervisor
-               )
+      # Start another supervisor with the same slot name
+      {:ok, pid} =
+        start_supervised({WalExSupervisor, other_app_opts},
+          restart: :temporary,
+          id: :other_app_supervisor
+        )
 
-      assert {:failed_to_start_child, WalEx.Replication.Supervisor, {:shutdown, error}} = error
-      assert {:failed_to_start_child, WalEx.Replication.Server, error} = error
-      assert %RuntimeError{message: "Durable slot already active"} = error
+      # Wait for the retry mechanism to complete
+      Process.sleep(10_000)
 
+      # Check that the other app is still running and waiting
+      assert Process.alive?(pid)
+
+      # The original slot should still be active
+      assert [^durable_slot] = pg_replication_slots(database_pid)
+
+      # Stop the first supervisor
       stop_supervised(:ok_supervisor)
-      # sleep to make sure that Postgres detect that the connection is closed
-      Process.sleep(1_000)
-      assert [^stopped_slot] = pg_replication_slots(database_pid)
 
-      start_supervised!({WalExSupervisor, durable_opts},
-        restart: :temporary,
-        id: :ok_supervisor
-      )
+      # Sleep to make sure that Postgres detects that the connection is closed
+      # and the second supervisor has time to activate the slot
+      Process.sleep(10_000)
 
+      # The slot should now be active under the second supervisor
       assert [^durable_slot] = pg_replication_slots(database_pid)
+
+      # Clean up
+      stop_supervised(:other_app_supervisor)
+
+      # Wait for the slot to become inactive after stopping all supervisors
+      Process.sleep(5_000)
+      [slot] = pg_replication_slots(database_pid)
+      assert slot == inactive_slot
     end
   end