Close connections from cancelled first checkouts

lib/ch.ex

@@ -234,8 +234,8 @@ defmodule Ch do
       NimblePool.checkout!(
         pool,
         :request,
-        fn {pid, _ref}, conn_or_template ->
-          with {:ok, conn} <- connect(conn_or_template, pid, deadline),
+        fn from, conn_or_template ->
+          with {:ok, conn} <- connect(conn_or_template, from, deadline),
                {:ok, conn, status, headers, data} <-
                  request(conn, "POST", path, headers, statement, deadline) do
             {{:ok, status, headers, data}, checkin(conn)}
@@ -283,6 +283,11 @@ defmodule Ch do
     {:ok, {:ok, conn}, conn, config}
   end
 
+  @impl NimblePool
+  def handle_update({:connected, %Mint.HTTP1{} = conn}, :template, config) do
+    {:ok, conn, config}
+  end
+
   @impl NimblePool
   def handle_checkin({:ok, conn}, _from, _prev, config) do
     {:ok, conn, config}
@@ -307,11 +312,13 @@ defmodule Ch do
     {:ok, config}
   end
 
-  defp connect({:template, scheme, host, port}, owner, deadline) do
+  defp connect({:template, scheme, host, port}, {owner, _ref} = from, deadline) do
     timeout = Ch.HTTP.to_timeout(deadline)
 
     case Mint.HTTP1.connect(scheme, host, port, mode: :passive, timeout: timeout) do
       {:ok, conn} ->
+        :ok = NimblePool.update(from, {:connected, conn})
+
         case Mint.HTTP1.controlling_process(conn, owner) do
           {:ok, _conn} = ok ->
             ok
@@ -326,7 +333,7 @@ defmodule Ch do
     end
   end
 
-  defp connect({:ok, _conn} = ok, _owner, _deadline), do: ok
+  defp connect({:ok, _conn} = ok, _from, _deadline), do: ok
 
   defp request(conn, method, path, headers, body, deadline) do
     result =

test/ch/faults_test.exs

@@ -19,6 +19,28 @@ defmodule Ch.FaultsTest do
     assert reason in [:econnrefused, :closed]
   end
 
+  test "closes a connection when its first checkout is cancelled", %{
+    listen: listen,
+    port: port
+  } do
+    {:ok, pool} =
+      Ch.start_link(
+        url: "http://localhost:#{port}",
+        pool_size: 1
+      )
+
+    {:ok, query} =
+      Task.start(fn ->
+        Ch.query(pool, "select sleep(1)", %{}, timeout: :infinity)
+      end)
+
+    {:ok, socket} = :gen_tcp.accept(listen)
+    assert_receive {:tcp, ^socket, _request}, to_timeout(second: 1)
+
+    Process.exit(query, :kill)
+    assert_receive {:tcp_closed, ^socket}, to_timeout(second: 1)
+  end
+
   test "removes a timed out connection and reconnects on the next query", ctx do
     %{port: port, listen: listen} = ctx
     {:ok, pool} = Ch.start_link(url: "http://localhost:#{port}")