fix(upstream): release closed connections and improve pool diagnostics

src/infra/network/upstream/bootstrap.rs

@@ -33,7 +33,7 @@ use crate::infra::network::upstream::conn::QuicConnection;
 use crate::infra::network::upstream::conn::{TcpConnection, UdpConnection};
 use crate::infra::network::upstream::pool::reuse::ReusePool;
 use crate::infra::network::upstream::pool::{
-    Connection, ConnectionBuilder, ConnectionPool, DeadlineOutcome, QueryDeadline,
+    Connection, ConnectionBuilder, ConnectionPool, DeadlineOutcome, PoolSize, QueryDeadline,
     QueryTimeoutPolicy,
 };
 use crate::infra::network::upstream::traits::Upstream;
@@ -102,7 +102,7 @@ pub(crate) struct BootstrapUpstream<C: Connection> {
     /// Upstream server domain name (for logging)
     server_name: String,
     /// Connection metadata (includes bootstrap config)
-    connection_info: ConnectionInfo,
+    connection_info: Arc<ConnectionInfo>,
     /// Bootstrap resolver for domain name resolution
     bootstrap: Arc<NameResolver>,
     /// Lock-free connection pool with current resolved IP and TTL deadline.
@@ -149,9 +149,11 @@ impl<C: Connection> BootstrapUpstream<C> {
         connection_info: ConnectionInfo,
         pool_factory: Box<dyn BootstrapPoolFactory<C>>,
     ) -> Self {
+        let connection_info = Arc::new(connection_info);
         let pool: Arc<dyn ConnectionPool<C>> = ReusePool::<C>::new(
-            0,
-            1,
+            connection_info.clone(),
+            connection_info.connection_type,
+            PoolSize::new(0, 1),
             ConnectionInfo::DEFAULT_CONN_IDLE_TIME,
             Box::new(DummyConnectionBuilder {}),
             QueryTimeoutPolicy::Close,
@@ -349,7 +351,12 @@ struct DummyConnectionBuilder {}
 
 #[async_trait]
 impl<C: Connection> ConnectionBuilder<C> for DummyConnectionBuilder {
-    async fn create_connection(&self, _conn_id: u16, _deadline: QueryDeadline) -> Result<Arc<C>> {
+    async fn create_connection(
+        &self,
+        _conn_id: u16,
+        _deadline: QueryDeadline,
+        _connection_info: Arc<ConnectionInfo>,
+    ) -> Result<Arc<C>> {
         Err(DnsError::protocol(
             "DummyConnectionBuilder cannot create connections (pool not yet initialized)",
         ))

src/infra/network/upstream/bootstrap_factory.rs

@@ -20,7 +20,9 @@ use crate::infra::network::upstream::conn::{
 };
 use crate::infra::network::upstream::pool::pipeline::PipelinePool;
 use crate::infra::network::upstream::pool::reuse::ReusePool;
-use crate::infra::network::upstream::pool::{Connection, ConnectionPool, QueryTimeoutPolicy};
+use crate::infra::network::upstream::pool::{
+    Connection, ConnectionPool, PoolSize, QueryTimeoutPolicy,
+};
 
 pub(super) trait BootstrapPoolFactory<C: Connection>: Debug + Send + Sync {
     fn create_pool(
@@ -39,11 +41,11 @@ impl BootstrapPoolFactory<UdpConnection> for UdpBootstrapPoolFactory {
         connection_info: &ConnectionInfo,
         ip: IpAddr,
     ) -> Arc<dyn ConnectionPool<UdpConnection>> {
-        let info = connection_info_with_ip(connection_info, ip);
+        let info = Arc::new(connection_info_with_ip(connection_info, ip));
         let builder = UdpConnectionBuilder::new(&info, pipeline_request_map_capacity());
         PipelinePool::new(
-            main_pool_min_conns(&info),
-            info.max_conns_or_default(),
+            info.clone(),
+            PoolSize::new(main_pool_min_conns(&info), info.max_conns_or_default()),
             ConnectionInfo::DEFAULT_MAX_CONNS_LOAD,
             info.idle_timeout,
             Box::new(builder),
@@ -62,12 +64,12 @@ impl BootstrapPoolFactory<TcpConnection> for TcpBootstrapPoolFactory {
         connection_info: &ConnectionInfo,
         ip: IpAddr,
     ) -> Arc<dyn ConnectionPool<TcpConnection>> {
-        let info = connection_info_with_ip(connection_info, ip);
+        let info = Arc::new(connection_info_with_ip(connection_info, ip));
         if info.enable_pipeline.unwrap_or(false) {
             let builder = TcpConnectionBuilder::new(&info, pipeline_request_map_capacity());
             PipelinePool::new(
-                main_pool_min_conns(&info),
-                info.max_conns_or_default(),
+                info.clone(),
+                PoolSize::new(main_pool_min_conns(&info), info.max_conns_or_default()),
                 ConnectionInfo::DEFAULT_MAX_CONNS_LOAD,
                 info.idle_timeout,
                 Box::new(builder),
@@ -77,8 +79,9 @@ impl BootstrapPoolFactory<TcpConnection> for TcpBootstrapPoolFactory {
         } else {
             let builder = TcpConnectionBuilder::new(&info, reuse_request_map_capacity());
             ReusePool::new(
-                main_pool_min_conns(&info),
-                info.max_conns_or_default(),
+                info.clone(),
+                info.connection_type,
+                PoolSize::new(main_pool_min_conns(&info), info.max_conns_or_default()),
                 info.idle_timeout,
                 Box::new(builder),
                 QueryTimeoutPolicy::Close,
@@ -99,11 +102,11 @@ impl BootstrapPoolFactory<QuicConnection> for QuicBootstrapPoolFactory {
         connection_info: &ConnectionInfo,
         ip: IpAddr,
     ) -> Arc<dyn ConnectionPool<QuicConnection>> {
-        let info = connection_info_with_ip(connection_info, ip);
+        let info = Arc::new(connection_info_with_ip(connection_info, ip));
         let builder = QuicConnectionBuilder::new(&info);
         PipelinePool::new(
-            main_pool_min_conns(&info),
-            info.max_conns_or_default(),
+            info.clone(),
+            PoolSize::new(main_pool_min_conns(&info), info.max_conns_or_default()),
             ConnectionInfo::DEFAULT_MAX_CONNS_LOAD,
             info.idle_timeout,
             Box::new(builder),
@@ -124,11 +127,11 @@ impl BootstrapPoolFactory<H2Connection> for H2BootstrapPoolFactory {
         connection_info: &ConnectionInfo,
         ip: IpAddr,
     ) -> Arc<dyn ConnectionPool<H2Connection>> {
-        let info = connection_info_with_ip(connection_info, ip);
+        let info = Arc::new(connection_info_with_ip(connection_info, ip));
         let builder = H2ConnectionBuilder::new(&info);
         PipelinePool::new(
-            main_pool_min_conns(&info),
-            info.max_conns_or_default(),
+            info.clone(),
+            PoolSize::new(main_pool_min_conns(&info), info.max_conns_or_default()),
             ConnectionInfo::DEFAULT_MAX_CONNS_LOAD,
             info.idle_timeout,
             Box::new(builder),
@@ -149,11 +152,11 @@ impl BootstrapPoolFactory<H3Connection> for H3BootstrapPoolFactory {
         connection_info: &ConnectionInfo,
         ip: IpAddr,
     ) -> Arc<dyn ConnectionPool<H3Connection>> {
-        let info = connection_info_with_ip(connection_info, ip);
+        let info = Arc::new(connection_info_with_ip(connection_info, ip));
         let builder = H3ConnectionBuilder::new(&info);
         PipelinePool::new(
-            main_pool_min_conns(&info),
-            info.max_conns_or_default(),
+            info.clone(),
+            PoolSize::new(main_pool_min_conns(&info), info.max_conns_or_default()),
             ConnectionInfo::DEFAULT_MAX_CONNS_LOAD,
             info.idle_timeout,
             Box::new(builder),

src/infra/network/upstream/builder.rs

@@ -1,6 +1,8 @@
 // SPDX-FileCopyrightText: 2025 Sven Shi
 // SPDX-License-Identifier: GPL-3.0-or-later
 
+use std::sync::Arc;
+
 use tracing::debug;
 
 use crate::infra::error::Result;
@@ -19,7 +21,9 @@ use crate::infra::network::upstream::conn::{
 };
 use crate::infra::network::upstream::pool::pipeline::PipelinePool;
 use crate::infra::network::upstream::pool::reuse::ReusePool;
-use crate::infra::network::upstream::pool::{Connection, ConnectionBuilder, QueryTimeoutPolicy};
+use crate::infra::network::upstream::pool::{
+    Connection, ConnectionBuilder, PoolSize, QueryTimeoutPolicy,
+};
 use crate::infra::network::upstream::pooled::{PooledUpstream, UdpTruncatedUpstream};
 use crate::infra::network::upstream::traits::Upstream;
 
@@ -37,13 +41,17 @@ impl UpstreamBuilder {
             let upstream: Box<dyn Upstream> = match connection_info.connection_type {
                 ConnectionType::UDP => {
                     debug!("Creating UDP upstream for {}", connection_info.raw_addr);
+                    let connection_info = Arc::new(connection_info);
                     let builder = UdpConnectionBuilder::new(
                         &connection_info,
                         pipeline_request_map_capacity(),
                     );
                     let main_pool = PipelinePool::new(
-                        main_pool_min_conns(&connection_info),
-                        connection_info.max_conns_or_default(),
+                        connection_info.clone(),
+                        PoolSize::new(
+                            main_pool_min_conns(&connection_info),
+                            connection_info.max_conns_or_default(),
+                        ),
                         ConnectionInfo::DEFAULT_MAX_CONNS_LOAD,
                         connection_info.idle_timeout,
                         Box::new(builder),
@@ -54,8 +62,12 @@ impl UpstreamBuilder {
                     let tcp_builder =
                         TcpConnectionBuilder::new(&connection_info, reuse_request_map_capacity());
                     let fallback_pool = ReusePool::new(
-                        udp_truncated_fallback_min_conns(),
-                        connection_info.max_conns_or_default(),
+                        connection_info.clone(),
+                        ConnectionType::TCP,
+                        PoolSize::new(
+                            udp_truncated_fallback_min_conns(),
+                            connection_info.max_conns_or_default(),
+                        ),
                         connection_info.idle_timeout,
                         Box::new(tcp_builder),
                         QueryTimeoutPolicy::Close,
@@ -262,12 +274,13 @@ pub(crate) fn create_pipeline_pool<C: Connection>(
     connection_info: ConnectionInfo,
     builder: Box<dyn ConnectionBuilder<C>>,
 ) -> PooledUpstream<C> {
+    let connection_info = Arc::new(connection_info);
     let timeout = connection_info.timeout;
     let min_size = main_pool_min_conns(&connection_info);
     PooledUpstream::<C> {
         pool: PipelinePool::new(
-            min_size,
-            connection_info.max_conns_or_default(),
+            connection_info.clone(),
+            PoolSize::new(min_size, connection_info.max_conns_or_default()),
             ConnectionInfo::DEFAULT_MAX_CONNS_LOAD,
             connection_info.idle_timeout,
             builder,
@@ -282,12 +295,14 @@ pub(crate) fn create_reuse_pool<C: Connection>(
     connection_info: ConnectionInfo,
     builder: Box<dyn ConnectionBuilder<C>>,
 ) -> PooledUpstream<C> {
+    let connection_info = Arc::new(connection_info);
     let timeout = connection_info.timeout;
     let min_size = main_pool_min_conns(&connection_info);
     PooledUpstream::<C> {
         pool: ReusePool::new(
-            min_size,
-            connection_info.max_conns_or_default(),
+            connection_info.clone(),
+            connection_info.connection_type,
+            PoolSize::new(min_size, connection_info.max_conns_or_default()),
             connection_info.idle_timeout,
             builder,
             QueryTimeoutPolicy::Close,

src/infra/network/upstream/conn/h2.rs

@@ -32,6 +32,7 @@ enum H2RecvError {
 
 #[derive(Debug)]
 pub struct H2Connection {
+    connection_info: Arc<ConnectionInfo>,
     id: u16,
     sender: SendRequest<Bytes>,
     using_count: AtomicU16,
@@ -47,7 +48,13 @@ impl Connection for H2Connection {
         if self.closed.swap(true, Ordering::AcqRel) {
             return;
         }
-        debug!(conn_id = self.id, "Closing DoH connection");
+        debug!(
+            conn_id = self.id,
+            upstream_tag = %self.connection_info.tag.as_deref().unwrap_or("<untagged>"),
+            upstream_host = %self.connection_info.server_name,
+            upstream_port = self.connection_info.port,
+            "Closing DoH connection"
+        );
         self.close_notify.notify_waiters();
     }
 
@@ -156,6 +163,7 @@ impl ConnectionBuilder<H2Connection> for H2ConnectionBuilder {
         &self,
         conn_id: u16,
         deadline: QueryDeadline,
+        connection_info: Arc<ConnectionInfo>,
     ) -> Result<Arc<H2Connection>> {
         let stream = match deadline
             .run(connect_tcp(
@@ -194,6 +202,7 @@ impl ConnectionBuilder<H2Connection> for H2ConnectionBuilder {
         };
 
         let h2_conn = Arc::new(H2Connection {
+            connection_info,
             id: conn_id,
             sender,
             closed: AtomicBool::new(false),

src/infra/network/upstream/conn/h3.rs

@@ -34,6 +34,7 @@ enum H3RecvError {
 }
 
 pub struct H3Connection {
+    connection_info: Arc<ConnectionInfo>,
     id: u16,
     sender: SendRequest<OpenStreams, Bytes>,
     using_count: AtomicU16,
@@ -54,7 +55,13 @@ impl Connection for H3Connection {
         if self.closed.swap(true, Ordering::AcqRel) {
             return;
         }
-        debug!(conn_id = self.id, "Closing H3 connection");
+        debug!(
+            conn_id = self.id,
+            upstream_tag = %self.connection_info.tag.as_deref().unwrap_or("<untagged>"),
+            upstream_host = %self.connection_info.server_name,
+            upstream_port = self.connection_info.port,
+            "Closing H3 connection"
+        );
         self.close_notify.notify_waiters();
     }
 
@@ -170,6 +177,7 @@ impl ConnectionBuilder<H3Connection> for H3ConnectionBuilder {
         &self,
         conn_id: u16,
         deadline: QueryDeadline,
+        connection_info: Arc<ConnectionInfo>,
     ) -> Result<Arc<H3Connection>> {
         let socket = connect_udp(UdpDialOptions::new(
             self.target.clone(),
@@ -201,6 +209,7 @@ impl ConnectionBuilder<H3Connection> for H3ConnectionBuilder {
         };
 
         let h3_conn = Arc::new(H3Connection {
+            connection_info,
             id: conn_id,
             sender: send_request,
             closed: AtomicBool::new(false),

src/infra/network/upstream/conn/mod.rs

@@ -3,7 +3,7 @@
 
 //! Protocol-specific upstream connection implementations.
 
-use std::sync::atomic::{AtomicU16, Ordering};
+use std::sync::atomic::{AtomicBool, AtomicU16, Ordering};
 #[cfg(any(feature = "upstream-doq", feature = "upstream-doh3"))]
 use std::time::Duration;
 
@@ -28,6 +28,17 @@ pub(crate) use quic::{QuicConnection, QuicConnectionBuilder};
 pub(crate) use tcp::{TcpConnection, TcpConnectionBuilder};
 pub(crate) use udp::{UdpConnection, UdpConnectionBuilder};
 
+/// Wait for a connection's persistent closed state without losing a wakeup.
+pub(super) async fn wait_for_close(closed: &AtomicBool, notify: &tokio::sync::Notify) {
+    // Create the future before checking state: notify_waiters wakes futures
+    // created before the notification, even if they have not been polled yet.
+    // The flag covers closure before this future was created.
+    let notified = notify.notified();
+    if !closed.load(Ordering::Acquire) {
+        notified.await;
+    }
+}
+
 /// RAII guard that decrements a connection's in-flight query counter on drop.
 ///
 /// Ensures `using_count` is always decremented even when the query future is

src/infra/network/upstream/conn/quic.rs

@@ -21,6 +21,7 @@ use crate::infra::network::upstream::{Connection, ConnectionInfo};
 use crate::proto::Message;
 
 pub struct QuicConnection {
+    connection_info: Arc<ConnectionInfo>,
     id: u16,
     transport: QuicTransport,
     using_count: AtomicU16,
@@ -47,6 +48,9 @@ impl Connection for QuicConnection {
         }
         debug!(
             conn_id = self.id,
+            upstream_tag = %self.connection_info.tag.as_deref().unwrap_or("<untagged>"),
+            upstream_host = %self.connection_info.server_name,
+            upstream_port = self.connection_info.port,
             "Closing QUIC connection, sending CONNECTION_CLOSE frame"
         );
         // Gracefully close the underlying QUIC connection with error code 0 (no
@@ -206,6 +210,7 @@ impl ConnectionBuilder<QuicConnection> for QuicConnectionBuilder {
         &self,
         conn_id: u16,
         deadline: QueryDeadline,
+        connection_info: Arc<ConnectionInfo>,
     ) -> Result<Arc<QuicConnection>> {
         let socket = connect_udp(UdpDialOptions::new(
             self.target.clone(),
@@ -235,6 +240,7 @@ impl ConnectionBuilder<QuicConnection> for QuicConnectionBuilder {
         );
 
         let quic_conn = Arc::new(QuicConnection {
+            connection_info,
             id: conn_id,
             transport: QuicTransport::new(quic_conn),
             closed: AtomicBool::new(false),