feat(sync-service): Terminate shape consumer processes after timeout

Closes #3253

Author: magnetised

.changeset/eighty-pandas-ring.md

@@ -0,0 +1,5 @@
+---
+'@core/sync-service': patch
+---
+
+Reduce memory usage by terminating consumer processes after the hibernation timeout

integration-tests/scripts/electric_dev.sh

@@ -2,8 +2,8 @@
 
 set -e
 
-SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &> /dev/null && pwd)
+SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)
 
 cd "$SCRIPT_DIR"/../../packages/sync-service
 
-ELECTRIC_STORAGE_DIR="$SCRIPT_DIR/../_storage" ELECTRIC_REPLICATION_STREAM_ID=integration iex -r "$SCRIPT_DIR/../test_utils/*.exs" -S mix
+ELECTRIC_STORAGE_DIR="$SCRIPT_DIR/../_storage" ELECTRIC_REPLICATION_STREAM_ID=integration iex -r "$SCRIPT_DIR/../test_utils/*.exs" "$@" -S mix

integration-tests/tests/_macros.luxinc

@@ -186,7 +186,7 @@
 
 [macro start_electric_script shell_name port env]
   [shell $shell_name]
-    !ELECTRIC_PORT=$port $env ../scripts/electric_dev.sh
+    !ELECTRIC_PORT=$port $env ../scripts/electric_dev.sh --no-color
 [endmacro]
 
 [macro stop_electric]

integration-tests/tests/shape-suspension-resumption.lux

