diff --git a/CHANGELOG.md b/CHANGELOG.md index c9120912..a89e232d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,8 @@ * Add SCRAM-SHA-256-PLUS support * Use reports for logging errors * Support arrays of domains defined over arrays + * Allow replication callbacks to pause the next receive batch without + disconnecting, resuming on ordinary callback returns. * Bug fixes * Do not UNLISTEN while disconnected diff --git a/lib/postgrex/protocol.ex b/lib/postgrex/protocol.ex index 2d629f74..5ce4d449 100644 --- a/lib/postgrex/protocol.ex +++ b/lib/postgrex/protocol.ex @@ -1288,57 +1288,64 @@ defmodule Postgrex.Protocol do msg_send(s, msgs, buffer) end - @spec handle_copy_recv(any, Keyword.t(), state) :: + @spec handle_copy_recv(any, pos_integer() | nil, state) :: {:ok, [binary | atom], state} | :unknown | {:error, Postgrex.Error.t(), state} | {:disconnect, %DBConnection.ConnectionError{}, state} def handle_copy_recv(msg, max_copies, s) do + handle_copy_recv(msg, max_copies, s, true) + end + + def handle_copy_recv(msg, max_copies, s, reactivate?) do case handle_socket(msg, s) do - {:data, data} -> handle_copy_recv(s, max_copies, [], 0, data) + {:data, data} -> handle_copy_recv(s, max_copies, [], 0, data, reactivate?) :ignore -> {:ok, [], s} :unknown -> :unknown disconnect -> disconnect end end - defp handle_copy_recv(s, max_copies, copies, max_copies, buffer) do - with {:ok, s} <- activate(s, buffer) do + defp handle_copy_recv(s, max_copies, copies, max_copies, buffer, reactivate?) do + with {:ok, s} <- copy_recv_buffer(s, buffer, reactivate?) do {:ok, Enum.reverse(copies), s} end end - defp handle_copy_recv(%{timeout: timeout} = s, max_copies, copies, ncopies, buffer) do + defp handle_copy_recv(%{timeout: timeout} = s, max_copies, copies, ncopies, buffer, reactivate?) do case msg_recv(s, timeout, buffer) do {:ok, msg_error(fields: fields), buffer} -> disconnect(s, Postgrex.Error.exception(postgres: fields), buffer) {:ok, msg_copy_data(data: data), <<>>} -> - with {:ok, s} <- activate(s, <<>>) do + with {:ok, s} <- copy_recv_buffer(s, <<>>, reactivate?) do {:ok, Enum.reverse([data | copies]), s} end {:ok, msg_copy_data(data: data), buffer} -> - handle_copy_recv(s, max_copies, [data | copies], ncopies + 1, buffer) + handle_copy_recv(s, max_copies, [data | copies], ncopies + 1, buffer, reactivate?) {:ok, msg_copy_done(), buffer} -> - handle_copy_recv(s, max_copies, copies, ncopies, buffer) + handle_copy_recv(s, max_copies, copies, ncopies, buffer, reactivate?) {:ok, msg_command_complete(), buffer} -> - handle_copy_recv(s, max_copies, copies, ncopies, buffer) + handle_copy_recv(s, max_copies, copies, ncopies, buffer, reactivate?) {:ok, msg_ready(status: postgres), buffer} -> s = %{s | postgres: postgres, buffer: buffer} {:ok, Enum.reverse([:copy_done | copies]), s} {:ok, _msg, buffer} -> - handle_copy_recv(s, max_copies, copies, ncopies, buffer) + handle_copy_recv(s, max_copies, copies, ncopies, buffer, reactivate?) {:disconnect, _, _} = dis -> dis end end + defp copy_recv_buffer(s, buffer, true), do: activate(s, buffer) + defp copy_recv_buffer(s, buffer, false), do: {:ok, %{s | buffer: buffer}} + @spec handle_streaming(String.t() | iolist(), state) :: {:ok, state} | {:error, Postgrex.Error.t(), state} diff --git a/lib/postgrex/replication_connection.ex b/lib/postgrex/replication_connection.ex index c6fc69d4..cab67d86 100644 --- a/lib/postgrex/replication_connection.ex +++ b/lib/postgrex/replication_connection.ex @@ -153,7 +153,9 @@ defmodule Postgrex.ReplicationConnection do state: nil, auto_reconnect: false, reconnect_backoff: 500, - streaming: nil + streaming: nil, + paused: false, + needs_checkin: false ## PUBLIC API ## @@ -227,10 +229,30 @@ defmodule Postgrex.ReplicationConnection do Replication messages may require explicit acknowledgement, which can be done by returning a list of binaries according to the replication protocol. + + Return `{:pause, state}` or `{:pause, replies, state}` to pause the next + receive batch. Messages already decoded in the current batch are still + passed to this callback, in order. Every callback returning `:noreply` + allows receiving again; return `:pause` from `handle_data/2`, `handle_info/2` + or `handle_call/3` to keep receiving paused while handling other events. + Use `max_messages: 1` to pause after each individual stream message. + Pausing is supported while a stream is active, not for the `:done` callback. + + The process remains responsive to calls and messages while paused. For logical + replication, the callback module must send periodic standby feedback using + `{:pause, replies, state}` to avoid the server's `wal_sender_timeout` while + remaining paused. Feedback must acknowledge only durably processed data. + Incoming keepalives are paused along with other stream messages. + + `:max_messages` bounds the number of decoded messages in a receive batch, not + its byte size. A partial wire message and the remainder of the socket buffer + may also be retained; this API is not a hard process-memory limit. """ @callback handle_data(binary | :done, state) :: {:noreply, state} | {:noreply, ack, state} + | {:pause, state} + | {:pause, ack, state} | {:query, query, state} | {:query, query, query_opts, state} | {:stream, query, stream_opts, state} @@ -242,6 +264,8 @@ defmodule Postgrex.ReplicationConnection do @callback handle_info(term, state) :: {:noreply, state} | {:noreply, ack, state} + | {:pause, state} + | {:pause, ack, state} | {:query, query, state} | {:query, query, query_opts, state} | {:stream, query, stream_opts, state} @@ -262,6 +286,8 @@ defmodule Postgrex.ReplicationConnection do @callback handle_call(term, :gen_statem.from(), state) :: {:noreply, state} | {:noreply, ack, state} + | {:pause, state} + | {:pause, ack, state} | {:query, query, state} | {:query, query, query_opts, state} | {:stream, query, stream_opts, state} @@ -518,7 +544,12 @@ defmodule Postgrex.ReplicationConnection do def handle_event(:internal, {:connect, _info}, @state, %{state: {mod, mod_state}} = s) do case Protocol.connect(opts()) do {:ok, protocol} -> - maybe_handle(mod, :handle_connect, [mod_state], %{s | protocol: protocol}) + maybe_handle(mod, :handle_connect, [mod_state], %{ + s + | protocol: protocol, + paused: false, + needs_checkin: false + }) {:error, reason} -> Logger.error( @@ -544,9 +575,12 @@ defmodule Postgrex.ReplicationConnection do end def handle_event(:info, msg, @state, %{protocol: protocol, streaming: streaming} = s) do - case Protocol.handle_copy_recv(msg, streaming, protocol) do + case Protocol.handle_copy_recv(msg, streaming, protocol, false) do + {:ok, [], ^protocol} -> + {:keep_state, s} + {:ok, copies, protocol} -> - handle_data(copies, %{s | protocol: protocol}) + handle_data(copies, %{s | protocol: protocol, needs_checkin: true}) :unknown -> %{state: {mod, mod_state}} = s @@ -578,7 +612,7 @@ defmodule Postgrex.ReplicationConnection do ]} end - defp handle_data([], s), do: {:keep_state, s} + defp handle_data([], s), do: maybe_checkin(s) defp handle_data([:copy_done | copies], %{state: {mod, mod_state}} = s) do with {:keep_state, s} <- @@ -604,15 +638,18 @@ defmodule Postgrex.ReplicationConnection do defp handle(mod, fun, args, from, %{streaming: streaming} = s) do case apply(mod, fun, args) do {:noreply, mod_state} -> - {:keep_state, %{s | state: {mod, mod_state}}} + finish_callback(fun, %{s | state: {mod, mod_state}, paused: false}) {:noreply, replies, mod_state} -> - s = %{s | state: {mod, mod_state}} + send_replies(replies, fun, %{s | state: {mod, mod_state}, paused: false}) - case Protocol.handle_copy_send(replies, s.protocol) do - :ok -> {:keep_state, s} - {error, reason, protocol} -> reconnect_or_stop(error, reason, protocol, s) - end + {:pause, mod_state} + when streaming != nil and fun in [:handle_data, :handle_info, :handle_call] -> + {:keep_state, %{s | state: {mod, mod_state}, paused: true}} + + {:pause, replies, mod_state} + when streaming != nil and fun in [:handle_data, :handle_info, :handle_call] -> + send_replies(replies, fun, %{s | state: {mod, mod_state}, paused: true}) {:stream, query, opts, mod_state} when streaming == nil -> s = %{s | state: {mod, mod_state}} @@ -620,7 +657,7 @@ defmodule Postgrex.ReplicationConnection do with {:ok, protocol} <- Protocol.handle_streaming(query, s.protocol), {:ok, protocol} <- Protocol.checkin(protocol) do - {:keep_state, %{s | protocol: protocol, streaming: max_messages}} + {:keep_state, %{s | protocol: protocol, streaming: max_messages, needs_checkin: false}} else {error_or_disconnect, reason, protocol} -> reconnect_or_stop(error_or_disconnect, reason, protocol, s) @@ -646,6 +683,27 @@ defmodule Postgrex.ReplicationConnection do end end + # Data callbacks must finish the decoded batch before reactivating the socket. + defp finish_callback(fun, s) when fun in [:handle_info, :handle_call], do: maybe_checkin(s) + defp finish_callback(_fun, s), do: {:keep_state, s} + + defp maybe_checkin(%{paused: false, needs_checkin: true, streaming: streaming} = s) + when streaming != nil do + case Protocol.checkin(s.protocol) do + {:ok, protocol} -> {:keep_state, %{s | protocol: protocol, needs_checkin: false}} + {error, reason, protocol} -> reconnect_or_stop(error, reason, protocol, s) + end + end + + defp maybe_checkin(s), do: {:keep_state, s} + + defp send_replies(replies, fun, s) do + case Protocol.handle_copy_send(replies, s.protocol) do + :ok -> finish_callback(fun, s) + {error, reason, protocol} -> reconnect_or_stop(error, reason, protocol, s) + end + end + defp handle_query(query, mod, from, s, mod_state, opts) do case Protocol.handle_simple(query, opts, s.protocol) do {:ok, results, protocol} when is_list(results) -> @@ -690,7 +748,9 @@ defmodule Postgrex.ReplicationConnection do ) {:keep_state, s} = maybe_handle(mod, :handle_disconnect, [mod_state], s) - {:keep_state, %{s | streaming: nil}, {:next_event, :internal, {:connect, :reconnect}}} + + {:keep_state, %{s | streaming: nil, paused: false, needs_checkin: false}, + {:next_event, :internal, {:connect, :reconnect}}} end defp pid_or_name do diff --git a/test/replication_backpressure_test.exs b/test/replication_backpressure_test.exs new file mode 100644 index 00000000..bf675424 --- /dev/null +++ b/test/replication_backpressure_test.exs @@ -0,0 +1,233 @@ +defmodule Postgrex.ReplicationBackpressureTest do + use ExUnit.Case, async: false + + alias Postgrex.ReplicationConnection, as: Connection + + @moduletag :logical_replication + @moduletag min_pg_version: "10.0" + + defmodule Receiver do + use Postgrex.ReplicationConnection + + def start_link(opts) do + {owner, opts} = Keyword.pop!(opts, :owner) + Connection.start_link(__MODULE__, owner, opts) + end + + @impl true + def init(owner), do: {:ok, owner} + + @impl true + def handle_connect(owner) do + send(owner, {:connected, self()}) + {:noreply, owner} + end + + @impl true + def handle_disconnect(owner), do: {:noreply, owner} + + @impl true + def handle_call({:query, sql}, from, owner), do: {:query, sql, {from, owner}} + + def handle_call({:stream, sql, batch_size}, from, owner) do + Connection.reply(from, :ok) + {:stream, sql, [max_messages: batch_size], owner} + end + + def handle_call(:ping, from, owner) do + Connection.reply(from, :pong) + {:pause, owner} + end + + def handle_call(:continue, from, owner) do + Connection.reply(from, :ok) + {:noreply, owner} + end + + @impl true + def handle_result(result, {from, owner}) do + Connection.reply(from, result) + {:noreply, owner} + end + + @impl true + def handle_data(:done, owner) do + send(owner, :done) + {:noreply, owner} + end + + def handle_data(<>, owner), do: {:noreply, owner} + + def handle_data("continue\n" = data, owner) do + send(owner, {:data, data}) + {:noreply, owner} + end + + def handle_data(data, owner) do + send(owner, {:data, data}) + {:pause, owner} + end + + @impl true + def handle_info(:resume, owner), do: {:noreply, owner} + def handle_info(:disconnect, _owner), do: {:disconnect, :test_disconnect} + + def handle_info(:resume_with_feedback, owner) do + {:pause, replies, owner} = handle_info(:heartbeat, owner) + {:noreply, replies, owner} + end + + def handle_info(:heartbeat, owner) do + clock = System.os_time(:microsecond) - 946_684_800_000_000 + {:pause, [<>], owner} + end + end + + for transport <- [:tcp, :ssl] do + @tag transport: transport + if transport == :ssl, do: @tag(:ssl) + + test "pauses and resumes ordered COPY data over #{transport}", %{transport: transport} do + opts = if transport == :ssl, do: [ssl: [verify: :verify_none]], else: [] + receiver = start_supervised!({Receiver, [owner: self(), database: "postgrex_test"] ++ opts}) + + assert :ok = + Connection.call( + receiver, + {:stream, "COPY (SELECT generate_series(1, 20)) TO STDOUT", 1} + ) + + for expected <- 1..20 do + assert_receive {:data, data} + assert data == "#{expected}\n" + assert Connection.call(receiver, :ping) == :pong + {:no_state, state} = :sys.get_state(receiver) + assert state.paused + refute Map.has_key?(state, :pending) + assert is_binary(state.protocol.buffer) + refute_receive {:data, _}, 10 + + if rem(expected, 2) == 0, + do: send(receiver, :resume), + else: assert(Connection.call(receiver, :continue) == :ok) + end + + assert_receive :done + assert [%Postgrex.Result{rows: [["1"]]}] = Connection.call(receiver, {:query, "SELECT 1"}) + end + end + + test "pause finishes the decoded batch before stopping the next receive" do + receiver = start_supervised!({Receiver, owner: self(), database: "postgrex_test"}) + + assert :ok = + Connection.call( + receiver, + {:stream, "COPY (SELECT generate_series(1, 20)) TO STDOUT", 3} + ) + + for batch <- Enum.chunk_every(Enum.to_list(1..20), 3) do + for expected <- batch do + assert_receive {:data, data} + assert data == "#{expected}\n" + end + + refute_receive {:data, _}, 20 + send(receiver, :resume) + end + + assert_receive :done + end + + test "an ordinary data return allows the next batch after a pause return" do + receiver = start_supervised!({Receiver, owner: self(), database: "postgrex_test"}) + sql = "COPY (SELECT unnest(ARRAY['pause', 'continue', 'third', 'fourth'])) TO STDOUT" + assert :ok = Connection.call(receiver, {:stream, sql, 2}) + + for expected <- ["pause", "continue", "third", "fourth"] do + assert_receive {:data, data} + assert data == expected <> "\n" + end + + refute_receive :done, 20 + send(receiver, :resume) + assert_receive :done + end + + @tag min_pg_version: "14.0" + test "paused logical replication keeps slot ownership and can send feedback" do + receiver = start_supervised!({Receiver, owner: self(), database: "postgrex_test"}) + slot = "postgrex_pause_#{System.unique_integer([:positive])}" + + assert [%Postgrex.Result{}] = + Connection.call( + receiver, + {:query, + "CREATE_REPLICATION_SLOT #{slot} TEMPORARY LOGICAL pgoutput NOEXPORT_SNAPSHOT"} + ) + + assert :ok = + Connection.call( + receiver, + {:stream, + "START_REPLICATION SLOT #{slot} LOGICAL 0/0 (proto_version '1', publication_names 'postgrex_example', messages 'true')", + 2} + ) + + writer = start_supervised!({Postgrex, database: "postgrex_test"}) + Postgrex.query!(writer, "SELECT pg_logical_emit_message(false, 'pause_test', 'first')", []) + assert_receive {:data, <>} + + [[backend_pid]] = + Postgrex.query!( + writer, + "SELECT active_pid FROM pg_replication_slots WHERE slot_name = $1", + [slot] + ).rows + + assert is_integer(backend_pid) + + send(receiver, :heartbeat) + assert Connection.call(receiver, :ping) == :pong + Postgrex.query!(writer, "SELECT pg_logical_emit_message(false, 'pause_test', 'second')", []) + refute_receive {:data, _}, 100 + + assert [[^backend_pid]] = + Postgrex.query!( + writer, + "SELECT active_pid FROM pg_replication_slots WHERE slot_name = $1", + [slot] + ).rows + + send(receiver, :resume_with_feedback) + assert_receive {:data, <>} + end + + @tag :capture_log + test "reconnection discards paused messages from the old connection" do + receiver = + start_supervised!( + {Receiver, owner: self(), database: "postgrex_test", auto_reconnect: true} + ) + + assert_receive {:connected, ^receiver} + + assert :ok = + Connection.call( + receiver, + {:stream, "COPY (SELECT generate_series(1, 20)) TO STDOUT", 1} + ) + + assert_receive {:data, "1\n"} + {:no_state, state} = :sys.get_state(receiver) + assert state.paused + + send(receiver, :disconnect) + assert_receive {:connected, ^receiver} + {:no_state, state} = :sys.get_state(receiver) + refute state.paused + refute Map.has_key?(state, :pending) + assert state.streaming == nil + assert [%Postgrex.Result{rows: [["1"]]}] = Connection.call(receiver, {:query, "SELECT 1"}) + end +end