From dada3902815d5660e2e79c3b99e1b561b9235b54 Mon Sep 17 00:00:00 2001
From: Hansie Odendaal <pluto@tari.com>
Date: Tue, 26 Nov 2024 14:00:23 +0200
Subject: [PATCH] Add connectivity timeout

Added a wallet connectivity timeout so long-winded failures can be managed timeously.
---
 .../wallet/src/connectivity_service/service.rs | 18 ++++++++++++++----
 comms/core/src/connectivity/error.rs           |  2 ++
 2 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/base_layer/wallet/src/connectivity_service/service.rs b/base_layer/wallet/src/connectivity_service/service.rs
index ab5ac039bf..b5159d4b05 100644
--- a/base_layer/wallet/src/connectivity_service/service.rs
+++ b/base_layer/wallet/src/connectivity_service/service.rs
@@ -34,7 +34,7 @@ use tari_core::base_node::{rpc::BaseNodeWalletRpcClient, sync::rpc::BaseNodeSync
 use tokio::{
     sync::{mpsc, oneshot, watch},
     time,
-    time::MissedTickBehavior,
+    time::{timeout, Duration as TokioDuration, MissedTickBehavior},
 };
 
 use crate::{
@@ -392,12 +392,22 @@ impl WalletConnectivityService {
     }
 
     async fn try_setup_rpc_pool(&mut self, peer_node_id: NodeId) -> Result<bool, WalletConnectivityError> {
-        let conn = match self.try_dial_peer(peer_node_id.clone()).await? {
-            Some(c) => c,
-            None => {
+        let dial_timeout = TokioDuration::from_secs(40);
+        let conn = match timeout(dial_timeout, self.try_dial_peer(peer_node_id.clone())).await {
+            Ok(Ok(Some(c))) => c,
+            Ok(Ok(None)) => {
                 warn!(target: LOG_TARGET, "Could not dial base node peer '{}'", peer_node_id);
                 return Ok(false);
             },
+            Ok(Err(e)) => return Err(e),
+            Err(_) => {
+                return Err(WalletConnectivityError::ConnectivityError(
+                    ConnectivityError::ClientCancelled(format!(
+                        "Could not connect to '{}' in {:?}",
+                        peer_node_id, dial_timeout
+                    )),
+                ));
+            },
         };
         debug!(
             target: LOG_TARGET,
diff --git a/comms/core/src/connectivity/error.rs b/comms/core/src/connectivity/error.rs
index acb554ac79..5dbfe2849d 100644
--- a/comms/core/src/connectivity/error.rs
+++ b/comms/core/src/connectivity/error.rs
@@ -43,6 +43,8 @@ pub enum ConnectivityError {
     OnlineWaitTimeout(usize),
     #[error("Pending dial was cancelled")]
     DialCancelled,
+    #[error("Client cancelled: '{0}'")]
+    ClientCancelled(String),
 }
 
 impl From<ConnectionManagerError> for ConnectivityError {