Connection fd callback before actual connection (#5494)

common/bandwidth-controller/src/lib.rs

@@ -105,26 +105,24 @@ impl<C, St: Storage> BandwidthController<C, St> {
     async fn get_aggregate_verification_key(
         &self,
         epoch_id: EpochId,
-        apis: &mut ApiClientsWrapper,
+        ecash_apis: &mut ApiClientsWrapper<'_, C>,
     ) -> Result<VerificationKeyAuth, BandwidthControllerError>
     where
         C: DkgQueryClient + Sync + Send,
         <St as Storage>::StorageError: Send + Sync + 'static,
     {
-        let ecash_apis = apis.get_or_init(epoch_id, &self.client).await?;
         get_aggregate_verification_key(&self.storage, epoch_id, ecash_apis).await
     }
 
     async fn get_coin_index_signatures(
         &self,
         epoch_id: EpochId,
-        apis: &mut ApiClientsWrapper,
+        ecash_apis: &mut ApiClientsWrapper<'_, C>,
     ) -> Result<Vec<AnnotatedCoinIndexSignature>, BandwidthControllerError>
     where
         C: DkgQueryClient + Sync + Send,
         <St as Storage>::StorageError: Send + Sync + 'static,
     {
-        let ecash_apis = apis.get_or_init(epoch_id, &self.client).await?;
         get_coin_index_signatures(&self.storage, epoch_id, ecash_apis).await
     }
 
@@ -132,13 +130,12 @@ impl<C, St: Storage> BandwidthController<C, St> {
         &self,
         epoch_id: EpochId,
         expiration_date: Date,
-        apis: &mut ApiClientsWrapper,
+        ecash_apis: &mut ApiClientsWrapper<'_, C>,
     ) -> Result<Vec<AnnotatedExpirationDateSignature>, BandwidthControllerError>
     where
         C: DkgQueryClient + Sync + Send,
         <St as Storage>::StorageError: Send + Sync + 'static,
     {
-        let ecash_apis = apis.get_or_init(epoch_id, &self.client).await?;
         get_expiration_date_signatures(&self.storage, epoch_id, expiration_date, ecash_apis).await
     }
 
@@ -154,7 +151,7 @@ impl<C, St: Storage> BandwidthController<C, St> {
     {
         let epoch_id = retrieved_ticketbook.ticketbook.epoch_id();
         let expiration_date = retrieved_ticketbook.ticketbook.expiration_date();
-        let mut api_clients = Default::default();
+        let mut api_clients = ApiClientsWrapper::new(&self.client, epoch_id);
 
         let verification_key = self
             .get_aggregate_verification_key(epoch_id, &mut api_clients)

common/bandwidth-controller/src/utils.rs

@@ -21,30 +21,67 @@ use rand::thread_rng;
 use std::fmt::Display;
 use std::future::Future;
 
-// it really doesn't need the RwLock because it's never moved across tasks,
+pub(crate) trait EcashClientsProvider {
-// but we need all the Send/Sync action
+    async fn try_get_ecash_clients(
-#[derive(Default)]
-pub(crate) struct ApiClientsWrapper(Option<Vec<EcashApiClient>>);
-
-impl ApiClientsWrapper {
-    pub(crate) async fn get_or_init<C>(
         &mut self,
+    ) -> Result<Vec<EcashApiClient>, BandwidthControllerError>;
+}
+
+impl EcashClientsProvider for Vec<EcashApiClient> {
+    async fn try_get_ecash_clients(
+        &mut self,
+    ) -> Result<Vec<EcashApiClient>, BandwidthControllerError> {
+        Ok(self.clone())
+    }
+}
+
+impl<C> EcashClientsProvider for &mut ApiClientsWrapper<'_, C>
+where
+    C: DkgQueryClient + Sync + Send,
+{
+    async fn try_get_ecash_clients(
+        &mut self,
+    ) -> Result<Vec<EcashApiClient>, BandwidthControllerError> {
+        self.clients().await
+    }
+}
+
+pub(crate) enum ApiClientsWrapper<'a, C> {
+    Uninitialised {
+        query_client: &'a C,
         epoch_id: EpochId,
-        dkg_client: &C,
+    },
-    ) -> Result<Vec<EcashApiClient>, BandwidthControllerError>
+    Cached {
+        clients: Vec<EcashApiClient>,
+    },
+}
+
+impl<'a, C> ApiClientsWrapper<'a, C> {
+    pub(crate) fn new(query_client: &'a C, epoch_id: EpochId) -> Self {
+        ApiClientsWrapper::Uninitialised {
+            query_client,
+            epoch_id,
+        }
+    }
+
+    async fn clients(&mut self) -> Result<Vec<EcashApiClient>, BandwidthControllerError>
     where
         C: DkgQueryClient + Sync + Send,
     {
-        if let Some(cached) = &self.0 {
+        match self {
-            return Ok(cached.clone());
+            ApiClientsWrapper::Uninitialised {
+                query_client,
+                epoch_id,
+            } => {
+                let clients = all_ecash_api_clients(*query_client, *epoch_id).await?;
+                *self = ApiClientsWrapper::Cached {
+                    clients: clients.clone(),
+                };
+
+                Ok(clients)
+            }
+            ApiClientsWrapper::Cached { clients } => Ok(clients.clone()),
         }
-
-        let clients = all_ecash_api_clients(dkg_client, epoch_id).await?;
-
-        // technically we don't have to be cloning all the clients here, but it's way simpler than
-        // dealing with locking and whatnot given the performance penalty is negligible
-        self.0 = Some(clients.clone());
-        Ok(clients)
     }
 }
 
@@ -76,7 +113,7 @@ where
 pub(crate) async fn get_aggregate_verification_key<St>(
     storage: &St,
     epoch_id: EpochId,
-    ecash_apis: Vec<EcashApiClient>,
+    mut ecash_apis: impl EcashClientsProvider,
 ) -> Result<VerificationKeyAuth, BandwidthControllerError>
 where
     St: Storage,
@@ -90,6 +127,8 @@ where
         return Ok(stored);
     };
 
+    let ecash_apis = ecash_apis.try_get_ecash_clients().await?;
+
     let master_vk = query_random_apis_until_success(
         ecash_apis,
         |api| async move { api.api_client.master_verification_key(Some(epoch_id)).await },
@@ -115,7 +154,7 @@ where
 pub(crate) async fn get_coin_index_signatures<St>(
     storage: &St,
     epoch_id: EpochId,
-    ecash_apis: Vec<EcashApiClient>,
+    mut ecash_apis: impl EcashClientsProvider,
 ) -> Result<Vec<AnnotatedCoinIndexSignature>, BandwidthControllerError>
 where
     St: Storage,
@@ -129,6 +168,8 @@ where
         return Ok(stored);
     };
 
+    let ecash_apis = ecash_apis.try_get_ecash_clients().await?;
+
     let index_sigs = query_random_apis_until_success(
         ecash_apis,
         |api| async move {
@@ -159,7 +200,7 @@ pub(crate) async fn get_expiration_date_signatures<St>(
     storage: &St,
     epoch_id: EpochId,
     expiration_date: Date,
-    ecash_apis: Vec<EcashApiClient>,
+    mut ecash_apis: impl EcashClientsProvider,
 ) -> Result<Vec<AnnotatedExpirationDateSignature>, BandwidthControllerError>
 where
     St: Storage,
@@ -173,6 +214,8 @@ where
         return Ok(stored);
     };
 
+    let ecash_apis = ecash_apis.try_get_ecash_clients().await?;
+
     let expiration_sigs = query_random_apis_until_success(
         ecash_apis,
         |api| async move {

common/client-libs/gateway-client/src/client/mod.rs

@@ -204,15 +204,15 @@ impl<C, St> GatewayClient<C, St> {
             "Attemting to establish connection to gateway at: {}",
             self.gateway_address
         );
-        let (ws_stream, _) = connect_async(&self.gateway_address).await?;
+        let (ws_stream, _) = connect_async(
+            &self.gateway_address,
+            #[cfg(unix)]
+            self.connection_fd_callback.clone(),
+        )
+        .await?;
 
         self.connection = SocketState::Available(Box::new(ws_stream));
 
-        #[cfg(unix)]
-        if let (Some(callback), Some(fd)) = (self.connection_fd_callback.as_ref(), self.ws_fd()) {
-            callback.as_ref()(fd);
-        }
-
         Ok(())
     }
 

common/client-libs/gateway-client/src/client/websockets.rs

@@ -1,6 +1,11 @@
 use crate::error::GatewayClientError;
 
 use nym_http_api_client::HickoryDnsResolver;
+#[cfg(unix)]
+use std::{
+    os::fd::{AsRawFd, RawFd},
+    sync::Arc,
+};
 use tokio::net::TcpStream;
 use tokio_tungstenite::{MaybeTlsStream, WebSocketStream};
 use tungstenite::handshake::client::Response;
@@ -11,7 +16,10 @@ use std::net::SocketAddr;
 #[cfg(not(target_arch = "wasm32"))]
 pub(crate) async fn connect_async(
     endpoint: &str,
+    #[cfg(unix)] connection_fd_callback: Option<Arc<dyn Fn(RawFd) + Send + Sync>>,
 ) -> Result<(WebSocketStream<MaybeTlsStream<TcpStream>>, Response), GatewayClientError> {
+    use tokio::net::TcpSocket;
+
     let resolver = HickoryDnsResolver::default();
     let uri =
         Url::parse(endpoint).map_err(|_| GatewayClientError::InvalidUrl(endpoint.to_owned()))?;
@@ -37,14 +45,41 @@ pub(crate) async fn connect_async(
         }
     };
 
-    let stream = TcpStream::connect(&sock_addrs[..]).await.map_err(|error| {
+    let mut stream = Err(GatewayClientError::NoEndpointForConnection {
-        GatewayClientError::NetworkConnectionFailed {
+        address: endpoint.to_owned(),
-            address: endpoint.to_owned(),
+    });
-            source: error.into(),
+    for sock_addr in sock_addrs {
+        let socket = if sock_addr.is_ipv4() {
+            TcpSocket::new_v4()
+        } else {
+            TcpSocket::new_v6()
         }
-    })?;
+        .map_err(|err| GatewayClientError::NetworkConnectionFailed {
+            address: endpoint.to_owned(),
+            source: err.into(),
+        })?;
 
-    tokio_tungstenite::client_async_tls(endpoint, stream)
+        #[cfg(unix)]
+        if let Some(callback) = connection_fd_callback.as_ref() {
+            callback.as_ref()(socket.as_raw_fd());
+        }
+
+        match socket.connect(sock_addr).await {
+            Ok(s) => {
+                stream = Ok(s);
+                break;
+            }
+            Err(err) => {
+                stream = Err(GatewayClientError::NetworkConnectionFailed {
+                    address: endpoint.to_owned(),
+                    source: err.into(),
+                });
+                continue;
+            }
+        }
+    }
+
+    tokio_tungstenite::client_async_tls(endpoint, stream?)
         .await
         .map_err(|error| GatewayClientError::NetworkConnectionFailed {
             address: endpoint.to_owned(),

common/client-libs/gateway-client/src/error.rs

@@ -43,6 +43,9 @@ pub enum GatewayClientError {
     #[error("connection failed: {address}: {source}")]
     NetworkConnectionFailed { address: String, source: WsError },
 
+    #[error("no socket address for endpoint: {address}")]
+    NoEndpointForConnection { address: String },
+
     #[error("Invalid URL: {0}")]
     InvalidUrl(String),
 

nym-wallet/Cargo.lock

@@ -165,6 +165,19 @@ version = "0.3.7"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "6b4930d2cb77ce62f89ee5d5289b4ac049559b1c45539271f5ed4fdc7db34545"
 
+[[package]]
+name = "async-compression"
+version = "0.4.18"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "df895a515f70646414f4b45c0b79082783b80552b373a68283012928df56f522"
+dependencies = [
+ "flate2",
+ "futures-core",
+ "memchr",
+ "pin-project-lite",
+ "tokio",
+]
+
 [[package]]
 name = "async-trait"
 version = "0.1.86"
@@ -560,9 +573,9 @@ dependencies = [
 
 [[package]]
 name = "celes"
-version = "2.4.0"
+version = "2.5.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "39b9a21273925d7cc9e8a9a5f068122341336813c607014f5ef64f82b6acba58"
+checksum = "54441489dce7026efc8f01d1aa996c23fa39dd615a953d0e934433a42f61dd30"
 dependencies = [
  "serde",
 ]
@@ -724,13 +737,12 @@ dependencies = [
 
 [[package]]
 name = "colored"
-version = "2.0.4"
+version = "2.2.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "2674ec482fbc38012cf31e6c42ba0177b431a0cb6f15fe40efa5aab1bda516f6"
+checksum = "117725a109d387c937a1533ce01b450cbde6b88abceea8473c4d7a85853cda3c"
 dependencies = [
- "is-terminal",
  "lazy_static",
- "windows-sys 0.48.0",
+ "windows-sys 0.52.0",
 ]
 
 [[package]]
@@ -3715,7 +3727,7 @@ dependencies = [
  "base64 0.22.1",
  "bip32",
  "bip39",
- "colored 2.0.4",
+ "colored 2.2.0",
  "cosmrs 0.21.1",
  "cosmwasm-std",
  "cw-controllers",
@@ -3823,7 +3835,7 @@ dependencies = [
  "base64 0.13.1",
  "bip39",
  "cfg-if",
- "colored 2.0.4",
+ "colored 2.2.0",
  "cosmrs 0.21.1",
  "cosmwasm-std",
  "dirs 4.0.0",
@@ -4758,6 +4770,7 @@ version = "0.12.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "566cafdd92868e0939d3fb961bd0dc25fcfaaed179291093b3d43e6b3150ea10"
 dependencies = [
+ "async-compression",
  "base64 0.22.1",
  "bytes",
  "encoding_rs",
@@ -4790,6 +4803,7 @@ dependencies = [
  "tokio",
  "tokio-native-tls",
  "tokio-rustls 0.25.0",
+ "tokio-util",
  "tower-service",
  "url",
  "wasm-bindgen",