feat(sync-service): repopulate log collectors state from the status ets (#3242)

Don't rely on consumer processes existing in order to initialise the
shape log collector -- all the information we need is in the shape
status ets.

We still need the consumer processes to register in order to synchronise
the filters etc with the existing shapes but we only need to do that if
the consumer was not in the original state.

This is handled by having an extra `action` param to the subscribe call.
Shape consumers being restored don't add their shape to the shape log
collector's internal filters while those being newly created do.

Author: magnetised

.changeset/seven-ladybugs-shave.md

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

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

@@ -75,15 +75,17 @@ defmodule Electric.Replication.ShapeLogCollector do
     :ok = GenServer.call(server, {:relation_msg, rel, trace_context}, :infinity)
   end
 
-  def subscribe(server_ref, shape_handle, shape) do
-    GenServer.call(server(server_ref), {:subscribe, shape_handle, shape})
+  def subscribe(server_ref, shape_handle, shape, action) when action in [:restore, :create] do
+    GenServer.call(server(server_ref), {:subscribe, shape_handle, shape, action})
   end
 
   def notify_flushed(server_ref, shape_handle, offset) do
     GenServer.cast(server(server_ref), {:writer_flushed, shape_handle, offset})
   end
 
   def init(opts) do
+    activate_mocked_functions_from_test_process()
+
     stack_id = opts.stack_id
 
     Process.set_label({:shape_log_collector, stack_id})
@@ -124,7 +126,32 @@ defmodule Electric.Replication.ShapeLogCollector do
           )
       })
 
-    {:ok, state}
+    {:ok, state, {:continue, :restore_shapes}}
+  end
+
+  def handle_continue(:restore_shapes, state) do
+    OpenTelemetry.with_span(
+      "shape_log_collector.restore_shapes",
+      [],
+      state.stack_id,
+      fn ->
+        {partitions, filter, layers} =
+          state.stack_id
+          |> Electric.ShapeCache.ShapeStatus.list_shapes()
+          |> Enum.reduce(
+            {state.partitions, state.filter, state.dependency_layers},
+            fn {shape_handle, shape}, {partitions, filter, layers} ->
+              {
+                Partitions.add_shape(partitions, shape_handle, shape),
+                Filter.add_shape(filter, shape_handle, shape),
+                DependencyLayers.add_dependency(layers, shape, shape_handle)
+              }
+            end
+          )
+
+        {:noreply, %{state | partitions: partitions, filter: filter, dependency_layers: layers}}
+      end
+    )
   end
 
   def handle_info({{:unsubscribe, shape_handle}, ref, :process, pid, _reason}, state) do
@@ -138,7 +165,7 @@ defmodule Electric.Replication.ShapeLogCollector do
     )
   end
 
