Fix notifications lost during idle ping

lib/postgrex/protocol.ex

@@ -313,12 +313,17 @@ defmodule Postgrex.Protocol do
   @spec ping(state) ::
           {:ok, state}
           | {:disconnect, Postgrex.Error.t() | %DBConnection.ConnectionError{}, state}
-  def ping(%{postgres: :transaction, transactions: :strict} = s) do
+  def ping(s), do: ping(s, [])
+
+  @spec ping(state, Keyword.t()) ::
+          {:ok, state}
+          | {:disconnect, Postgrex.Error.t() | %DBConnection.ConnectionError{}, state}
+  def ping(%{postgres: :transaction, transactions: :strict} = s, _opts) do
     sync_error(s, :transaction)
   end
 
-  def ping(%{buffer: buffer} = s) do
-    status = new_status([], mode: :transaction)
+  def ping(%{buffer: buffer} = s, opts) do
+    status = new_status(opts, mode: :transaction)
     s = %{s | buffer: nil}
 
     case msg_send(s, msg_sync(), buffer) do

lib/postgrex/simple_connection.ex

@@ -397,8 +397,15 @@ defmodule Postgrex.SimpleConnection do
     handle(mod, :handle_call, [msg, callback_from, mod_state], from, state)
   end
 
-  def handle_event(:timeout, nil, @state, %{protocol: protocol} = state) do
-    case Protocol.ping(protocol) do
+  def handle_event(
+        :timeout,
+        nil,
+        @state,
+        %{protocol: protocol, state: {mod, mod_state}} = state
+      ) do
+    opts = [notify: &mod.notify(&1, &2, mod_state)]
+
+    case Protocol.ping(protocol, opts) do
       {:ok, protocol} ->
         {:keep_state, %{state | protocol: protocol}, {:timeout, state.idle_interval, nil}}
 

test/protocol_test.exs

@@ -34,6 +34,29 @@ defmodule Postgrex.ProtocolTest do
     assert_receive {:sent, <<?Q, _size::32, "START_REPLICATION", 0>>}
   end
 
+  test "ping handles notifications received before ready" do
+    responses =
+      IO.iodata_to_binary([
+        backend_message(?A, [<<123::32>>, "events", 0, "ready", 0]),
+        backend_message(?Z, [?I])
+      ])
+
+    state = %Protocol{
+      sock: {Socket, self()},
+      buffer: responses,
+      postgres: :idle,
+      transactions: :naive,
+      messages: []
+    }
+
+    notify = fn channel, payload -> send(self(), {:notification, channel, payload}) end
+
+    assert {:ok, state} = Protocol.ping(state, notify: notify)
+    assert state.buffer == ""
+    assert_receive {:notification, "events", "ready"}
+    assert_receive {:sent, <<?S, 4::32>>}
+  end
+
   defp backend_message(type, data) do
     [type, <<IO.iodata_length(data) + 4::32>>, data]
   end

test/simple_connection_test.exs

@@ -1,7 +1,8 @@
 defmodule SimpleConnectionTest do
   use ExUnit.Case, async: true
 
-  alias Postgrex.SimpleConnection, as: SC
+  alias Postgrex.{Protocol, SimpleConnection}
+  alias SimpleConnection, as: SC
 
   defmodule Conn do
     @behaviour Postgrex.SimpleConnection
@@ -60,6 +61,13 @@ defmodule SimpleConnectionTest do
     end
   end
 
+  defmodule Socket do
+    def send(pid, data) do
+      Kernel.send(pid, {:sent, IO.iodata_to_binary(data)})
+      :ok
+    end
+  end
+
   @opts [database: "postgrex_test", sync_connect: true, auto_reconnect: false]
 
   setup context do
@@ -130,6 +138,37 @@ defmodule SimpleConnectionTest do
     end
   end
 
+  describe "idle ping" do
+    test "relays notifications received while pinging" do
+      responses =
+        IO.iodata_to_binary([
+          backend_message(?A, [<<123::32>>, "events", 0, "ready", 0]),
+          backend_message(?Z, [?I])
+        ])
+
+      protocol = %Protocol{
+        sock: {Socket, self()},
+        buffer: responses,
+        postgres: :idle,
+        transactions: :naive,
+        messages: []
+      }
+
+      state = %SC{
+        idle_interval: 10,
+        protocol: protocol,
+        state: {Conn, %{pid: self()}}
+      }
+
+      assert {:keep_state, state, {:timeout, 10, nil}} =
+               SC.handle_event(:timeout, nil, :no_state, state)
+
+      assert state.protocol.buffer == ""
+      assert_receive {"events", "ready"}
+      assert_receive {:sent, <<?S, 4::32>>}
+    end
+  end
+
   describe "auto-reconnect" do
     @tag opts: [auto_reconnect: true]
     test "disconnect and connect handlers are invoked on reconnection", context do
@@ -188,4 +227,8 @@ defmodule SimpleConnectionTest do
     {:gen_tcp, sock} = state.protocol.sock
     :gen_tcp.shutdown(sock, :read_write)
   end
+
+  defp backend_message(type, data) do
+    [type, <<IO.iodata_length(data) + 4::32>>, data]
+  end
 end