Ensure pool connections terminate when the pool manager shuts down

Closes elixir-ecto/postgrex#763.

Author: josevalim

lib/db_connection/connection_pool.ex

@@ -18,6 +18,7 @@ defmodule DBConnection.ConnectionPool do
   @queue_interval 2000
   @idle_interval 1000
   @time_unit 1000
+  @watcher_ref {__MODULE__, :watcher_ref}
 
   @doc false
   def start_link({mod, opts}) do
@@ -46,6 +47,7 @@ defmodule DBConnection.ConnectionPool do
 
   @impl GenServer
   def init({mod, opts}) do
+    Process.flag(:trap_exit, true)
     DBConnection.register_as_pool(mod)
 
     queue = :ets.new(__MODULE__.Queue, [:protected, :ordered_set, decentralized_counters: true])
@@ -64,7 +66,8 @@ defmodule DBConnection.ConnectionPool do
       end
 
     ts = {nil, max_lifetime}
-    {:ok, _} = DBConnection.ConnectionPool.Pool.start_supervised(queue, mod, opts)
+    {:ok, watcher_ref} = DBConnection.ConnectionPool.Pool.start_supervised(queue, mod, opts)
+    Process.put(@watcher_ref, watcher_ref)
     target = Keyword.get(opts, :queue_target, @queue_target)
     interval = Keyword.get(opts, :queue_interval, @queue_interval)
     idle_interval = Keyword.get(opts, :idle_interval, @idle_interval)
@@ -229,6 +232,19 @@ defmodule DBConnection.ConnectionPool do
     drop_idle(past_in_native, limit, status, queue, codel, ts)
   end
 
+  def handle_info({:EXIT, _pid, reason}, state) do
+    {:stop, reason, state}
+  end
+
+  @impl GenServer
+  def terminate(_reason, _state) do
+    if watcher_ref = Process.get(@watcher_ref) do
+      DBConnection.ConnectionPool.Pool.stop_supervised(watcher_ref)
+    end
+
+    :ok
+  end
+
   defp drop_idle(past_in_native, limit, status, queue, codel, ts) do
     with true <- status == :ready and limit > 0,
          {queued_in_native, holder} = key when queued_in_native <= past_in_native <-

lib/db_connection/connection_pool/pool.ex

@@ -9,6 +9,13 @@ defmodule DBConnection.ConnectionPool.Pool do
     )
   end
 
+  def stop_supervised(ref) do
+    case DBConnection.Watcher.unwatch(ref) do
+      pid when is_pid(pid) -> wait_until_down(pid)
+      nil -> :ok
+    end
+  end
+
   def start_link(arg) do
     Supervisor.start_link(__MODULE__, arg)
   end
@@ -30,4 +37,12 @@ defmodule DBConnection.ConnectionPool.Pool do
     child_opts = [id: {mod, owner, id}] ++ Keyword.take(opts, [:shutdown])
     DBConnection.Connection.child_spec(mod, [pool_index: id] ++ opts, owner, tag, child_opts)
   end
+
+  defp wait_until_down(pid) do
+    ref = Process.monitor(pid)
+
+    receive do
+      {:DOWN, ^ref, _, _, _} -> :ok
+    end
+  end
 end

lib/db_connection/watcher.ex

@@ -12,6 +12,10 @@ defmodule DBConnection.Watcher do
     GenServer.call(@name, {:watch, supervisor, args}, :infinity)
   end
 
+  def unwatch(ref) do
+    GenServer.call(@name, {:unwatch, ref}, :infinity)
+  end
+
   @impl true
   def init(:ok) do
     Process.flag(:trap_exit, true)
@@ -27,33 +31,39 @@ defmodule DBConnection.Watcher do
         started_ref = Process.monitor(started_pid)
         caller_refs = Map.put(caller_refs, caller_ref, {supervisor, started_pid, started_ref})
         started_refs = Map.put(started_refs, started_ref, {caller_pid, caller_ref})
-        {:reply, {:ok, started_pid}, {caller_refs, started_refs}}
+        {:reply, {:ok, caller_ref}, {caller_refs, started_refs}}
 
       other ->
         {:reply, other, {caller_refs, started_refs}}
     end
   end
 
+  @impl true
+  def handle_call({:unwatch, ref}, _from, {caller_refs, started_refs}) do
+    case Map.pop(caller_refs, ref) do
+      {{_supervisor, started_pid, started_ref}, caller_refs} ->
+        terminate_started(started_pid)
+        Process.demonitor(ref, [:flush])
+        Process.demonitor(started_ref, [:flush])
+        {:reply, started_pid, {caller_refs, Map.delete(started_refs, started_ref)}}
+
+      {nil, caller_refs} ->
+        {:reply, nil, {caller_refs, started_refs}}
+    end
+  end
+
   @impl true
   def handle_info({:DOWN, ref, _, _, _}, {caller_refs, started_refs}) do
     case caller_refs do
       %{^ref => {_supervisor, started_pid, started_ref}} ->
-        try do
-          :sys.terminate(started_pid, :shutdown, :infinity)
-        catch
-          :exit, {:noproc, {:sys, :terminate, _}} -> :ok
-          :exit, {:shutdown, {:sys, :terminate, _}} -> :ok
-        end
-
+        terminate_started(started_pid)
+        Process.demonitor(started_ref, [:flush])
         caller_refs = Map.delete(caller_refs, ref)
-        started_refs = Map.put(started_refs, started_ref, {nil, nil})
+        started_refs = Map.delete(started_refs, started_ref)
         {:noreply, {caller_refs, started_refs}}
 
       %{} ->
         case started_refs do
-          %{^ref => {nil, nil}} ->
-            {:noreply, {caller_refs, Map.delete(started_refs, ref)}}
-
           %{^ref => {caller_pid, caller_ref}} ->
             Process.demonitor(caller_ref, [:flush])
             Process.exit(caller_pid, :kill)
@@ -66,9 +76,18 @@ defmodule DBConnection.Watcher do
     {:noreply, state}
   end
 
+  defp terminate_started(started_pid) do
+    try do
+      :sys.terminate(started_pid, :shutdown, :infinity)
+    catch
+      :exit, {:noproc, {:sys, :terminate, _}} -> :ok
+      :exit, {:shutdown, {:sys, :terminate, _}} -> :ok
+    end
+  end
+
   @impl true
   def terminate(_, {_, started_refs}) do
-    for {_, {caller_pid, _}} when caller_pid != nil <- started_refs do
+    for {_, {caller_pid, _}} <- started_refs do
       Process.exit(caller_pid, :kill)
     end
 

test/db_connection/watcher_test.exs

@@ -47,7 +47,7 @@ defmodule DBConnection.WatcherTest do
     # Trigger pool termination
     stop_supervised!(:pool1)
 
-    {:ok, agent2} = A.start_link([{:ok, :state}])
+    {:ok, agent2} = A.start_link([{:ok, :state}, {:ok, :state}])
 
     task =
       Task.async(fn ->