use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;
use deadpool::Runtime;
use deadpool::managed::{self, Metrics, Pool, RecycleError, RecycleResult};
use crate::error::{Error, Result};
use crate::tcp::connect::{self, ConnectKind};
use crate::tcp::connection_actor::{ActorConfig, ConnectionActor, ConnectionHandle};
use crate::tcp::handshake::HandshakeConfig;
use crate::tcp::retry::RetryPolicy;
pub(crate) const DEFAULT_MAX_SIZE: usize = 8;
pub(crate) const DEFAULT_ACQUIRE_TIMEOUT: Duration = Duration::from_secs(30);
pub(crate) const DEFAULT_CREATE_TIMEOUT: Duration = Duration::from_secs(10);
pub(crate) const DEFAULT_RECYCLE_TIMEOUT: Duration = Duration::from_secs(5);
#[non_exhaustive]
pub struct TcpConnectionManager {
pub endpoints: Vec<String>,
pub next: Arc<AtomicUsize>,
pub kind: ConnectKind,
pub config: HandshakeConfig,
pub max_lifetime: Option<Duration>,
pub read_timeout: Option<Duration>,
}
impl managed::Manager for TcpConnectionManager {
type Type = ConnectionHandle;
type Error = Error;
async fn create(&self) -> Result<ConnectionHandle> {
let n = self.endpoints.len();
if n == 0 {
return Err(Error::Connect(
"tcp: no endpoints configured for this pool".to_string(),
));
}
let start = self.next.fetch_add(1, Ordering::Relaxed) % n;
let mut last_err: Option<Error> = None;
for i in 0..n {
let endpoint = &self.endpoints[(start + i) % n];
let resolved = match tokio::net::lookup_host(endpoint).await {
Ok(mut addrs) => addrs.next().ok_or_else(|| {
Error::Connect(format!("tcp: {endpoint:?} resolved to no addresses"))
}),
Err(e) => Err(Error::Connect(format!(
"tcp: cannot resolve {endpoint:?}: {e}"
))),
};
let addr = match resolved {
Ok(addr) => addr,
Err(e) => {
last_err = Some(e);
continue;
}
};
match connect::open_handshaken(addr, &self.kind, &self.config).await {
Ok((stream, hello)) => {
return Ok(ConnectionActor::spawn_with_config(
stream,
hello,
ActorConfig {
read_timeout: self.read_timeout,
quota_key: self.config.quota_key.clone(),
},
));
}
Err(e) => last_err = Some(e),
}
}
Err(last_err.unwrap_or_else(|| {
Error::Connect("tcp: no endpoints configured for this pool".to_string())
}))
}
async fn recycle(&self, obj: &mut ConnectionHandle, metrics: &Metrics) -> RecycleResult<Error> {
match recycle_refusal(self.max_lifetime, metrics.created.elapsed(), obj.is_alive()) {
Some(reason) => {
close_refused(obj).await;
Err(RecycleError::message(reason))
}
None => Ok(()),
}
}
}
fn recycle_refusal(
max_lifetime: Option<Duration>,
created_elapsed: Duration,
is_alive: bool,
) -> Option<&'static str> {
if max_lifetime.is_some_and(|max| created_elapsed >= max) {
return Some("connection exceeded max_lifetime");
}
(!is_alive).then_some("connection poisoned")
}
async fn close_refused(obj: &ConnectionHandle) {
if let Err(e) = obj.clone().close().await {
tracing::warn!(
target: "clickhouse::tcp",
error = %e,
"refused connection did not shut down cleanly"
);
}
}
pub type NativePool = Pool<TcpConnectionManager>;
#[derive(Clone, Copy, Debug)]
#[non_exhaustive]
pub struct PoolConfig {
pub max_size: usize,
pub acquire_timeout: Option<Duration>,
pub create_timeout: Option<Duration>,
pub recycle_timeout: Option<Duration>,
pub max_lifetime: Option<Duration>,
pub read_timeout: Option<Duration>,
}
impl Default for PoolConfig {
fn default() -> Self {
Self {
max_size: DEFAULT_MAX_SIZE,
acquire_timeout: Some(DEFAULT_ACQUIRE_TIMEOUT),
create_timeout: Some(DEFAULT_CREATE_TIMEOUT),
recycle_timeout: Some(DEFAULT_RECYCLE_TIMEOUT),
max_lifetime: None,
read_timeout: None,
}
}
}
#[derive(Clone, Debug, Default)]
#[non_exhaustive]
pub struct TcpClientConfig {
pub endpoints: Vec<String>,
pub kind: ConnectKindConfig,
pub handshake: HandshakeConfig,
pub pool: PoolConfig,
pub retry: Option<RetryPolicy>,
}
#[derive(Clone, Debug, Default)]
#[non_exhaustive]
pub enum ConnectKindConfig {
#[default]
Plain,
Tls {
server_name: String,
},
}
#[cfg(feature = "tls")]
#[non_exhaustive]
pub enum TcpTls {
NotConfigured,
Resolved(std::sync::Arc<tokio_rustls::rustls::ClientConfig>),
ConfiguredButFailed,
}
#[cfg_attr(not(feature = "tls"), allow(clippy::needless_pass_by_value))]
pub fn build_pool(
endpoints: Vec<String>,
kind: ConnectKindConfig,
config: HandshakeConfig,
pool_cfg: PoolConfig,
#[cfg(feature = "tls")] tls: TcpTls,
) -> Result<NativePool> {
assert!(
!endpoints.is_empty(),
"build_pool requires a non-empty endpoint list"
);
let kind = match kind {
ConnectKindConfig::Plain => ConnectKind::Plain,
#[cfg(feature = "tls")]
ConnectKindConfig::Tls { server_name } => match tls {
TcpTls::Resolved(config) => ConnectKind::Tls {
server_name,
config,
},
TcpTls::NotConfigured => {
let config = crate::tls::build_client_config(&crate::tls::TlsConfigSource::Trust(
crate::tls::TlsTrust::default(),
))?;
ConnectKind::Tls {
server_name,
config,
}
}
TcpTls::ConfiguredButFailed => ConnectKind::TlsFailClosed,
},
#[cfg(not(feature = "tls"))]
ConnectKindConfig::Tls { .. } => {
return Err(Error::Custom(
"tcp: TLS requested but the `tls` feature is not enabled".into(),
));
}
};
let manager = TcpConnectionManager {
endpoints,
next: Arc::new(AtomicUsize::new(0)),
kind,
config,
max_lifetime: pool_cfg.max_lifetime,
read_timeout: pool_cfg.read_timeout,
};
Pool::builder(manager)
.max_size(pool_cfg.max_size)
.wait_timeout(pool_cfg.acquire_timeout)
.create_timeout(pool_cfg.create_timeout)
.recycle_timeout(pool_cfg.recycle_timeout)
.runtime(Runtime::Tokio1)
.build()
.map_err(|e| Error::Custom(format!("tcp: failed to build pool: {e}")))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::tcp::connection_actor::ConnectionActor;
use crate::tcp::protocol::{DBMS_TCP_PROTOCOL_VERSION, ServerHello};
use crate::tcp::transport::MaybeTlsStream;
use tokio::net::{TcpListener, TcpStream};
async fn fabricate_handle() -> ConnectionHandle {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let connect_fut = TcpStream::connect(addr);
let accept_fut = async { listener.accept().await.unwrap().0 };
let (client, _server) = tokio::join!(connect_fut, accept_fut);
let client = client.unwrap();
let _ = client.set_nodelay(true);
let stream = MaybeTlsStream::Plain(client);
let hello = ServerHello {
server_name: "mock".to_string(),
version: (1, 0, 0),
revision: DBMS_TCP_PROTOCOL_VERSION,
timezone: None,
display_name: None,
};
ConnectionActor::spawn(stream, hello)
}
struct TestManager {
creates: Arc<AtomicUsize>,
max_lifetime: Option<Duration>,
}
impl TestManager {
fn new() -> Self {
Self {
creates: Arc::new(AtomicUsize::new(0)),
max_lifetime: None,
}
}
}
impl managed::Manager for TestManager {
type Type = ConnectionHandle;
type Error = Error;
async fn create(&self) -> Result<ConnectionHandle> {
self.creates.fetch_add(1, Ordering::AcqRel);
Ok(fabricate_handle().await)
}
fn recycle(
&self,
obj: &mut ConnectionHandle,
metrics: &Metrics,
) -> impl std::future::Future<Output = RecycleResult<Error>> {
std::future::ready(
match recycle_refusal(self.max_lifetime, metrics.created.elapsed(), obj.is_alive())
{
Some(reason) => Err(RecycleError::message(reason)),
None => Ok(()),
},
)
}
}
fn build_test_pool(mgr: TestManager, cfg: PoolConfig) -> Pool<TestManager> {
Pool::<TestManager>::builder(mgr)
.max_size(cfg.max_size)
.wait_timeout(cfg.acquire_timeout)
.create_timeout(cfg.create_timeout)
.recycle_timeout(cfg.recycle_timeout)
.runtime(Runtime::Tokio1)
.build()
.expect("test pool builds")
}
#[tokio::test]
async fn pool_recycle_drops_poisoned_handle() {
let mgr = TestManager::new();
let creates = Arc::clone(&mgr.creates);
let pool = build_test_pool(mgr, PoolConfig::default());
let first = pool.get().await.expect("acquire");
assert_eq!(creates.load(Ordering::Acquire), 1);
first.poison();
assert!(!first.is_alive());
drop(first);
let _second = pool.get().await.expect("acquire again");
assert_eq!(
creates.load(Ordering::Acquire),
2,
"poisoned handle should have been dropped and a fresh one created"
);
}
#[tokio::test]
async fn pool_recycle_keeps_alive_handle() {
let mgr = TestManager::new();
let creates = Arc::clone(&mgr.creates);
let pool = build_test_pool(mgr, PoolConfig::default());
let first = pool.get().await.expect("acquire");
assert_eq!(creates.load(Ordering::Acquire), 1);
assert!(first.is_alive());
drop(first);
let _second = pool.get().await.expect("acquire again");
assert_eq!(
creates.load(Ordering::Acquire),
1,
"alive handle should have been reused, not re-created"
);
}
#[tokio::test]
async fn pool_recycle_drops_aged_connection() {
let creates = Arc::new(AtomicUsize::new(0));
let mgr = TestManager {
creates: Arc::clone(&creates),
max_lifetime: Some(Duration::from_millis(10)),
};
let pool = build_test_pool(mgr, PoolConfig::default());
let first = pool.get().await.expect("acquire");
assert_eq!(creates.load(Ordering::Acquire), 1);
assert!(first.is_alive(), "fresh handle is alive");
drop(first);
tokio::time::sleep(Duration::from_millis(30)).await;
let _second = pool.get().await.expect("acquire again");
assert_eq!(
creates.load(Ordering::Acquire),
2,
"connection older than max_lifetime should be dropped and re-created"
);
}
#[test]
fn recycle_refusal_covers_age_then_liveness() {
let ms = Duration::from_millis;
let cases: [(Option<Duration>, Duration, bool, Option<&str>); 6] = [
(None, ms(10_000), true, None),
(None, ms(0), false, Some("connection poisoned")),
(Some(ms(10)), ms(9), true, None),
(
Some(ms(10)),
ms(10),
true,
Some("connection exceeded max_lifetime"),
),
(
Some(ms(10)),
ms(11),
false,
Some("connection exceeded max_lifetime"),
),
(Some(ms(10)), ms(9), false, Some("connection poisoned")),
];
for (max, elapsed, alive, want) in cases {
assert_eq!(
recycle_refusal(max, elapsed, alive),
want,
"max={max:?} elapsed={elapsed:?} alive={alive}"
);
}
}
#[tokio::test]
async fn pool_max_size_blocks_excess_acquires() {
let cfg = PoolConfig {
max_size: 2,
acquire_timeout: None,
..PoolConfig::default()
};
let mgr = TestManager::new();
let pool = build_test_pool(mgr, cfg);
let a = pool.get().await.expect("a");
let b = pool.get().await.expect("b");
let pool_clone = pool.clone();
let third = tokio::spawn(async move { pool_clone.get().await.is_ok() });
tokio::time::sleep(Duration::from_millis(50)).await;
assert!(
!third.is_finished(),
"third acquire should block while pool is saturated"
);
drop(a);
let acquired = tokio::time::timeout(Duration::from_secs(2), third)
.await
.expect("third acquire should unblock once a slot is freed")
.expect("join");
assert!(acquired, "third acquire should succeed after a drop");
drop(b);
}
#[tokio::test]
async fn create_timeout_surfaces_as_a_pool_error() {
struct StalledManager;
impl managed::Manager for StalledManager {
type Type = ConnectionHandle;
type Error = Error;
async fn create(&self) -> Result<ConnectionHandle> {
std::future::pending::<()>().await;
unreachable!("the create timeout fires first")
}
fn recycle(
&self,
_: &mut ConnectionHandle,
_: &Metrics,
) -> impl std::future::Future<Output = RecycleResult<Error>> {
std::future::ready(Ok(()))
}
}
let pool = Pool::<StalledManager>::builder(StalledManager)
.max_size(1)
.create_timeout(Some(Duration::from_millis(20)))
.runtime(Runtime::Tokio1)
.build()
.expect("test pool builds");
let Ok(acquired) = tokio::time::timeout(Duration::from_secs(2), pool.get()).await else {
panic!("acquire hung past the create timeout")
};
let Err(err) = acquired else {
panic!("create never completes; acquire must fail")
};
let mapped = crate::tcp::retry::map_pool_error(err);
assert!(
matches!(mapped, Error::Transient(_)),
"a create timeout must map to a transient error, got {mapped:?}"
);
assert!(crate::tcp::retry::is_retriable_transport(&mapped));
}
#[tokio::test]
async fn pool_acquire_timeout_returns_err() {
let cfg = PoolConfig {
max_size: 1,
acquire_timeout: Some(Duration::from_millis(10)),
..PoolConfig::default()
};
let mgr = TestManager::new();
let pool = build_test_pool(mgr, cfg);
let _held = pool.get().await.expect("first acquire");
let start = tokio::time::Instant::now();
let result = tokio::time::timeout(Duration::from_millis(200), pool.get()).await;
let elapsed = start.elapsed();
let Ok(acquired) = result else {
panic!("acquire hung past the wait_timeout bound")
};
let Err(_refused) = acquired else {
panic!("acquire should not have succeeded with one held slot")
};
assert!(
elapsed < Duration::from_millis(200),
"acquire should have errored quickly, took {elapsed:?}"
);
}
use crate::tcp::connect::ConnectKind;
use crate::tcp::mock::serve_one_handshake;
use deadpool::managed::Manager as _;
async fn spawn_mock_endpoint(
idx: usize,
accepts: usize,
order: Arc<std::sync::Mutex<Vec<usize>>>,
) -> String {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap().to_string();
tokio::spawn(async move {
for _ in 0..accepts {
let Ok((mut sock, _)) = listener.accept().await else {
break;
};
order.lock().unwrap().push(idx);
serve_one_handshake(&mut sock).await;
}
});
addr
}
fn manager_over(endpoints: Vec<String>) -> TcpConnectionManager {
TcpConnectionManager {
endpoints,
next: Arc::new(AtomicUsize::new(0)),
kind: ConnectKind::Plain,
config: HandshakeConfig::default(),
max_lifetime: None,
read_timeout: None,
}
}
#[tokio::test]
async fn create_round_robins_across_endpoints() {
let order = Arc::new(std::sync::Mutex::new(Vec::new()));
let mut endpoints = Vec::new();
for idx in 0..3 {
endpoints.push(spawn_mock_endpoint(idx, 1, Arc::clone(&order)).await);
}
let mgr = manager_over(endpoints);
for _ in 0..3 {
mgr.create()
.await
.expect("create connects to the next endpoint");
}
let observed = order.lock().unwrap().clone();
assert_eq!(
observed,
vec![0, 1, 2],
"three creates should round-robin across the three endpoints in order"
);
}
#[tokio::test]
async fn create_fails_over_to_next_endpoint() {
let dead = TcpListener::bind("127.0.0.1:0").await.unwrap();
let refusing = dead.local_addr().unwrap().to_string();
drop(dead);
let order = Arc::new(std::sync::Mutex::new(Vec::new()));
let accepting = spawn_mock_endpoint(1, 1, Arc::clone(&order)).await;
let mgr = manager_over(vec![refusing, accepting]);
mgr.create()
.await
.expect("create should fail over to the accepting endpoint");
let observed = order.lock().unwrap().clone();
assert_eq!(
observed,
vec![1],
"create should have connected only to the accepting endpoint after skipping the refusing one"
);
}
#[tokio::test]
async fn create_fails_over_past_an_unresolvable_endpoint() {
let order = Arc::new(std::sync::Mutex::new(Vec::new()));
let accepting = spawn_mock_endpoint(1, 1, Arc::clone(&order)).await;
let mgr = manager_over(vec!["no-such-host.invalid:9000".to_string(), accepting]);
mgr.create()
.await
.expect("an unresolvable endpoint must be skipped, not fatal");
assert_eq!(
order.lock().unwrap().clone(),
vec![1],
"create should have reached the accepting endpoint"
);
}
#[tokio::test]
async fn create_surfaces_the_resolver_error_when_no_endpoint_resolves() {
let mgr = manager_over(vec![
"no-such-host-a.invalid:9000".to_string(),
"no-such-host-b.invalid:9000".to_string(),
]);
let Err(err) = mgr.create().await else {
panic!("no endpoint resolves; create must not connect")
};
let msg = err.to_string();
assert!(
msg.contains("no-such-host-b.invalid"),
"expected the last endpoint's resolver error, got: {msg}"
);
}
#[tokio::test]
async fn create_over_no_endpoints_errors_rather_than_panicking() {
let mgr = manager_over(Vec::new());
let Err(err) = mgr.create().await else {
panic!("an empty endpoint list cannot produce a connection")
};
assert!(err.to_string().contains("no endpoints configured"), "{err}");
}
#[tokio::test]
async fn create_surfaces_error_when_all_endpoints_refuse() {
let mut refusing = Vec::new();
for _ in 0..2 {
let dead = TcpListener::bind("127.0.0.1:0").await.unwrap();
refusing.push(dead.local_addr().unwrap().to_string());
drop(dead);
}
let mgr = manager_over(refusing);
let Err(err) = mgr.create().await else {
panic!("all endpoints refuse; create must not connect")
};
assert!(
crate::tcp::retry::is_retriable_transport(&err),
"an all-endpoints-refused connect error should be retriable, got {err:?}"
);
}
#[cfg(feature = "tls")]
#[tokio::test]
async fn tls_pool_fails_closed_when_configured_trust_unresolved() {
let pool = build_pool(
vec!["127.0.0.1:9440".to_string()],
ConnectKindConfig::Tls {
server_name: "example.invalid".to_string(),
},
HandshakeConfig::default(),
PoolConfig::default(),
TcpTls::ConfiguredButFailed,
)
.expect("configured-but-failed trust still BUILDS a (fail-closed) pool");
let Err(err) = pool.get().await else {
panic!("a fail-closed TLS pool must refuse to connect")
};
let msg = format!("{err}");
assert!(
msg.contains("refusing to connect with default trust"),
"expected fail-closed connect error, got: {msg}"
);
}
}