Add DBConnection.disconnect_all/3 (#240)

Author: josevalim

integration_test/connection_pool/disconnect_all_test.exs

@@ -0,0 +1,83 @@
+defmodule TestPoolDisconnectAll do
+  use ExUnit.Case, async: true
+
+  alias TestPool, as: P
+  alias TestAgent, as: A
+  alias TestQuery, as: Q
+  alias TestResult, as: R
+
+  test "disconnect on checkin" do
+    stack = [
+      {:ok, :state},
+      {:ok, %Q{}, %R{}, :new_state1},
+      {:ok, %Q{}, %R{}, :new_state2},
+      {:ok, :dead_state},
+      {:ok, :final_state},
+      {:ok, %Q{}, %R{}, :final_state1},
+      {:ok, %Q{}, %R{}, :final_state2}
+    ]
+
+    {:ok, agent} = A.start_link(stack)
+
+    opts = [agent: agent, parent: self()]
+    {:ok, pool} = P.start_link(opts)
+    assert P.execute(pool, %Q{}, [:param]) == {:ok, %Q{}, %R{}}
+
+    P.disconnect_all(pool, 0)
+    assert P.execute(pool, %Q{}, [:param]) == {:ok, %Q{}, %R{}}
+    assert P.execute(pool, %Q{}, [:param]) == {:ok, %Q{}, %R{}}
+    assert P.execute(pool, %Q{}, [:param]) == {:ok, %Q{}, %R{}}
+
+    err = %DBConnection.ConnectionError{message: "disconnect_all requested", severity: :info}
+
+    assert [
+             connect: [_],
+             handle_execute: [_, _, _, :state],
+             handle_execute: [_, _, _, :new_state1],
+             disconnect: [^err, :new_state2],
+             connect: [_],
+             handle_execute: [_, _, _, :final_state],
+             handle_execute: [_, _, _, :final_state1]
+           ] = A.record(agent)
+  end
+
+  @tag :idle_interval
+  test "disconnect on ping" do
+    parent = self()
+
+    stack = [
+      {:ok, :state},
+      {:ok, %Q{}, %R{}, :new_state},
+      fn _, _ ->
+        send(parent, :disconnecting)
+        {:ok, :dead_state}
+      end,
+      {:ok, :final_state},
+      {:ok, %Q{}, %R{}, :final_state1},
+      {:ok, %Q{}, %R{}, :final_state2},
+    ]
+
+    {:ok, agent} = A.start_link(stack)
+
+    opts = [agent: agent, parent: self(), idle_interval: 50]
+    {:ok, pool} = P.start_link(opts)
+    assert P.execute(pool, %Q{}, [:param]) == {:ok, %Q{}, %R{}}
+
+    # High intervals do not affect ping because those are always disconnected.
+    P.disconnect_all(pool, 45_000)
+    assert_receive :disconnecting
+    assert P.execute(pool, %Q{}, [:param]) == {:ok, %Q{}, %R{}}
+    assert P.execute(pool, %Q{}, [:param]) == {:ok, %Q{}, %R{}}
+
+    err = %DBConnection.ConnectionError{message: "disconnect_all requested", severity: :info}
+
+    assert [
+             connect: [_],
+             handle_execute: [_, _, _, :state],
+             disconnect: [^err, :new_state],
+             connect: [_],
+             handle_execute: [_, _, _, :final_state],
+             handle_execute: [_, _, _, :final_state1]
+           ] = A.record(agent)
+  end
+end

integration_test/ownership/disconnect_all_test.exs

@@ -0,0 +1,49 @@
+defmodule TestOwnershipDisconnectAll do
+  use ExUnit.Case, async: true
+
+  alias TestPool, as: P
+  alias TestAgent, as: A
+  alias TestQuery, as: Q
+  alias TestResult, as: R
+  alias DBConnection.Ownership
+
+  test "disconnect on checkin" do
+    stack = [
+      {:ok, :state},
+      {:ok, %Q{}, %R{}, :new_state1},
+      {:ok, %Q{}, %R{}, :new_state2},
+      {:ok, %Q{}, %R{}, :new_state3},
+      {:ok, :dead_state},
+      {:ok, :final_state},
+      {:ok, %Q{}, %R{}, :final_state1},
+      {:ok, %Q{}, %R{}, :final_state2}
+    ]
+
+    {:ok, agent} = A.start_link(stack)
+
+    opts = [agent: agent, parent: self()]
+    {:ok, pool} = P.start_link(opts)
+    assert P.execute(pool, %Q{}, [:param]) == {:ok, %Q{}, %R{}}
+
+    P.disconnect_all(pool, 0)
+    assert P.execute(pool, %Q{}, [:param]) == {:ok, %Q{}, %R{}}
+    assert P.execute(pool, %Q{}, [:param]) == {:ok, %Q{}, %R{}}
+
+    assert Ownership.ownership_checkin(pool, []) == :ok
+    assert P.execute(pool, %Q{}, [:param]) == {:ok, %Q{}, %R{}}
+    assert P.execute(pool, %Q{}, [:param]) == {:ok, %Q{}, %R{}}
+
+    err = %DBConnection.ConnectionError{message: "disconnect_all requested", severity: :info}
+
+    assert [
+             connect: [_],
+             handle_execute: [_, _, _, :state],
+             handle_execute: [_, _, _, :new_state1],
+             handle_execute: [_, _, _, :new_state2],
+             disconnect: [^err, :new_state3],
+             connect: [_],
+             handle_execute: [_, _, _, :final_state],
+             handle_execute: [_, _, _, :final_state1]
+           ] = A.record(agent)
+  end
+end

lib/db_connection.ex

@@ -369,7 +369,11 @@ defmodule DBConnection do
       Defaults to 1000ms.
     * `:queue_target` and `:queue_interval` - See "Queue config" below
     * `:max_restarts` and `:max_seconds` - Configures the `:max_restarts` and
-      `:max_seconds` for the connection pool supervisor (see the `Supervisor` docs)
+      `:max_seconds` for the connection pool supervisor (see the `Supervisor` docs).
+      Typically speaking the connection process doesn't terminate, except due to
+      faults in DBConnection. However, if backoff has been disabled, then they
+      also terminate whenever a connection is disconnected (for instance, due to
+      client or server errors)
     * `:show_sensitive_data_on_connection_error` - By default, `DBConnection`
       hides all information during connection errors to avoid leaking credentials
       or other sensitive information. You can set this option if you wish to
@@ -454,6 +458,29 @@ defmodule DBConnection do
     pool.child_spec({conn_mod, opts})
   end
 
+  @doc """
+  Forces all connections in the pool to disconnect within the given interval.
+
+  Once this function is called, the pool will disconnect all of its connections
+  as they are checked in or as they are pinged. Checked in connections will be
+  randomly disconnected within the given time interval. Pinged connections are
+  immediately disconnected - as they are idle (according to `:idle_interval`).
+
+  If the connection has a backoff configured (which is the case by default),
+  disconnecting means an attempt at a new connection will be done immediately
+  after, without starting a new process for each connection. However, if backoff
+  has been disabled, the connection process will terminate. In such cases,
+  disconnecting all connections may cause the pool supervisor to restart
+  depending on the max_restarts/max_seconds configuration of the pool,
+  so you will want to set those carefully.
+  """
+  @spec disconnect_all(conn, non_neg_integer, opts :: Keyword.t()) :: :ok
+  def disconnect_all(conn, interval, opts \\ []) when interval >= 0 do
+    pool = Keyword.get(opts, :pool, DBConnection.ConnectionPool)
+    interval = System.convert_time_unit(interval, :millisecond, :native)
+    pool.disconnect_all(conn, interval, opts)
+  end
+
   @doc """
   Prepare a query with a database connection for later execution.