From 07ac754634f6cbeff22323c93fdecebb0fef04d2 Mon Sep 17 00:00:00 2001
From: link2xt <link2xt@testrun.org>
Date: Sun, 25 Aug 2024 07:45:40 +0000
Subject: [PATCH] feat: parallelize connection attempts

---
 src/imap/client.rs | 70 +++++++++++++++++++++++++++++-----------------
 1 file changed, 44 insertions(+), 26 deletions(-)

diff --git a/src/imap/client.rs b/src/imap/client.rs
index 74d64c166e..eb8ca61271 100644
--- a/src/imap/client.rs
+++ b/src/imap/client.rs
@@ -106,6 +106,38 @@ impl Client {
         Ok(Session::new(session, capabilities))
     }
 
+    async fn connection_attempt(
+        context: &Context,
+        host: &str,
+        security: ConnectionSecurity,
+        resolved_addr: SocketAddr,
+        strict_tls: bool,
+        load_cache: bool,
+    ) -> Result<Self> {
+        let res = match security {
+            ConnectionSecurity::Tls => Client::connect_secure(resolved_addr, host, strict_tls).await,
+            ConnectionSecurity::Starttls => {
+                Client::connect_starttls(resolved_addr, host, strict_tls).await
+            }
+            ConnectionSecurity::Plain => Client::connect_insecure(resolved_addr).await,
+        };
+        match res {
+            Ok(client) => {
+                let ip_addr = resolved_addr.ip().to_string();
+                let port = resolved_addr.port();
+                if load_cache {
+                    update_connect_timestamp(context, host, &ip_addr).await?;
+                }
+                update_connection_history(context, "imap", host, port, &ip_addr, time()).await?;
+                Ok(client)
+            }
+            Err(err) => {
+                warn!(context, "Failed to connect to {resolved_addr}: {err:#}.");
+                Err(err)
+            }
+        }
+    }
+
     pub async fn connect(
         context: &Context,
         socks5_config: Option<Socks5Config>,
@@ -131,40 +163,26 @@ impl Client {
             };
             Ok(client)
         } else {
-            let mut first_error = None;
             let load_cache = match security {
                 ConnectionSecurity::Tls | ConnectionSecurity::Starttls => strict_tls,
                 ConnectionSecurity::Plain => false,
             };
+
+            let mut connection_attempts = Vec::new();
             for resolved_addr in
                 lookup_host_with_cache(context, host, port, "imap", load_cache).await?
             {
-                let res = match security {
-                    ConnectionSecurity::Tls => {
-                        Client::connect_secure(resolved_addr, host, strict_tls).await
-                    }
-                    ConnectionSecurity::Starttls => {
-                        Client::connect_starttls(resolved_addr, host, strict_tls).await
-                    }
-                    ConnectionSecurity::Plain => Client::connect_insecure(resolved_addr).await,
-                };
-                match res {
-                    Ok(client) => {
-                        let ip_addr = resolved_addr.ip().to_string();
-                        if load_cache {
-                            update_connect_timestamp(context, host, &ip_addr).await?;
-                        }
-                        update_connection_history(context, "imap", host, port, &ip_addr, time())
-                            .await?;
-                        return Ok(client);
-                    }
-                    Err(err) => {
-                        warn!(context, "Failed to connect to {resolved_addr}: {err:#}.");
-                        first_error.get_or_insert(err);
-                    }
-                }
+                let fut = Box::pin(Self::connection_attempt(context, host, security, resolved_addr, strict_tls, load_cache));
+                connection_attempts.push(fut);
             }
-            Err(first_error.unwrap_or_else(|| format_err!("no DNS resolution results for {host}")))
+
+            if connection_attempts.is_empty() {
+                return Err(format_err!("no DNS resolution results for {host}"));
+            }
+
+            connection_attempts.reverse();
+            let (res, _remaining_futures) = futures::future::select_ok(connection_attempts).await?;
+            Ok(res)
         }
     }