feat: Restore PublicationManager state from ShapeStatus ETS on startup (#3275)

## Edited by @msfstef 

Closes #3245

1. On bootup, pub man filters state is recovered directly from
`ShapeStatus`
2. It immediately goes into a `:continue` clause to update the
configuration
3. Keeps track of a `restore_complete?` flag as well as
`restore_waiters` from the new `wait_for_restore` sync API
3a. I decided to add waiters as I want the update to be able to fail due
to e.g. network errors and still accept requests and handle them without
crashing the shape subsystem, in the spirit of previous work.
4a. If restore succeeds, we reply to the waiters
4b. If restore fails, we go into a refresh loop with a shorter timeout
than our regular refresh timeout until restore completes
5. We `wait_for_restore` in the `ShapeCache` before we mark consumers as
ready, meaning we block there until we manage to recover the
publication.

The last step is still coupling and blocking ShapeCache to PubMan but in
the spirit of maintaining the current behaviour but avoiding the
onslaught of recovery messages from all shapes I think this is a good
first step, and @magnetised feel free to move things around in your work
to decouple things.

The main thing to note is that _we should not_ start streaming changes
until we know for sure that the publication has actually been restored -
but I believe the shape subsystem should be able to handle existing
shape data while we wait for that to happen.

I've also made it so that we only call `add_shape` to the publication
manager if a snapshot _has not_ been started, so that most restored
shapes don't end up calling the pubman (@magnetised was this the
intention? if not, I can let you sort it out in your own PRs)

---------

Co-authored-by: msfstef <[email protected]>

Author: magnetised

.changeset/seven-ladybugs-shave.md

@@ -1,5 +1,5 @@
 ---
-"@core/sync-service": patch
+'@core/sync-service': patch
 ---
 
 Restore ShapeLogCollector's state from the ShapeStatus ETS table at startup

.changeset/sharp-worms-give.md

@@ -0,0 +1,5 @@
+---
+'@core/sync-service': patch
+---
+
+Restore shapes in `PublicationManager` via the `ShapeStatus` ETS to avoid message congestion.

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

@@ -1,10 +1,19 @@
 defmodule Electric.Replication.PublicationManager do
-  @moduledoc false
+  @moduledoc """
+  Manages a PostgreSQL publication for a given Electric stack, tracking shapes
+  and ensuring that the publication configuration matches the required set of
+  relations that need to be published for the shapes to function correctly.
+
+  Includes periodic checks of the publication to ensure that it remains valid,
+  and expires any shapes that are no longer valid due to schema changes or
+  permission issues.
+  """
   use GenServer
 
   alias Electric.Postgres.Configuration
   alias Electric.ShapeCache.ShapeCleaner
   alias Electric.Shapes.Shape
+  alias Electric.Telemetry.OpenTelemetry
   alias Electric.Utils
 
   require Logger
@@ -14,6 +23,7 @@ defmodule Electric.Replication.PublicationManager do
   @callback name(binary() | Keyword.t()) :: term()
   @callback add_shape(shape_handle(), Shape.t(), Keyword.t()) :: :ok
   @callback remove_shape(shape_handle(), Keyword.t()) :: :ok
+  @callback wait_for_restore(Keyword.t()) :: :ok
 
   defstruct [
     :stack_id,
@@ -23,13 +33,16 @@ defmodule Electric.Replication.PublicationManager do
     :can_alter_publication?,
     :manual_table_publishing?,
     :publication_refresh_period,
+    :restore_retry_timeout,
     relation_ref_counts: %{},
     prepared_relation_filters: MapSet.new(),
     committed_relation_filters: MapSet.new(),
     tracked_shape_handles: %{},
     waiters: %{},
     scheduled_updated_ref: nil,
-    next_update_forced?: false
+    next_update_forced?: false,
+    restore_waiters: [],
+    restore_complete?: false
   ]
 
   @type relation_filters() :: MapSet.t(Electric.oid_relation())
