diff --git a/CHANGELOG.md b/CHANGELOG.md index c9120912..f343c174 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 stream delivery without disconnecting, + then resume buffered messages in order while continuing to send feedback. * 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..b0e27f91 100644 --- a/lib/postgrex/replication_connection.ex +++ b/lib/postgrex/replication_connection.ex @@ -153,7 +153,10 @@ defmodule Postgrex.ReplicationConnection do state: nil, auto_reconnect: false, reconnect_backoff: 500, - streaming: nil + streaming: nil, + paused: false, + pending: [], + needs_checkin: false ## PUBLIC API ## @@ -227,10 +230,35 @@ 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}` from this callback to + pause delivery after the current message. Already received messages remain + ordered and the socket is not reactivated until `handle_info/2` or + `handle_call/3` returns `{:resume, state}` or `{:resume, replies, state}`. + Pausing is supported while a stream is active, not for the `:done` callback. + + A paused stream is still active. Returning `{:query, ...}` while streaming, + including while paused, is rejected by the existing stream-in-progress + guard, the query is not sent or queued. Queries needed to process replication + messages must use a separate database connection. Queries on this connection + are supported before streaming starts or after it finishes, including from + `handle_data(:done, state)`. + + The process remains responsive to calls and messages while paused. For logical + replication, the callback module must send periodic standby feedback using + `{:noreply, replies, state}` to avoid the server's `wal_sender_timeout`. 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 +270,8 @@ defmodule Postgrex.ReplicationConnection do @callback handle_info(term, state) :: {:noreply, state} | {:noreply, ack, state} + | {:resume, state} + | {:resume, ack, state} | {:query, query, state} | {:query, query, query_opts, state} | {:stream, query, stream_opts, state} @@ -262,6 +292,8 @@ defmodule Postgrex.ReplicationConnection do @callback handle_call(term, :gen_statem.from(), state) :: {:noreply, state} | {:noreply, ack, state} + | {:resume, state} + | {:resume, ack, state} | {:query, query, state} | {:query, query, query_opts, state} | {:stream, query, stream_opts, state} @@ -518,7 +550,13 @@ 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, + pending: [], + needs_checkin: false + }) {:error, reason} -> Logger.error( @@ -543,10 +581,17 @@ defmodule Postgrex.ReplicationConnection do handle(mod, :handle_call, [msg, from, mod_state], from, s) end + def handle_event(:internal, :resume, @state, s) do + handle_data(s.pending, %{s | pending: []}) + 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,6 +623,19 @@ defmodule Postgrex.ReplicationConnection do ]} end + defp handle_data(copies, %{paused: true} = s), do: {:keep_state, %{s | pending: copies}} + + defp handle_data( + [], + %{needs_checkin: true, streaming: streaming, protocol: %{buffer: buffer}} = s + ) + when streaming != nil and is_binary(buffer) 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 handle_data([], s), do: {:keep_state, s} defp handle_data([:copy_done | copies], %{state: {mod, mod_state}} = s) do @@ -614,13 +672,26 @@ defmodule Postgrex.ReplicationConnection do {error, reason, protocol} -> reconnect_or_stop(error, reason, protocol, s) end + {:pause, mod_state} when fun == :handle_data and streaming != nil -> + {:keep_state, %{s | state: {mod, mod_state}, paused: true}} + + {:pause, replies, mod_state} when fun == :handle_data and streaming != nil -> + s = %{s | state: {mod, mod_state}, paused: true} + send_replies(replies, s, []) + + {:resume, mod_state} when fun in [:handle_info, :handle_call] -> + resume(%{s | state: {mod, mod_state}}, []) + + {:resume, replies, mod_state} when fun in [:handle_info, :handle_call] -> + resume(%{s | state: {mod, mod_state}}, replies) + {:stream, query, opts, mod_state} when streaming == nil -> s = %{s | state: {mod, mod_state}} max_messages = opts[:max_messages] || @max_messages 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 +717,20 @@ defmodule Postgrex.ReplicationConnection do end end + defp resume(%{paused: true} = s, replies) do + send_replies(replies, %{s | paused: false}, [{:next_event, :internal, :resume}]) + end + + defp resume(s, replies), do: send_replies(replies, s, []) + + defp send_replies(replies, s, actions) do + case Protocol.handle_copy_send(replies, s.protocol) do + :ok when actions == [] -> {:keep_state, s} + :ok -> {:keep_state, s, actions} + {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 +775,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, pending: [], needs_checkin: false}, + {:next_event, :internal, {:connect, :reconnect}}} end defp pid_or_name do diff --git a/test/query_test.exs b/test/query_test.exs index c3702dc2..64120a7b 100644 --- a/test/query_test.exs +++ b/test/query_test.exs @@ -2180,6 +2180,7 @@ defmodule QueryTest do end test "terminate backend during query returns FATAL error", context do + Process.flag(:trap_exit, true) assert {:ok, pid} = P.start_link([idle_interval: 10] ++ context[:options]) %Postgrex.Result{connection_id: connection_id} = Postgrex.query!(pid, "SELECT 42", []) @@ -2204,6 +2205,8 @@ defmodule QueryTest do assert {:error, %Postgrex.Error{postgres: %{code: :admin_shutdown, severity: "FATAL"}}} = Task.await(task, 5000) + + assert_receive {:EXIT, ^pid, :killed}, 5000 after :erlang.trace_pattern({Postgrex.Protocol, :recv_bind, :_}, false, []) :erlang.trace(:all, false, [:call]) diff --git a/test/replication_backpressure_test.exs b/test/replication_backpressure_test.exs new file mode 100644 index 00000000..ab23e843 --- /dev/null +++ b/test/replication_backpressure_test.exs @@ -0,0 +1,180 @@ +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_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) + {: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(data, owner) do + send(owner, {:data, data}) + {:pause, owner} + end + + @impl true + def handle_info(:resume, owner), do: {:resume, owner} + def handle_info(:disconnect, _owner), do: {:disconnect, :test_disconnect} + + def handle_info(:resume_with_feedback, owner) do + {:noreply, replies, owner} = handle_info(:heartbeat, owner) + {:resume, replies, owner} + end + + def handle_info(:heartbeat, owner) do + clock = System.os_time(:microsecond) - 946_684_800_000_000 + {:noreply, [<>], 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", 3} + ) + + 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 + assert length(state.pending) <= 2 + assert is_binary(state.protocol.buffer) + refute_receive {:data, _}, 10 + send(receiver, :resume) + end + + assert_receive :done + assert [%Postgrex.Result{rows: [["1"]]}] = Connection.call(receiver, {:query, "SELECT 1"}) + end + 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", 3} + ) + + 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 + assert state.pending == [] + assert state.streaming == nil + assert [%Postgrex.Result{rows: [["1"]]}] = Connection.call(receiver, {:query, "SELECT 1"}) + end +end diff --git a/test/replication_connection_test.exs b/test/replication_connection_test.exs index 19453fa7..4111e8fc 100644 --- a/test/replication_connection_test.exs +++ b/test/replication_connection_test.exs @@ -83,6 +83,10 @@ defmodule ReplicationTest do {:stream, query, [], pid} end + def handle_info(:request_keepalive, pid) do + {:noreply, [<>], pid} + end + def handle_info(_, pid) do {:noreply, pid} end @@ -239,7 +243,9 @@ defmodule ReplicationTest do @tag :capture_log test "resumes streaming after reconnect", context do assert_receive {:connect, i1} + # Request a reply instead of relying on an unsolicited server keepalive. start_replication(context.repl) + send(context.repl, :request_keepalive) assert_receive <>, @timeout :sys.suspend(context.repl) @@ -250,6 +256,7 @@ defmodule ReplicationTest do assert_receive {:disconnect, i2} when i1 < i2, @timeout assert_receive {:connect, i3} when i2 < i3, @timeout start_replication(context.repl) + send(context.repl, :request_keepalive) assert_receive <>, @timeout end end @@ -319,6 +326,7 @@ defmodule ReplicationTest do test "on replication", context do start_replication(context.repl) + send(context.repl, :request_keepalive) assert_receive <>, @timeout end @@ -363,6 +371,7 @@ defmodule ReplicationTest do assert_receive {"1\t1\n", i2} when i1 < i2, @timeout assert_receive {:done, i3} when i2 < i3, @timeout # Prior to allowing one stream to start after another, this would fail + send(context.repl, :request_keepalive) assert_receive <>, @timeout end end