chore(sync-service): Recover shapes without holding supervisor up (#3293)

This is a small part of the shape subsystem work. This PR recovers
shapes without holding it's supervisor up. Eventually this will mean we
can recover shapes asynchronously while the database connection is being
established.

The effect of this currently is that the ExpiryManager will start before
the shapes have been restores so this PR also makes the ExpiryManager
check if the stack is ready before expiring.

I've also used this opportunity to re-write the ExpiryManager tests that
were a mess (I wrote them)

The other process this affects is the SchemaReconciler, it may now
reconcile before the shapes have been restored but given that it loads
its own data from PersistentKV I don't think this is an issue (but
please let me know if you know otherwise!)

Author: robacourt

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

@@ -220,28 +220,24 @@ defmodule Electric.ShapeCache do
       snapshot_timeout_to_first_data: opts.snapshot_timeout_to_first_data
     }
 
+    {:ok, state, {:continue, {:recover_shapes, opts.recover_shape_timeout}}}
+  end
+
+  @impl GenServer
+  def handle_continue({:recover_shapes, recover_shape_timeout}, state) do
     {last_processed_lsn, total_recovered, total_failed_to_recover} =
-      recover_shapes(state, opts.recover_shape_timeout)
+      recover_shapes(state, recover_shape_timeout)
 
     # Empirical evidence shows that after recovering 50K shapes ShapeStatusOwner and ShapeCache
     # each take up 200+MB of memory. Explicitly running garbage collection for both immediately
     # takes that down to 4-5MB.
     :erlang.garbage_collect()
 
