From 3cbef8f0e14ef5a987975e62a6a4044e7723c98c Mon Sep 17 00:00:00 2001
From: Nick Steel <nick@nsteel.co.uk>
Date: Fri, 20 Sep 2024 00:29:44 +0100
Subject: [PATCH] core: retry session connection

---
 core/src/session.rs | 38 +++++++++++++++++++++++++-------------
 1 file changed, 25 insertions(+), 13 deletions(-)

diff --git a/core/src/session.rs b/core/src/session.rs
index 3944ce83e..5f7063a34 100644
--- a/core/src/session.rs
+++ b/core/src/session.rs
@@ -144,7 +144,7 @@ impl Session {
 
     async fn connect_inner(
         &self,
-        access_point: SocketAddress,
+        access_point: &SocketAddress,
         credentials: Credentials,
     ) -> Result<(Credentials, Transport), Error> {
         let mut transport = connection::connect(
@@ -187,19 +187,30 @@ impl Session {
         credentials: Credentials,
         store_credentials: bool,
     ) -> Result<(), Error> {
-        let (reusable_credentials, transport) = loop {
+        const MAX_RETRIES: u8 = 2;
+        let (reusable_credentials, transport) = 'ap: loop {
             let ap = self.apresolver().resolve("accesspoint").await?;
             info!("Connecting to AP \"{}:{}\"", ap.0, ap.1);
-            match self.connect_inner(ap, credentials.clone()).await {
-                Ok(ct) => break ct,
-                Err(e) => {
-                    if let Some(AuthenticationError::LoginFailed(ErrorCode::TryAnotherAP)) =
-                        e.error.downcast_ref::<AuthenticationError>()
-                    {
-                        warn!("Instructed to try another access point...");
-                        continue;
-                    } else {
-                        return Err(e);
+            let mut num_retries = 0;
+            loop {
+                match self.connect_inner(&ap, credentials.clone()).await {
+                    Ok(ct) => break 'ap ct,
+                    Err(e) => {
+                        if let Some(AuthenticationError::LoginFailed(ErrorCode::TryAnotherAP)) =
+                            e.error.downcast_ref::<AuthenticationError>()
+                        {
+                            warn!("Instructed to try another access point...");
+                            continue 'ap;
+                        } else {
+                            debug!("Connection failed: {e}");
+                            if num_retries == MAX_RETRIES {
+                                warn!("Retries exceeded, try another access point...");
+                                continue 'ap;
+                            } else {
+                                num_retries += 1;
+                                trace!("Retry access point...");
+                            }
+                        }
                     }
                 }
             }
@@ -567,7 +578,7 @@ impl Session {
     }
 
     pub fn shutdown(&self) {
-        debug!("Invalidating session");
+        debug!("Shutdown: Invalidating session");
         self.0.data.write().invalid = true;
         self.mercury().shutdown();
         self.channel().shutdown();
@@ -624,6 +635,7 @@ where
                     return Poll::Ready(Ok(()));
                 }
                 Some(Err(e)) => {
+                    error!("Connection to server closed.");
                     session.shutdown();
                     return Poll::Ready(Err(e));
                 }