@@ -47,7 +60,10 @@ defmodule Electric.Replication.PublicationManager do
            can_alter_publication?: boolean(),
            manual_table_publishing?: boolean(),
            publication_refresh_period: non_neg_integer(),
-           next_update_forced?: boolean()
+           next_update_forced?: boolean(),
+           restore_waiters: [GenServer.from()],
+           restore_complete?: boolean(),
+           restore_retry_timeout: non_neg_integer()
          }
 
   # The default debounce timeout is 0, which means that the publication update
@@ -56,6 +72,9 @@ defmodule Electric.Replication.PublicationManager do
   # windows to aggregate shape filter updates
   @default_debounce_timeout 0
 
+  # The default retry timeout in case of failed restore attempts
+  @default_restore_retry_timeout 1_000
+
   @name_schema_tuple {:tuple, [:atom, :atom, :any]}
   @genserver_name_schema {:or, [:atom, @name_schema_tuple]}
   @schema NimbleOptions.new!(
@@ -67,7 +86,12 @@ defmodule Electric.Replication.PublicationManager do
             manual_table_publishing?: [type: :boolean, required: false, default: false],
             update_debounce_timeout: [type: :timeout, default: @default_debounce_timeout],
             server: [type: :any, required: false],
-            refresh_period: [type: :pos_integer, required: false, default: 60_000]
+            refresh_period: [type: :pos_integer, required: false, default: 60_000],
+            restore_retry_timeout: [
+              type: :pos_integer,
+              required: false,
+              default: @default_restore_retry_timeout
+            ]
           )
 
   @behaviour __MODULE__
@@ -83,26 +107,34 @@ defmodule Electric.Replication.PublicationManager do
   end
 
   @impl __MODULE__
-  def add_shape(shape_id, shape, opts \\ []) do
+  def add_shape(shape_handle, shape, opts \\ []) do
     server = Access.get(opts, :server, name(opts))
     oid_relation = get_oid_relation_from_shape(shape)
 
-    case GenServer.call(server, {:add_shape, shape_id, oid_relation}) do
+    case GenServer.call(server, {:add_shape, shape_handle, oid_relation}) do
       :ok -> :ok
       {:error, err} -> raise err
     end
   end
 
   @impl __MODULE__
-  def remove_shape(shape_id, opts \\ []) do
+  def remove_shape(shape_handle, opts \\ []) do
     server = Access.get(opts, :server, name(opts))
 
-    case GenServer.call(server, {:remove_shape, shape_id}) do
+    case GenServer.call(server, {:remove_shape, shape_handle}) do
       :ok -> :ok
       {:error, err} -> raise err
     end
   end
 
+  @impl __MODULE__
+  def wait_for_restore(opts \\ []) do
+    server = Access.get(opts, :server, name(opts))
+
+    GenServer.call(server, :wait_for_restore, Keyword.get(opts, :timeout, :infinity))
+    :ok
+  end
+
   def start_link(opts) do
     with {:ok, opts} <- NimbleOptions.validate(opts, @schema) do
       stack_id = Keyword.fetch!(opts, :stack_id)
@@ -133,10 +165,39 @@ defmodule Electric.Replication.PublicationManager do
       db_pool: opts.db_pool,
       can_alter_publication?: opts.can_alter_publication?,
       manual_table_publishing?: opts.manual_table_publishing?,
-      publication_refresh_period: opts.refresh_period
+      publication_refresh_period: opts.refresh_period,
+      restore_retry_timeout: opts.restore_retry_timeout
     }
 
