From 69a51db7d22b970f52a84f97ce069f1c3af139fe Mon Sep 17 00:00:00 2001 From: joshdchang Date: Tue, 23 Jun 2026 13:34:39 -0400 Subject: [PATCH] perf(sync-service): batch cold shape startup work --- .../lib/electric/postgres/configuration.ex | 68 +++++ .../postgres/inspector/direct_inspector.ex | 69 ++++- .../postgres/inspector/ets_inspector.ex | 257 +++++++++++++++--- .../replication/publication_manager.ex | 1 + .../publication_manager/configurator.ex | 72 ++++- .../publication_manager/relation_tracker.ex | 93 ++++++- .../electric/postgres/configuration_test.exs | 32 +++ .../postgres/inspector/ets_inspector_test.exs | 106 +++++++- .../replication/publication_manager_test.exs | 149 ++++++++++ 9 files changed, 775 insertions(+), 72 deletions(-) diff --git a/packages/sync-service/lib/electric/postgres/configuration.ex b/packages/sync-service/lib/electric/postgres/configuration.ex index a25121fcb9..8dc8a63f4b 100644 --- a/packages/sync-service/lib/electric/postgres/configuration.ex +++ b/packages/sync-service/lib/electric/postgres/configuration.ex @@ -128,6 +128,34 @@ defmodule Electric.Postgres.Configuration do ) end + @spec configure_tables_for_replication( + Postgrex.conn(), + String.t(), + [Electric.oid_relation()], + [Electric.oid_relation()], + timeout() + ) :: :ok | {:error, term()} + def configure_tables_for_replication( + conn, + publication_name, + additions, + replica_identity_relations, + timeout \\ @default_action_timeout + ) do + additions = Enum.sort(additions) + replica_identity_relations = Enum.sort(replica_identity_relations) + + run_in_transaction( + conn, + fn conn -> + add_tables_to_publication!(conn, publication_name, additions) + set_tables_replica_identity_full!(conn, replica_identity_relations) + :ok + end, + timeout + ) + end + @spec drop_table_from_publication( Postgrex.conn(), String.t(), @@ -190,6 +218,46 @@ defmodule Electric.Postgres.Configuration do end end + defp add_tables_to_publication!(_conn, _publication_name, []), do: :ok + + defp add_tables_to_publication!(conn, publication_name, additions) do + tables = + additions + |> Enum.map(fn {_oid, relation} -> Utils.relation_to_sql(relation) end) + |> Enum.join(", ") + + Postgrex.query!( + conn, + "ALTER PUBLICATION #{Utils.quote_name(publication_name)} ADD TABLE #{tables}", + [] + ) + end + + defp set_tables_replica_identity_full!(_conn, []), do: :ok + + defp set_tables_replica_identity_full!(conn, relations) do + statements = + Enum.map_join(relations, "\n", fn {_oid, relation} -> + " ALTER TABLE #{Utils.relation_to_sql(relation)} REPLICA IDENTITY FULL;" + end) + + delimiter = dollar_quote_delimiter(statements) + + Postgrex.query!( + conn, + "DO #{delimiter}\nBEGIN\n#{statements}\nEND\n#{delimiter};", + [] + ) + end + + defp dollar_quote_delimiter(contents) do + delimiter = "$electric_#{System.unique_integer([:positive])}$" + + if String.contains?(contents, delimiter), + do: dollar_quote_delimiter(contents), + else: delimiter + end + @spec exec_set_replica_identity_full(Postgrex.conn(), String.t()) :: :ok | {:error, term()} defp exec_set_replica_identity_full(conn, table) do Postgrex.query!(conn, "SAVEPOINT before_replica", []) diff --git a/packages/sync-service/lib/electric/postgres/inspector/direct_inspector.ex b/packages/sync-service/lib/electric/postgres/inspector/direct_inspector.ex index 9f228057ac..d68f129b67 100644 --- a/packages/sync-service/lib/electric/postgres/inspector/direct_inspector.ex +++ b/packages/sync-service/lib/electric/postgres/inspector/direct_inspector.ex @@ -15,7 +15,7 @@ defmodule Electric.Postgres.Inspector.DirectInspector do @spec load_relation_oid(Electric.relation(), conn :: Postgrex.conn()) :: {:ok, Electric.oid_relation()} | :table_not_found | {:error, String.t()} def load_relation_oid({schema, table}, conn) do - query = load_relation_query(@oid_from_schema_table_name_subquery) + query = load_relation_query("pc.oid = #{@oid_from_schema_table_name_subquery}") case do_load_relation(conn, query, [schema, table]) do {:ok, []} -> @@ -39,7 +39,7 @@ defmodule Electric.Postgres.Inspector.DirectInspector do | :table_not_found | {:error, String.t() | :connection_not_available} def normalize_and_load_relation_info({schema, table}, conn) do - query = load_relation_query(@oid_from_schema_table_name_subquery) + query = load_relation_query("pc.oid = #{@oid_from_schema_table_name_subquery}") case do_load_relation(conn, query, [schema, table]) do {:ok, []} -> @@ -53,13 +53,35 @@ defmodule Electric.Postgres.Inspector.DirectInspector do end end + @doc """ + Normalizes and loads multiple relations in one catalog query. + + Missing relations are omitted from the returned list so the caller can map + them back to `:table_not_found` without issuing another query. + """ + @spec normalize_and_load_relations_info([Electric.relation()], Postgrex.conn()) :: + {:ok, [Inspector.relation_info()]} | {:error, String.t() | :connection_not_available} + def normalize_and_load_relations_info([], _conn), do: {:ok, []} + + def normalize_and_load_relations_info(relations, conn) when is_list(relations) do + {schemas, tables} = Enum.unzip(relations) + + query = + load_relation_query( + "(pn.nspname, pc.relname) IN " <> + "(SELECT * FROM unnest($1::text[], $2::text[]))" + ) + + do_load_relation(conn, query, [schemas, tables]) + end + @impl Electric.Postgres.Inspector @spec load_relation_info(Electric.relation_id(), conn :: Postgrex.conn()) :: {:ok, Inspector.relation_info()} | :table_not_found | {:error, String.t() | :connection_not_available} def load_relation_info(oid, conn) when is_relation_id(oid) do - query = load_relation_query("$1::oid") + query = load_relation_query("pc.oid = $1::oid") case do_load_relation(conn, query, [oid]) do {:ok, []} -> @@ -74,7 +96,7 @@ defmodule Electric.Postgres.Inspector.DirectInspector do end def load_relations_by_oids(oids, conn) when is_list(oids) do - query = load_relation_query("ANY ($1::oid[])") + query = load_relation_query("pc.oid = ANY ($1::oid[])") do_load_relation(conn, query, [oids]) end @@ -95,11 +117,28 @@ defmodule Electric.Postgres.Inspector.DirectInspector do {:ok, relations} {:error, err} -> - {:error, Exception.message(err)} + {:error, normalize_query_error(err)} + end + end + + @doc false + # Convert an in-band query error into the inspector error contract. + # Connection-class errors (e.g. "ssl recv: closed") are returned by + # Postgrex as `{:error, %DBConnection.ConnectionError{}}` rather than + # raised, so without this they would bypass the `:connection_not_available` + # mapping that callers rely on for retries. + @spec normalize_query_error(Exception.t()) :: String.t() | :connection_not_available + def normalize_query_error(%DBConnection.ConnectionError{} = err) do + if Electric.DbConnectionError.from_error(err).type != :unknown do + :connection_not_available + else + Exception.message(err) end end - defp load_relation_query(match) do + def normalize_query_error(err), do: Exception.message(err) + + defp load_relation_query(condition) do # partitions can live in other namespaces from the parent/root table, so we # need to keep track of them [ @@ -123,8 +162,7 @@ defmodule Electric.Postgres.Inspector.DirectInspector do WHERE pc.relkind IN ('r', 'p') AND """, - "pc.oid = ", - match + condition ] end @@ -187,7 +225,7 @@ defmodule Electric.Postgres.Inspector.DirectInspector do {:ok, rows} {:error, err} -> - {:error, err} + {:error, normalize_query_error(err)} end end @@ -204,6 +242,19 @@ defmodule Electric.Postgres.Inspector.DirectInspector do end end + def load_column_info_by_oids(oids, conn) when is_list(oids) do + query = """ + #{@column_info_query_base} + WHERE pg_class.oid = ANY ($1::oid[]) + ORDER BY pg_class.oid, attnum + """ + + case do_query_column_info(conn, query, [oids]) do + {:ok, rows} -> {:ok, Enum.group_by(rows, & &1.relation_id)} + {:error, reason} -> {:error, normalize_query_error(reason)} + end + end + defp do_query_column_info(conn, query, params) do with {:ok, %{rows: rows, columns: columns}} <- Postgrex.query(conn, query, params) do columns = Enum.map(columns, &String.to_atom/1) diff --git a/packages/sync-service/lib/electric/postgres/inspector/ets_inspector.ex b/packages/sync-service/lib/electric/postgres/inspector/ets_inspector.ex index 1adf3262de..75df46a1b0 100644 --- a/packages/sync-service/lib/electric/postgres/inspector/ets_inspector.ex +++ b/packages/sync-service/lib/electric/postgres/inspector/ets_inspector.ex @@ -35,6 +35,8 @@ defmodule Electric.Postgres.Inspector.EtsInspector do # table. They are bounded by this interval rather than living forever, even # though the negative key space is client-controlled (distinct table names). @negative_cache_sweep_interval_ms 60_000 + @default_fetch_batch_window_ms 10 + @default_max_fetch_batch_size 64 alias Electric.Postgres.Inspector alias Electric.Telemetry.OpenTelemetry @@ -149,6 +151,16 @@ defmodule Electric.Postgres.Inspector.EtsInspector do ]) persistence_key = "#{opts.stack_id}:ets_inspector_state" + fetch_batch_window_ms = Map.get(opts, :fetch_batch_window_ms, @default_fetch_batch_window_ms) + max_fetch_batch_size = Map.get(opts, :max_fetch_batch_size, @default_max_fetch_batch_size) + + if not (is_integer(fetch_batch_window_ms) and fetch_batch_window_ms >= 0) do + raise ArgumentError, "fetch_batch_window_ms must be a non-negative integer" + end + + if not (is_integer(max_fetch_batch_size) and max_fetch_batch_size > 0) do + raise ArgumentError, "max_fetch_batch_size must be a positive integer" + end state = %{ @@ -160,6 +172,11 @@ defmodule Electric.Postgres.Inspector.EtsInspector do task_sup: task_supervisor_name(opts.stack_id), in_flight: %{}, in_flight_refs: %{}, + pending_relation_keys: :queue.new(), + relation_batch_ref: nil, + relation_batch_timer: nil, + fetch_batch_window_ms: fetch_batch_window_ms, + max_fetch_batch_size: max_fetch_batch_size, negative_cache_ttl_ms: Map.get(opts, :negative_cache_ttl_ms, @default_negative_cache_ttl_ms) } @@ -236,10 +253,14 @@ defmodule Electric.Postgres.Inspector.EtsInspector do {nil, state} -> {:noreply, state} - {{key, entry}, state} -> - state = apply_fill_result(state, key, result) - reply_waiters(state, key, result, entry.waiters) - {:noreply, state} + {entries, state} -> + state = apply_fill_results(state, result) + + for {key, entry} <- entries do + reply_waiters(state, key, Map.fetch!(result, key), entry.waiters) + end + + {:noreply, finish_relation_batch(state, ref)} end end @@ -248,16 +269,46 @@ defmodule Electric.Postgres.Inspector.EtsInspector do {nil, state} -> {:noreply, state} - {{key, entry}, state} -> - Logger.warning( - "EtsInspector fill worker for #{inspect(key)} exited before replying: #{inspect(reason)}" - ) + {entries, state} -> + Logger.warning("EtsInspector fill worker exited before replying: #{inspect(reason)}") - for {from, _reader} <- entry.waiters do + for {_key, entry} <- entries, + {from, _reader} <- entry.waiters do GenServer.reply(from, {:error, :connection_not_available}) end + {:noreply, finish_relation_batch(state, ref)} + end + end + + def handle_info(:flush_relation_batch, state) do + {keys, pending_relation_keys} = + take_pending_keys(state.pending_relation_keys, state.max_fetch_batch_size) + + state = %{state | relation_batch_timer: nil, pending_relation_keys: pending_relation_keys} + + case keys do + [] -> {:noreply, state} + + keys -> + %{ref: ref} = + Task.Supervisor.async_nolink(state.task_sup, fn -> + fetch_relation_batch(keys, state.pg_pool, state.stack_id) + end) + + in_flight = + Enum.reduce(keys, state.in_flight, fn key, in_flight -> + Map.update!(in_flight, key, &Map.put(&1, :ref, ref)) + end) + + {:noreply, + %{ + state + | in_flight: in_flight, + in_flight_refs: Map.put(state.in_flight_refs, ref, keys), + relation_batch_ref: ref + }} end end @@ -284,43 +335,90 @@ defmodule Electric.Postgres.Inspector.EtsInspector do Process.send_after(self(), :sweep_negative_cache, @negative_cache_sweep_interval_ms) end - # Coalesce concurrent loads of the same key: the first waiter spawns one - # supervised worker to do the DB lookup; later waiters for the same key just - # park their `from` and are all answered when the worker reports back. This - # keeps the inspector's mailbox population proportional to unique in-flight - # keys rather than to in-flight requests, and avoids head-of-line blocking on - # a single slow lookup. `in_flight_refs` maps the worker's monitor ref back to - # its key so completion/crash handling is O(1). + # Coalesce concurrent loads of the same key. Relation-name misses also wait + # for a short batching window so one catalog transaction can fill many cold + # entries without consuming one pool checkout per table. defp enqueue_waiter(state, key, waiter) do case Map.fetch(state.in_flight, key) do {:ok, entry} -> entry = %{entry | waiters: [waiter | entry.waiters]} %{state | in_flight: Map.put(state.in_flight, key, entry)} - :error -> - %{ref: ref} = - Task.Supervisor.async_nolink(state.task_sup, fn -> - fetch_for_key(key, state.pg_pool, state.stack_id) - end) - - entry = %{waiters: [waiter], ref: ref} - - %{ + :error when is_tuple(key) and elem(key, 0) == :rel -> + state = %{ state - | in_flight: Map.put(state.in_flight, key, entry), - in_flight_refs: Map.put(state.in_flight_refs, ref, key) + | in_flight: Map.put(state.in_flight, key, %{waiters: [waiter], ref: nil}), + pending_relation_keys: :queue.in(key, state.pending_relation_keys) } + + schedule_relation_batch(state) + + :error -> + start_single_fetch(state, key, waiter) end end + defp start_single_fetch(state, key, waiter) do + %{ref: ref} = + Task.Supervisor.async_nolink(state.task_sup, fn -> + %{key => fetch_for_key(key, state.pg_pool, state.stack_id)} + end) + + entry = %{waiters: [waiter], ref: ref} + + %{ + state + | in_flight: Map.put(state.in_flight, key, entry), + in_flight_refs: Map.put(state.in_flight_refs, ref, [key]) + } + end + defp pop_in_flight_by_ref(state, ref) do case Map.pop(state.in_flight_refs, ref) do {nil, _refs} -> {nil, state} - {key, in_flight_refs} -> - {entry, in_flight} = Map.pop(state.in_flight, key) - {{key, entry}, %{state | in_flight: in_flight, in_flight_refs: in_flight_refs}} + {keys, in_flight_refs} -> + {entries, in_flight} = + Enum.map_reduce(keys, state.in_flight, fn key, in_flight -> + {entry, in_flight} = Map.pop(in_flight, key) + {{key, entry}, in_flight} + end) + + {entries, %{state | in_flight: in_flight, in_flight_refs: in_flight_refs}} + end + end + + defp schedule_relation_batch(state) do + if is_nil(state.relation_batch_ref) and is_nil(state.relation_batch_timer) do + timer = Process.send_after(self(), :flush_relation_batch, state.fetch_batch_window_ms) + %{state | relation_batch_timer: timer} + else + state + end + end + + defp finish_relation_batch(%{relation_batch_ref: ref} = state, ref) do + state + |> Map.put(:relation_batch_ref, nil) + |> schedule_relation_batch_if_pending() + end + + defp finish_relation_batch(state, _ref), do: state + + defp schedule_relation_batch_if_pending(state) do + if :queue.is_empty(state.pending_relation_keys), + do: state, + else: schedule_relation_batch(state) + end + + defp take_pending_keys(queue, limit), do: take_pending_keys(queue, limit, []) + defp take_pending_keys(queue, 0, keys), do: {Enum.reverse(keys), queue} + + defp take_pending_keys(queue, remaining, keys) do + case :queue.out(queue) do + {:empty, queue} -> {Enum.reverse(keys), queue} + {{:value, key}, queue} -> take_pending_keys(queue, remaining - 1, [key | keys]) end end @@ -342,6 +440,74 @@ defmodule Electric.Postgres.Inspector.EtsInspector do ) end + defp fetch_relation_batch(keys, pool, stack_id) do + key_type = if length(keys) == 1, do: "relation", else: "relation_batch" + + OpenTelemetry.with_span( + "inspector.fetch_db", + [{"inspector.key_type", key_type}, {"inspector.batch_size", length(keys)}], + stack_id, + fn -> + results = do_fetch_relation_batch(keys, pool) + + OpenTelemetry.add_span_attributes(%{ + "inspector.result" => fetch_batch_outcome(results) + }) + + results + end + ) + end + + defp do_fetch_relation_batch(keys, pool) do + relations = Enum.map(keys, fn {:rel, relation} -> relation end) + + result = + wrap_in_db_errors(fn -> + Postgrex.transaction( + pool, + fn conn -> + with {:ok, relation_infos} <- + DirectInspector.normalize_and_load_relations_info(relations, conn), + {:ok, columns_by_oid} <- + relation_infos + |> Enum.map(& &1.relation_id) + |> DirectInspector.load_column_info_by_oids(conn) do + relation_infos = Map.new(relation_infos, &{&1.relation, &1}) + + Map.new(keys, fn {:rel, relation} = key -> + case Map.fetch(relation_infos, relation) do + {:ok, relation_info} -> + columns = Map.get(columns_by_oid, relation_info.relation_id, []) + + if columns == [], + do: {key, {:ok, :table_not_found}}, + else: {key, {:ok, {relation_info, columns}}} + + :error -> + {key, {:ok, :table_not_found}} + end + end) + else + {:error, reason} -> Postgrex.rollback(conn, reason) + end + end, + timeout: @fetch_db_timeout + ) + end) + + case result do + {:ok, results} -> results + {:error, reason} -> Map.new(keys, &{&1, {:error, reason}}) + end + end + + defp fetch_batch_outcome(results) do + if Enum.any?(results, fn {_key, result} -> match?({:error, _}, result) end), + do: "error", + else: "ok" + end + defp do_fetch_for_key({:rel, rel}, pool), do: fetch_from_db(rel, pool) defp do_fetch_for_key({:oid, oid}, pool), do: fetch_from_db(oid, pool) @@ -357,24 +523,25 @@ defmodule Electric.Postgres.Inspector.EtsInspector do defp fetch_outcome({:ok, _}), do: "ok" defp fetch_outcome({:error, _}), do: "error" - defp apply_fill_result(state, :supported_features, {:ok, features}) do - state |> store_supported_features(features) |> persist_data() - state - end + defp apply_fill_results(state, results) do + {state, persist?} = + Enum.reduce(results, {state, false}, fn + {:supported_features, {:ok, features}}, {state, _persist?} -> + {store_supported_features(state, features), true} - defp apply_fill_result(state, _key, {:ok, {rel, cols}}) do - state |> store_relation_info(rel, cols) |> persist_data() - state - end + {_key, {:ok, {rel, cols}}}, {state, _persist?} -> + {store_relation_info(state, rel, cols), true} + + {key, {:ok, :table_not_found}}, {state, persist?} -> + {put_negative_cache(state, key, :table_not_found), persist?} - # Terminal negative results (table-not-found / DB error) are cached for a short - # TTL so a burst against a failing key drains the mailbox instead of refilling - # it; the client reads this cache and short-circuits before messaging us again. - defp apply_fill_result(state, key, {:ok, :table_not_found}), - do: put_negative_cache(state, key, :table_not_found) + {key, {:error, reason}}, {state, persist?} -> + {put_negative_cache(state, key, {:error, reason}), persist?} + end) - defp apply_fill_result(state, key, {:error, reason}), - do: put_negative_cache(state, key, {:error, reason}) + if persist?, do: persist_data(state) + state + end defp reply_waiters(state, key, result, waiters) do for {from, reader} <- waiters do diff --git a/packages/sync-service/lib/electric/replication/publication_manager.ex b/packages/sync-service/lib/electric/replication/publication_manager.ex index 5a1fa361a9..7b1a1e6d0a 100644 --- a/packages/sync-service/lib/electric/replication/publication_manager.ex +++ b/packages/sync-service/lib/electric/replication/publication_manager.ex @@ -14,6 +14,7 @@ defmodule Electric.Replication.PublicationManager do defdelegate name(opts), to: __MODULE__.RelationTracker defdelegate add_shape(stack_id, shape_handle, shape), to: __MODULE__.RelationTracker + defdelegate register_shapes(stack_id, shapes), to: __MODULE__.RelationTracker defdelegate remove_shape(stack_id, shape_handle), to: __MODULE__.RelationTracker defdelegate wait_for_restore(stack_id, opts \\ []), to: __MODULE__.RelationTracker end diff --git a/packages/sync-service/lib/electric/replication/publication_manager/configurator.ex b/packages/sync-service/lib/electric/replication/publication_manager/configurator.ex index 6fd8cefd33..5e9d976a23 100644 --- a/packages/sync-service/lib/electric/replication/publication_manager/configurator.ex +++ b/packages/sync-service/lib/electric/replication/publication_manager/configurator.ex @@ -213,10 +213,58 @@ defmodule Electric.Replication.PublicationManager.Configurator do {:noreply, state} end - def handle_continue({:perform_relation_actions, [{action, filter} | rest]}, state) do + def handle_continue({:perform_relation_actions, relation_actions}, state) do + case addition_batch(relation_actions) do + {:ok, additions, replica_identity_relations} -> + result = + OpenTelemetry.with_span( + "publication_manager.update_publication_relations", + [ + action: "batch_add", + relation_count: length(additions), + replica_identity_count: length(replica_identity_relations) + ], + state.stack_id, + fn -> + Configuration.configure_tables_for_replication( + state.db_pool, + state.publication_name, + additions, + replica_identity_relations + ) + end + ) + + case result do + :ok -> + notify_filters_result(additions, {:ok, :configured}, state) + {:noreply, state, {:continue, {:perform_relation_actions, []}}} + + {:error, error} -> + Logger.debug( + "Batch publication update failed; retrying relations individually: #{inspect(error)}" + ) + + {:noreply, state, + {:continue, {:perform_relation_actions_individually, relation_actions}}} + end + + :not_batchable -> + {:noreply, state, {:continue, {:perform_relation_actions_individually, relation_actions}}} + end + end + + def handle_continue({:perform_relation_actions_individually, []}, state) do + {:noreply, state, {:continue, {:perform_relation_actions, []}}} + end + + def handle_continue( + {:perform_relation_actions_individually, [{action, filter} | rest]}, + state + ) do {_oid, rel} = filter - res = + result = OpenTelemetry.with_span( "publication_manager.update_publication_relation", [action: inspect(action), relation: Utils.relation_to_sql(rel)], @@ -224,9 +272,9 @@ defmodule Electric.Replication.PublicationManager.Configurator do fn -> do_publication_update(action, filter, state) end ) - notify_filter_result(filter, res, state) + notify_filter_result(filter, result, state) - {:noreply, state, {:continue, {:perform_relation_actions, rest}}} + {:noreply, state, {:continue, {:perform_relation_actions_individually, rest}}} end @spec check_publication_status(state()) :: @@ -309,6 +357,22 @@ defmodule Electric.Replication.PublicationManager.Configurator do end end + defp addition_batch(relation_actions) do + if length(relation_actions) > 1 and + Enum.all?(relation_actions, fn {action, _filter} -> + action in [:add, :add_and_configure] + end) do + additions = Enum.map(relation_actions, &elem(&1, 1)) + + replica_identity_relations = + for {:add_and_configure, filter} <- relation_actions, do: filter + + {:ok, additions, replica_identity_relations} + else + :not_batchable + end + end + defp handle_publication_fatally_misconfigured(err, %{manual_table_publishing?: true}), do: Logger.warning(""" diff --git a/packages/sync-service/lib/electric/replication/publication_manager/relation_tracker.ex b/packages/sync-service/lib/electric/replication/publication_manager/relation_tracker.ex index 8c1d41e07f..214e5f46dc 100644 --- a/packages/sync-service/lib/electric/replication/publication_manager/relation_tracker.ex +++ b/packages/sync-service/lib/electric/replication/publication_manager/relation_tracker.ex @@ -29,6 +29,9 @@ defmodule Electric.Replication.PublicationManager.RelationTracker do prepared_relation_filters: MapSet.new(), submitted_relation_filters: MapSet.new(), committed_relation_filters: MapSet.new(), + # The relations covered by the configuration request the Configurator is + # currently working on, or nil when none is outstanding. + in_flight_relation_filters: nil, waiters: %{}, # start with optimistic assumption about what the # publication supports (altering and generated columns) @@ -48,6 +51,7 @@ defmodule Electric.Replication.PublicationManager.RelationTracker do prepared_relation_filters: internal_relation_filters(), submitted_relation_filters: internal_relation_filters(), committed_relation_filters: internal_relation_filters(), + in_flight_relation_filters: internal_relation_filters() | nil, waiters: %{Electric.relation_id() => [waiter(), ...]}, publication_name: String.t(), publishes_generated_columns?: boolean(), @@ -62,7 +66,22 @@ defmodule Electric.Replication.PublicationManager.RelationTracker do def add_shape(stack_id, shape_handle, shape) do pub_filter = get_publication_filter_from_shape(shape) - case GenServer.call(name(stack_id), {:add_shape, shape_handle, pub_filter}) do + case GenServer.call(name(stack_id), {:add_shape, shape_handle, pub_filter}, :infinity) do + :ok -> :ok + {:error, err} -> raise err + end + end + + @spec register_shapes(stack_id(), [{shape_handle(), Electric.Shapes.Shape.t()}]) :: :ok + def register_shapes(_stack_id, []), do: :ok + + def register_shapes(stack_id, shapes) do + shapes = + Enum.map(shapes, fn {shape_handle, shape} -> + {shape_handle, get_publication_filter_from_shape(shape)} + end) + + case GenServer.call(name(stack_id), {:register_shapes, shapes}, :infinity) do :ok -> :ok {:error, err} -> raise err end @@ -196,9 +215,7 @@ defmodule Electric.Replication.PublicationManager.RelationTracker do # if the relation is already committed AND part of the last made # update submission, we can consider it ready - relation_ready? = - MapSet.member?(state.submitted_relation_filters, oid) and - MapSet.member?(state.committed_relation_filters, oid) + relation_ready? = relation_ready?(oid, state) state = add_shape_to_publication_filters(shape_handle, publication_filter, state) state = update_publication_if_necessary(state) @@ -228,6 +245,16 @@ defmodule Electric.Replication.PublicationManager.RelationTracker do end end + def handle_call({:register_shapes, shapes}, _from, state) do + state = + Enum.reduce(shapes, state, fn {shape_handle, publication_filter}, state -> + add_shape_to_publication_filters(shape_handle, publication_filter, state) + end) + |> update_publication_if_necessary() + + {:reply, :ok, state, state.publication_refresh_period} + end + def handle_call({:remove_shape, shape_handle}, _from, state) do state = remove_shape_from_publication_filters(shape_handle, state) state = update_publication_if_necessary(state) @@ -265,16 +292,20 @@ defmodule Electric.Replication.PublicationManager.RelationTracker do def handle_cast({:relation_configuration_result, {oid, _rel}, {:ok, :dropped}}, state) do new_committed_filters = MapSet.delete(state.committed_relation_filters, oid) - {:noreply, %{state | committed_relation_filters: new_committed_filters}, - state.publication_refresh_period} + state = + clear_in_flight_if_committed(%{state | committed_relation_filters: new_committed_filters}) + + {:noreply, state, state.publication_refresh_period} end def handle_cast({:relation_configuration_result, {oid, _} = oid_rel, {:ok, :configured}}, state) do state = reply_to_relation_waiters(oid_rel, :ok, state) new_committed_filters = MapSet.put(state.committed_relation_filters, oid) - {:noreply, %{state | committed_relation_filters: new_committed_filters}, - state.publication_refresh_period} + state = + clear_in_flight_if_committed(%{state | committed_relation_filters: new_committed_filters}) + + {:noreply, state, state.publication_refresh_period} end def handle_cast({:relation_configuration_result, oid_rel, {:error, error}}, state) do @@ -296,11 +327,13 @@ defmodule Electric.Replication.PublicationManager.RelationTracker do ) end - {:noreply, state, state.publication_refresh_period} + {:noreply, %{state | in_flight_relation_filters: nil}, state.publication_refresh_period} end def handle_cast({:configuration_error, {:error, error}}, state) do - {:noreply, reply_to_all_waiters({:error, error}, state), state.publication_refresh_period} + state = reply_to_all_waiters({:error, error}, state) + state = %{state | in_flight_relation_filters: nil} + {:noreply, state, state.publication_refresh_period} end @impl true @@ -328,7 +361,16 @@ defmodule Electric.Replication.PublicationManager.RelationTracker do expand_oids(state.prepared_relation_filters, state) ) - %{state | submitted_relation_filters: state.prepared_relation_filters} + %{ + state + | submitted_relation_filters: state.prepared_relation_filters, + in_flight_relation_filters: state.prepared_relation_filters + } + end + + defp relation_ready?(oid, state) do + MapSet.member?(state.submitted_relation_filters, oid) and + MapSet.member?(state.committed_relation_filters, oid) end @spec expand_oids(MapSet.t(Electric.relation_id()), state()) :: @@ -345,9 +387,34 @@ defmodule Electric.Replication.PublicationManager.RelationTracker do defp update_needed?(%__MODULE__{ prepared_relation_filters: prepared, submitted_relation_filters: submitted, - committed_relation_filters: committed + committed_relation_filters: committed, + in_flight_relation_filters: in_flight }) do - not MapSet.equal?(prepared, submitted) or not MapSet.equal?(submitted, committed) + prepared_changed? = not MapSet.equal?(prepared, submitted) + + # "submitted but not yet committed" is the legitimate "the last request + # didn't fully go through, retry it" signal. But we must not act on it while + # a request for the same relations is still outstanding, otherwise every + # add_shape/remove_shape in that window would send another duplicate request. + # If a request dies without ever reporting back, the publication_refresh_period + # inactivity timeout is the fallback that retries. + retry_needed? = is_nil(in_flight) and not MapSet.equal?(submitted, committed) + + prepared_changed? or retry_needed? + end + + @spec clear_in_flight_if_committed(state()) :: state() + defp clear_in_flight_if_committed( + %__MODULE__{ + committed_relation_filters: committed, + in_flight_relation_filters: in_flight + } = state + ) do + if not is_nil(in_flight) and MapSet.equal?(committed, in_flight) do + %{state | in_flight_relation_filters: nil} + else + state + end end @spec add_shape_to_publication_filters(shape_handle(), publication_filter(), state()) :: state() diff --git a/packages/sync-service/test/electric/postgres/configuration_test.exs b/packages/sync-service/test/electric/postgres/configuration_test.exs index ea336eca19..069fc50de8 100644 --- a/packages/sync-service/test/electric/postgres/configuration_test.exs +++ b/packages/sync-service/test/electric/postgres/configuration_test.exs @@ -69,6 +69,38 @@ defmodule Electric.Postgres.ConfigurationTest do assert list_tables_in_publication(conn, publication) == [{"public", "items"}] end + test "configures multiple tables for replication as one batch", %{ + pool: conn, + publication_name: publication + } do + Postgrex.query!( + conn, + "CREATE TABLE other_items (id UUID PRIMARY KEY, value TEXT NOT NULL)", + [] + ) + + relations = + for relation <- [{"public", "items"}, {"public", "other_items"}] do + {get_table_oid(conn, relation), relation} + end + + assert :ok = + Configuration.configure_tables_for_replication( + conn, + publication, + relations, + relations + ) + + assert list_tables_in_publication(conn, publication) == [ + {"public", "items"}, + {"public", "other_items"} + ] + + assert get_table_identity(conn, {"public", "items"}) == "f" + assert get_table_identity(conn, {"public", "other_items"}) == "f" + end + test "drops table from publication", %{ pool: conn, publication_name: publication diff --git a/packages/sync-service/test/electric/postgres/inspector/ets_inspector_test.exs b/packages/sync-service/test/electric/postgres/inspector/ets_inspector_test.exs index a93b1e2f5b..324b348f60 100644 --- a/packages/sync-service/test/electric/postgres/inspector/ets_inspector_test.exs +++ b/packages/sync-service/test/electric/postgres/inspector/ets_inspector_test.exs @@ -743,7 +743,7 @@ defmodule Electric.Postgres.Inspector.EtsInspectorTest do negative_cache_ttl_ms: 50} ) - %{opts: [stack_id: ctx.stack_id, server: server]} + %{opts: [stack_id: ctx.stack_id, server: server], pool: ctx.db_conn} end test "re-attempts the DB after the negative TTL expires", %{opts: opts, db_conn: pool} do @@ -784,6 +784,110 @@ defmodule Electric.Postgres.Inspector.EtsInspectorTest do end end + describe "batched cold-cache lookups" do + setup :with_shared_db + setup :in_transaction + setup :with_stack_id_from_test + setup [:with_persistent_kv, :with_basic_tables] + + setup ctx do + start_supervised!({Task.Supervisor, name: EtsInspector.task_supervisor_name(ctx.stack_id)}) + + server = + start_supervised!( + {EtsInspector, + stack_id: ctx.stack_id, + pool: ctx.db_conn, + persistent_kv: ctx.persistent_kv, + fetch_batch_window_ms: 50} + ) + + %{opts: [stack_id: ctx.stack_id, server: server]} + end + + test "starts one fill worker for a burst of distinct relations", %{opts: opts} do + test_pid = self() + Repatch.allow(test_pid, opts[:server]) + + Repatch.patch(Task.Supervisor, :async_nolink, [mode: :shared], fn supervisor, fun -> + task = Repatch.real(Task.Supervisor.async_nolink(supervisor, fun)) + Repatch.allow(test_pid, task.pid) + send(test_pid, :fill_worker_started) + task + end) + + tasks = + for suffix <- 1..3 do + Task.async(fn -> + EtsInspector.load_relation_oid({"public", "missing_#{suffix}"}, opts) + end) + end + + assert_receive :fill_worker_started, 200 + refute_receive :fill_worker_started, 100 + + Enum.each(tasks, &Task.shutdown(&1, :brutal_kill)) + end + + test "loads a relation burst in one catalog transaction", %{opts: opts} do + test_pid = self() + + Repatch.patch(Postgrex, :transaction, [mode: :shared], fn _pool, fun, _db_opts -> + send(test_pid, :db_transaction) + {:ok, fun.(:batch_connection)} + end) + + Repatch.patch( + Electric.Postgres.Inspector.DirectInspector, + :normalize_and_load_relations_info, + [mode: :shared], + fn relations, :batch_connection -> + send(test_pid, {:relations_loaded, relations}) + + {:ok, + Enum.map(relations, fn {"public", table} = relation -> + oid = %{"alpha" => 100, "beta" => 101, "gamma" => 102} |> Map.fetch!(table) + %{relation: relation, relation_id: oid, kind: :ordinary_table} + end)} + end + ) + + Repatch.patch( + Electric.Postgres.Inspector.DirectInspector, + :load_column_info_by_oids, + [mode: :shared], + fn oids, :batch_connection -> + send(test_pid, {:columns_loaded, oids}) + {:ok, Map.new(oids, &{&1, [%{relation_id: &1, name: "id"}]})} + end + ) + + allow_fill_workers(opts[:server]) + relations = [{"public", "alpha"}, {"public", "beta"}, {"public", "gamma"}] + + results = + relations + |> Enum.map(fn relation -> + Task.async(fn -> EtsInspector.load_relation_oid(relation, opts) end) + end) + |> Task.await_many(1_000) + + assert Enum.sort(results) == + Enum.sort([ + {:ok, {100, {"public", "alpha"}}}, + {:ok, {101, {"public", "beta"}}}, + {:ok, {102, {"public", "gamma"}}} + ]) + + assert_receive :db_transaction + refute_receive :db_transaction + assert_receive {:relations_loaded, loaded_relations} + assert Enum.sort(loaded_relations) == Enum.sort(relations) + assert_receive {:columns_loaded, loaded_oids} + assert Enum.sort(loaded_oids) == [100, 101, 102] + end + end + describe "telemetry" do setup :with_shared_db setup :in_transaction diff --git a/packages/sync-service/test/electric/replication/publication_manager_test.exs b/packages/sync-service/test/electric/replication/publication_manager_test.exs index 6158fe1cb0..0548b5af2f 100644 --- a/packages/sync-service/test/electric/replication/publication_manager_test.exs +++ b/packages/sync-service/test/electric/replication/publication_manager_test.exs @@ -14,6 +14,7 @@ defmodule Electric.Replication.PublicationManagerTest do require Repatch alias Electric.Utils + alias Electric.Postgres.Configuration alias Electric.Replication.Eval.Expr alias Electric.Replication.PublicationManager @@ -175,6 +176,154 @@ defmodule Electric.Replication.PublicationManagerTest do assert_pub_tables(ctx, [ctx.relation, alt_relation]) end + @tag update_debounce_timeout: 50 + test "batches a burst of new relations", ctx do + Postgrex.query!( + ctx.pool, + "CREATE TABLE other_table (id UUID PRIMARY KEY, value TEXT NOT NULL)", + [] + ) + + alt_relation = {"public", "other_table"} + alt_relation_oid = lookup_relation_oid(ctx.pool, alt_relation) + test_pid = self() + + Repatch.patch( + Configuration, + :configure_tables_for_replication, + [mode: :shared], + fn conn, publication, additions, replica_identity_relations -> + send(test_pid, {:batch_configuration, additions, replica_identity_relations}) + + Repatch.real( + Configuration.configure_tables_for_replication( + conn, + publication, + additions, + replica_identity_relations + ) + ) + end + ) + + tasks = [ + Task.async(fn -> + shape = generate_shape(ctx.relation_with_oid, @where_clause_1) + PublicationManager.add_shape(ctx.stack_id, @shape_handle_1, shape) + end), + Task.async(fn -> + shape = generate_shape({alt_relation_oid, alt_relation}, @where_clause_1) + PublicationManager.add_shape(ctx.stack_id, @shape_handle_2, shape) + end) + ] + + assert_receive {:batch_configuration, additions, replica_identity_relations}, 1_000 + + assert MapSet.new(additions) == + MapSet.new([ctx.relation_with_oid, {alt_relation_oid, alt_relation}]) + + assert MapSet.new(replica_identity_relations) == MapSet.new(additions) + assert Task.await_many(tasks) == [:ok, :ok] + assert_pub_tables(ctx, [ctx.relation, alt_relation]) + end + + @tag update_debounce_timeout: 50 + test "registers a complete shape tree in one publication batch", ctx do + Postgrex.query!( + ctx.pool, + "CREATE TABLE other_table (id UUID PRIMARY KEY, value TEXT NOT NULL)", + [] + ) + + alt_relation = {"public", "other_table"} + alt_relation_oid = lookup_relation_oid(ctx.pool, alt_relation) + test_pid = self() + + Repatch.patch( + Configuration, + :configure_tables_for_replication, + [mode: :shared], + fn conn, publication, additions, replica_identity_relations -> + send(test_pid, {:batch_configuration, additions, replica_identity_relations}) + + Repatch.real( + Configuration.configure_tables_for_replication( + conn, + publication, + additions, + replica_identity_relations + ) + ) + end + ) + + shapes = [ + {@shape_handle_1, generate_shape(ctx.relation_with_oid, @where_clause_1)}, + {@shape_handle_2, generate_shape({alt_relation_oid, alt_relation}, @where_clause_1)} + ] + + assert :ok == PublicationManager.register_shapes(ctx.stack_id, shapes) + + assert_receive {:batch_configuration, additions, replica_identity_relations}, 1_000 + + assert MapSet.new(additions) == + MapSet.new([ctx.relation_with_oid, {alt_relation_oid, alt_relation}]) + + assert MapSet.new(replica_identity_relations) == MapSet.new(additions) + assert_pub_tables(ctx, [ctx.relation, alt_relation]) + end + + @tag update_debounce_timeout: 200 + test "does not starve publication updates while registrations continue", ctx do + Postgrex.query!( + ctx.pool, + "CREATE TABLE other_table (id UUID PRIMARY KEY, value TEXT NOT NULL)", + [] + ) + + alt_relation = {"public", "other_table"} + alt_relation_oid = lookup_relation_oid(ctx.pool, alt_relation) + test_pid = self() + + Repatch.patch( + Configuration, + :configure_tables_for_replication, + [mode: :shared], + fn conn, publication, additions, replica_identity_relations -> + send(test_pid, {:batch_configuration, additions, replica_identity_relations}) + + Repatch.real( + Configuration.configure_tables_for_replication( + conn, + publication, + additions, + replica_identity_relations + ) + ) + end + ) + + assert :ok == + PublicationManager.register_shapes(ctx.stack_id, [ + {@shape_handle_1, generate_shape(ctx.relation_with_oid, @where_clause_1)} + ]) + + Process.sleep(100) + + assert :ok == + PublicationManager.register_shapes(ctx.stack_id, [ + {@shape_handle_2, + generate_shape({alt_relation_oid, alt_relation}, @where_clause_1)} + ]) + + assert_receive {:batch_configuration, additions, replica_identity_relations}, 150 + + assert MapSet.new(additions) == + MapSet.new([ctx.relation_with_oid, {alt_relation_oid, alt_relation}]) + + assert MapSet.new(replica_identity_relations) == MapSet.new(additions) + end + @tag update_debounce_timeout: 100 test "doesn't update when adding same relation again", ctx do notify_alter_queries()