From b43c69af47b97e61f39a523ce65e016df9b00f8b Mon Sep 17 00:00:00 2001
From: Paolo Barbolini <paolo@paolo565.org>
Date: Tue, 6 Oct 2020 09:56:36 +0200
Subject: [PATCH] Allow changing defaults for the connection pool configuration

---
 src/transport/smtp/mod.rs       |  5 ++-
 src/transport/smtp/pool.rs      | 70 ++++++++++++++++++++++++++++++++-
 src/transport/smtp/transport.rs | 35 ++++++++++-------
 3 files changed, 93 insertions(+), 17 deletions(-)

diff --git a/src/transport/smtp/mod.rs b/src/transport/smtp/mod.rs
index 3d7bec6..5bc89c1 100644
--- a/src/transport/smtp/mod.rs
+++ b/src/transport/smtp/mod.rs
@@ -158,6 +158,8 @@
 pub use self::async_transport::{
     AsyncSmtpConnector, AsyncSmtpTransport, AsyncSmtpTransportBuilder, Tokio02Connector,
 };
+#[cfg(feature = "r2d2")]
+pub use self::pool::PoolConfig;
 pub(crate) use self::transport::SmtpClient;
 pub use self::{
     error::Error,
@@ -182,8 +184,7 @@ pub mod commands;
 mod error;
 pub mod extension;
 #[cfg(feature = "r2d2")]
-#[cfg_attr(docsrs, doc(cfg(feature = "r2d2")))]
-pub mod pool;
+mod pool;
 pub mod response;
 mod transport;
 pub mod util;
diff --git a/src/transport/smtp/pool.rs b/src/transport/smtp/pool.rs
index 6f01d0c..92fae3a 100644
--- a/src/transport/smtp/pool.rs
+++ b/src/transport/smtp/pool.rs
@@ -1,5 +1,73 @@
+use std::time::Duration;
+
 use crate::transport::smtp::{client::SmtpConnection, error::Error, SmtpClient};
-use r2d2::ManageConnection;
+
+use r2d2::{ManageConnection, Pool};
+
+/// Configuration for a connection pool
+#[derive(Debug, Clone)]
+#[allow(missing_copy_implementations)]
+#[cfg_attr(docsrs, doc(cfg(feature = "r2d2")))]
+pub struct PoolConfig {
+    min_idle: u32,
+    max_size: u32,
+    connection_timeout: Duration,
+    idle_timeout: Duration,
+}
+
+impl PoolConfig {
+    /// Minimum number of idle connections
+    ///
+    /// Defaults to `0`
+    pub fn min_idle(mut self, min_idle: u32) -> Self {
+        self.min_idle = min_idle;
+        self
+    }
+
+    /// Maximum number of pooled connections
+    ///
+    /// Defaults to `10`
+    pub fn max_size(mut self, max_size: u32) -> Self {
+        self.min_idle = max_size;
+        self
+    }
+
+    /// Connection timeout
+    ///
+    /// Defaults to `30 seconds`
+    pub fn connection_timeout(mut self, connection_timeout: Duration) -> Self {
+        self.connection_timeout = connection_timeout;
+        self
+    }
+
+    /// Connection idle timeout
+    ///
+    /// Defaults to `60 seconds`
+    pub fn idle_timeout(mut self, idle_timeout: Duration) -> Self {
+        self.idle_timeout = idle_timeout;
+        self
+    }
+
+    pub(crate) fn build<C: ManageConnection>(&self, client: C) -> Pool<C> {
+        Pool::builder()
+            .min_idle(Some(self.min_idle))
+            .max_size(self.max_size)
+            .connection_timeout(self.connection_timeout)
+            .idle_timeout(Some(self.idle_timeout))
+            .build_unchecked(client)
+    }
+}
+
+impl Default for PoolConfig {
+    fn default() -> Self {
+        Self {
+            min_idle: 0,
+            max_size: 10,
+            connection_timeout: Duration::from_secs(30),
+            idle_timeout: Duration::from_secs(60),
+        }
+    }
+}
 
 impl ManageConnection for SmtpClient {
     type Connection = SmtpConnection;
diff --git a/src/transport/smtp/transport.rs b/src/transport/smtp/transport.rs
index 7f7d5b0..1e1ee05 100644
--- a/src/transport/smtp/transport.rs
+++ b/src/transport/smtp/transport.rs
@@ -3,6 +3,8 @@ use std::time::Duration;
 #[cfg(feature = "r2d2")]
 use r2d2::Pool;
 
+#[cfg(feature = "r2d2")]
+use super::PoolConfig;
 use super::{ClientId, Credentials, Error, Mechanism, Response, SmtpConnection, SmtpInfo};
 #[cfg(any(feature = "native-tls", feature = "rustls-tls"))]
 use super::{Tls, TlsParameters, SUBMISSIONS_PORT, SUBMISSION_PORT};
@@ -95,7 +97,12 @@ impl SmtpTransport {
     pub fn builder_dangerous<T: Into<String>>(server: T) -> SmtpTransportBuilder {
         let mut new = SmtpInfo::default();
         new.server = server.into();
-        SmtpTransportBuilder { info: new }
+
+        SmtpTransportBuilder {
+            info: new,
+            #[cfg(feature = "r2d2")]
+            pool_config: PoolConfig::default(),
+        }
     }
 }
 
@@ -104,6 +111,8 @@ impl SmtpTransport {
 #[derive(Clone)]
 pub struct SmtpTransportBuilder {
     info: SmtpInfo,
+    #[cfg(feature = "r2d2")]
+    pool_config: PoolConfig,
 }
 
 /// Builder for the SMTP `SmtpTransport`
@@ -145,27 +154,25 @@ impl SmtpTransportBuilder {
         self
     }
 
-    /// Build the client
-    fn build_client(self) -> SmtpClient {
-        SmtpClient { info: self.info }
+    /// Use a custom configuration for the connection pool
+    ///
+    /// Defaults can be found at [`PoolConfig`]
+    #[cfg(feature = "r2d2")]
+    #[cfg_attr(docsrs, doc(cfg(feature = "r2d2")))]
+    pub fn pool_config(mut self, pool_config: PoolConfig) -> Self {
+        self.pool_config = pool_config;
+        self
     }
 
     /// Build the transport
     ///
     /// If the `r2d2` feature is enabled an `Arc` wrapped pool is be created.
-    /// Defaults:
-    ///
-    /// * 60 seconds idle timeout
-    /// * 30 minutes max connection lifetime
-    /// * max pool size of 10 connections
+    /// Defaults can be found at [`PoolConfig`]
     pub fn build(self) -> SmtpTransport {
-        let client = self.build_client();
+        let client = SmtpClient { info: self.info };
         SmtpTransport {
             #[cfg(feature = "r2d2")]
-            inner: Pool::builder()
-                .min_idle(Some(0))
-                .idle_timeout(Some(Duration::from_secs(60)))
-                .build_unchecked(client),
+            inner: self.pool_config.build(client),
             #[cfg(not(feature = "r2d2"))]
             inner: client,
         }