-    {:ok, state, state.publication_refresh_period}
+    {:ok, state, {:continue, :restore_relations}}
+  end
+
+  @impl true
+  def handle_continue(:restore_relations, state) do
+    OpenTelemetry.with_span(
+      "publication_manager.restore_relations",
+      [],
+      state.stack_id,
+      fn ->
+        state =
+          state.stack_id
+          |> Electric.ShapeCache.ShapeStatus.list_shapes()
+          |> Enum.reduce(
+            state,
+            fn {shape_handle, shape}, state ->
+              rel_key = get_oid_relation_from_shape(shape)
+              do_update_relation_filters_with_shape(shape_handle, rel_key, :add, state)
+            end
+          )
+
+        state =
+          if update_needed?(state),
+            do: schedule_update_publication(0, true, state),
+            else: mark_restore_complete(state)
+
+        {:noreply, state, refresh_timeout(state)}
+      end
+    )
   end
 
   @impl true
@@ -146,9 +207,9 @@ defmodule Electric.Replication.PublicationManager do
 
     if not relation_tracked?(oid_rel, state) do
       state = add_waiter(from, oid_rel, state)
-      {:noreply, state}
+      {:noreply, state, refresh_timeout(state)}
     else
-      {:reply, :ok, state, state.publication_refresh_period}
+      {:reply, :ok, state, refresh_timeout(state)}
     end
   end
 
@@ -160,7 +221,16 @@ defmodule Electric.Replication.PublicationManager do
     # reconcile the publication, otherwise you run into issues where only the last
     # removal fails and all others succeed. No removal guarantees anything about
     # the state of the publication.
-    {:reply, :ok, state, state.publication_refresh_period}
+    {:reply, :ok, state, refresh_timeout(state)}
+  end
+
+  def handle_call(:wait_for_restore, from, state) do
+    if state.restore_complete? do
+      {:reply, :ok, state, refresh_timeout(state)}
+    else
+      state = %{state | restore_waiters: [from | state.restore_waiters]}
+      {:noreply, state, refresh_timeout(state)}
+    end
   end
 
   @impl true
@@ -170,33 +240,45 @@ defmodule Electric.Replication.PublicationManager do
     state = %{state | scheduled_updated_ref: nil}
 
     state =
-      case check_publication_status(state) do
-        {:ok, state} ->
-          case configure_publication(state) do
+      OpenTelemetry.with_span(
+        "publication_manager.update_publication",
+        [
+          is_restore: not state.restore_complete?
+        ],
+        state.stack_id,
+        fn ->
+          case check_publication_status(state) do
             {:ok, state} ->
-              reply_to_all_waiters(:ok, state)
+              case configure_publication(state) do
+                {:ok, state} ->
+                  state = mark_restore_complete(state)
+                  reply_to_all_waiters(:ok, state)
 
-            {:ok, relations_configured, state} ->
-              handle_publication_update_result(relations_configured, state)
+                {:ok, relations_configured, state} ->
+                  state = mark_restore_complete(state)
+                  handle_publication_update_result(relations_configured, state)
+
+                {:error, err, state} ->
+                  reply_to_all_waiters({:error, err}, state)
+              end
 
             {:error, err, state} ->
+              Logger.warning("Failed to confirm publication status: #{inspect(err)}")
               reply_to_all_waiters({:error, err}, state)
           end
-
-        {:error, err, state} ->
-          Logger.warning("Failed to confirm publication status: #{inspect(err)}")
-          reply_to_all_waiters({:error, err}, state)
-      end
+        end
+      )
 
     # Schedule a forced refresh to happen periodically unless there's an explicit call to
     # update the publication that happens sooner.
-    {:noreply, state, state.publication_refresh_period}
+    {:noreply, state, refresh_timeout(state)}
   end
 
   def handle_info(:timeout, state) do
     case Electric.StatusMonitor.status(state.stack_id) do
       %{conn: :up} ->
-        handle_info(:update_publication, %{state | next_update_forced?: true})
+        state = schedule_update_publication(0, true, state)
+        {:noreply, state, refresh_timeout(state)}
 
       status ->
         Logger.debug("Publication update skipped due to inactive stack: #{inspect(status)}")
