fix(sync-service): Fix issue where replication traffic slows down shape restoration (#2840)

Fixes #2767 by not starting the replication client streaming until the
consumers are ready. This avoids the need for the consumers to buffer
messages until ready, which also reduces memory consumption and uses one
less feature of GenStage making it easier to replace.

Author: robacourt

.changeset/loud-eagles-end.md

@@ -0,0 +1,5 @@
+---
+"@core/sync-service": patch
+---
+
+Fix issue where replication traffic slows down shape restoration

packages/sync-service/lib/electric/connection/manager.ex

@@ -64,6 +64,7 @@ defmodule Electric.Connection.Manager do
             | {:start_replication_client, :configuring_connection}
             | :start_connection_pool
             | :start_shapes_supervisor
+            | :waiting_for_consumers
             | {:start_replication_client, :start_streaming}
             # Steps of the :running phase:
             | :waiting_for_streaming_confirmation
@@ -193,6 +194,10 @@ defmodule Electric.Connection.Manager do
     GenServer.cast(manager, :lock_connection_started)
   end
 
+  def consumers_ready(stack_id) do
+    GenServer.cast(name(stack_id), :consumers_ready)
+  end
+
   def exclusive_connection_lock_acquisition_failed(manager, error) do
     GenServer.cast(manager, {:exclusive_connection_lock_acquisition_failed, error})
   end
@@ -476,11 +481,11 @@ defmodule Electric.Connection.Manager do
     state = %State{
       state
       | shape_log_collector_pid: log_collector_pid,
-        current_step: {:start_replication_client, :start_streaming},
+        current_step: :waiting_for_consumers,
         purge_all_shapes?: false
     }
 
-    {:noreply, state, {:continue, :start_streaming}}
+    {:noreply, state}
   end
 
   def handle_continue(
@@ -877,6 +882,17 @@ defmodule Electric.Connection.Manager do
     end
   end
 
+  def handle_cast(
+        :consumers_ready,
+        %State{
+          current_phase: :connection_setup,
+          current_step: :waiting_for_consumers
+        } = state
+      ) do
+    state = %State{state | current_step: {:start_replication_client, :start_streaming}}
+    {:noreply, state, {:continue, :start_streaming}}
+  end
+
   def handle_cast(
         :replication_client_streamed_first_message,
         %State{current_phase: :running, current_step: :waiting_for_streaming_confirmation} = state

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

@@ -33,10 +33,10 @@ defmodule Electric.Replication.ShapeLogCollector do
     Electric.ProcessRegistry.name(stack_id, __MODULE__)
   end
 
-  def start_processing(server, last_processed_lsn) do
+  def set_last_processed_lsn(server, last_processed_lsn) do
     # Allow 60s for this call as it may need to wait for thousands of restored shapes
     # to subscribe before it returns.
-    GenStage.call(server, {:start_processing, last_processed_lsn}, 60_000)
+    GenStage.call(server, {:set_last_processed_lsn, last_processed_lsn}, 60_000)
   end
 
   # use `GenStage.call/2` here to make the event processing synchronous.
@@ -87,10 +87,7 @@ defmodule Electric.Replication.ShapeLogCollector do
         tracked_relations: tracker_state
       })
 
-    # start in demand: :accumulate mode so that the ShapeCache is able to start
-    # all active consumers before we start sending transactions
-    {:producer, state,
-     dispatcher: {Electric.Shapes.Dispatcher, inspector: state.inspector}, demand: :accumulate}
+    {:producer, state, dispatcher: {Electric.Shapes.Dispatcher, inspector: state.inspector}}
   end
 
   def handle_subscribe(:consumer, _opts, from, state) do
@@ -131,9 +128,8 @@ defmodule Electric.Replication.ShapeLogCollector do
     {:noreply, [], remove_subscription(from, state)}
   end
 
-  def handle_call({:start_processing, lsn}, _from, state) do
+  def handle_call({:set_last_processed_lsn, lsn}, _from, state) do
     LsnTracker.init(lsn, state.stack_id)
-    GenStage.demand(self(), :forward)
     Electric.StatusMonitor.mark_shape_log_collector_ready(state.stack_id, self())
     {:reply, :ok, [], Map.put(state, :last_processed_lsn, lsn)}
   end

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

@@ -266,7 +266,9 @@ defmodule Electric.ShapeCache do
 
   @impl GenServer
   def handle_continue({:consumers_ready, last_processed_lsn}, state) do
-    ShapeLogCollector.start_processing(state.log_producer, last_processed_lsn)
+    ShapeLogCollector.set_last_processed_lsn(state.log_producer, last_processed_lsn)
+    Electric.Connection.Manager.consumers_ready(state.stack_id)
+
     {:noreply, state}
   end
 

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

@@ -352,7 +352,7 @@ defmodule Electric.Replication.ShapeLogCollectorTest do
     prev_lsn = Lsn.increment(start_lsn, -1)
     next_lsn = Lsn.increment(start_lsn, +1)
 
-    ShapeLogCollector.start_processing(pid, start_lsn)
+    ShapeLogCollector.set_last_processed_lsn(pid, start_lsn)
 
     assert start_lsn == LsnTracker.get_last_processed_lsn(ctx.stack_id)
 

packages/sync-service/test/electric/shapes/consumer_test.exs

@@ -118,7 +118,7 @@ defmodule Electric.Shapes.ConsumerTest do
           inspector: @base_inspector
         )
 
-      ShapeLogCollector.start_processing(producer, Lsn.from_integer(0))
+      ShapeLogCollector.set_last_processed_lsn(producer, Lsn.from_integer(0))
 
       consumers =
         for {shape_handle, shape} <- ctx.shapes do