Add pause and resume for replication streams

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

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}

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

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])