-  def handle_call({:subscribe, shape_handle, shape}, {pid, _ref}, state) do
+  def handle_call({:subscribe, shape_handle, shape, action}, {pid, _ref}, state) do
     OpenTelemetry.with_span(
       "shape_log_collector.subscribe",
       [shape_handle: shape_handle],
@@ -148,14 +175,23 @@ defmodule Electric.Replication.ShapeLogCollector do
         from = {pid, ref}
 
         state =
-          %{
-            state
-            | partitions: Partitions.add_shape(state.partitions, shape_handle, shape),
-              filter: Filter.add_shape(state.filter, shape_handle, shape),
-              pids_by_shape_handle: Map.put(state.pids_by_shape_handle, shape_handle, pid),
-              dependency_layers:
-                DependencyLayers.add_dependency(state.dependency_layers, shape, shape_handle)
-          }
+          case action do
+            :restore ->
+              # Once we move consumer monitoring out of this process,
+              # subscribing with action :restore will be a no-op that we can
+              # filter in the `subscribe/4` function
+              state
+
+            :create ->
+              %{
+                state
+                | partitions: Partitions.add_shape(state.partitions, shape_handle, shape),
+                  filter: Filter.add_shape(state.filter, shape_handle, shape),
+                  dependency_layers:
+                    DependencyLayers.add_dependency(state.dependency_layers, shape, shape_handle)
+              }
+          end
+          |> Map.update!(:pids_by_shape_handle, &Map.put(&1, shape_handle, pid))
           |> Map.update!(:subscriptions, fn {count, set} ->
             {count + 1, MapSet.put(set, from)}
           end)
@@ -401,4 +437,12 @@ defmodule Electric.Replication.ShapeLogCollector do
   defp server(stack_id) when is_binary(stack_id), do: name(stack_id)
   defp server({:via, _, _} = name), do: name
   defp server(pid) when is_pid(pid), do: pid
+
+  if Mix.env() == :test do
+    def activate_mocked_functions_from_test_process do
+      Support.TestUtils.activate_mocked_functions_for_module(__MODULE__)
+    end
+  else
+    def activate_mocked_functions_from_test_process, do: :noop
+  end
 end

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

@@ -302,7 +302,7 @@ defmodule Electric.ShapeCache do
   end
 
   defp start_and_recover_shape(shape_handle, shape, state) do
-    case start_shape(shape_handle, shape, state) do
+    case start_shape(shape_handle, shape, state, nil, :restore) do
       :ok ->
         consumer = Shapes.Consumer.name(state.stack_id, shape_handle)
         # This `initial_state` is a GenServer call, so we're blocked until consumer is ready
@@ -354,15 +354,15 @@ defmodule Electric.ShapeCache do
 
       Logger.info("Creating new shape for #{inspect(shape)} with handle #{shape_handle}")
 
-      :ok = start_shape(shape_handle, shape, state, otel_ctx)
+      :ok = start_shape(shape_handle, shape, state, otel_ctx, :create)
 
       # In this branch of `if`, we're guaranteed to have a newly started shape, so we can be sure about it's
       # "latest offset" because it'll be in the snapshotting stage
       {shape_handle, LogOffset.last_before_real_offsets()}
     end
   end
 
-  defp start_shape(shape_handle, shape, state, otel_ctx \\ nil) do
+  defp start_shape(shape_handle, shape, state, otel_ctx, action) do
     case Electric.Shapes.DynamicConsumerSupervisor.start_shape_consumer(
            state.consumer_supervisor,
            stack_id: state.stack_id,
@@ -376,7 +376,8 @@ defmodule Electric.ShapeCache do
            db_pool: state.db_pool,
            hibernate_after: state.shape_hibernate_after,
            otel_ctx: otel_ctx,
-           snapshot_timeout_to_first_data: state.snapshot_timeout_to_first_data
+           snapshot_timeout_to_first_data: state.snapshot_timeout_to_first_data,
+           action: action
          ) do
       {:ok, _supervisor_pid} ->
         :ok

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

@@ -76,6 +76,8 @@ defmodule Electric.Shapes.Consumer do
     Logger.metadata(metadata)
     Electric.Telemetry.Sentry.set_tags_context(metadata)
 
+    {action, config} = Map.pop!(config, :action)
+
     state =
       Map.merge(config, %{
         snapshot_started: false,
@@ -91,11 +93,11 @@ defmodule Electric.Shapes.Consumer do
         shape_status_mod: Map.get(config, :shape_status_mod) || Electric.ShapeCache.ShapeStatus
       })
 
-    {:ok, state, {:continue, :init_storage}}
+    {:ok, state, {:continue, {:init_consumer, action}}}
   end
 
   @impl GenServer
-  def handle_continue(:init_storage, state) do
+  def handle_continue({:init_consumer, action}, state) do
     %{
       storage: storage,
       shape_status_mod: shape_status_mod
@@ -135,7 +137,7 @@ defmodule Electric.Shapes.Consumer do
       Materializer.subscribe(state.stack_id, shape_handle)
     end
 
-    ShapeLogCollector.subscribe(state.stack_id, state.shape_handle, state.shape)
+    ShapeLogCollector.subscribe(state.stack_id, state.shape_handle, state.shape, action)
 
     Logger.debug("Writer for #{state.shape_handle} initialized")
 

packages/sync-service/lib/electric/shapes/consumer_supervisor.ex

@@ -21,6 +21,7 @@ defmodule Electric.Shapes.ConsumerSupervisor do
               default: :timer.seconds(30)
             ],
             hibernate_after: [type: :integer, required: true],
+            action: [type: {:in, [:restore, :create]}, default: :create],
             otel_ctx: [type: :any, required: false]
           )
 
@@ -109,7 +110,8 @@ defmodule Electric.Shapes.ConsumerSupervisor do
          shape_handle: shape_handle,
          shape_status_mod: Map.get(config, :shape_status_mod),
          stack_id: config.stack_id,
-         storage: shape_storage
+         storage: shape_storage,
+         action: config.action
        }}
     ]
 

packages/sync-service/test/electric/replication/shape_log_collector_test.exs