@@ -247,22 +329,7 @@ defmodule Electric.Replication.PublicationManager do
       state = %{state | next_update_forced?: false}
 
       try do
-        relations_configured =
-          if can_update? do
-            Configuration.configure_publication!(
-              state.db_pool,
-              state.publication_name,
-              state.prepared_relation_filters
-            )
-          else
-            Configuration.validate_publication_configuration!(
-              state.db_pool,
-              state.publication_name,
-              state.prepared_relation_filters
-            )
-          end
-
-        {:ok, relations_configured, state}
+        do_configure_publication!(state)
       rescue
         err ->
           Logger.warning("Failed to #{key_word} publication: #{inspect(err)}")
@@ -275,6 +342,25 @@ defmodule Electric.Replication.PublicationManager do
     end
   end
 
+  defp do_configure_publication!(state) do
+    relations_configured =
+      if can_update_publication?(state) do
+        Configuration.configure_publication!(
+          state.db_pool,
+          state.publication_name,
+          state.prepared_relation_filters
+        )
+      else
+        Configuration.validate_publication_configuration!(
+          state.db_pool,
+          state.publication_name,
+          state.prepared_relation_filters
+        )
+      end
+
+    {:ok, relations_configured, state}
+  end
+
   defguardp is_known_publication_error(error)
             when is_exception(error) and
                    (is_struct(error, Electric.DbConfigurationError) or
@@ -349,6 +435,13 @@ defmodule Electric.Replication.PublicationManager do
        ),
        do: %{state | next_update_forced?: forced? or state.next_update_forced?}
 
+  defp mark_restore_complete(%{restore_complete?: true} = state), do: state
+
+  defp mark_restore_complete(state) do
+    for waiter <- state.restore_waiters, do: GenServer.reply(waiter, :ok)
+    %{state | restore_complete?: true, restore_waiters: []}
+  end
+
   defp relation_tracked?(oid_rel, %__MODULE__{committed_relation_filters: committed}) do
     MapSet.member?(committed, oid_rel)
   end
@@ -367,6 +460,9 @@ defmodule Electric.Replication.PublicationManager do
     can_alter and not manual
   end
 
+  defp refresh_timeout(%{restore_complete?: false, restore_retry_timeout: timeout}), do: timeout
+  defp refresh_timeout(%{publication_refresh_period: period}), do: period
+
   defguardp is_tracking_shape_handle?(shape_handle, state)
             when is_map_key(state.tracked_shape_handles, shape_handle)
 

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

@@ -239,6 +239,9 @@ defmodule Electric.ShapeCache do
         {:consumers_ready, last_processed_lsn, total_recovered, total_failed_to_recover},
         state
       ) do
+    {pub_man, pub_man_opts} = state.publication_manager
+    pub_man.wait_for_restore(pub_man_opts)
+
     ShapeLogCollector.set_last_processed_lsn(state.stack_id, last_processed_lsn)
 
     Electric.Connection.Manager.consumers_ready(

packages/sync-service/lib/electric/shapes/consumer/snapshotter.ex

@@ -67,16 +67,16 @@ defmodule Electric.Shapes.Consumer.Snapshotter do
             stack_id,
             fn ->
               try do
-                OpenTelemetry.with_span(
-                  "shape_snapshot.prepare_tables",
-                  shape_attrs(shape_handle, shape),
-                  stack_id,
-                  fn ->
-                    publication_manager.add_shape(shape_handle, shape, publication_manager_opts)
-                  end
-                )
-
                 if not Storage.snapshot_started?(state.storage) do
+                  OpenTelemetry.with_span(
+                    "shape_snapshot.prepare_tables",
+                    shape_attrs(shape_handle, shape),
+                    stack_id,
+                    fn ->
+                      publication_manager.add_shape(shape_handle, shape, publication_manager_opts)
+                    end
+                  )
+
                   start_streaming_snapshot_from_db(
                     consumer,
                     shape_handle,