Skip to content

Commit

Permalink
added test start fn to client_pool: uses specified gw for all clients
Browse files Browse the repository at this point in the history
  • Loading branch information
mfahampshire committed Jan 3, 2025
1 parent 376eca6 commit 93f063b
Showing 1 changed file with 47 additions and 0 deletions.
47 changes: 47 additions & 0 deletions sdk/rust/nym-sdk/src/client_pool/mixnet_client_pool.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,9 @@
// Copyright 2024 - Nym Technologies SA <contact@nymtech.net>
// SPDX-License-Identifier: GPL-3.0-only

use crate::mixnet::{MixnetClient, MixnetClientBuilder, NymNetworkDetails};
use anyhow::Result;
use nym_crypto::asymmetric::ed25519;
use std::fmt;
use std::sync::Arc;
use tokio::sync::RwLock;
Expand Down Expand Up @@ -108,6 +112,49 @@ impl ClientPool {
}
}

// Even though this is basically start() with an extra param since I think this
// will only be used for testing scenarios, and I didn't want to unnecessarily add
// another param to the function that will be used elsewhere, hence this is its own fn
pub async fn start_with_specified_gateway(&self, gateway: ed25519::PublicKey) -> Result<()> {
loop {
let spawned_clients = self.clients.read().await.len();
let addresses = self;
debug!(
"Currently spawned clients: {}: {:?}",
spawned_clients, addresses
);
if self.cancel_token.is_cancelled() {
break Ok(());
}
if spawned_clients >= self.client_pool_reserve_number {
debug!("Got enough clients already: sleeping");
} else {
info!(
"Clients in reserve = {}, reserve amount = {}, spawning new client",
spawned_clients, self.client_pool_reserve_number
);
let client = loop {
let net = NymNetworkDetails::new_from_env();
match MixnetClientBuilder::new_ephemeral()
.network_details(net)
.request_gateway(gateway.to_string())
.build()?
.connect_to_mixnet()
.await
{
Ok(client) => break client,
Err(err) => {
warn!("Error creating client: {:?}, will retry in 100ms", err);
tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
}
}
};
self.clients.write().await.push(Arc::new(client));
}
tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
}
}

pub async fn disconnect_pool(&self) {
info!("Triggering Client Pool disconnect");
self.cancel_token.cancel();
Expand Down

0 comments on commit 93f063b

Please sign in to comment.