@@ -0,0 +1,92 @@
+[doc Verify that consumer processes suspend and resume correctly]
+
+[include _macros.luxinc]
+
+[global pg_container_name=db-connection-scaledown__pg]
+[global database_url=postgresql://electric_test:password@localhost:$pg_host_port/electric?sslmode=disable]
+
+###
+
+# Start Postgres and create an additional role that will be used by Electric in this test
+[invoke setup_pg_with_shell_name \
+  "pg" \
+  "-e INIT_DB_SQL=\"\
+    CREATE ROLE electric_test LOGIN PASSWORD 'password' REPLICATION;\
+    GRANT CREATE ON DATABASE electric TO electric_test\"" \
+  "" \
+  "-v $(realpath ../scripts/init_db.sh):/docker-entrypoint-initdb.d/initdb-init_db.sh"]
+
+# Create a table for subsequent shape requests
+[invoke start_psql]
+
+[shell psql]
+  !create table items(val text);
+  ??CREATE
+
+  !insert into items values ('1'), ('2');
+  ??INSERT
+
+  !alter table items owner to electric_test;
+  ??ALTER TABLE
+
+# Start Electric and wait for it to finish initialization.
+# Set a very small hibernation timeout so consumers will shutdown quickly
+[invoke setup_electric_with_env "ELECTRIC_SHAPE_HIBERNATE_AFTER=200ms"]
+[shell electric]
+  [timeout 10]
+  ??[debug] Replication client started streaming
+
+# Start a live shape request
+[shell client]
+  [invoke curl_shape "http://localhost:3000/v1/shape?table=items&offset=-1"]
+
+  ??HTTP/1.1 200 OK
+
+  ?electric-handle: ([\d-]+)
+  [global handle=$1]
+
+  ??electric-offset: 0_0
+
+  ??[{"headers":{"operation":"insert","relation":["public","items"]},"key":"\"public\".\"items\"/\"1\"","value":{"val":"1"}},\
+     {"headers":{"operation":"insert","relation":["public","items"]},"key":"\"public\".\"items\"/\"2\"","value":{"val":"2"}},\
+     {"headers":{"control":"snapshot-end"
+
+  [invoke curl_shape "http://localhost:3000/v1/shape?table=items&handle=$handle&offset=0_0&live"]
+
+# Check that the consumer process for the given handle has suspended itself
+[shell electric]
+  [sleep 1]
+
+  ??[debug] Suspending consumer $handle
+  !{:active_consumer_count, Electric.Shapes.ConsumerRegistry.active_consumer_count("single_stack")}
+  ?{:active_consumer_count, 0}
+
+# Verify that writes to the handle re-create the consumer process
+[shell psql]
+  !insert into items values ('3');
+  ??INSERT
+
+
+[shell electric]
+  ??[info] Started consumer for existing handle $handle
+  !{:active_consumer_count, Electric.Shapes.ConsumerRegistry.active_consumer_count("single_stack")}
+  ?{:active_consumer_count, 1}
+
+[shell client]
+  [timeout 1]
+
+  [invoke curl_shape "http://localhost:3000/v1/shape?table=items&handle=$handle&offset=0_0&live"]
+
+  ??HTTP/1.1 200
+
+  ??transfer-encoding: chunked
+  ?electric-offset: ([0-9]+)_0
+  [local global_last_seen_lsn=$1]
+
+  ??"headers":{"last":true,"lsn":"$global_last_seen_lsn","op_position":0,"operation":"insert","relation":["public","items"],"txids":
+  ??"key":"\"public\".\"items\"/\"3\"","value":{"val":"3"}
+  ??"headers":{"control":"up-to-date","global_last_seen_lsn":"$global_last_seen_lsn"}
+
+
+[cleanup]
+  [invoke teardown]

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

@@ -15,9 +15,11 @@ defmodule Electric.ShapeCache.ShapeCleaner do
   @type stack_id() :: Electric.stack_id()
 
   @shutdown_cleanup {:shutdown, :cleanup}
+  @shutdown_suspend {:shutdown, :suspend}
 
   # Public API
   def consumer_cleanup_reason, do: @shutdown_cleanup
+  def consumer_suspend_reason, do: @shutdown_suspend
 
   @spec remove_shapes(stack_id(), [shape_handle()], term()) :: :ok | {:error, term()}
   def remove_shapes(stack_id, shape_handles, reason \\ @shutdown_cleanup)
@@ -76,6 +78,8 @@ defmodule Electric.ShapeCache.ShapeCleaner do
     end)
   end
 
+  @type reason() :: {:shutdown, :cleanup} | {:shutdown, :suspend} | term()
+  @spec handle_writer_termination(stack_id(), shape_handle(), reason()) :: :removed | :ok
   def handle_writer_termination(stack_id, shape_handle, @shutdown_cleanup) do
     Logger.info("Removing shape #{inspect(shape_handle)}")
 
@@ -84,6 +88,13 @@ defmodule Electric.ShapeCache.ShapeCleaner do
     :removed
   end
 
+  def handle_writer_termination(stack_id, shape_handle, @shutdown_suspend) do
+    # deregister the consumer without removing it from the rest of the system
+    # the next time a txn comes in matching this consumer it will be re-started
+    # by the consumer registry as per any other lazily loaded consumer
+    Electric.Shapes.ConsumerRegistry.remove_consumer(shape_handle, stack_id)
+  end
+
   def handle_writer_termination(_stack_id, _shape_handle, reason)
       when reason in [:normal, :killed, :shutdown] or
              (is_tuple(reason) and elem(reason, 0) == :shutdown) do

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

@@ -362,9 +362,23 @@ defmodule Electric.Shapes.Consumer do
   end
 
   def handle_info(:timeout, state) do
-    state = %{state | writer: ShapeCache.Storage.hibernate(state.writer)}
+    # we can only suspend (terminate) the consumer process if
+    # 1. we're not waiting for snapshot information
+    # 2. we are not part of a subquery dependency tree, that is either
+    #   a. we have no dependent shapes
+    #   b. we don't have a materializer subscribed
+    can_suspend? =
+      state.snapshot_started and Enum.empty?(state.shape.shape_dependencies_handles) and
+        not state.materializer_subscribed?
+
+    if can_suspend? do
+      Logger.debug(fn -> ["Suspending consumer ", to_string(state.shape_handle)] end)
+      {:stop, ShapeCleaner.consumer_suspend_reason(), state}
+    else
+      state = %{state | writer: ShapeCache.Storage.hibernate(state.writer)}
 
-    {:noreply, state, :hibernate}
+      {:noreply, state, :hibernate}
+    end
   end
 
   @impl GenServer

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

@@ -17,6 +17,8 @@ defmodule Electric.Shapes.ConsumerRegistry do
           stack_id: stack_id()
         }
 
+  @consumer_suspend_reason Electric.ShapeCache.ShapeCleaner.consumer_suspend_reason()
+
   def name(stack_id, shape_handle) when is_stack_id(stack_id) and is_shape_handle(shape_handle) do
     {:via, __MODULE__, {stack_id, shape_handle}}
   end
@@ -67,22 +69,36 @@ defmodule Electric.Shapes.ConsumerRegistry do
   end
 
   @spec publish([shape_handle()], term(), t()) :: :ok
+  def publish([], _event, _registry_state) do
+    :ok
+  end
+
   def publish(shape_handles, event, registry_state) do
     %{table: table} = registry_state
 
     shape_handles
-    |> Enum.flat_map(fn handle ->
-      (consumer_pid(handle, table) || start_consumer!(handle, registry_state))
-      |> List.wrap()
+    |> Enum.map(fn handle ->
+      {handle, consumer_pid(handle, table) || start_consumer!(handle, registry_state)}
     end)
     |> broadcast(event)
+    |> publish(event, registry_state)
   end
 
   @spec remove_consumer(shape_handle(), t()) :: :ok
   def remove_consumer(shape_handle, %__MODULE__{table: table}) do
+    do_remove_consumer(shape_handle, table)
+  end
+
+  @spec remove_consumer(shape_handle(), stack_id()) :: :ok
+  def remove_consumer(shape_handle, stack_id) when is_stack_id(stack_id) do
+    do_remove_consumer(shape_handle, ets_name(stack_id))
+  end
+
+  @spec do_remove_consumer(shape_handle(), :ets.table()) :: :ok
+  defp do_remove_consumer(shape_handle, table) when is_atom(table) or is_reference(table) do
     :ets.delete(table, shape_handle)
 
-    Logger.debug(fn -> "Stopped and removed consumer #{shape_handle}" end)
+    Logger.debug(fn -> "Removed consumer #{shape_handle}" end)
 
     :ok
   end
@@ -96,26 +112,39 @@ defmodule Electric.Shapes.ConsumerRegistry do
   There is no timeout so if the GenServers do not respond or die, this
   function will block indefinitely.
   """
-  @spec broadcast([pid()], term()) :: :ok
-  def broadcast(pids, message) do
+  @spec broadcast([{shape_handle(), pid()}], term()) :: [shape_handle()]
+  def broadcast(handle_pids, message) do
     # Based on OTP GenServer.call, see:
     # https://github.com/erlang/otp/blob/090c308d7c925e154240685174addaa516ea2f69/lib/stdlib/src/gen.erl#L243
-    pids
-    |> Enum.map(fn pid ->
+    handle_pids
+    |> Enum.map(fn {handle, pid} ->
       ref = Process.monitor(pid)
       send(pid, {:"$gen_call", {self(), ref}, message})
-      ref
+      {handle, ref}
     end)
-    |> Enum.each(fn ref ->
+    |> Enum.flat_map(fn {handle, ref} ->
       receive do
         {^ref, _reply} ->
           Process.demonitor(ref, [:flush])
-          :ok
+          []
+
+        {:DOWN, ^ref, _, _, @consumer_suspend_reason} ->
+          # Catch the race condition where a consumer is in the act of
+          # suspending as the txn arrives by retrying those handles (which will
+          # start a new consumer instance).
+          [handle]
 
         {:DOWN, ^ref, _, _, _reason} ->
-          :ok
+          []
       end
     end)
+    |> tap(fn
+      [] ->
+        :ok
+
+      suspended_handles ->
+        Logger.debug(fn -> ["Re-trying suspended shape handles ", inspect(suspended_handles)] end)
+    end)
   end
 
   defp consumer_pid(handle, table) do