-    # 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.
-    {:ok, state,
-     {:continue, {:consumers_ready, last_processed_lsn, total_recovered, total_failed_to_recover}}}
-  end
-
-  @impl GenServer
-  def handle_continue(
-        {: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)
 
+    # 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)
 
     Electric.Connection.Manager.consumers_ready(

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

@@ -2,6 +2,7 @@ defmodule Electric.ShapeCache.ExpiryManager do
   use GenServer
 
   alias Electric.ShapeCache.ShapeStatus
+  alias Electric.StatusMonitor
   alias Electric.Telemetry.OpenTelemetry
 
   require Logger
@@ -60,10 +61,19 @@ defmodule Electric.ShapeCache.ExpiryManager do
   defp maybe_expire_shapes(%{max_shapes: nil}), do: :ok
 
   defp maybe_expire_shapes(%{max_shapes: max_shapes} = state) do
-    shape_count = shape_count(state)
-
-    if shape_count > max_shapes do
-      expire_shapes(shape_count, state)
+    case StatusMonitor.status(state.stack_id) do
+      %{shape: :up} ->
+        shape_count = shape_count(state)
+
+        if shape_count > max_shapes do
+          expire_shapes(shape_count, state)
+        end
+
+      status ->
+        # We do not expire shapes if the stack is not active since this may mean that
+        # shapes have not fully restored yet and we don't want to expire while restoring
+        # as this may cause race conditions.
+        Logger.debug("Expiry check skipped due to inactive stack: #{inspect(status)}")
     end
   end
 

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

@@ -84,21 +84,17 @@ defmodule Electric.ShapeCache.ShapeStatus do
   @snapshot_started :snapshot_started
 
   @impl true
-  def initialise(stack_ref, storage) do
+  def initialise(stack_ref, storage \\ nil) do
     last_used_table = shape_last_used_table(stack_ref)
     :ets.new(last_used_table, [:named_table, :public, :ordered_set])
 
     meta_table = shape_meta_table(stack_ref)
 
-    case load_table_backup(meta_table, storage) do
-      {:ok, ^meta_table, path} ->
-        Logger.info("Loaded shape status from backup at #{path}")
-        :ok
-
-      _ ->
-        Logger.debug("No shape status backup loaded, creating new table #{meta_table}")
-        :ets.new(meta_table, [:named_table, :public, :ordered_set])
-        load(meta_table, last_used_table, storage)
+    if storage do
+      load_from_storage(stack_ref, storage)
+    else
+      :ets.new(meta_table, [:named_table, :public, :ordered_set])
+      :ok
     end
   end
 
@@ -416,6 +412,22 @@ defmodule Electric.ShapeCache.ShapeStatus do
     end
   end
 
+  defp load_from_storage(stack_ref, storage) do
+    last_used_table = shape_last_used_table(stack_ref)
+    meta_table = shape_meta_table(stack_ref)
+
+    case load_table_backup(meta_table, storage) do
+      {:ok, ^meta_table, path} ->
+        Logger.info("Loaded shape status from backup at #{path}")
+        :ok
+
+      _ ->
+        Logger.debug("No shape status backup loaded, creating new table #{meta_table}")
+        :ets.new(meta_table, [:named_table, :public, :ordered_set])
+        load(meta_table, last_used_table, storage)
+    end
+  end
+
   defp load_table_backup(meta_table, storage) do
     case backup_dir(storage) do
       nil ->

packages/sync-service/test/electric/shape_cache/expiry_manager_test.exs

@@ -1,138 +1,79 @@
 defmodule Electric.ExpiryManagerTest do
   use ExUnit.Case, async: true
-  use Support.Mock
   use Repatch.ExUnit
 
-  alias Electric.Replication.Changes
-  alias Electric.Replication.LogOffset
-  alias Electric.ShapeCache
   alias Electric.ShapeCache.ExpiryManager
-  alias Electric.ShapeCache.Storage
+  alias Electric.ShapeCache.ShapeCleaner
+  alias Electric.ShapeCache.ShapeStatus
   alias Electric.Shapes.Shape
+  alias Support.RepatchExt
 
-  import Mox
   import Support.ComponentSetup
   import Support.TestUtils
 
-  @stub_inspector Support.StubInspector.new(
-                    tables: [{1, {"public", "items"}}],
-                    columns: [
-                      %{
-                        name: "id",
-                        type: "int8",
-                        type_id: {20, 1},
-                        pk_position: 0,
-                        is_generated: false
-                      },
-                      %{name: "value", type: "text", type_id: {25, 1}, is_generated: false}
-                    ]
-                  )
-  @shape Shape.new!("items", inspector: @stub_inspector)
-
-  # {xmin, xmax, xip_list}
-  @pg_snapshot_xmin_10 {10, 11, [10]}
-
-  @moduletag :tmp_dir
-
-  defmodule TempPubManager do
-    def add_shape(_handle, _, opts) do
-      send(opts[:test_pid], {:called, :prepare_tables_fn})
-    end
-  end
-
-  setup :verify_on_exit!
-
-  setup do
-    %{inspector: @stub_inspector, pool: nil}
-  end
-
   setup [
-    :with_persistent_kv,
     :with_stack_id_from_test,
-    :with_async_deleter,
-    :with_pure_file_storage,
-    :with_shape_status,
-    :with_shape_cleaner,
-    :with_status_monitor,
-    :with_shape_monitor,
-    :with_log_chunking,
-    :with_registry,
-    :with_shape_log_collector,
-    :with_noop_publication_manager
+    :with_status_monitor
   ]
 
-  test "expires shapes if shape count has gone over max_shapes", ctx do
-    Support.TestUtils.patch_snapshotter(fn parent, shape_handle, _shape, %{storage: storage} ->
-      GenServer.cast(parent, {:pg_snapshot_known, shape_handle, @pg_snapshot_xmin_10})
-      Storage.make_new_snapshot!([["test"]], storage)
-      GenServer.cast(parent, {:snapshot_started, shape_handle})
-    end)
-
-    %{shape_cache_opts: opts} = with_shape_cache(ctx)
-
-    start_supervised!(
-      {ExpiryManager, max_shapes: 1, expiry_batch_size: 1, period: 10, stack_id: ctx.stack_id}
-    )
-
-    {shape_handle, _} = ShapeCache.get_or_create_shape_handle(@shape, opts)
-    assert :started = ShapeCache.await_snapshot_start(shape_handle, opts)
+  @max_shapes 10
 
-    consumer_ref =
-      Electric.Shapes.Consumer.whereis(ctx.stack_id, shape_handle)
-      |> Process.monitor()
+  setup %{stack_id: stack_id} do
+    ShapeStatus.initialise(stack_id)
 
-    storage = Storage.for_shape(shape_handle, ctx.storage)
-    writer = Storage.init_writer!(storage, @shape)
-
-    Storage.append_to_log!(
-      changes_to_log_items([
-        %Changes.NewRecord{
-          relation: {"public", "items"},
-          record: %{"id" => "1", "value" => "Alice"},
-          log_offset: LogOffset.new(Electric.Postgres.Lsn.from_integer(1000), 0)
-        }
-      ]),
-      writer
-    )
+    expiry_manager =
+      start_supervised!(
+        {ExpiryManager,
+         max_shapes: @max_shapes, expiry_batch_size: 1, period: 1, stack_id: stack_id}
+      )
 
-    assert Storage.snapshot_started?(storage)
+    Repatch.patch(ShapeCleaner, :remove_shape, [mode: :shared], fn shape_handle,
+                                                                   stack_id: stack_id ->
+      ShapeStatus.remove_shape(stack_id, shape_handle)
+    end)
 
-    assert Enum.count(Storage.get_log_stream(LogOffset.last_before_real_offsets(), storage)) ==
-             1
+    Repatch.allow(self(), expiry_manager)
+    %{expiry_manager: expiry_manager}
+  end
 
-    {new_shape_handle, _} =
-      ShapeCache.get_or_create_shape_handle(%{@shape | where: "1 == 1"}, opts)
+  describe "when stack is active" do
+    setup :set_status_to_active
 
-    assert :started = ShapeCache.await_snapshot_start(new_shape_handle, opts)
+    test "expires shapes if shape count has gone over max_shapes", ctx do
+      for i <- 1..(@max_shapes + 1) do
+        ShapeStatus.add_shape(ctx.stack_id, create_shape(i))
+      end
 
-    assert_receive {:DOWN, ^consumer_ref, :process, _pid, {:shutdown, :cleanup}}
+      assert RepatchExt.called_within_ms?(ShapeCleaner, :remove_shape, 2, 50, ctx.expiry_manager)
+    end
 
-    assert :ok = await_for_storage_to_raise(storage)
+    test "does not expires shapes if shape count has not gone over max_shapes", ctx do
+      for i <- 1..@max_shapes do
+        ShapeStatus.add_shape(ctx.stack_id, create_shape(i))
+      end
 
-    {shape_handle2, _} = ShapeCache.get_or_create_shape_handle(@shape, opts)
-    assert shape_handle != shape_handle2
-    assert :started = ShapeCache.await_snapshot_start(shape_handle2, opts)
+      refute RepatchExt.called_within_ms?(ShapeCleaner, :remove_shape, 2, 50, ctx.expiry_manager)
+    end
   end
 
-  defp await_for_storage_to_raise(storage, timeout \\ 5_000)
+  describe "when stack is not active" do
+    test "does not expires shapes even if shape count has gone over max_shapes", ctx do
+      for i <- 1..(@max_shapes + 1) do
+        ShapeStatus.add_shape(ctx.stack_id, create_shape(i))
+      end
 
-  defp await_for_storage_to_raise(_storage, timeout) when timeout <= 0 do
-    raise "Storage did not raise Storage.Error in time"
+      refute RepatchExt.called_within_ms?(ShapeCleaner, :remove_shape, 2, 50, ctx.expiry_manager)
+    end
   end
 
-  defp await_for_storage_to_raise(storage, timeout) do
-    try do
-      start_time = System.monotonic_time()
-      Stream.run(Storage.get_log_stream(LogOffset.before_all(), storage))
-      Process.sleep(50)
-      elapsed = System.monotonic_time() - start_time
+  @inspector Support.StubInspector.new(
+               tables: ["t1"],
+               columns: [
+                 %{name: "id", type: "int8", pk_position: 0}
+               ]
+             )
 
-      await_for_storage_to_raise(
-        storage,
-        timeout - System.convert_time_unit(elapsed, :native, :millisecond)
-      )
-    rescue
-      Storage.Error -> :ok
-    end
+  defp create_shape(id) do
+    Shape.new!("t1", where: "id = #{id}", inspector: @inspector)
   end
 end

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

@@ -1,19 +1,19 @@
 defmodule Support.RepatchExt do
-  def called_within_ms?(module, function, args, ms) do
-    called_within_ms?(module, function, args, ms, System.monotonic_time(:millisecond))
+  def called_within_ms?(module, function, args, ms, by \\ :any) do
+    called_within_ms?(module, function, args, ms, by, System.monotonic_time(:millisecond))
   end
 
-  defp called_within_ms?(module, function, args, ms, start_time) do
+  defp called_within_ms?(module, function, args, ms, by, start_time) do
     cond do
       System.monotonic_time(:millisecond) - start_time > ms ->
         false
 
-      Repatch.called?(module, function, args, by: :any) ->
+      Repatch.called?(module, function, args, by: by) ->
         true
 
       true ->
         Process.sleep(1)
-        called_within_ms?(module, function, args, ms, start_time)
+        called_within_ms?(module, function, args, ms, by, start_time)
     end
   end
 end