Add Redix.PubSub.ping/2

Closes #274 (because it supersedes it).

Author: whatyouhide

lib/redix/pubsub.ex

@@ -180,6 +180,7 @@ defmodule Redix.PubSub do
   @type subscriber() :: pid() | port() | atom() | {atom(), node()}
   @type connection() :: GenServer.server()
 
+  alias Redix.ConnectionError
   alias Redix.StartOptions
 
   @doc """
@@ -329,6 +330,33 @@ defmodule Redix.PubSub do
     :gen_statem.call(conn, {:subscribe, List.wrap(channels), subscriber})
   end
 
+  @doc """
+  Sends a `PING` to the Redis server and waits for a `PONG` reply.
+
+  Upon a successful reply from the server, returns `:ok`. If no reply is received
+  within the given timeout, returns `:error`.
+
+  This is useful for periodic health checks of the underlying connection. `PING` is
+  one of the few commands that Redis allows on a pub/sub connection.
+
+  ## Examples
+
+      iex> Redix.PubSub.ping(pubsub)
+      :ok
+
+      iex> Redix.PubSub.ping(pubsub, 1000)
+      :ok
+
+  """
+  @doc since: "1.6.0"
+  @spec ping(connection(), timeout()) :: :ok | {:error, ConnectionError.t()}
+  def ping(conn, timeout \\ 5000)
+      when (is_integer(timeout) and timeout >= 0) or timeout == :infinity do
+    :gen_statem.call(conn, :ping, timeout)
+  catch
+    :exit, {:timeout, _} -> {:error, %ConnectionError{reason: :timeout}}
+  end
+
   @doc """
   Subscribes `subscriber` to the given pattern or list of patterns.
 

lib/redix/pubsub/connection.ex

@@ -15,7 +15,8 @@ defmodule Redix.PubSub.Connection do
     :connected_address,
     :client_id,
     subscriptions: %{},
-    monitors: %{}
+    monitors: %{},
+    ping_callers: :queue.new()
   ]
 
   @backoff_exponent 1.5
@@ -187,6 +188,10 @@ defmodule Redix.PubSub.Connection do
     {:keep_state, data}
   end
 
+  def disconnected({:call, from}, :ping, _data) do
+    {:keep_state_and_data, {:reply, from, {:error, %ConnectionError{reason: :closed}}}}
+  end
+
   def disconnected({:call, from}, :get_client_id, _data) do
     reply = {:error, %ConnectionError{reason: :closed}}
     {:keep_state_and_data, {:reply, from, reply}}
@@ -234,6 +239,18 @@ defmodule Redix.PubSub.Connection do
     end
   end
 
+  def connected({:call, from}, :ping, data) do
+    data = %{data | ping_callers: :queue.in(from, data.ping_callers)}
+
+    case data.transport.send(data.socket, Protocol.pack(["PING"])) do
+      :ok ->
+        {:keep_state, data}
+
+      {:error, reason} ->
+        disconnect(data, reason, _handle_disconnection? = true)
+    end
+  end
+
   def connected({:call, from}, :get_client_id, data) do
     reply =
       if id = data.client_id do
@@ -345,6 +362,13 @@ defmodule Redix.PubSub.Connection do
     end
   end
 
+  # When a Redis connection is in pub/sub mode (subscribed to at least one
+  # channel), Redis wraps PING responses as pub/sub messages: ["pong", ""]. When
+  # not subscribed to anything, Redis returns a plain +PONG\r\n simple string,
+  # which parses to "PONG".
+  defp handle_pubsub_msg(data, ["pong", _message]), do: handle_pong(data)
+  defp handle_pubsub_msg(data, "PONG"), do: handle_pong(data)
+
   defp handle_pubsub_msg(data, ["message", channel, payload]) do
     properties = %{channel: channel, payload: payload}
     handle_pubsub_message_with_payload(data, {:channel, channel}, :message, properties)
@@ -369,6 +393,31 @@ defmodule Redix.PubSub.Connection do
     {:ok, data}
   end
 
+  # Ping handling.
+
+  defp handle_pong(data) do
+    {popped, data} = get_and_update_in(data.ping_callers, &:queue.out/1)
+
+    # :empty is not supposed to happen but we're not gonna crash the process
+    # if it does.
+    case popped do
+      {:value, from} -> :gen_statem.reply(from, :ok)
+      :empty -> :ok
+    end
+
+    {:ok, data}
+  end
+
+  defp drain_ping_callers(data) do
+    error = {:error, %ConnectionError{reason: :closed}}
+
+    data.ping_callers
+    |> :queue.to_list()
+    |> Enum.each(&:gen_statem.reply(&1, error))
+
+    %{data | ping_callers: :queue.new()}
+  end
+
   # Subscribing.
 
   defp subscribe_pid_to_targets(data, operation, targets, pid) do
@@ -622,6 +671,7 @@ defmodule Redix.PubSub.Connection do
   end
 
   def disconnect(data, reason, handle_disconnection?) do
+    data = drain_ping_callers(data)
     {next_backoff, data} = next_backoff(data)
 
     if data.socket do

test/redix/pubsub_test.exs

@@ -365,6 +365,50 @@ defmodule Redix.PubSubTest do
     refute new_state.continuation
   end
 
+  describe "ping/2" do
+    test "returns :ok when connected", %{pubsub: pubsub} do
+      assert PubSub.ping(pubsub) == :ok
+      assert PubSub.ping(pubsub, _with_custom_timeout = 10_000) == :ok
+      assert PubSub.ping(pubsub, _with_custom_timeout = :infinity) == :ok
+    end
+
+    test "returns :ok when subscribed to channels", %{pubsub: pubsub} do
+      assert {:ok, ref} = PubSub.subscribe(pubsub, "foo", self())
+      assert_receive {:redix_pubsub, ^pubsub, ^ref, :subscribed, %{channel: "foo"}}
+      assert PubSub.ping(pubsub) == :ok
+    end
+
+    test "returns :error when disconnected" do
+      {:ok, pubsub} = PubSub.start_link(port: 9999)
+      assert PubSub.ping(pubsub) == {:error, %ConnectionError{reason: :closed}}
+    end
+
+    @tag :capture_log
+    test "returns :error when connection goes down", %{pubsub: pubsub, conn: conn} do
+      assert {:ok, ref} = PubSub.subscribe(pubsub, "foo", self())
+      assert_receive {:redix_pubsub, ^pubsub, ^ref, :subscribed, %{channel: "foo"}}
+
+      Redix.command!(conn, ~w(CLIENT KILL TYPE pubsub))
+      assert_receive {:redix_pubsub, ^pubsub, ^ref, :disconnected, _properties}
+
+      assert PubSub.ping(pubsub) == {:error, %ConnectionError{reason: :closed}}
+    end
+
+    test "returns :error when timeout expires", %{pubsub: pubsub} do
+      assert PubSub.ping(pubsub, 0) == {:error, %ConnectionError{reason: :timeout}}
+    end
+
+    test "multiple concurrent pings return :ok", %{pubsub: pubsub} do
+      tasks =
+        for _ <- 1..5 do
+          Task.async(fn -> PubSub.ping(pubsub) end)
+        end
+
+      results = Task.await_many(tasks)
+      assert Enum.all?(results, &(&1 == :ok))
+    end
+  end
+
   defp wait_until_passes(timeout, fun) when timeout <= 0 do
     fun.()
   end