feat: LSN tracker persistence through replication slot confirmed flush position (#3462)

Closes #3448

This is part of a larger piece of work described in
#3442

The initial idea was to have a separate, periodic persister of the last
processed LSN, but as I was working through the code I realised that we
might not even need that.

Since now consumers are idempotent in the transactions they handle (see
#3452), even the our global
last processed LSN is slightly behind any given shape that is ok.

Instead of persisting it, we actually recover it from the replication
slot's confirmed flush LSN.
- If the slot is new, we reset the last processed LSN to the new slot's
LSN since all shapes will be invalidated
- If the slot is an existing one, we initialise it to the confirmed
flush LSN
- If the reconnection to the slot happens at runtime (e.g. because of
network partition), we preserve the current last processed LSN
- If the reconnection happens at startup, we initialise the LSN from the
slot. This value might be _slightly_ behind if shapes processed
transactions that they did not manage to acknowledge back, but we can
tolerate this.

This way, shape recovery does not have to read all shapes' last
processed LSN during startup.

If we need stronger guarantees of the global processed lsn being the
max, we can also attempt to set it everytime we lazily recover an
existing shape (if the shape's max lsn is larger than the current max
processed lsn). Happy to include this in this PR or a next one.

Author: msfstef

.changeset/polite-moons-compete.md

@@ -0,0 +1,5 @@
+---
+'@core/sync-service': patch
+---
+
+Replace max LSN recovery from on-disk shapes with direct read from replication slot flushed LSN.

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

@@ -2,32 +2,47 @@ defmodule Electric.LsnTracker do
   alias Electric.Postgres.Lsn
 
   # this function is idempotent to avoid problems in tests
-  def create_table(stack_id) do
+  @spec initialize(Electric.stack_id()) :: :ok
+  def initialize(stack_id) do
     table = table(stack_id)
 
     case :ets.info(table, :id) do
       :undefined ->
         :ets.new(table, [:public, :named_table])
+        :ok
 
       ref when is_reference(ref) ->
         :ok
     end
   end
 
-  @spec set_last_processed_lsn(Lsn.t() | non_neg_integer(), String.t()) :: :ok
-  def set_last_processed_lsn(lsn, stack_id) when is_struct(lsn, Lsn) do
+  @spec set_last_processed_lsn(Electric.stack_id(), Lsn.t() | non_neg_integer()) :: :ok
+  def set_last_processed_lsn(stack_id, lsn) when is_struct(lsn, Lsn) do
     stack_id
     |> table()
     |> :ets.insert({:last_processed_lsn, lsn})
 
     :ok
   end
 
-  def set_last_processed_lsn(lsn, stack_id) when is_integer(lsn) do
-    set_last_processed_lsn(Lsn.from_integer(lsn), stack_id)
+  def set_last_processed_lsn(stack_id, lsn) when is_integer(lsn) do
+    set_last_processed_lsn(stack_id, Lsn.from_integer(lsn))
+  end
+
+  @spec initialize_last_processed_lsn(Electric.stack_id(), Lsn.t()) :: :ok
+  def initialize_last_processed_lsn(stack_id, lsn) when is_struct(lsn, Lsn) do
+    stack_id
+    |> table()
+    |> :ets.insert_new({:last_processed_lsn, lsn})
+
+    :ok
+  end
+
+  def initialize_last_processed_lsn(stack_id, lsn) when is_integer(lsn) do
+    initialize_last_processed_lsn(stack_id, Lsn.from_integer(lsn))
   end
 
-  @spec get_last_processed_lsn(String.t()) :: Lsn.t()
+  @spec get_last_processed_lsn(Electric.stack_id()) :: Lsn.t()
   def get_last_processed_lsn(stack_id) do
     [last_processed_lsn: lsn] =
       stack_id

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

@@ -24,6 +24,7 @@ defmodule Electric.Postgres.ReplicationClient do
           | :check_if_publication_exists
           | :drop_slot
           | :create_slot
+          | :query_slot_flushed_lsn
           | :set_display_setting
           | :ready_to_stream
           | :start_streaming
@@ -246,8 +247,15 @@ defmodule Electric.Postgres.ReplicationClient do
       end
     end
 
-    if current_step == :create_slot and extra_info == :created_new_slot,
-      do: notify_created_new_slot(state)
+    # for new slots, always reset the last processed LSN
+    if current_step == :create_slot and extra_info == :created_new_slot do
+      Electric.LsnTracker.set_last_processed_lsn(state.stack_id, state.flushed_wal)
+      notify_created_new_slot(state)
+    end
+
+    # for existing slots, populate the last processed LSN if not present
+    if current_step == :query_slot_flushed_lsn,
+      do: Electric.LsnTracker.initialize_last_processed_lsn(state.stack_id, state.flushed_wal)
 
     if next_step == :ready_to_stream,
       do: notify_ready_to_stream(state)

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

@@ -9,6 +9,7 @@ defmodule Electric.Postgres.ReplicationClient.ConnectionSetup do
   """
   alias Electric.Utils
   alias Electric.Postgres.ReplicationClient.State
+  alias Electric.Postgres.Lsn
 
   require Logger
 
@@ -263,8 +264,9 @@ defmodule Electric.Postgres.ReplicationClient.ConnectionSetup do
     } = result
 
     Logger.debug("Created new slot at lsn=#{lsn_str}")
+    lsn = lsn_str |> Lsn.from_string() |> Lsn.to_integer()
 
-    {:created_new_slot, state}
+    {:created_new_slot, %{state | flushed_wal: lsn}}
   end
 
   defp create_slot_result(%Postgrex.Error{} = error, state) do
@@ -284,6 +286,32 @@ defmodule Electric.Postgres.ReplicationClient.ConnectionSetup do
 
   ###
 
+  defp query_slot_flushed_lsn_query(%State{slot_name: slot_name} = state) do
+    Logger.debug("ReplicationClient step: query_slot_flushed_lsn")
+
+    query = """
+      SELECT confirmed_flush_lsn
+      FROM pg_replication_slots
+      WHERE slot_name = #{Utils.quote_string(slot_name)}
+    """
+
+    {:query, query, state}
+  end
+
+  defp query_slot_flushed_lsn_result([%Postgrex.Result{} = result], state) do
+    %{rows: [[lsn_str]]} = result
+    Logger.debug("Queried existing slot flushed lsn=#{lsn_str}")
+    lsn = lsn_str |> Lsn.from_string() |> Lsn.to_integer()
+    %{state | flushed_wal: lsn}
+  end
+
+  defp query_slot_flushed_lsn_result(%Postgrex.Error{} = error, _state) do
+    # Unexpected error, fail loudly.
+    raise error
+  end
+
+  ###
+
   defp set_display_setting_query(%{display_settings: [query | rest]} = state) do
     Logger.debug("ReplicationClient step: set_display_setting")
     {:query, query, %{state | display_settings: rest}}
@@ -363,9 +391,15 @@ defmodule Electric.Postgres.ReplicationClient.ConnectionSetup do
   defp next_step(%{step: :drop_slot}),
     do: :create_slot
 
+  defp next_step(%{step: :create_slot, flushed_wal: 0}),
+    do: :query_slot_flushed_lsn
+
   defp next_step(%{step: :create_slot}),
     do: :set_display_setting
 
+  defp next_step(%{step: :query_slot_flushed_lsn}),
+    do: :set_display_setting
+
   defp next_step(%{step: :set_display_setting, display_settings: queries}) when queries != [],
     do: :set_display_setting
 
@@ -392,6 +426,7 @@ defmodule Electric.Postgres.ReplicationClient.ConnectionSetup do
 
   defp query_for_step(:drop_slot, state), do: drop_slot_query(state)
   defp query_for_step(:create_slot, state), do: create_slot_query(state)
+  defp query_for_step(:query_slot_flushed_lsn, state), do: query_slot_flushed_lsn_query(state)
   defp query_for_step(:set_display_setting, state), do: set_display_setting_query(state)
   defp query_for_step(:ready_to_stream, state), do: ready_to_stream(state)
   defp query_for_step(:start_streaming, state), do: start_replication_slot_query(state)
@@ -424,6 +459,9 @@ defmodule Electric.Postgres.ReplicationClient.ConnectionSetup do
   defp dispatch_query_result(:create_slot, result, state),
     do: create_slot_result(result, state)
 
+  defp dispatch_query_result(:query_slot_flushed_lsn, result, state),
+    do: query_slot_flushed_lsn_result(result, state)
+
   defp dispatch_query_result(:set_display_setting, result, state),
     do: set_display_setting_result(result, state)
 end

packages/sync-service/lib/electric/replication/shape_log_collector.ex

@@ -47,8 +47,8 @@ defmodule Electric.Replication.ShapeLogCollector do
     Electric.ProcessRegistry.name(stack_id, __MODULE__)
   end
 
-  def set_last_processed_lsn(server_ref, last_processed_lsn) do
-    GenServer.call(server(server_ref), {:set_last_processed_lsn, last_processed_lsn})
+  def mark_as_ready(server_ref) do
+    GenServer.call(server(server_ref), :mark_as_ready)
   end
 
   # use `GenServer.call/2` here to make the event processing synchronous.
@@ -212,8 +212,8 @@ defmodule Electric.Replication.ShapeLogCollector do
     )
   end
 
-  def handle_call({:set_last_processed_lsn, lsn}, _from, state) do
-    LsnTracker.set_last_processed_lsn(lsn, state.stack_id)
+  def handle_call(:mark_as_ready, _from, state) do
+    lsn = LsnTracker.get_last_processed_lsn(state.stack_id)
     Electric.StatusMonitor.mark_shape_log_collector_ready(state.stack_id, self())
     {:reply, :ok, Map.put(state, :last_processed_lsn, lsn)}
   end
@@ -377,7 +377,7 @@ defmodule Electric.Replication.ShapeLogCollector do
 
     OpenTelemetry.start_interval("shape_log_collector.set_last_processed_lsn")
 
-    LsnTracker.set_last_processed_lsn(state.last_processed_lsn, state.stack_id)
+    LsnTracker.set_last_processed_lsn(state.stack_id, state.last_processed_lsn)
 
     flush_tracker =
       if is_struct(event, Transaction) do

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

@@ -27,11 +27,9 @@ end
 defmodule Electric.ShapeCache do
   use GenServer
 
-  alias Electric.Postgres.Lsn
   alias Electric.Replication.LogOffset
   alias Electric.Replication.ShapeLogCollector
   alias Electric.ShapeCache.ShapeStatus
-  alias Electric.ShapeCache
   alias Electric.Shapes
   alias Electric.ShapeCache.ShapeCleaner
   alias Electric.Shapes.Shape
@@ -195,29 +193,30 @@ defmodule Electric.ShapeCache do
       subscription: nil
     }
 
-    {:ok, state, {:continue, :recover_shapes}}
+    {:ok, state, {:continue, :wait_for_restore}}
   end
 
   @impl GenServer
-  def handle_continue(:recover_shapes, state) do
+  def handle_continue(:wait_for_restore, state) do
     start_time = System.monotonic_time()
-    {last_processed_lsn, total_recovered, total_failed_to_recover} = recover_shapes(state)
+
+    total_recovered = ShapeStatus.count_shapes(state.stack_id)
 
     Electric.Replication.PublicationManager.wait_for_restore(state.stack_id)
 
     # Let ShapeLogCollector that it can start processing after finishing this function so that
     # we're subscribed to the producer before it starts forwarding its demand.
-    ShapeLogCollector.set_last_processed_lsn(state.stack_id, last_processed_lsn)
+    ShapeLogCollector.mark_as_ready(state.stack_id)
 
     duration = System.monotonic_time() - start_time
 
     Logger.notice(
-      "Consumers ready in #{System.convert_time_unit(duration, :native, :millisecond)}ms (#{total_recovered} shapes, #{total_failed_to_recover} failed to recover)"
+      "Consumers ready in #{System.convert_time_unit(duration, :native, :millisecond)}ms (#{total_recovered} shapes)"
     )
 
     Electric.Telemetry.OpenTelemetry.execute(
       [:electric, :connection, :consumers_ready],
-      %{duration: duration, total: total_recovered, failed_to_recover: total_failed_to_recover},
+      %{duration: duration, total: total_recovered},
       %{stack_id: state.stack_id}
     )
 
@@ -255,55 +254,6 @@ defmodule Electric.ShapeCache do
     end
   end
 
-  defp recover_shapes(%{stack_id: stack_id} = _state) do
-    import Electric.Postgres.Lsn, only: [is_larger: 2]
-
-    start_time = System.monotonic_time()
-    storage = ShapeCache.Storage.for_stack(stack_id)
-    all_handles_and_shapes = ShapeStatus.list_shapes(stack_id)
-
-    {max_lsn, total_recovered} =
-      all_handles_and_shapes
-      |> Task.async_stream(
-        fn {shape_handle, shape} ->
-          shape_storage = ShapeCache.Storage.for_shape(shape_handle, storage)
-
-          case ShapeCache.Storage.get_current_position(shape_storage) do
-            {:ok, latest_offset, _pg_snapshot} ->
-              {shape_handle, LogOffset.extract_lsn(latest_offset)}
-
-            {:error, reason} ->
-              Logger.error([
-                "shape #{inspect(shape)} (#{inspect(shape_handle)})",
-                " returned error from get_current_position: #{inspect(reason)}"
-              ])
-
-              ShapeCleaner.remove_shape(stack_id, shape_handle)
-
-              {shape_handle, :error}
-          end
-        end,
-        ordered: false
-      )
-      |> Enum.reduce({Lsn.from_integer(0), 0}, fn
-        {:ok, {_handle, :error}}, acc -> acc
-        {:ok, {_handle, lsn}}, {max, recovered} when is_larger(lsn, max) -> {lsn, recovered + 1}
-        _, {max, recovered} -> {max, recovered + 1}
-      end)
-
-    total_failed_to_recover = length(all_handles_and_shapes) - total_recovered
-
-    duration = System.monotonic_time() - start_time
-
-    Logger.info([
-      "Restored LSN position #{max_lsn} in",
-      " #{System.convert_time_unit(duration, :native, :millisecond)}ms",
-      " (#{total_recovered} shapes, #{total_failed_to_recover} failed to recover)"
-    ])
-
-    {max_lsn, total_recovered, total_failed_to_recover}
-  end
-
   defp maybe_create_shape(shape, otel_ctx, %{stack_id: stack_id} = state) do
     if shape_state = ShapeStatus.get_existing_shape(stack_id, shape) do
       shape_state

packages/sync-service/lib/electric/shape_cache/shape_status_owner.ex

@@ -41,8 +41,7 @@ defmodule Electric.ShapeCache.ShapeStatusOwner do
     Electric.Telemetry.Sentry.set_tags_context(stack_id: stack_id)
 
     :ok = ShapeStatus.initialize_from_storage(stack_id, config.storage)
-
-    Electric.LsnTracker.create_table(stack_id)
+    :ok = Electric.LsnTracker.initialize(stack_id)
 
     {:ok, %{stack_id: stack_id, backup_dir: ShapeStatus.backup_dir(config.storage)}}
   end

packages/sync-service/test/electric/lsn_tracker_test.exs

@@ -9,28 +9,68 @@ defmodule Electric.LsnTrackerTest do
 
   describe "get_last_processed_lsn/1" do
     setup ctx do
-      LsnTracker.create_table(ctx.stack_id)
+      LsnTracker.initialize(ctx.stack_id)
       :ok
     end
 
     test "returns inital lsn", %{stack_id: stack_id} do
       lsn = Lsn.from_integer(7)
-      LsnTracker.set_last_processed_lsn(lsn, stack_id)
+      LsnTracker.set_last_processed_lsn(stack_id, lsn)
 
       assert LsnTracker.get_last_processed_lsn(stack_id) == lsn
     end
 
     test "returns last set lsn", %{stack_id: stack_id} do
       lsn = Lsn.from_integer(7)
-      LsnTracker.set_last_processed_lsn(lsn, stack_id)
+      LsnTracker.set_last_processed_lsn(stack_id, lsn)
 
       lsn = Lsn.from_integer(77)
-      LsnTracker.set_last_processed_lsn(lsn, stack_id)
+      LsnTracker.set_last_processed_lsn(stack_id, lsn)
 
       lsn = Lsn.from_integer(111)
-      LsnTracker.set_last_processed_lsn(lsn, stack_id)
+      LsnTracker.set_last_processed_lsn(stack_id, lsn)
 
       assert LsnTracker.get_last_processed_lsn(stack_id) == lsn
     end
   end
+
+  describe "initialize_last_processed_lsn/2" do
+    setup ctx do
+      LsnTracker.initialize(ctx.stack_id)
+      :ok
+    end
+
+    test "sets lsn when not previously set", %{stack_id: stack_id} do
+      lsn = Lsn.from_integer(42)
+      assert :ok = LsnTracker.initialize_last_processed_lsn(stack_id, lsn)
+      assert LsnTracker.get_last_processed_lsn(stack_id) == lsn
+    end
+
+    test "accepts integer and converts to Lsn", %{stack_id: stack_id} do
+      assert :ok = LsnTracker.initialize_last_processed_lsn(stack_id, 100)
+      assert LsnTracker.get_last_processed_lsn(stack_id) == Lsn.from_integer(100)
+    end
+
+    test "does not overwrite existing lsn", %{stack_id: stack_id} do
+      initial_lsn = Lsn.from_integer(50)
+      LsnTracker.set_last_processed_lsn(stack_id, initial_lsn)
+
+      new_lsn = Lsn.from_integer(100)
+      assert :ok = LsnTracker.initialize_last_processed_lsn(stack_id, new_lsn)
+
+      # Should still be the initial LSN, not the new one
+      assert LsnTracker.get_last_processed_lsn(stack_id) == initial_lsn
+    end
+
+    test "is idempotent - can be called multiple times", %{stack_id: stack_id} do
+      lsn1 = Lsn.from_integer(10)
+      lsn2 = Lsn.from_integer(20)
+
+      assert :ok = LsnTracker.initialize_last_processed_lsn(stack_id, lsn1)
+      assert :ok = LsnTracker.initialize_last_processed_lsn(stack_id, lsn2)
+
+      # First call wins
+      assert LsnTracker.get_last_processed_lsn(stack_id) == lsn1
+    end
+  end
 end