diff --git a/filtermail/src/smtp_responses.rs b/filtermail/src/smtp_responses.rs index e12fba2a..b290e1a7 100644 --- a/filtermail/src/smtp_responses.rs +++ b/filtermail/src/smtp_responses.rs @@ -4,4 +4,4 @@ pub const OK_SMTP_250: &str = "250 OK (SMTP)"; pub const ENCRYPTION_NEEDED_523: &str = "523 Encryption Needed: Invalid Unencrypted Mail"; pub const LOCAL_ERROR_451: &str = "451 Local error"; -pub const WORKER_BUSY_421: &str = "421 Worker for this destination is busy"; +pub const TRANSPORT_BUSY_421: &str = "421 Transport busy"; diff --git a/filtermail/src/snapshots/filtermail__transport__tests__smtp_send_mail_defer.snap b/filtermail/src/snapshots/filtermail__transport__tests__smtp_send_mail_defer.snap index 9dd43b2f..01244c79 100644 --- a/filtermail/src/snapshots/filtermail__transport__tests__smtp_send_mail_defer.snap +++ b/filtermail/src/snapshots/filtermail__transport__tests__smtp_send_mail_defer.snap @@ -20,7 +20,7 @@ TRANSACTION 1 > message > . < 250 OK (SMTP) -421 Worker for this destination is busy +421 Transport busy [filtermail-transport -> destination A] < 220 filtermail SMTP diff --git a/filtermail/src/snapshots/filtermail__transport__tests__smtp_send_mail_worker_cap_exceeded.snap b/filtermail/src/snapshots/filtermail__transport__tests__smtp_send_mail_worker_cap_exceeded.snap new file mode 100644 index 00000000..4c37501d --- /dev/null +++ b/filtermail/src/snapshots/filtermail__transport__tests__smtp_send_mail_worker_cap_exceeded.snap @@ -0,0 +1,51 @@ +--- +source: src/transport.rs +expression: "format!(\"TRANSACTION 1\\r\\n\\\n [postfix -> filtermail-transport]\\r\\n{record_postfix_1}\\r\\n\\\n [filtermail-transport -> destination A]\\r\\n{record_filtermail_1}\\r\\n\\r\\n\\\n TRANSACTION 2\\r\\n\\\n [postfix -> filtermail-transport]\\r\\n{record_postfix_2}\")" +--- +TRANSACTION 1 +[postfix -> filtermail-transport] +< 220 filtermail SMTP +> LHLO postfix +< 250-filtermail +250-8BITMIME +250 OK +> MAIL FROM: +< 250 OK +> RCPT TO: +< 250 OK +> DATA +< 354 End data with . +> message +> . +< 250 OK (SMTP) + +[filtermail-transport -> destination A] +< 220 filtermail SMTP +> EHLO example.org +< 250-filtermail +250-8BITMIME +250 OK +> MAIL FROM: +< 250 OK +> RCPT TO: +< 250 OK +> DATA +< 354 End data with . +> message +. +< 250 OK + + +TRANSACTION 2 +[postfix -> filtermail-transport] +< 220 filtermail SMTP +> LHLO postfix +< 250-filtermail +250-8BITMIME +250 OK +> MAIL FROM: +< 250 OK +> RCPT TO: +< 250 OK +> DATA +< 421 Transport busy diff --git a/filtermail/src/transport.rs b/filtermail/src/transport.rs index 0a3997c1..cd79b8a3 100644 --- a/filtermail/src/transport.rs +++ b/filtermail/src/transport.rs @@ -2,7 +2,7 @@ mod https_client; mod worker; use crate::config::Config; -use crate::smtp_responses::{LOCAL_ERROR_451, WORKER_BUSY_421}; +use crate::smtp_responses::{LOCAL_ERROR_451, TRANSPORT_BUSY_421}; use crate::smtp_server::{SmtpHandler, Transaction}; use crate::tcp::{TcpConnect, TcpStreamTrait}; use crate::utils::AddressDomain; @@ -32,12 +32,17 @@ where Ok(Self { workers }) } - /// Same as [`Self::new`], but lets you set worker queue size. + /// Same as [`Self::new`], but lets you set worker queue size and pool capacity. /// /// Only used for tests. #[cfg(test)] - pub fn with_queue_size(config: Config, queue_size: usize) -> Result { - let workers = WorkerPool::with_queue_size(config, queue_size)?; + pub fn with_queue_size_and_pool_capacity( + config: Config, + queue_size: usize, + pool_capacity: usize, + ) -> Result { + let workers = + WorkerPool::with_queue_size_and_pool_capacity(config, queue_size, pool_capacity)?; Ok(Self { workers }) } @@ -94,7 +99,7 @@ where // deferred mails to be constantly retried. if transaction.state.permits.is_empty() { - return Err(WORKER_BUSY_421.to_string()); + return Err(TRANSPORT_BUSY_421.to_string()); } Ok(()) @@ -133,11 +138,10 @@ where if let Some(permit) = transaction.state.permits.remove(rcpt_domain) { let (message, receiver) = WorkerMessage::new(domain_envelope); permit.send(message); - // todo: receiver timeout? transactions.spawn(receiver).id() } else { transactions - .spawn(async move { Ok(Err(WORKER_BUSY_421.to_string())) }) + .spawn(async move { Ok(Err(TRANSPORT_BUSY_421.to_string())) }) .id() }; task_id_domain_map.insert(receiver_task_id, rcpt_domain); @@ -152,11 +156,14 @@ where let smtp_response = match result { Ok((_, Ok(Ok(resp)))) | Ok((_, Ok(Err(resp)))) => resp, - Ok((_, Err(e))) => { + Ok((_, Err(_))) => { log::error!( - "Worker task failed while delivering to {}: {e}", + "Delivery to {} failed due to a dead worker!", domain.map(AsRef::as_ref).unwrap_or("") ); + if let Some(domain) = domain { + self.workers.cleanup(domain); + } LOCAL_ERROR_451.to_string() } Err(e) => { @@ -243,9 +250,16 @@ mod tests { /// Spawns filtermail-transport. /// /// Returns a pointer to the underlying handler. - fn spawn_filtermail_transport() -> TestResult>> { + fn spawn_filtermail_transport( + queue_size: usize, + pool_capacity: usize, + ) -> TestResult>> { let config = Config::default(); - let transport = Arc::new(TransportHandler::with_queue_size(config.clone(), 1)?); + let transport = Arc::new(TransportHandler::with_queue_size_and_pool_capacity( + config.clone(), + queue_size, + pool_capacity, + )?); tokio::spawn(run_smtp_server( &FILTERMAIL_ADDR, transport.clone(), @@ -302,8 +316,11 @@ mod tests { #[rstest] #[tokio::test] async fn test_rcpt_to_and_start_data(addrs1: Vec, addrs2: Vec) -> TestResult { - let transport_handler = - TransportHandler::::with_queue_size(Config::default(), 1)?; + let transport_handler = TransportHandler::::with_queue_size_and_pool_capacity( + Config::default(), + 1, + 500, + )?; let domain1 = AddressDomain::from_str(addrs1.first().unwrap())?; let domain2 = AddressDomain::from_str(addrs2.first().unwrap())?; @@ -358,7 +375,7 @@ mod tests { #[tokio::test] async fn test_smtp_send_mail() -> TestResult { let mut remote_mta = spawn_mock_mta()?; - spawn_filtermail_transport()?; + spawn_filtermail_transport(1, 2)?; let envelope = Envelope { mail_from: "sender@here".to_string(), @@ -398,7 +415,7 @@ mod tests { #[tokio::test] async fn test_smtp_send_mail_defer() -> TestResult { let mut remote_mta = spawn_mock_mta()?; - let transport = spawn_filtermail_transport()?; + let transport = spawn_filtermail_transport(1, 500)?; let mut envelope = Envelope { mail_from: "sender@here".to_string(), @@ -438,4 +455,48 @@ mod tests { Ok(()) } + + #[rstest] + #[serial] + #[tokio::test] + async fn test_smtp_send_mail_worker_cap_exceeded() -> TestResult { + let mut remote_mta = spawn_mock_mta()?; + spawn_filtermail_transport(30, 1)?; + + let envelope_1 = Envelope { + mail_from: "sender@here".to_string(), + rcpt_to: vec!["a1@localhost".to_string()], + data: "message\r\n".as_bytes().to_vec(), + }; + + let envelope_2 = Envelope { + mail_from: "sender@here".to_string(), + rcpt_to: vec!["b1@[127.0.0.1]".to_string()], + data: "message\r\n".as_bytes().to_vec(), + }; + + // 1 + let (record_postfix_1, record_filtermail_1) = { + let record_postfix = lmtp_send(&envelope_1).await?; + let record_filtermail = remote_mta.recv().await.unwrap(); + tokio::time::sleep(Duration::from_secs(1)).await; + assert!(remote_mta.is_empty()); + (record_postfix, record_filtermail) + }; + + // 2 + let record_postfix_2 = lmtp_send(&envelope_2).await?; + tokio::time::sleep(Duration::from_secs(1)).await; + assert!(remote_mta.is_empty()); + + insta::assert_snapshot!(format!( + "TRANSACTION 1\r\n\ + [postfix -> filtermail-transport]\r\n{record_postfix_1}\r\n\ + [filtermail-transport -> destination A]\r\n{record_filtermail_1}\r\n\r\n\ + TRANSACTION 2\r\n\ + [postfix -> filtermail-transport]\r\n{record_postfix_2}" + )); + + Ok(()) + } } diff --git a/filtermail/src/transport/worker.rs b/filtermail/src/transport/worker.rs index 0924589c..36c670e7 100644 --- a/filtermail/src/transport/worker.rs +++ b/filtermail/src/transport/worker.rs @@ -17,6 +17,7 @@ use tokio::sync::mpsc::OwnedPermit; use tokio::sync::{mpsc, oneshot}; use tokio::task; use tokio::task::JoinHandle; +use tokio::time::timeout; use tokio_rustls::rustls; #[cfg(not(test))] @@ -35,16 +36,25 @@ const SMTP_SKIP_TLS: bool = true; /// all new messages will be immediately deferred. const PER_DESTINATION_QUEUE_SIZE: usize = 30; +/// Max number of [`Worker`]s operating at the same time. +const MAX_WORKERS: usize = 500; + +/// How long a worker can stay idle. +/// +/// Exceeding this value will cause a worker shutdown. +const WORKER_KEEPALIVE_DURATION: Duration = Duration::from_secs(60); + type SMTPResponse = Result; pub struct WorkerPool { - inner: RwLock>>, + inner: Arc>>>, client_hostname: String, smtp_connection_pool: Arc>, mxdeliv_unsupported_hosts: Arc>, monitor_handle: JoinHandle<()>, dns_resolver: Arc, queue_size: usize, + capacity: usize, } impl WorkerPool @@ -72,48 +82,61 @@ where mxdeliv_unsupported_hosts: mxdeliv_cache, monitor_handle, queue_size: PER_DESTINATION_QUEUE_SIZE, + capacity: MAX_WORKERS, }) } - /// Same as [`Self::new`], but lets you set the size of the queue. + /// Same as [`Self::new`], but lets you set the size of the queue and pool capacity. /// /// Used only for tests. #[cfg(test)] - pub fn with_queue_size(config: Config, queue_size: usize) -> Result { + pub fn with_queue_size_and_pool_capacity( + config: Config, + queue_size: usize, + pool_capacity: usize, + ) -> Result { let mut this = Self::new(config)?; this.queue_size = queue_size; + this.capacity = pool_capacity; Ok(this) } - fn get_or_create_worker(&self, destination: &AddressDomain) -> Arc { + /// Gets a worker for provided `destination`, + /// spawning a new one if required. + /// + /// Returns [`None`], if operation requires spawning a new worker, + /// but pool already operates at maximum worker capacity. + fn get_or_create_worker(&self, destination: &AddressDomain) -> Option> { // NOTE: these locks are blocking, but critical section here is quite small and // shouldn't cause issues in async code. // NOTE: read() returns a guard that is dropped before the match statement. // This must be ensured or else, the write() line would cause a deadlock. let worker = { - let mut worker = { - let map = self.inner.read(); - map.get(destination).cloned() - }; - // Remove (and re-create) worker if it finished/crashed. - // In reality, this should never happen. - if let Some(w) = &worker - && w.handle.is_finished() - { - log::error!("Worker for destination {destination} crashed! Restarting...",); - worker = None; - { - let mut map = self.inner.write(); - map.remove(destination); - } - }; - worker + let map = self.inner.read(); + map.get(destination).cloned() }; match worker { - Some(worker) => worker, + Some(worker) => Some(worker), None => { // Worker for this destination wasn't spawned yet. + let mut map = self.inner.write(); + + // we check if it wasn't spawned in the meantime first + if let Some(worker) = map.get(destination) + && !worker.handle.is_finished() + { + return Some(worker.clone()); + } + + if map.len() >= self.capacity { + log::warn!( + "Worker pool operating at maximum capacity! \ + Messages to new destinations will be deferred." + ); + return None; + } + let (tx, rx) = mpsc::channel(self.queue_size); let handle = tokio::spawn(Worker::run( destination.clone(), @@ -122,24 +145,47 @@ where self.smtp_connection_pool.clone(), self.mxdeliv_unsupported_hosts.clone(), self.dns_resolver.clone(), + self.inner.clone(), )); log::trace!("Worker {} spawned", handle.id()); let worker = Arc::new(Worker { tx, handle }); - self.inner - .write() - .insert(destination.clone(), worker.clone()); - worker + map.insert(destination.clone(), worker.clone()); + Some(worker) } } } /// Tries to get an [`OwnedPermit`] to the worker for specified destination. /// - /// Returns [`None`] if the worker's queue is full. + /// Returns [`None`] if: + /// + /// - the worker's queue is full, + /// - operation requires spawning a new worker, + /// but pool already operates at maximum worker capacity. pub fn get_permit(&self, destination: &AddressDomain) -> Option> { - let worker = self.get_or_create_worker(destination); - worker.tx.clone().try_reserve_owned().ok() + if let Some(worker) = self.get_or_create_worker(destination) { + return worker.tx.clone().try_reserve_owned().ok(); + } + None + } + + /// Informs pool that a worker has exited unexpectedly. + /// + /// Such worker will be removed from the pool, + /// and re-created on the next [`Self::get_permit`] call. + pub fn cleanup(&self, destination: &AddressDomain) { + let mut map = self.inner.write(); + // we first check in case other task already removed/re-created it + if let Some(worker) = map.get(destination) + && worker.handle.is_finished() + { + log::info!( + "Removing a dead worker {} for destination {destination}", + worker.handle.id() + ); + map.remove(destination); + } } } @@ -169,6 +215,7 @@ impl Worker { smtp_connection_pool: Arc>, mxdeliv_unsupported_hosts: Arc>, dns_resolver: Arc, + worker_pool: Arc>>>, ) -> Result<(), crate::error::Error> where S: TcpStreamTrait + TcpConnect, @@ -182,7 +229,7 @@ impl Worker { let tls_resumption_store = Arc::new(rustls::client::ClientSessionMemoryCache::new(256)); let https_client = HttpsClient::new(tls_resumption_store.clone())?; - while let Some(message) = rx.recv().await { + while let Ok(Some(message)) = timeout(WORKER_KEEPALIVE_DURATION, rx.recv()).await { log::trace!( "Worker {worker_id} received a message from {}", message.envelope.mail_from @@ -205,6 +252,9 @@ impl Worker { }; } + log::info!("Worker {worker_id} for domain {destination} shutting down..."); + worker_pool.write().remove(&destination); + Ok(()) }