From 10098b6b68efa83b3c213a08b3a933b552aea9a7 Mon Sep 17 00:00:00 2001
From: ruslandoga <ruslandoga+gh@icloud.com>
Date: Mon, 3 Aug 2026 15:56:30 +0300
Subject: [PATCH 1/2] fix cancelled first checkout connection leak

---
 lib/ch.ex               | 15 +++++++++++----
 test/ch/faults_test.exs | 22 ++++++++++++++++++++++
 2 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/lib/ch.ex b/lib/ch.ex
index 472bfae1..f177aa75 100644
--- a/lib/ch.ex
+++ b/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 =
diff --git a/test/ch/faults_test.exs b/test/ch/faults_test.exs
index 2f43dc62..8c9b9ae2 100644
--- a/test/ch/faults_test.exs
+++ b/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 1", %{}, timeout: :infinity)
+      end)
+
+    {:ok, socket} = :gen_tcp.accept(listen)
+    assert_receive {:tcp, ^socket, _request}, 1_000
+
+    Process.exit(query, :kill)
+    assert_receive {:tcp_closed, ^socket}, 1_000
+  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}")

From 4cd3d194839886ed8bf082c194b8f7128900412b Mon Sep 17 00:00:00 2001
From: ruslandoga <ruslandoga+gh@icloud.com>
Date: Mon, 3 Aug 2026 15:57:49 +0300
Subject: [PATCH 2/2] refine cancelled checkout regression test

---
 test/ch/faults_test.exs | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/test/ch/faults_test.exs b/test/ch/faults_test.exs
index 8c9b9ae2..a7d504b2 100644
--- a/test/ch/faults_test.exs
+++ b/test/ch/faults_test.exs
@@ -31,14 +31,14 @@ defmodule Ch.FaultsTest do
 
     {:ok, query} =
       Task.start(fn ->
-        Ch.query(pool, "select 1", %{}, timeout: :infinity)
+        Ch.query(pool, "select sleep(1)", %{}, timeout: :infinity)
       end)
 
     {:ok, socket} = :gen_tcp.accept(listen)
-    assert_receive {:tcp, ^socket, _request}, 1_000
+    assert_receive {:tcp, ^socket, _request}, to_timeout(second: 1)
 
     Process.exit(query, :kill)
-    assert_receive {:tcp_closed, ^socket}, 1_000
+    assert_receive {:tcp_closed, ^socket}, to_timeout(second: 1)
   end
 
   test "removes a timed out connection and reconnects on the next query", ctx do