feat: remove full table info from shape definitions to lower mem (#2454)

## Reasons

We have a process setup where we have 2 layers of supervisors which
store entire child spec of their children as state, plus the
`ShapeLogCollector` that the shape consumers register with, plus the
consumers themselves -- all these processes happen to have shapes in
their state. _Arguably, we can inject shapes into consumers bypassing
supervisors, but that's a different PR._ These shapes, however, happened
to know entire schema of the table they were defined against, which was
duplicated a lot. This caused a lot of unnecessary memory usage when the
table schema was sufficiently large.

So, in numbers, for a 75 column table, without any transactions going
through (just shape creation)
Baseline, no shapes: 71kB
- Current, 300 shapes
	- Total (for shape-related processes): 64.1MB
	- Total overhead in shared processes: 28MB 
	- Total overhead in shared per shape: 93kB
	- Per-shape process total: 35.9MB
	- Per-shape 119kB
- With this PR, 300 shapes
	- Total (for shape-related processes): 21MB
	- Total overhead in shared processes: 7.3MB
	- Total overhead in shared per shape: 24kB
	- Per-shape process total: 13MB
	- Per-shape 45kB

So for a sufficiently large table, this PR drops used memory a lot, 60%
in this particular case. Further optimizations of state are possible and
we'll do them, but this is the lowest hanging fruit.

## Changes

Shapes now don't store their table's entire schema. That's OK when
processing transactions or querying the DB, but that's not enough to
invalidate the shape when the schema changes, and it can change while
Electric is down as well.

To address this, now the burden of figuring out whether a
schema-breaking change has happened from the incoming relations now lies
upon `ShapeLogCollector`. It tracks all seen relation messages (which
should be relatively rare unless PG is constantly disconnecting), and
marks the columns that have changed within this relation as compared to
a previously seen version. It also saves these last-seen relations when
all shapes have reacted to it and invalidated themselves.

Author: icehaunter

.changeset/big-candles-approve.md

@@ -0,0 +1,5 @@
+---
+"@core/sync-service": patch
+---
+
+feat: lower the per-shape memory usage, especially for very large tables

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

@@ -177,13 +177,14 @@ defmodule Electric.Replication.Changes do
 
   defmodule Relation do
     @derive Jason.Encoder
-    defstruct [:id, :schema, :table, :columns]
+    defstruct [:id, :schema, :table, :columns, affected_columns: []]
 
     @type t() :: %__MODULE__{
             id: Changes.relation_id(),
             schema: Changes.db_identifier(),
             table: Changes.db_identifier(),
-            columns: [Column.t()]
+            columns: [Column.t()],
+            affected_columns: [Changes.db_identifier()]
           }
   end
 

packages/sync-service/lib/electric/replication/eval/expr.ex

@@ -3,6 +3,7 @@ defmodule Electric.Replication.Eval.Expr do
   Parsed expression, available for evaluation using the runner
   """
 
+  alias Electric.Replication.Eval.Parser
   alias Electric.Replication.Eval.Env
 
   defstruct [:query, :eval, :used_refs, :returns]
@@ -37,4 +38,41 @@ defmodule Electric.Replication.Eval.Expr do
     |> Enum.filter(&match?({[_], _}, &1))
     |> Enum.map(fn {[key], _} -> key end)
   end
+
+  @doc false
+  @spec to_json_safe(t()) :: map()
+  def to_json_safe(%__MODULE__{query: query, used_refs: used_refs}) do
+    %{
+      version: 1,
+      query: query,
+      used_refs:
+        Enum.map(used_refs, fn
+          {k, v} when is_tuple(v) -> [k, Tuple.to_list(v)]
+          {k, v} -> [k, v]
+        end)
+    }
+  end
+
+  @doc false
+  @spec from_json_safe(map()) :: {:ok, t()} | {:error, String.t()}
+  def from_json_safe(%{"version" => 1, "query" => query, "used_refs" => refs}) do
+    refs =
+      Map.new(refs, fn
+        [k, v] when is_list(v) -> {k, List.to_tuple(v)}
+        [k, v] when is_binary(v) -> {k, String.to_existing_atom(v)}
+      end)
+
+    Parser.parse_and_validate_expression(query, refs: refs)
+  end
+
+  def from_json_safe(_),
+    do: {:error, "Incorrect serialized format: keys must be `version`, `query`, `used_refs`"}
+end
+
+defimpl Jason.Encoder, for: Electric.Replication.Eval.Expr do
+  def encode(expr, opts) do
+    expr
+    |> Electric.Replication.Eval.Expr.to_json_safe()
+    |> Jason.Encode.map(opts)
+  end
 end

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

@@ -1,6 +1,7 @@
 defmodule Electric.Replication.PersistentReplicationState do
   alias Electric.PersistentKV
   alias Electric.Postgres.Lsn
+  alias Electric.Replication.Changes
   require Logger
 
   @type opts() :: [
@@ -29,9 +30,33 @@ defmodule Electric.Replication.PersistentReplicationState do
     |> Lsn.from_integer()
   end
 
+  @base_tracked_relations %{
+    table_to_id: %{},
+    id_to_table_info: %{}
+  }
+
+  @type tracked_relations :: %{
+          table_to_id: %{{String.t(), String.t()} => Changes.relation_id()},
+          id_to_table_info: %{Changes.relation_id() => Changes.Relation.t()}
+        }
+
+  @spec set_tracked_relations(tracked_relations, opts()) :: :ok
+  def set_tracked_relations(tracked_relations, opts) do
+    set("tracked_relations", tracked_relations, opts)
+  end
+
+  @spec get_tracked_relations(opts()) :: tracked_relations()
+  def get_tracked_relations(opts) do
+    case get("tracked_relations", opts) do
+      {:ok, tracked_relations} -> tracked_relations
+      {:error, :not_found} -> @base_tracked_relations
+    end
+  end
+
   @spec reset(opts()) :: :ok
   def reset(opts) do
     set(@last_processed_lsn_key, 0, opts)
+    set_tracked_relations(@base_tracked_relations, opts)
   end
 
   @spec set(String.t(), any(), opts()) :: :ok

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

@@ -150,6 +150,9 @@ defmodule Electric.Replication.PublicationManager do
 
   @impl true
   def init(opts) do
+    Logger.metadata(stack_id: Access.fetch!(opts, :stack_id))
+    Process.set_label({:publication_manager, Access.fetch!(opts, :stack_id)})
+
     state = %__MODULE__{
       relation_filter_counters: %{},
       prepared_relation_filters: %{},
@@ -437,7 +440,7 @@ defmodule Electric.Replication.PublicationManager do
   end
 
   @spec get_selected_columns_for_shape(Shape.t()) :: MapSet.t(String.t() | nil)
-  defp get_selected_columns_for_shape(%Shape{where: _, selected_columns: nil}),
+  defp get_selected_columns_for_shape(%Shape{where: _, flags: %{selects_all_columns: true}}),
     do: MapSet.new([nil])
 
   defp get_selected_columns_for_shape(%Shape{where: nil, selected_columns: columns}),
@@ -453,20 +456,8 @@ defmodule Electric.Replication.PublicationManager do
           MapSet.t(Electric.Replication.Eval.Expr.t() | nil)
   defp get_where_clauses_for_shape(%Shape{where: nil}), do: MapSet.new([nil])
   # TODO: flatten where clauses by splitting top level ANDs
-  defp get_where_clauses_for_shape(%Shape{
-         where: where,
-         table_info: table_info,
-         root_table: root_table
-       }) do
-    unqualified_refs = Expr.unqualified_refs(where)
-
-    table_info[root_table].columns
-    |> Enum.filter(&(&1.name in unqualified_refs))
-    |> Enum.any?(fn
-      %{type_kind: kind} when kind in [:enum, :domain, :composite] -> true
-      _ -> false
-    end)
-    |> if do
+  defp get_where_clauses_for_shape(%Shape{where: where, flags: flags}) do
+    if Map.get(flags, :non_primitive_columns_in_where, false) do
       MapSet.new([nil])
     else
       MapSet.new([where])

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

@@ -6,6 +6,7 @@ defmodule Electric.Replication.ShapeLogCollector do
   use GenStage
 
   require Electric.Postgres.Lsn
+  alias Electric.Replication.ShapeLogCollector.AffectedColumns
   alias Electric.Postgres.Lsn
   alias Electric.Replication.PersistentReplicationState
   alias Electric.Postgres.Inspector
@@ -67,13 +68,19 @@ defmodule Electric.Replication.ShapeLogCollector do
       persistent_kv: opts.persistent_kv
     ]
 
+    {:ok, tracker_state} =
+      persistent_replication_data_opts
+      |> PersistentReplicationState.get_tracked_relations()
+      |> AffectedColumns.init()
+
     state =
       Map.merge(opts, %{
         producer: nil,
         subscriptions: {0, MapSet.new()},
         persistent_replication_data_opts: persistent_replication_data_opts,
         last_seen_lsn:
-          PersistentReplicationState.get_last_processed_lsn(persistent_replication_data_opts)
+          PersistentReplicationState.get_last_processed_lsn(persistent_replication_data_opts),
+        tracked_relations: tracker_state
       })
 
     # start in demand: :accumulate mode so that the ShapeCache is able to start
@@ -110,6 +117,12 @@ defmodule Electric.Replication.ShapeLogCollector do
         state.persistent_replication_data_opts
       )
 
+    :ok =
+      PersistentReplicationState.set_tracked_relations(
+        state.tracked_relations,
+        state.persistent_replication_data_opts
+      )
+
     {:noreply, [], %{state | producer: nil}}
   end
 
@@ -181,19 +194,29 @@ defmodule Electric.Replication.ShapeLogCollector do
     {:noreply, [txn], %{state | producer: from} |> put_last_seen_lsn(txn.lsn)}
   end
 
-  defp handle_relation(rel, _from, %{subscriptions: {0, _}} = state) do
-    Logger.debug(fn ->
-      "Dropping relation message for #{inspect(rel.schema)}.#{inspect(rel.table)}: no active consumers"
-    end)
+  defp handle_relation(rel, from, state) do
+    OpenTelemetry.add_span_attributes("rel.is_dropped": false)
 
-    OpenTelemetry.add_span_attributes("rel.is_dropped": true)
+    {updated_rel, tracker_state} =
+      AffectedColumns.transform_relation(rel, state.tracked_relations)
 
-    {:reply, :ok, [], state}
-  end
+    case state do
+      %{subscriptions: {0, _}} ->
+        Logger.debug(fn ->
+          "Dropping relation message for #{inspect(rel.schema)}.#{inspect(rel.table)}: no active consumers"
+        end)
 
-  defp handle_relation(rel, from, state) do
-    OpenTelemetry.add_span_attributes("rel.is_dropped": false)
-    {:noreply, [rel], %{state | producer: from}}
+        :ok =
+          PersistentReplicationState.set_tracked_relations(
+            state.tracked_relations,
+            state.persistent_replication_data_opts
+          )
+
+        {:reply, :ok, [], %{state | tracked_relations: tracker_state}}
+
+      _ ->
+        {:noreply, [updated_rel], %{state | producer: from, tracked_relations: tracker_state}}
+    end
   end
 
   defp drop_transaction(txn, state) do

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

@@ -0,0 +1,87 @@
+defmodule Electric.Replication.ShapeLogCollector.AffectedColumns do
+  @moduledoc false
+
+  require Logger
+  alias Electric.Replication.Changes.Relation
+
+  def init(%{id_to_table_info: id_to_table_info, table_to_id: table_to_id})
+      when is_map(id_to_table_info) and is_map(table_to_id) do
+    {:ok, %{id_to_table_info: id_to_table_info, table_to_id: table_to_id}}
+  end
+
+  def transform_relation(
+        %Relation{schema: schema, table: table, id: id} = rel,
+        %{
+          id_to_table_info: id_to_table_info,
+          table_to_id: table_to_id
+        } = state
+      ) do
+    schema_table = {schema, table}
+
+    existing_id = Map.get(table_to_id, schema_table)
+    existing_rel = Map.get(id_to_table_info, id)
+
+    case {existing_id, existing_rel} do
+      # New relation, register it
+      {nil, nil} ->
+        {rel, add_relation(state, id, rel)}
+
+      # Relation identity matches known, let's compare columns
+      {^id, %Relation{schema: ^schema, table: ^table}} ->
+        case find_differing_columns(existing_rel, rel) do
+          # No (noticable) changes to the relation, continue as-is
+          [] ->
+            {rel, state}
+
+          affected_cols ->
+            updated_rel = %{rel | affected_columns: affected_cols}
+            {updated_rel, add_relation(state, id, rel)}
+        end
+
+      # Some part of identity changed, update the state and pass it through
+      {_, _} ->
+        Logger.debug(fn ->
+          "Relation identity changed: #{existing_id}/#{inspect(existing_rel)} -> #{inspect(rel)}"
+        end)
+
+        {rel,
+         state
+         |> delete_tracked_relation(schema_table_key(existing_rel), existing_id)
+         |> add_relation(id, rel)}
+    end
+  end
+
+  defp schema_table_key(%Relation{schema: schema, table: table}), do: {schema, table}
+  defp schema_table_key(nil), do: nil
+
+  defp add_relation(state, id, rel) do
+    state
+    |> put_in([:table_to_id, schema_table_key(rel)], id)
+    |> put_in([:id_to_table_info, id], rel)
+  end
+
+  defp delete_tracked_relation(state, schema_table, id) do
+    state
+    |> update_in([:table_to_id], &Map.delete(&1, schema_table))
+    |> update_in([:id_to_table_info], &Map.delete(&1, id))
+  end
+
+  defp find_differing_columns(%Relation{columns: old_cols}, %Relation{columns: new_cols})
+       when old_cols == new_cols,
+       do: []
+
+  defp find_differing_columns(%Relation{columns: old_cols}, %Relation{columns: new_cols}) do
+    (old_cols ++ new_cols)
+    |> Enum.reduce(%{}, fn
+      %{name: name, type_oid: type_oid}, acc when is_map_key(acc, name) ->
+        # We're seeing column with this name for a second time, so we can remove it from the diff if type oid is the same
+        if acc[name] == type_oid, do: Map.delete(acc, name), else: acc
+
+      %{name: name, type_oid: type_oid}, acc ->
+        # If we're seeing column with this name for a first time, it'll either stay if it's present only in one set,
+        # or be deleted if seen again
+        Map.put(acc, name, type_oid)
+    end)
+    |> Map.keys()
+  end
+end

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

@@ -167,7 +167,7 @@ defmodule Electric.ShapeCache.FileStorage do
 
           with {:ok, shape_def_encoded} <- File.read(shape_def_path),
                {:ok, shape_def_json} <- Jason.decode(shape_def_encoded),
-               shape = Electric.Shapes.Shape.from_json_safe!(shape_def_json) do
+               {:ok, shape} <- Electric.Shapes.Shape.from_json_safe(shape_def_json) do
             Map.put(acc, shape_handle, shape)
           else
             # if the shape definition file cannot be read/decoded, just ignore it

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

@@ -1,4 +1,5 @@
 defmodule Electric.Shapes.Api do
+  alias Electric.Postgres.Inspector
   alias Electric.Replication.LogOffset
   alias Electric.Shapes
   alias Electric.Telemetry.OpenTelemetry
@@ -624,18 +625,25 @@ defmodule Electric.Shapes.Api do
     apply(encoder, type, [message])
   end
 
-  def schema(%Request{params: params}) do
-    schema(params)
-  end
-
-  def schema(%{shape_definition: %Shapes.Shape{} = shape}) do
-    shape.table_info
-    |> Map.fetch!(shape.root_table)
-    |> Map.fetch!(:columns)
-    |> Electric.Schema.from_column_info(shape.selected_columns)
+  def schema(%Response{
+        api: %Api{inspector: inspector},
+        shape_definition: %Shapes.Shape{} = shape
+      }) do
+    # This technically does double work because we've already fetched this info to build the shape,
+    # but that's not a big deal as it's all ETS backed. This also has an added benefit that
+    # if table schema changes in a way that doesn't invalidate the shape or we can't detect
+    # (e.g. column nullability changes but the type remains the same), we might return the new
+    # version if it's invalidated in ETS or server is restarted.
+    case Inspector.load_column_info(shape.root_table, inspector) do
+      {:ok, columns} ->
+        Electric.Schema.from_column_info(columns, shape.selected_columns)
+
+      :table_not_found ->
+        nil
+    end
   end
 
-  def schema(_) do
+  def schema(_req) do
     nil
   end
 

packages/sync-service/lib/electric/shapes/api/response.ex

@@ -54,6 +54,7 @@ defmodule Electric.Shapes.Api.Response do
       args
       |> Keyword.put_new(:status, 400)
       |> Keyword.put(:body, error_body(api, message, args))
+      |> Keyword.put(:api, api)
 
     struct(__MODULE__, opts)
   end
@@ -64,6 +65,7 @@ defmodule Electric.Shapes.Api.Response do
       |> Keyword.put_new(:status, 400)
       |> Keyword.put(:body, error_body(request, message, args))
       |> Keyword.put(:shape_definition, request.params.shape_definition)
+      |> Keyword.put(:api, request.api)
 
     struct(__MODULE__, opts)
   end