perf(smtp-client): Cache connections (#117)

This commit is contained in:
Jagoda Estera Ślązak
2026-04-17 11:37:05 +02:00
committed by GitHub
parent 45eaa2f659
commit dee00ebdcf
6 changed files with 277 additions and 62 deletions
+137
View File
@@ -67,6 +67,35 @@ dependencies = [
"windows-sys 0.61.2",
]
[[package]]
name = "async-io"
version = "2.6.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "456b8a8feb6f42d237746d4b3e9a178494627745c3c56c6ea55d92ba50d026fc"
dependencies = [
"autocfg",
"cfg-if",
"concurrent-queue",
"futures-io",
"futures-lite",
"parking",
"polling",
"rustix",
"slab",
"windows-sys 0.61.2",
]
[[package]]
name = "async-lock"
version = "3.4.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "290f7f2596bd5b78a9fec8088ccd89180d7f9f55b94b0576823bbbdc72ee8311"
dependencies = [
"event-listener",
"event-listener-strategy",
"pin-project-lite",
]
[[package]]
name = "async-trait"
version = "0.1.89"
@@ -155,6 +184,15 @@ version = "1.0.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b05b61dc5112cbb17e4b6cd61790d9845d13888356391624cbe7e41efeac1e75"
[[package]]
name = "concurrent-queue"
version = "2.5.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4ca0197aee26d1ae37445ee532fefce43251d24cc7c166799f4d46817f1d3973"
dependencies = [
"crossbeam-utils",
]
[[package]]
name = "const-oid"
version = "0.9.6"
@@ -383,6 +421,33 @@ dependencies = [
"windows-sys 0.61.2",
]
[[package]]
name = "event-listener"
version = "5.4.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e13b66accf52311f30a0db42147dadea9850cb48cd070028831ae5f5d4b856ab"
dependencies = [
"concurrent-queue",
"parking",
"pin-project-lite",
]
[[package]]
name = "event-listener-strategy"
version = "0.5.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8be9f3dfaaffdae2972880079a491a1a8bb7cbed0b8dd7a347f668b4150a3b93"
dependencies = [
"event-listener",
"pin-project-lite",
]
[[package]]
name = "fastrand"
version = "2.4.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9f1f227452a390804cdb637b74a86990f2a7d7ba4b7d5693aac9b4dd6defd8d6"
[[package]]
name = "fiat-crypto"
version = "0.2.9"
@@ -402,6 +467,7 @@ dependencies = [
"lru",
"mailparse",
"parking_lot",
"retainer",
"rstest",
"serde",
"serini",
@@ -456,6 +522,19 @@ version = "0.3.31"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9e5c1b78ca4aae1ac06c48a526a655760685149f0d465d21f37abfe57ce075c6"
[[package]]
name = "futures-lite"
version = "2.6.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f78e10609fe0e0b3f4157ffab1876319b5b0db102a2c60dc4626306dc46b44ad"
dependencies = [
"fastrand",
"futures-core",
"futures-io",
"parking",
"pin-project-lite",
]
[[package]]
name = "futures-macro"
version = "0.3.31"
@@ -587,6 +666,12 @@ version = "0.5.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea"
[[package]]
name = "hermit-abi"
version = "0.5.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c"
[[package]]
name = "hickory-proto"
version = "0.25.2"
@@ -837,6 +922,12 @@ version = "0.2.16"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981"
[[package]]
name = "linux-raw-sys"
version = "0.12.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53"
[[package]]
name = "litemap"
version = "0.8.1"
@@ -986,6 +1077,12 @@ version = "1.70.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe"
[[package]]
name = "parking"
version = "2.2.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f38d5652c16fde515bb1ecef450ab0f6a219d619a7274976324d5e377f7dceba"
[[package]]
name = "parking_lot"
version = "0.12.5"
@@ -1057,6 +1154,20 @@ dependencies = [
"spki",
]
[[package]]
name = "polling"
version = "3.11.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5d0e4f59085d47d8241c88ead0f274e8a0cb551f3625263c05eb8dd897c34218"
dependencies = [
"cfg-if",
"concurrent-queue",
"hermit-abi",
"pin-project-lite",
"rustix",
"windows-sys 0.61.2",
]
[[package]]
name = "portable-atomic"
version = "1.13.0"
@@ -1267,6 +1378,19 @@ version = "0.7.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1e061d1b48cb8d38042de4ae0a7a6401009d6143dc80d2e2d6f31f0bdd6470c7"
[[package]]
name = "retainer"
version = "0.4.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7b071fe646a2ab077f74656a4602c16528829c1fafa81946c5e88eaeccf08d5b"
dependencies = [
"async-io",
"async-lock",
"futures-lite",
"log",
"rand 0.9.2",
]
[[package]]
name = "ring"
version = "0.17.14"
@@ -1339,6 +1463,19 @@ dependencies = [
"semver",
]
[[package]]
name = "rustix"
version = "1.1.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190"
dependencies = [
"bitflags",
"errno",
"libc",
"linux-raw-sys",
"windows-sys 0.61.2",
]
[[package]]
name = "rustls"
version = "0.23.38"
+1
View File
@@ -30,6 +30,7 @@ tokio-rustls = { version = "0.26.4", default-features = false, features = [
] }
webpki-roots = "1.0.6"
tokio-io-timeout = "1.2.1"
retainer = "0.4.0"
[dev-dependencies]
rstest = "0.26.1"
+4
View File
@@ -4,6 +4,7 @@ use crate::ENCRYPTION_NEEDED_523;
use crate::config::Config;
use crate::dkim_verifier::DkimVerifier;
use crate::message::{check_encrypted, is_securejoin};
use crate::smtp_client::SmtpConnectionPool;
pub use crate::smtp_server::Envelope;
use crate::smtp_server::SmtpHandler;
use crate::utils::{AddressDomain, build_resolver, extract_address, log_eml};
@@ -19,6 +20,7 @@ pub struct IncomingBeforeQueueHandler {
dns_resolver: Arc<TokioResolver>,
dkim_verifier: DkimVerifier,
skip_dkim: bool,
smtp_connection_pool: Arc<SmtpConnectionPool>,
}
impl IncomingBeforeQueueHandler {
@@ -29,6 +31,7 @@ impl IncomingBeforeQueueHandler {
dns_resolver: dns_resolver.clone(),
dkim_verifier: DkimVerifier::new(dns_resolver),
skip_dkim,
smtp_connection_pool: SmtpConnectionPool::new(),
})
}
@@ -152,6 +155,7 @@ impl SmtpHandler for IncomingBeforeQueueHandler {
&hostname,
None,
self.dns_resolver.clone(),
self.smtp_connection_pool.clone(),
)
.await
.map_err(|e| {
+4
View File
@@ -3,6 +3,7 @@
use crate::ENCRYPTION_NEEDED_523;
use crate::config::Config;
use crate::message::{check_encrypted, is_securejoin, recipient_matches_passthrough};
use crate::smtp_client::SmtpConnectionPool;
pub use crate::smtp_server::Envelope;
use crate::smtp_server::SmtpHandler;
use crate::utils::{build_resolver, extract_address};
@@ -17,6 +18,7 @@ pub struct OutgoingBeforeQueueHandler {
config: Config,
dns_resolver: Arc<TokioResolver>,
send_rate_limiter: DefaultKeyedRateLimiter<String>,
smtp_connection_pool: Arc<SmtpConnectionPool>,
}
impl OutgoingBeforeQueueHandler {
@@ -28,6 +30,7 @@ impl OutgoingBeforeQueueHandler {
config,
dns_resolver,
send_rate_limiter: RateLimiter::keyed(quota),
smtp_connection_pool: SmtpConnectionPool::new(),
})
}
}
@@ -150,6 +153,7 @@ impl SmtpHandler for OutgoingBeforeQueueHandler {
&hostname,
None,
self.dns_resolver.clone(),
self.smtp_connection_pool.clone(),
)
.await
.map_err(|e| {
+125 -61
View File
@@ -7,13 +7,62 @@ use std::time::Duration;
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufStream};
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
use tokio::net::TcpStream;
use tokio::task::JoinSet;
use tokio::task::{JoinHandle, JoinSet};
use tokio_io_timeout::TimeoutStream;
use tokio_rustls::rustls::client::ClientSessionMemoryCache;
/// A connection pool for SMTP connections, keyed by (address, port).
///
/// Connections are cached for up to 100 seconds of idle time.
///
/// Only a single connection is cached per address/port pair.
pub struct SmtpConnectionPool {
pool: Arc<retainer::Cache<(String, u16), BufStream<SmtpStream>>>,
monitor_handle: JoinHandle<()>,
}
impl SmtpConnectionPool {
/// Creates a new connection pool and starts the cache monitoring task.
pub fn new() -> Arc<Self> {
let pool = Arc::new(retainer::Cache::new());
let pool_clone = pool.clone();
let monitor_handle =
tokio::spawn(async move { pool_clone.monitor(4, 0.25, Duration::from_secs(10)).await });
Arc::new(Self {
pool,
monitor_handle,
})
}
/// Takes a connection from the pool for the given address and port, if available.
pub async fn take(&self, address: &str, port: u16) -> Option<BufStream<SmtpStream>> {
self.pool.remove(&(address.to_string(), port)).await
}
/// Puts a connection into the pool for the given address and port, with a 100s timeout.
pub async fn put(&self, address: &str, port: u16, stream: BufStream<SmtpStream>) {
// similarly to postfix default -> 100s max idle time.
self.pool
.insert(
(address.to_string(), port),
stream,
Duration::from_secs(100),
)
.await;
}
}
impl Drop for SmtpConnectionPool {
fn drop(&mut self) {
self.monitor_handle.abort();
}
}
/// A [`TcpStream`] used for SMTP communication.
#[expect(clippy::large_enum_variant)]
enum SmtpStream {
pub enum SmtpStream {
/// A plain TCP stream.
Plain(Pin<Box<TimeoutStream<TcpStream>>>),
/// A TLS-encrypted stream.
@@ -144,6 +193,8 @@ async fn establish_tcp_connection(
///
/// If `address` is a domain that resolves to multiple IP addresses,
/// all will be tried in parallel and the first successful connection will be used.
///
/// `pool` is used to reuse existing connections to the same address and port, if available.
pub async fn send(
address: &str,
port: u16,
@@ -151,23 +202,30 @@ pub async fn send(
client_hostname: &str,
tls_config: Option<TlsConfig>,
dns_resolver: Arc<TokioResolver>,
pool: Arc<SmtpConnectionPool>,
) -> Result<(), crate::error::Error> {
let stream = establish_tcp_connection(address, port, dns_resolver).await?;
let (mut buf_stream, reused) = if let Some(stream) = pool.take(address, port).await {
log::debug!("Reusing existing connection to {address}:{port}",);
if tls_config.is_some() {
// This should never happen,
// assert to make sure we never accidentally use a plain connection while expecting TLS.
assert!(
matches!(stream.get_ref(), SmtpStream::Tls(_)),
"Expected TLS stream from pool, but got plain stream."
);
}
(stream, true)
} else {
let stream = establish_tcp_connection(address, port, dns_resolver).await?;
log::debug!("Successfully connected to {}", stream.peer_addr()?);
(BufStream::new(SmtpStream::plain(stream)), false)
};
log::debug!(
"SMTP client: successfully connected to {}",
stream.peer_addr()?
);
let mut buf_stream = BufStream::new(SmtpStream::plain(stream));
let mut response = String::new();
macro_rules! smtp_write {
($command: expr) => {
log::trace!(
"SMTP client: sending: {}",
String::from_utf8_lossy($command)
);
log::trace!("Sending: {}", String::from_utf8_lossy($command));
buf_stream.write_all($command).await?;
buf_stream.flush().await?;
};
@@ -203,61 +261,65 @@ pub async fn send(
};
}
// Read initial greeting
smtp_read!("initial greeting", "220");
if tls_config.is_some() {
smtp_cmd!(
format!("EHLO {client_hostname}\r\n").as_bytes(),
"EHLO",
"250"
);
if reused {
smtp_cmd!(b"RSET\r\n", "RSET on reused connection", "250");
} else {
smtp_cmd!(
format!("HELO {client_hostname}\r\n").as_bytes(),
"HELO",
"250"
);
};
// Read initial greeting
smtp_read!("initial greeting", "220");
// STARTTLS
if let Some(tls_config) = tls_config {
if !response.to_uppercase().contains("STARTTLS") {
// TLS was requested, but server doesn't support STARTTLS.
return Err(crate::error::Error::MailSend {
context: "STARTTLS".to_string(),
raw_smtp_answer: response.clone(),
});
}
log::trace!("Initiating STARTTLS...");
smtp_cmd!(b"STARTTLS\r\n", "STARTTLS", "220");
let stream = buf_stream.into_inner();
let raw_tcp = match stream {
SmtpStream::Plain(s) => s,
SmtpStream::Tls(_) => {
unreachable!("This is the first and only place we upgrade to TLS.")
}
if tls_config.is_some() {
smtp_cmd!(
format!("EHLO {client_hostname}\r\n").as_bytes(),
"EHLO",
"250"
);
} else {
smtp_cmd!(
format!("HELO {client_hostname}\r\n").as_bytes(),
"HELO",
"250"
);
};
let tls_stream = crate::tls::wrap_rustls(
address,
raw_tcp,
tls_config.session_cache,
tls_config.allow_invalid_cert,
)
.await?;
// STARTTLS
if let Some(tls_config) = tls_config {
if !response.to_uppercase().contains("STARTTLS") {
// TLS was requested, but server doesn't support STARTTLS.
return Err(crate::error::Error::MailSend {
context: "STARTTLS".to_string(),
raw_smtp_answer: response.clone(),
});
}
let smtp_stream = SmtpStream::Tls(tls_stream);
log::trace!("Initiating STARTTLS...");
smtp_cmd!(b"STARTTLS\r\n", "STARTTLS", "220");
buf_stream = BufStream::new(smtp_stream);
let stream = buf_stream.into_inner();
let raw_tcp = match stream {
SmtpStream::Plain(s) => s,
SmtpStream::Tls(_) => {
unreachable!("This is the first and only place we upgrade to TLS.")
}
};
smtp_cmd!(
format!("EHLO {client_hostname}\r\n").as_bytes(),
"EHLO after STARTTLS",
"250"
);
let tls_stream = crate::tls::wrap_rustls(
address,
raw_tcp,
tls_config.session_cache,
tls_config.allow_invalid_cert,
)
.await?;
let smtp_stream = SmtpStream::Tls(tls_stream);
buf_stream = BufStream::new(smtp_stream);
smtp_cmd!(
format!("EHLO {client_hostname}\r\n").as_bytes(),
"EHLO after STARTTLS",
"250"
);
}
}
// MAIL FROM
@@ -282,5 +344,7 @@ pub async fn send(
smtp_write!(b".\r\n");
smtp_read!("end of DATA", "250");
pool.put(address, port, buf_stream).await;
Ok(())
}
+6 -1
View File
@@ -1,5 +1,5 @@
use crate::config::Config;
use crate::smtp_client::TlsConfig;
use crate::smtp_client::{SmtpConnectionPool, TlsConfig};
use crate::smtp_server::{Envelope, SmtpHandler};
use crate::utils::{AddressDomain, build_resolver};
use async_trait::async_trait;
@@ -14,6 +14,7 @@ pub struct TransportHandler {
config: Config,
dns_resolver: Arc<TokioResolver>,
tls_resumption_store: Arc<rustls::client::ClientSessionMemoryCache>,
smtp_connection_pool: Arc<SmtpConnectionPool>,
}
impl TransportHandler {
@@ -24,12 +25,14 @@ impl TransportHandler {
config,
dns_resolver,
tls_resumption_store,
smtp_connection_pool: SmtpConnectionPool::new(),
})
}
/// Handles a single email transaction for a single recipient domain.
async fn handle_single_domain(
tls_resumption_store: Arc<rustls::client::ClientSessionMemoryCache>,
smtp_connection_pool: Arc<SmtpConnectionPool>,
dns_resolver: Arc<TokioResolver>,
domain: AddressDomain,
envelope: Envelope,
@@ -97,6 +100,7 @@ impl TransportHandler {
&client_hostname,
tls_config.clone(),
dns_resolver.clone(),
smtp_connection_pool.clone(),
)
.await
{
@@ -179,6 +183,7 @@ impl SmtpHandler for TransportHandler {
let task_id = transactions
.spawn(Self::handle_single_domain(
self.tls_resumption_store.clone(),
self.smtp_connection_pool.clone(),
self.dns_resolver.clone(),
rcpt_domain.clone(),
domain_envelope,