Add query timeout option on ReplicationConnection (#748)

Author: edgurgel

lib/postgrex/protocol.ex

@@ -1153,36 +1153,37 @@ defmodule Postgrex.Protocol do
 
   ## replication/notifications
 
-  @spec handle_simple(String.t() | iolist(), state) ::
+  @spec handle_simple(String.t() | iolist(), Keyword.t(), state) ::
           {:ok, [Postgrex.Result.t()], state}
           | {:error, Postgrex.Error.t(), state}
           | {:disconnect, %DBConnection.ConnectionError{}, state}
   def handle_simple(statement, opts \\ [], %{buffer: buffer} = s) do
     status = new_status(opts, mode: :transaction)
     msgs = [msg_query(statement: statement)]
+    timeout = Keyword.get(opts, :timeout, :infinity)
 
     case msg_send(%{s | buffer: nil}, msgs, buffer) do
       :ok ->
-        recv_simple(s, status, [], [], [], buffer)
+        recv_simple(s, status, [], [], [], buffer, timeout)
 
       {:disconnect, err, s} ->
         {:disconnect, err, s}
     end
   end
 
-  defp recv_simple(s, status, results, columns, rows, buffer) do
-    case msg_recv(s, :infinity, buffer) do
+  defp recv_simple(s, status, results, columns, rows, buffer, timeout) do
+    case msg_recv(s, timeout, buffer) do
       {:ok, msg_row_desc(fields: fields), buffer} ->
         columns = column_names(fields)
-        recv_simple(s, status, results, columns, rows, buffer)
+        recv_simple(s, status, results, columns, rows, buffer, timeout)
 
       {:ok, msg_data_row(values: values), buffer} ->
         row = Types.decode_simple(values, s.types)
-        recv_simple(s, status, results, columns, [row | rows], buffer)
+        recv_simple(s, status, results, columns, [row | rows], buffer, timeout)
 
       {:ok, msg_command_complete(tag: tag), buffer} ->
         {result, s} = done(s, columns, Enum.reverse(rows), [tag])
-        recv_simple(s, status, [result | results], [], [], buffer)
+        recv_simple(s, status, [result | results], [], [], buffer, timeout)
 
       {:ok, msg_error(fields: fields), buffer} ->
         err = Postgrex.Error.exception(postgres: fields)
@@ -1194,7 +1195,7 @@ defmodule Postgrex.Protocol do
 
       {:ok, msg, buffer} ->
         s = handle_msg(s, status, msg)
-        recv_simple(s, status, results, columns, rows, buffer)
+        recv_simple(s, status, results, columns, rows, buffer, timeout)
 
       {:disconnect, _, _} = dis ->
         dis

lib/postgrex/replication_connection.ex

@@ -172,6 +172,17 @@ defmodule Postgrex.ReplicationConnection do
 
   """
   @type stream_opts :: [max_messages: pos_integer]
+
+  @query_timeout :infinity
+
+  @typedoc """
+  The following options configure querying:
+
+    * `:timeout` - Query request timeout (default: `#{@query_timeout}`);
+
+  """
+  @type query_opts :: [timeout: timeout]
+
   @max_lsn_component_size 8
   @max_uint64 18_446_744_073_709_551_615
   @max_messages 500
@@ -192,6 +203,7 @@ defmodule Postgrex.ReplicationConnection do
               {:noreply, state}
               | {:noreply, ack, state}
               | {:query, query, state}
+              | {:query, query, query_opts, state}
               | {:stream, query, stream_opts, state}
               | {:disconnect, reason}
 
@@ -220,6 +232,7 @@ defmodule Postgrex.ReplicationConnection do
               {:noreply, state}
               | {:noreply, ack, state}
               | {:query, query, state}
+              | {:query, query, query_opts, state}
               | {:stream, query, stream_opts, state}
               | {:disconnect, reason}
 
@@ -230,6 +243,7 @@ defmodule Postgrex.ReplicationConnection do
               {:noreply, state}
               | {:noreply, ack, state}
               | {:query, query, state}
+              | {:query, query, query_opts, state}
               | {:stream, query, stream_opts, state}
               | {:disconnect, reason}
 
@@ -249,17 +263,19 @@ defmodule Postgrex.ReplicationConnection do
               {:noreply, state}
               | {:noreply, ack, state}
               | {:query, query, state}
+              | {:query, query, query_opts, state}
               | {:stream, query, stream_opts, state}
               | {:disconnect, reason}
 
   @doc """
   Callback for `:query` outputs.
 
-  If any callback returns `{:query, iodata, state}`,
-  then this callback will be immediately called with
-  the result of the query. Please note that even though
-  replication connections use the simple query protocol,
-  Postgres currently limits them to single command queries.
+  If any callback returns `{:query, iodata, state}` or
+  `{:query, iodata, opts, state}`, then this callback will
+  be immediately called with the result of the query.
+  Please note that even though replication connections use
+  the simple query protocol, Postgres currently limits them to
+  single command queries.
   Due to this constraint, this callback will be passed
   either a list with a single successful query result or
   an error.
@@ -268,6 +284,7 @@ defmodule Postgrex.ReplicationConnection do
               {:noreply, state}
               | {:noreply, ack, state}
               | {:query, query, state}
+              | {:query, query, query_opts, state}
               | {:stream, query, stream_opts, state}
               | {:disconnect, reason}
 
@@ -586,25 +603,35 @@ defmodule Postgrex.ReplicationConnection do
         stream_in_progress(:stream, mod, mod_state, from, s)
 
       {:query, query, mod_state} when streaming == nil ->
-        case Protocol.handle_simple(query, [], s.protocol) do
-          {:ok, results, protocol} when is_list(results) ->
-            handle(mod, :handle_result, [results, mod_state], from, %{s | protocol: protocol})
-
-          {:error, %Postgrex.Error{} = error, protocol} ->
-            handle(mod, :handle_result, [error, mod_state], from, %{s | protocol: protocol})
+        handle_query(query, mod, from, s, mod_state, timeout: @query_timeout)
 
-          {:disconnect, reason, protocol} ->
-            reconnect_or_stop(:disconnect, reason, protocol, %{s | state: {mod, mod_state}})
-        end
+      {:query, query, opts, mod_state} when streaming == nil ->
+        handle_query(query, mod, from, s, mod_state, opts)
 
       {:query, _query, mod_state} ->
         stream_in_progress(:query, mod, mod_state, from, s)
 
+      {:query, _query, _opts, mod_state} ->
+        stream_in_progress(:query, mod, mod_state, from, s)
+
       {:disconnect, reason} ->
         reconnect_or_stop(:disconnect, reason, s.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) ->
+        handle(mod, :handle_result, [results, mod_state], from, %{s | protocol: protocol})
+
+      {:error, %Postgrex.Error{} = error, protocol} ->
+        handle(mod, :handle_result, [error, mod_state], from, %{s | protocol: protocol})
+
+      {:disconnect, reason, protocol} ->
+        reconnect_or_stop(:disconnect, reason, protocol, %{s | state: {mod, mod_state}})
+    end
+  end
+
   defp stream_in_progress(command, mod, mod_state, from, s) do
     Logger.warning("received #{command} while stream is already in progress")
     from && reply(from, {__MODULE__, :stream_in_progress})

test/replication_connection_test.exs

@@ -93,11 +93,14 @@ defmodule ReplicationTest do
       {:noreply, pid}
     end
 
-    @impl true
     def handle_call({:query, query}, from, pid) do
       {:query, query, {from, pid}}
     end
 
+    def handle_call({:query, query, %{timeout: timeout}}, from, pid) do
+      {:query, query, [timeout: timeout], {from, pid}}
+    end
+
     # This is part of the "stream_continuation" test and handles call that
     # triggers that chain of events.
     def handle_call(
@@ -108,7 +111,6 @@ defmodule ReplicationTest do
       {:query, query, Map.merge(opts, %{from: from, pid: pid})}
     end
 
-    @impl true
     def handle_call({:disconnect, reason}, _, _) do
       {:disconnect, reason}
     end
@@ -167,6 +169,24 @@ defmodule ReplicationTest do
       assert {:error, %Postgrex.Error{}} = PR.call(context.repl, {:query, "SELCT"})
     end
 
+    @tag :capture_log
+    test "on timeout", context do
+      Process.flag(:trap_exit, true)
+      repl_ref = Process.monitor(context.repl)
+
+      {_pid, call_ref} =
+        spawn_monitor(fn ->
+          assert {:error, %Postgrex.Error{}} =
+                   PR.call(context.repl, {:query, "SELECT pg_sleep(0.1);", %{timeout: 10}})
+        end)
+
+      assert_receive {:DOWN, ^call_ref, _, _,
+                      {%DBConnection.ConnectionError{message: "tcp async_recv: timeout"}, _}}
+
+      assert_receive {:DOWN, ^repl_ref, _, _,
+                      %DBConnection.ConnectionError{message: "tcp async_recv: timeout"}}
+    end
+
     @tag :capture_log
     test "on disconnect", context do
       Process.flag(:trap_exit, true)