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

.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,79 @@
+[doc Verify that consumer processes suspend and resume correctly]
+
+[include _macros.luxinc]
+
+[global pg_container_name=consumer-suspend-resume__pg]
+
+###
+
+[invoke setup_pg]
+
+# 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
+
+# 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 ELECTRIC_SHAPE_SUSPEND_CONSUMER=true ELECTRIC_FEATURE_FLAGS=suspend_consumers"]
+[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]
+
+  ??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/config/runtime.exs

@@ -172,6 +172,8 @@ call_home_telemetry_url =
 shape_hibernate_after =
   env!("ELECTRIC_SHAPE_HIBERNATE_AFTER", &Electric.Config.parse_human_readable_time!/1, nil)
 
+shape_enable_suspend? = env!("ELECTRIC_SHAPE_SUSPEND_CONSUMER", :boolean, nil)
+
 system_metrics_poll_interval =
   env!(
     "ELECTRIC_SYSTEM_METRICS_POLL_INTERVAL",
@@ -234,6 +236,7 @@ config :electric,
   max_txn_size: env!("ELECTRIC_EXPERIMENTAL_MAX_TXN_SIZE", :integer, nil),
   service_port: env!("ELECTRIC_PORT", :integer, nil),
   shape_hibernate_after: shape_hibernate_after,
+  shape_enable_suspend?: shape_enable_suspend?,
   storage_dir: storage_dir,
   storage: storage_spec,
   cleanup_interval_ms:

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

@@ -38,7 +38,7 @@ defmodule Electric.Config do
 
   @build_env Mix.env()
 
-  @known_feature_flags ~w[allow_subqueries]
+  @known_feature_flags ~w[allow_subqueries suspend_consumers]
 
   @defaults [
     ## Database
@@ -99,6 +99,9 @@ defmodule Electric.Config do
     telemetry_long_message_queue_disable_threshold: 0,
     ## Memory
     shape_hibernate_after: :timer.seconds(30),
+    # Should we terminate consumer processes after `shape_hibernate_after` ms
+    # or just hibernate them?
+    shape_enable_suspend?: true,
     ## Performance tweaks
     publication_alter_debounce_ms: 0,
     ## Misc

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,33 @@ 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. Consumer suspend has been enabled in the stack config
+    # 2. we're not waiting for snapshot information
+    # 3. 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
+
+    if consumer_suspend_enabled?(state) and consumer_can_suspend?(state) 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}
+    end
+  end
+
+  defp consumer_suspend_enabled?(%{stack_id: stack_id}) do
+    # TODO: remove additional feature flag check once feature is stabilised
+    Electric.StackConfig.lookup(stack_id, :shape_enable_suspend?, true) and
+      "suspend_consumers" in Electric.StackConfig.lookup(stack_id, :feature_flags, [])
+  end
 
-    {:noreply, state, :hibernate}
+  defp consumer_can_suspend?(state) do
+    state.snapshot_started and Enum.empty?(state.shape.shape_dependencies_handles) and
+      not state.materializer_subscribed?
   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
@@ -150,7 +179,12 @@ defmodule Electric.Shapes.ConsumerRegistry do
 
   @doc false
   def registry_table(stack_id) do
-    :ets.new(ets_name(stack_id), [:public, :named_table, write_concurrency: :auto])
+    :ets.new(ets_name(stack_id), [
+      :public,
+      :named_table,
+      write_concurrency: :auto,
+      read_concurrency: true
+    ])
   end
 
   def new(stack_id, opts \\ []) when is_binary(stack_id) do

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

@@ -23,7 +23,9 @@ defmodule Electric.StackConfig do
     [
       snapshot_timeout_to_first_data: :timer.seconds(30),
       shape_hibernate_after: Electric.Config.default(:shape_hibernate_after),
-      chunk_bytes_threshold: Electric.ShapeCache.LogChunker.default_chunk_size_threshold()
+      shape_enable_suspend?: Electric.Config.default(:shape_enable_suspend?),
+      chunk_bytes_threshold: Electric.ShapeCache.LogChunker.default_chunk_size_threshold(),
+      feature_flags: []
     ]
   end
 

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

@@ -124,6 +124,10 @@ defmodule Electric.StackSupervisor do
                        type: :integer,
                        default: Electric.Config.default(:shape_hibernate_after)
                      ],
+                     shape_enable_suspend?: [
+                       type: :boolean,
+                       default: Electric.Config.default(:shape_enable_suspend?)
+                     ],
                      snapshot_timeout_to_first_data: [
                        type: :pos_integer,
                        default: Electric.Config.default(:snapshot_timeout_to_first_data)
@@ -302,6 +306,7 @@ defmodule Electric.StackSupervisor do
     shape_changes_registry_name = registry_name(stack_id)
 
     shape_hibernate_after = Keyword.fetch!(config.tweaks, :shape_hibernate_after)
+    shape_enable_suspend? = Keyword.fetch!(config.tweaks, :shape_enable_suspend?)
 
     shape_cache_opts = [
       stack_id: stack_id
@@ -361,7 +366,9 @@ defmodule Electric.StackSupervisor do
              chunk_bytes_threshold: config.chunk_bytes_threshold,
              snapshot_timeout_to_first_data: config.tweaks[:snapshot_timeout_to_first_data],
              inspector: inspector,
-             shape_hibernate_after: shape_hibernate_after
+             shape_hibernate_after: shape_hibernate_after,
+             shape_enable_suspend?: shape_enable_suspend?,
+             feature_flags: Map.get(config, :feature_flags, [])
            ]},
           {Electric.AsyncDeleter,
            stack_id: stack_id,