@@ -44,14 +44,25 @@ defmodule Electric.Replication.ShapeLogCollectorTest do
   @shape_handle "the-shape-handle"
 
   def setup_log_collector(ctx) do
+    %{stack_id: stack_id} = ctx
     # Start a test Registry
     registry_name = Module.concat(__MODULE__, Registry)
     start_link_supervised!({Registry, keys: :duplicate, name: registry_name})
 
+    existing_shapes = Map.get(ctx, :restore_shapes, [])
+
+    Repatch.patch(Electric.ShapeCache.ShapeStatus, :list_shapes, [mode: :shared], fn ^stack_id ->
+      existing_shapes
+    end)
+
+    Support.TestUtils.activate_mocks_for_descendant_procs(ShapeLogCollector)
+
+    inspector = Map.get(ctx, :inspector, {Mock.Inspector, elem(@inspector, 1)})
+
     # Start the ShapeLogCollector process
     opts = [
-      stack_id: ctx.stack_id,
-      inspector: {Mock.Inspector, elem(@inspector, 1)},
+      stack_id: stack_id,
+      inspector: inspector,
       persistent_kv: ctx.persistent_kv
     ]
 
@@ -61,16 +72,14 @@ defmodule Electric.Replication.ShapeLogCollectorTest do
       :ok
     end)
 
-    Repatch.allow(self(), pid)
-
     shape_cache_opts =
       [
         storage: {Mock.Storage, []},
         chunk_bytes_threshold: Electric.ShapeCache.LogChunker.default_chunk_size_threshold(),
         inspector: {Mock.Inspector, elem(@inspector, 1)},
         publication_manager: ctx.publication_manager,
-        stack_id: ctx.stack_id,
-        consumer_supervisor: Electric.Shapes.DynamicConsumerSupervisor.name(ctx.stack_id),
+        stack_id: stack_id,
+        consumer_supervisor: Electric.Shapes.DynamicConsumerSupervisor.name(stack_id),
         registry: registry_name
       ]
 
@@ -79,6 +88,44 @@ defmodule Electric.Replication.ShapeLogCollectorTest do
     %{server: pid, registry: registry_name, shape_cache: shape_cache_pid}
   end
 
+  describe "shape restoration" do
+    setup :setup_log_collector
+
+    @tag restore_shapes: [{@shape_handle, @shape}], inspector: @inspector
+    test "populates the filter, partitions and layers from the shape_status table", ctx do
+      parent = self()
+
+      xmin = 100
+      lsn = Lsn.from_string("0/10")
+      last_log_offset = LogOffset.new(lsn, 0)
+
+      consumer =
+        start_link_supervised!(
+          {Support.TransactionConsumer,
+           [
+             id: 1,
+             parent: parent,
+             producer: ctx.server,
+             shape: @shape,
+             shape_handle: @shape_handle,
+             action: :restore
+           ]}
+        )
+
+      txn =
+        %Transaction{xid: xmin, lsn: lsn, last_log_offset: last_log_offset}
+        |> Transaction.prepend_change(%Changes.NewRecord{
+          relation: {"public", "test_table"},
+          record: %{"id" => "2", "name" => "foo"}
+        })
+
+      assert :ok = ShapeLogCollector.store_transaction(txn, ctx.server)
+
+      xids = Support.TransactionConsumer.assert_consume([{1, consumer}], [txn])
+      assert xids == [xmin]
+    end
+  end
+
   describe "store_transaction/2" do
     setup :setup_log_collector
 

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

@@ -975,8 +975,8 @@ defmodule Electric.ShapeCacheTest do
 
       context =
         context
-        |> Map.merge(with_shape_log_collector(context))
         |> Map.merge(with_shape_status(context))
+        |> Map.merge(with_shape_log_collector(context))
 
       with_shape_cache(context, opts)
     end

packages/sync-service/test/support/transaction_consumer.ex

@@ -41,9 +41,11 @@ defmodule Support.TransactionConsumer do
     {:ok, producer} = Keyword.fetch(opts, :producer)
     {:ok, parent} = Keyword.fetch(opts, :parent)
     {:ok, id} = Keyword.fetch(opts, :id)
+    action = Keyword.get(opts, :action, :create)
     shape = Keyword.fetch!(opts, :shape)
     shape_handle = Keyword.fetch!(opts, :shape_handle)
-    Electric.Replication.ShapeLogCollector.subscribe(producer, shape_handle, shape)
+
+    Electric.Replication.ShapeLogCollector.subscribe(producer, shape_handle, shape, action)
 
     {:ok, {id, parent}}
   end