diff --git a/src/infra/network/upstream/bootstrap.rs b/src/infra/network/upstream/bootstrap.rs index 1a558d5d..3a04deab 100644 --- a/src/infra/network/upstream/bootstrap.rs +++ b/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 { /// Upstream server domain name (for logging) server_name: String, /// Connection metadata (includes bootstrap config) - connection_info: ConnectionInfo, + connection_info: Arc, /// Bootstrap resolver for domain name resolution bootstrap: Arc, /// Lock-free connection pool with current resolved IP and TTL deadline. @@ -149,9 +149,11 @@ impl BootstrapUpstream { connection_info: ConnectionInfo, pool_factory: Box>, ) -> Self { + let connection_info = Arc::new(connection_info); let pool: Arc> = ReusePool::::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 ConnectionBuilder for DummyConnectionBuilder { - async fn create_connection(&self, _conn_id: u16, _deadline: QueryDeadline) -> Result> { + async fn create_connection( + &self, + _conn_id: u16, + _deadline: QueryDeadline, + _connection_info: Arc, + ) -> Result> { Err(DnsError::protocol( "DummyConnectionBuilder cannot create connections (pool not yet initialized)", )) diff --git a/src/infra/network/upstream/bootstrap_factory.rs b/src/infra/network/upstream/bootstrap_factory.rs index b9919547..6db11f70 100644 --- a/src/infra/network/upstream/bootstrap_factory.rs +++ b/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: Debug + Send + Sync { fn create_pool( @@ -39,11 +41,11 @@ impl BootstrapPoolFactory for UdpBootstrapPoolFactory { connection_info: &ConnectionInfo, ip: IpAddr, ) -> Arc> { - 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 for TcpBootstrapPoolFactory { connection_info: &ConnectionInfo, ip: IpAddr, ) -> Arc> { - 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 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 for QuicBootstrapPoolFactory { connection_info: &ConnectionInfo, ip: IpAddr, ) -> Arc> { - 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 for H2BootstrapPoolFactory { connection_info: &ConnectionInfo, ip: IpAddr, ) -> Arc> { - 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 for H3BootstrapPoolFactory { connection_info: &ConnectionInfo, ip: IpAddr, ) -> Arc> { - 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), diff --git a/src/infra/network/upstream/builder.rs b/src/infra/network/upstream/builder.rs index 3e638a4e..4fac7e3c 100644 --- a/src/infra/network/upstream/builder.rs +++ b/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 = 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( connection_info: ConnectionInfo, builder: Box>, ) -> PooledUpstream { + let connection_info = Arc::new(connection_info); let timeout = connection_info.timeout; let min_size = main_pool_min_conns(&connection_info); PooledUpstream:: { 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( connection_info: ConnectionInfo, builder: Box>, ) -> PooledUpstream { + let connection_info = Arc::new(connection_info); let timeout = connection_info.timeout; let min_size = main_pool_min_conns(&connection_info); PooledUpstream:: { 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, diff --git a/src/infra/network/upstream/conn/h2.rs b/src/infra/network/upstream/conn/h2.rs index ec2416cf..2369a474 100644 --- a/src/infra/network/upstream/conn/h2.rs +++ b/src/infra/network/upstream/conn/h2.rs @@ -32,6 +32,7 @@ enum H2RecvError { #[derive(Debug)] pub struct H2Connection { + connection_info: Arc, id: u16, sender: SendRequest, 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(""), + 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 for H2ConnectionBuilder { &self, conn_id: u16, deadline: QueryDeadline, + connection_info: Arc, ) -> Result> { let stream = match deadline .run(connect_tcp( @@ -194,6 +202,7 @@ impl ConnectionBuilder for H2ConnectionBuilder { }; let h2_conn = Arc::new(H2Connection { + connection_info, id: conn_id, sender, closed: AtomicBool::new(false), diff --git a/src/infra/network/upstream/conn/h3.rs b/src/infra/network/upstream/conn/h3.rs index b607ffdb..5cb2bb2a 100644 --- a/src/infra/network/upstream/conn/h3.rs +++ b/src/infra/network/upstream/conn/h3.rs @@ -34,6 +34,7 @@ enum H3RecvError { } pub struct H3Connection { + connection_info: Arc, id: u16, sender: SendRequest, 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(""), + 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 for H3ConnectionBuilder { &self, conn_id: u16, deadline: QueryDeadline, + connection_info: Arc, ) -> Result> { let socket = connect_udp(UdpDialOptions::new( self.target.clone(), @@ -201,6 +209,7 @@ impl ConnectionBuilder for H3ConnectionBuilder { }; let h3_conn = Arc::new(H3Connection { + connection_info, id: conn_id, sender: send_request, closed: AtomicBool::new(false), diff --git a/src/infra/network/upstream/conn/mod.rs b/src/infra/network/upstream/conn/mod.rs index 1fdc4920..fcb59e7f 100644 --- a/src/infra/network/upstream/conn/mod.rs +++ b/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 diff --git a/src/infra/network/upstream/conn/quic.rs b/src/infra/network/upstream/conn/quic.rs index 7af00fac..583b3fc1 100644 --- a/src/infra/network/upstream/conn/quic.rs +++ b/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, 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(""), + 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 for QuicConnectionBuilder { &self, conn_id: u16, deadline: QueryDeadline, + connection_info: Arc, ) -> Result> { let socket = connect_udp(UdpDialOptions::new( self.target.clone(), @@ -235,6 +240,7 @@ impl ConnectionBuilder for QuicConnectionBuilder { ); let quic_conn = Arc::new(QuicConnection { + connection_info, id: conn_id, transport: QuicTransport::new(quic_conn), closed: AtomicBool::new(false), diff --git a/src/infra/network/upstream/conn/tcp.rs b/src/infra/network/upstream/conn/tcp.rs index cd1eedc9..79421561 100644 --- a/src/infra/network/upstream/conn/tcp.rs +++ b/src/infra/network/upstream/conn/tcp.rs @@ -11,6 +11,7 @@ use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender, unbounded_channel}; use tokio::sync::{Notify, oneshot}; use tracing::{debug, error, trace, warn}; +use super::wait_for_close; use crate::infra::clock::AppClock; use crate::infra::error::{DnsError, Result}; #[cfg(feature = "upstream-dot")] @@ -34,6 +35,7 @@ use crate::proto::Message; /// asynchronous DNS queries and concurrent request tracking. #[derive(Debug)] pub struct TcpConnection { + connection_info: Arc, /// Unique connection ID for logging/tracing. id: u16, /// Sender for the unbounded outgoing TCP message channel. @@ -76,6 +78,9 @@ impl Connection for TcpConnection { let cleared = self.request_map.clear(); debug!( conn_id = self.id, + upstream_tag = %self.connection_info.tag.as_deref().unwrap_or(""), + upstream_host = %self.connection_info.server_name, + upstream_port = self.connection_info.port, canceled_queries = cleared, "Initiating TCP connection close sequence" ); @@ -185,12 +190,18 @@ impl TcpConnection { /// # Arguments /// * `conn_id` - Unique connection identifier for logging and debugging /// * `sender` - Unbounded channel for queuing outbound DNS messages - fn new(conn_id: u16, sender: UnboundedSender, request_map_capacity: u16) -> Self { + fn new( + conn_id: u16, + sender: UnboundedSender, + request_map_capacity: u16, + connection_info: Arc, + ) -> Self { debug!( conn_id, "Initialized TCP connection wrapper with async I/O tasks" ); Self { + connection_info, id: conn_id, sender, close_notify: Notify::new(), @@ -213,19 +224,28 @@ impl TcpConnection { mut writer: TcpTransportWriter, mut receiver: UnboundedReceiver, ) { - let mut closing = false; debug!( conn_id = self.id, "TCP sender task started, ready to transmit queued messages" ); - while !closing { + loop { select! { - Some(queued) = receiver.recv() => { - if let Err(e) = writer - .write_message_with_id(&queued.message, queued.query_id) - .await - { + biased; + _ = wait_for_close(&self.closed, &self.close_notify) => break, + queued = receiver.recv() => { + let Some(queued) = queued else { + self.close(); + break; + }; + // Closing retires the stream, so abandoning a partial + // frame is safe and must also interrupt blocked writes. + let result = select! { + biased; + _ = wait_for_close(&self.closed, &self.close_notify) => break, + result = writer.write_message_with_id(&queued.message, queued.query_id) => result, + }; + if let Err(e) = result { error!( conn_id = self.id, error = ?e, @@ -233,15 +253,9 @@ impl TcpConnection { ); self.writeable.store(false, Ordering::Release); self.close(); + break; } } - _ = self.close_notify.notified() => { - debug!( - conn_id = self.id, - "TCP sender received close notification, shutting down stream" - ); - closing = true; - } } } @@ -263,23 +277,15 @@ impl TcpConnection { self: Arc, mut reader: TcpTransportReader, ) { - let mut closing = false; debug!( conn_id = self.id, "TCP listener task started, waiting for DNS responses" ); loop { - if closing && self.request_map.is_empty() { - debug!(conn_id = self.id, "TCP listener exiting (no more requests)"); - break; - } - if self.closed.load(Ordering::Acquire) { - debug!(conn_id = self.id, "TCP listener detected closed connection"); - break; - } - select! { + biased; + _ = wait_for_close(&self.closed, &self.close_notify) => break, res = reader.read_message() => { match res { Ok(msg) => { @@ -311,15 +317,6 @@ impl TcpConnection { } } } - _ = self.close_notify.notified() => { - closing = true; - debug!( - conn_id = self.id, - pending_queries = self.request_map.size(), - "TCP listener received close notification, draining remaining responses" - ); - continue; - } } } @@ -380,6 +377,7 @@ impl ConnectionBuilder for TcpConnectionBuilder { &self, conn_id: u16, deadline: QueryDeadline, + connection_info: Arc, ) -> Result> { let stream = match deadline .run(connect_tcp( @@ -402,7 +400,8 @@ impl ConnectionBuilder for TcpConnectionBuilder { ); let (sender, receiver) = unbounded_channel(); - let connection = TcpConnection::new(conn_id, sender, self.request_map_capacity); + let connection = + TcpConnection::new(conn_id, sender, self.request_map_capacity, connection_info); let arc = Arc::new(connection); if self.tls_enabled { @@ -463,6 +462,96 @@ mod tests { #[cfg(feature = "upstream-dot")] use crate::infra::network::upstream::ConnectionType; + fn test_connection_info() -> Arc { + Arc::new(ConnectionInfo::with_addr("tcp://127.0.0.1:53").unwrap()) + } + + #[tokio::test] + async fn test_sender_exits_and_releases_connection_after_write_failure() { + AppClock::start(); + let (sender, receiver) = unbounded_channel(); + let connection = Arc::new(TcpConnection::new(1, sender, 1, test_connection_info())); + let weak = Arc::downgrade(&connection); + let (stream, peer) = tokio::io::duplex(64); + drop(peer); + connection + .sender + .send(QueuedQuery { + message: Message::new(), + query_id: 0, + }) + .unwrap(); + + tokio::time::timeout( + Duration::from_secs(1), + connection.send_dns_request(TcpTransportWriter::new(stream), receiver), + ) + .await + .expect("sender must exit after a failed write"); + assert!(weak.upgrade().is_none(), "sender retained the connection"); + } + + #[tokio::test] + async fn test_close_stops_tcp_tasks_before_start_and_while_idle() { + AppClock::start(); + for start_tasks in [false, true] { + let (sender, receiver) = unbounded_channel(); + let connection = Arc::new(TcpConnection::new(1, sender, 1, test_connection_info())); + let weak = Arc::downgrade(&connection); + let (stream, _peer) = tokio::io::duplex(64); + let (reader, writer) = tokio::io::split(stream); + let mut tasks = Box::pin(async { + tokio::join!( + connection + .clone() + .listen_dns_response(TcpTransportReader::new(reader)), + connection + .clone() + .send_dns_request(TcpTransportWriter::new(writer), receiver), + ); + }); + if start_tasks { + assert!(futures::poll!(&mut tasks).is_pending()); + } + connection.close(); + connection.close(); + tokio::time::timeout(Duration::from_secs(1), tasks) + .await + .expect("both TCP tasks must observe close"); + drop(connection); + assert!(weak.upgrade().is_none()); + } + } + + #[tokio::test] + async fn test_close_interrupts_blocked_tcp_write() { + AppClock::start(); + let (sender, receiver) = unbounded_channel(); + let connection = Arc::new(TcpConnection::new(1, sender, 1, test_connection_info())); + let weak = Arc::downgrade(&connection); + // One byte of capacity forces the DNS frame write to remain pending. + let (stream, _peer) = tokio::io::duplex(1); + connection + .sender + .send(QueuedQuery { + message: Message::new(), + query_id: 0, + }) + .unwrap(); + let mut writer = Box::pin( + connection + .clone() + .send_dns_request(TcpTransportWriter::new(stream), receiver), + ); + assert!(futures::poll!(&mut writer).is_pending()); + connection.close(); + drop(connection); + tokio::time::timeout(Duration::from_secs(1), writer) + .await + .expect("close must cancel a blocked write"); + assert!(weak.upgrade().is_none()); + } + #[cfg(feature = "upstream-dot")] #[test] fn test_builder_new_marks_dot_connections_as_tls_enabled() { @@ -485,7 +574,12 @@ mod tests { async fn test_query_returns_error_when_connection_is_closed() { AppClock::start(); let (sender, _receiver) = unbounded_channel(); - let connection = TcpConnection::new(7, sender, DEFAULT_REQUEST_MAP_CAPACITY); + let connection = TcpConnection::new( + 7, + sender, + DEFAULT_REQUEST_MAP_CAPACITY, + test_connection_info(), + ); connection.close(); let result = connection @@ -504,7 +598,12 @@ mod tests { async fn test_query_removes_request_when_cancelled() { AppClock::start(); let (sender, _receiver) = unbounded_channel(); - let connection = TcpConnection::new(8, sender, DEFAULT_REQUEST_MAP_CAPACITY); + let connection = TcpConnection::new( + 8, + sender, + DEFAULT_REQUEST_MAP_CAPACITY, + test_connection_info(), + ); let result = tokio::time::timeout( Duration::from_millis(10), diff --git a/src/infra/network/upstream/conn/udp.rs b/src/infra/network/upstream/conn/udp.rs index 9938e2b2..49f97088 100644 --- a/src/infra/network/upstream/conn/udp.rs +++ b/src/infra/network/upstream/conn/udp.rs @@ -13,6 +13,7 @@ use tokio::sync::{Notify, oneshot}; use tokio::time::timeout; use tracing::{debug, error, trace, warn}; +use super::wait_for_close; use crate::infra::clock::AppClock; use crate::infra::error::{DnsError, Result}; use crate::infra::network::dial::{DialTarget, SocketOptions, UdpDialOptions, connect_udp}; @@ -29,6 +30,7 @@ const UDP_RECV_BUFFER_SIZE: usize = 8_196; /// of request IDs to response channels for asynchronous query handling. #[derive(Debug)] pub struct UdpConnection { + connection_info: Arc, /// Unique connection ID (for debugging/tracing) id: u16, /// The underlying UDP transport bound to a local address @@ -69,6 +71,9 @@ impl Connection for UdpConnection { let cleared = self.request_map.clear(); debug!( conn_id = self.id, + upstream_tag = %self.connection_info.tag.as_deref().unwrap_or(""), + upstream_host = %self.connection_info.server_name, + upstream_port = self.connection_info.port, canceled_queries = cleared, "Closing UDP connection and signaling listener task" ); @@ -194,8 +199,14 @@ impl UdpConnection { /// # Arguments /// * `conn_id` - Unique connection identifier for logging /// * `socket` - Pre-configured UDP socket connected to remote server - fn new(conn_id: u16, socket: UdpSocket, request_map_capacity: u16) -> UdpConnection { + fn new( + conn_id: u16, + socket: UdpSocket, + request_map_capacity: u16, + connection_info: Arc, + ) -> UdpConnection { Self { + connection_info, id: conn_id, transport: UdpTransport::new(socket), close_notify: Notify::new(), @@ -213,11 +224,10 @@ impl UdpConnection { /// the connection closes. /// /// # Buffer Size - /// Uses 4KB buffer which is sufficient for most DNS responses. + /// Uses an 8KB buffer which is sufficient for most DNS responses. /// Larger responses would typically use TCP (with TC bit set). async fn listen_dns_response(self: Arc) { let mut buf = vec![0u8; UDP_RECV_BUFFER_SIZE]; - let mut closing = false; debug!( conn_id = self.id, @@ -225,12 +235,9 @@ impl UdpConnection { ); loop { - if closing && self.request_map.is_empty() { - debug!(conn_id = self.id, "Listener exiting (connection dropped)"); - break; - } - select! { + biased; + _ = wait_for_close(&self.closed, &self.close_notify) => break, recv = self.transport.read_message(&mut buf) => { match recv { Ok(msg) => { @@ -253,19 +260,16 @@ impl UdpConnection { } Err(e) => { if self.closed.load(Ordering::Acquire) { - closing = true; // graceful shutdown path - continue; + break; } warn!(conn_id = self.id, err = %e, "UDP listener error"); continue; } } } - _ = self.close_notify.notified() => { - closing = true; - } } } + debug!(conn_id = self.id, "UDP listener task terminated"); } } @@ -311,6 +315,7 @@ impl ConnectionBuilder for UdpConnectionBuilder { &self, conn_id: u16, _deadline: QueryDeadline, + connection_info: Arc, ) -> Result> { let socket = connect_udp(UdpDialOptions::new( self.target.clone(), @@ -328,6 +333,7 @@ impl ConnectionBuilder for UdpConnectionBuilder { conn_id, UdpSocket::from_std(socket)?, self.request_map_capacity, + connection_info, ); let arc = Arc::new(connection); @@ -343,6 +349,41 @@ mod tests { use super::*; use crate::infra::network::upstream::ConnectionType; + fn test_connection_info() -> Arc { + Arc::new(ConnectionInfo::with_addr("udp://127.0.0.1:53").unwrap()) + } + + #[tokio::test] + async fn test_close_stops_udp_listener_before_start_and_while_idle() { + AppClock::start(); + for start_listener in [false, true] { + let socket = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let connection = Arc::new(UdpConnection::new(1, socket, 1, test_connection_info())); + let weak = Arc::downgrade(&connection); + let (sender, receiver) = oneshot::channel(); + let guard = connection.request_map.store(sender).unwrap(); + let mut listener = Box::pin(connection.clone().listen_dns_response()); + if start_listener { + assert!(futures::poll!(&mut listener).is_pending()); + } + connection.close(); + connection.close(); + assert_eq!(connection.using_count(), 0); + assert!( + tokio::time::timeout(Duration::from_secs(1), receiver) + .await + .expect("close must cancel pending queries") + .is_err() + ); + drop(guard); + drop(connection); + tokio::time::timeout(Duration::from_secs(1), listener) + .await + .expect("UDP listener must observe close"); + assert!(weak.upgrade().is_none(), "listener retained the connection"); + } + } + #[test] fn test_builder_new_copies_connection_info_fields() { let mut connection_info = diff --git a/src/infra/network/upstream/pool/mod.rs b/src/infra/network/upstream/pool/mod.rs index dcb56797..1f8dd970 100644 --- a/src/infra/network/upstream/pool/mod.rs +++ b/src/infra/network/upstream/pool/mod.rs @@ -43,6 +43,7 @@ use tokio::task::yield_now; use crate::infra::error::Result; pub use crate::infra::network::deadline::{DeadlineOutcome, QueryDeadline}; +use crate::infra::network::upstream::config::ConnectionInfo; use crate::infra::task as task_center; use crate::proto::Message; @@ -60,6 +61,19 @@ pub enum QueryTimeoutPolicy { Close, } +/// Minimum and maximum connections configured for a pool. +#[derive(Debug, Clone, Copy)] +pub(crate) struct PoolSize { + pub(crate) min: usize, + pub(crate) max: usize, +} + +impl PoolSize { + pub(crate) const fn new(min: usize, max: usize) -> Self { + Self { min, max } + } +} + /// Connection trait - represents a single persistent connection to an upstream /// DNS server /// @@ -108,7 +122,12 @@ pub trait ConnectionBuilder: Send + Sync + Debug + 'static { /// # Returns /// Arc-wrapped connection on success, or error if connection establishment /// fails - async fn create_connection(&self, conn_id: u16, deadline: QueryDeadline) -> Result>; + async fn create_connection( + &self, + conn_id: u16, + deadline: QueryDeadline, + connection_info: Arc, + ) -> Result>; } /// Connection pool trait - manages a pool of connections for load balancing diff --git a/src/infra/network/upstream/pool/pipeline.rs b/src/infra/network/upstream/pool/pipeline.rs index 47b3d713..19caa20f 100644 --- a/src/infra/network/upstream/pool/pipeline.rs +++ b/src/infra/network/upstream/pool/pipeline.rs @@ -13,9 +13,10 @@ use tracing::{debug, warn}; use crate::infra::clock::AppClock; use crate::infra::error::Result; +use crate::infra::network::upstream::config::ConnectionInfo; use crate::infra::network::upstream::pool::{ Connection, ConnectionBuilder, ConnectionPool, DeadlineOutcome, ManagedMaintenanceTask, - QueryDeadline, QueryTimeoutPolicy, start_maintenance, + PoolSize, QueryDeadline, QueryTimeoutPolicy, start_maintenance, }; use crate::infra::task as task_center; use crate::proto::Message; @@ -28,6 +29,7 @@ const SLOT_CLOSED: u8 = 2; #[derive(Debug)] pub struct PipelinePool { + connection_info: Arc, /// Round-robin index for load balancing across slots. index: AtomicUsize, /// List of connection slots (lock-free with ArcSwap). @@ -147,9 +149,13 @@ impl ConnectionPool for PipelinePool { if !close_after_swap.is_empty() { debug!( - "Pipeline pool maintenance: removed {} slots, {} active", - close_after_swap.len(), - new_len + upstream_tag = %self.connection_info.tag.as_deref().unwrap_or(""), + upstream_host = %self.connection_info.server_name, + target_ip = ?self.connection_info.remote_ip, + upstream_port = self.connection_info.port, + removed_slots = close_after_swap.len(), + active_slots = new_len, + "Pipeline pool maintenance removed connections" ); self.release_notified.notify_waiters(); } @@ -167,15 +173,29 @@ impl ConnectionPool for PipelinePool { impl PipelinePool { pub fn new( - min_size: usize, - max_size: usize, + connection_info: Arc, + size: PoolSize, max_load: u16, idle_time: Duration, connection_builder: Box>, timeout_policy: QueryTimeoutPolicy, connect_timeout: Duration, ) -> Arc> { + let PoolSize { + min: min_size, + max: max_size, + } = size; + debug!( + upstream_tag = %connection_info.tag.as_deref().unwrap_or(""), + upstream_host = %connection_info.server_name, + target_ip = ?connection_info.remote_ip, + upstream_port = connection_info.port, + min_size, + max_size, + "Creating PipelinePool" + ); let pool = Arc::new(Self { + connection_info, index: AtomicUsize::new(0), slots: ArcSwap::from_pointee(Vec::new()), reserved_slots: AtomicUsize::new(0), @@ -318,7 +338,11 @@ impl PipelinePool { ) -> Result>>> { let id = self.next_id.fetch_add(1, Ordering::Relaxed); match deadline - .run(self.connection_builder.create_connection(id, deadline)) + .run(self.connection_builder.create_connection( + id, + deadline, + self.connection_info.clone(), + )) .await { DeadlineOutcome::Completed(Ok(conn)) => { @@ -724,6 +748,7 @@ mod tests { &self, _conn_id: u16, _deadline: QueryDeadline, + _connection_info: Arc, ) -> Result> { self.planned .lock() @@ -747,6 +772,7 @@ mod tests { .map(|conn| Arc::new(PipelineSlot::new(conn))) .collect(); PipelinePool { + connection_info: Arc::new(ConnectionInfo::with_addr("udp://127.0.0.1:53").unwrap()), index: AtomicUsize::new(0), slots: ArcSwap::from_pointee(slots), reserved_slots: AtomicUsize::new(0), diff --git a/src/infra/network/upstream/pool/reuse.rs b/src/infra/network/upstream/pool/reuse.rs index 6d8c447e..0a31a93f 100644 --- a/src/infra/network/upstream/pool/reuse.rs +++ b/src/infra/network/upstream/pool/reuse.rs @@ -13,9 +13,10 @@ use tracing::{debug, info, warn}; use crate::infra::clock::AppClock; use crate::infra::error::Result; +use crate::infra::network::upstream::config::{ConnectionInfo, ConnectionType}; use crate::infra::network::upstream::pool::{ Connection, ConnectionBuilder, ConnectionPool, DeadlineOutcome, ManagedMaintenanceTask, - QueryDeadline, QueryTimeoutPolicy, start_maintenance, + PoolSize, QueryDeadline, QueryTimeoutPolicy, start_maintenance, }; use crate::infra::task as task_center; use crate::proto::Message; @@ -36,6 +37,8 @@ fn close_conns(conns: &[Arc]) { /// - Thread-safe, designed for async DNS request handling #[derive(Debug)] pub struct ReusePool { + connection_info: Arc, + transport: ConnectionType, /// Queue holding idle connections connections: ArrayQueue>, /// Number of active connections in use or queued @@ -183,20 +186,33 @@ impl ConnectionPool for ReusePool { impl ReusePool { /// Create a new reusable connection pool - pub fn new( - min_size: usize, - max_size: usize, + pub(crate) fn new( + connection_info: Arc, + transport: ConnectionType, + size: PoolSize, idle_time: Duration, connection_builder: Box>, timeout_policy: QueryTimeoutPolicy, connect_timeout: Duration, ) -> Arc> { + let PoolSize { + min: min_size, + max: max_size, + } = size; info!( - "Creating ReusePool (min_size={}, max_size={})", - min_size, max_size + upstream_tag = %connection_info.tag.as_deref().unwrap_or(""), + upstream_host = %connection_info.server_name, + target_ip = ?connection_info.remote_ip, + upstream_port = connection_info.port, + transport = ?transport, + min_size, + max_size, + "Creating ReusePool" ); let pool = Arc::new(Self { + connection_info, + transport, connections: ArrayQueue::new(max_size), min_size, max_size, @@ -388,7 +404,11 @@ impl ReusePool { ) -> Result> { let id = self.next_id.fetch_add(1, Ordering::Relaxed); match deadline - .run(self.connection_builder.create_connection(id, deadline)) + .run(self.connection_builder.create_connection( + id, + deadline, + self.connection_info.clone(), + )) .await { DeadlineOutcome::Completed(Ok(conn)) => { @@ -471,7 +491,22 @@ impl<'a, C: Connection> BorrowedConnection<'a, C> { impl Drop for BorrowedConnection<'_, C> { fn drop(&mut self) { if let Some(conn) = self.conn.take() { - warn!("Borrowed reuse-pool connection dropped before release, closing it"); + // Concurrent forwarding intentionally cancels losing queries. + // Drop cannot distinguish that from other early exits or unwinding. + debug!( + upstream_tag = %self.pool.connection_info.tag.as_deref().unwrap_or(""), + upstream_host = %self.pool.connection_info.server_name, + target_ip = ?self.pool.connection_info.remote_ip, + upstream_port = self.pool.connection_info.port, + transport = ?self.pool.transport, + connection_type = std::any::type_name::(), + connection_available = conn.available(), + inflight_queries = conn.using_count(), + pool_active_before_close = self.pool.active_count.load(Ordering::Relaxed), + pool_idle = self.pool.connections.len(), + pool_max = self.pool.max_size, + "Closing borrowed upstream connection after query cancellation or early drop" + ); self.pool.close_active_connection(conn); } } @@ -503,8 +538,12 @@ impl Drop for ReusePool { #[cfg(test)] mod tests { use std::collections::VecDeque; - use std::sync::Mutex; + use std::io; use std::sync::atomic::{AtomicBool, AtomicU64}; + use std::sync::{Arc, Mutex}; + + use tracing::dispatcher::Dispatch; + use tracing_subscriber::fmt::MakeWriter; use super::*; use crate::infra::error::{DnsError, Result}; @@ -593,6 +632,7 @@ mod tests { &self, _conn_id: u16, _deadline: QueryDeadline, + _connection_info: Arc, ) -> Result> { self.planned .lock() @@ -610,6 +650,8 @@ mod tests { ) -> ReusePool { AppClock::start(); ReusePool { + connection_info: Arc::new(ConnectionInfo::with_addr("tcp://127.0.0.1:53").unwrap()), + transport: ConnectionType::TCP, connections: ArrayQueue::new(max_size.max(1)), active_count: AtomicUsize::new(0), max_size, @@ -624,6 +666,66 @@ mod tests { } } + #[derive(Clone, Default)] + struct CapturedLogs(Arc>>); + + impl<'a> MakeWriter<'a> for CapturedLogs { + type Writer = Self; + + fn make_writer(&'a self) -> Self::Writer { + self.clone() + } + } + + impl io::Write for CapturedLogs { + fn write(&mut self, buf: &[u8]) -> io::Result { + self.0.lock().unwrap().extend_from_slice(buf); + Ok(buf.len()) + } + + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } + } + + #[test] + fn test_early_drop_logs_upstream_identity_without_url_credentials() { + let mut info = ConnectionInfo::with_addr("tcp://user:secret@resolver.example:5353") + .expect("upstream URL should parse"); + info.tag = Some("resolver_primary".to_string()); + info.remote_ip = Some("192.0.2.53".parse().unwrap()); + + let mut pool = make_pool(0, 1, 10, MockBuilder::new(vec![])); + pool.connection_info = Arc::new(info); + pool.active_count.store(1, Ordering::Relaxed); + let logs = CapturedLogs::default(); + let subscriber = tracing_subscriber::fmt() + .with_ansi(false) + .without_time() + .with_max_level(tracing::Level::DEBUG) + .with_writer(logs.clone()) + .finish(); + + tracing::dispatcher::with_default(&Dispatch::new(subscriber), || { + drop(BorrowedConnection::new( + &pool, + Arc::new(MockConnection::new(true, 0, 0)), + )); + }); + + let output = String::from_utf8(logs.0.lock().unwrap().clone()).unwrap(); + assert!(output.contains("DEBUG")); + assert!(output.contains("upstream_tag=resolver_primary"), "{output}"); + assert!( + output.contains("upstream_host=resolver.example"), + "{output}" + ); + assert!(output.contains("target_ip=Some(192.0.2.53)"), "{output}"); + assert!(output.contains("upstream_port=5353"), "{output}"); + assert!(output.contains("transport=TCP"), "{output}"); + assert!(!output.contains("secret")); + } + #[test] fn test_close_conns_closes_every_connection_once() { let first = Arc::new(MockConnection::new(true, 0, 0)); @@ -793,6 +895,35 @@ mod tests { assert_eq!(pool.connections.len(), 1); } + #[tokio::test] + async fn test_cancelled_query_closes_connection_and_releases_pool_capacity() { + let replacement = Arc::new(MockConnection::new(true, 0, 0)); + let pool = make_pool(0, 1, 10, MockBuilder::new(vec![Ok(replacement.clone())])); + let conn = + Arc::new(MockConnection::new(true, 0, 0).with_query_delay(Duration::from_secs(60))); + pool.connections.push(conn.clone()).unwrap(); + pool.active_count.store(1, Ordering::Relaxed); + + let mut query = + Box::pin(pool.query(Message::new(), QueryDeadline::new(Duration::from_secs(30)))); + assert!(futures::poll!(&mut query).is_pending()); + assert_eq!(conn.query_calls(), 1); + drop(query); + + assert_eq!(conn.close_calls(), 1); + assert!(!conn.available()); + assert!(pool.connections.is_empty()); + assert_eq!(pool.active_count.load(Ordering::Relaxed), 0); + + let borrowed = pool + .get(QueryDeadline::new(Duration::from_secs(1))) + .await + .expect("cancelled query must free capacity for a replacement"); + assert!(std::ptr::eq(borrowed.connection(), replacement.as_ref())); + borrowed.release(); + assert_eq!(pool.active_count.load(Ordering::Relaxed), 1); + } + #[tokio::test] async fn test_query_closes_borrowed_connection_when_deadline_expires() { let pool = make_pool(0, 1, 10, MockBuilder::new(vec![])); diff --git a/src/infra/network/upstream/pooled.rs b/src/infra/network/upstream/pooled.rs index 54bdf704..36ef7ce7 100644 --- a/src/infra/network/upstream/pooled.rs +++ b/src/infra/network/upstream/pooled.rs @@ -22,7 +22,7 @@ use crate::proto::Message; #[derive(Debug)] pub(crate) struct PooledUpstream { /// Connection metadata (remote address, port, etc.) - pub(crate) connection_info: ConnectionInfo, + pub(crate) connection_info: Arc, /// Connection pool for load balancing and connection reuse pub(crate) pool: Arc>, } @@ -59,7 +59,7 @@ impl Upstream for PooledUpstream { #[derive(Debug)] pub(crate) struct UdpTruncatedUpstream { /// Connection configuration (includes timeout) - pub(crate) connection_info: ConnectionInfo, + pub(crate) connection_info: Arc, /// Primary UDP connection pool (fast path) pub(crate) main_pool: Arc>, /// Fallback TCP connection pool (used when UDP response is truncated) diff --git a/src/infra/network/upstream/tests.rs b/src/infra/network/upstream/tests.rs index c2809d2c..0b68ff44 100644 --- a/src/infra/network/upstream/tests.rs +++ b/src/infra/network/upstream/tests.rs @@ -86,6 +86,7 @@ impl ConnectionBuilder for NoopConnectionBuilder { &self, _conn_id: u16, _deadline: QueryDeadline, + _connection_info: Arc, ) -> Result> { Ok(Arc::new(NoopConnection { available: AtomicBool::new(true), @@ -751,7 +752,7 @@ async fn test_pooled_upstream_lets_pool_handle_deadline_expiry() { ConnectionInfo::with_addr("tcp://127.0.0.1").expect("upstream should parse"); connection_info.timeout = Duration::from_millis(10); let upstream = PooledUpstream:: { - connection_info, + connection_info: Arc::new(connection_info), pool: Arc::new(DeadlineHandlingPool { handled_timeout: handled_timeout.clone(), }),