feat(transport): Worker eviction (#185)

Adds max worker capacity of 500 to the worker pool,
and an automatic shutdown of idle workers.
Messages that would cause the capacity to be exceeded,
are deferred.

Additionally, ensures that the same worker
is not spawned by two tasks at the same time.

Signed-off-by: Jagoda Ślązak <jslazak@jslazak.com>

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";

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

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:<sender@here>
+< 250 OK
+> RCPT TO:<a1@localhost>
+< 250 OK
+> DATA
+< 354 End data with <CR><LF>.<CR><LF>
+> message
+> .
+< 250 OK (SMTP)
+
+[filtermail-transport -> destination A]
+< 220 filtermail SMTP
+> EHLO example.org
+< 250-filtermail
+250-8BITMIME
+250 OK
+> MAIL FROM:<sender@here>
+< 250 OK
+> RCPT TO:<a1@localhost>
+< 250 OK
+> DATA
+< 354 End data with <CR><LF>.<CR><LF>
+> message
+.
+< 250 OK
+
+
+TRANSACTION 2
+[postfix -> filtermail-transport]
+< 220 filtermail SMTP
+> LHLO postfix
+< 250-filtermail
+250-8BITMIME
+250 OK
+> MAIL FROM:<sender@here>
+< 250 OK
+> RCPT TO:<b1@[127.0.0.1]>
+< 250 OK
+> DATA
+< 421 Transport busy

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<Self, crate::error::Error> {
+    pub fn with_queue_size_and_pool_capacity(
-        let workers = WorkerPool::with_queue_size(config, queue_size)?;
+        config: Config,
+        queue_size: usize,
+        pool_capacity: usize,
+    ) -> Result<Self, crate::error::Error> {
+        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("<unknown>")
                     );
+                    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<Arc<TransportHandler<TcpStream>>> {
+    fn spawn_filtermail_transport(
+        queue_size: usize,
+        pool_capacity: usize,
+    ) -> TestResult<Arc<TransportHandler<TcpStream>>> {
         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<String>, addrs2: Vec<String>) -> TestResult {
-        let transport_handler =
+        let transport_handler = TransportHandler::<TcpStream>::with_queue_size_and_pool_capacity(
-            TransportHandler::<TcpStream>::with_queue_size(Config::default(), 1)?;
+            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(())
+    }
 }

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<String, String>;
 
 pub struct WorkerPool<S: TcpConnect> {
-    inner: RwLock<BTreeMap<AddressDomain, Arc<Worker>>>,
+    inner: Arc<RwLock<BTreeMap<AddressDomain, Arc<Worker>>>>,
     client_hostname: String,
     smtp_connection_pool: Arc<SmtpConnectionPool<S>>,
     mxdeliv_unsupported_hosts: Arc<retainer::Cache<String, ()>>,
     monitor_handle: JoinHandle<()>,
     dns_resolver: Arc<TokioResolver>,
     queue_size: usize,
+    capacity: usize,
 }
 
 impl<S> WorkerPool<S>
@@ -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<Self, crate::error::Error> {
+    pub fn with_queue_size_and_pool_capacity(
+        config: Config,
+        queue_size: usize,
+        pool_capacity: usize,
+    ) -> Result<Self, crate::error::Error> {
         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<Worker> {
+    /// 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<Arc<Worker>> {
         // 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();
-                let map = self.inner.read();
+            map.get(destination).cloned()
-                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
         };
 
         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
+                map.insert(destination.clone(), worker.clone());
-                    .write()
+                Some(worker)
-                    .insert(destination.clone(), worker.clone());
-                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<OwnedPermit<WorkerMessage>> {
-        let worker = self.get_or_create_worker(destination);
+        if let Some(worker) = self.get_or_create_worker(destination) {
-        worker.tx.clone().try_reserve_owned().ok()
+            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<SmtpConnectionPool<S>>,
         mxdeliv_unsupported_hosts: Arc<retainer::Cache<String, ()>>,
         dns_resolver: Arc<TokioResolver>,
+        worker_pool: Arc<RwLock<BTreeMap<AddressDomain, Arc<Worker>>>>,
     ) -> 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(())
     }