GRIN/nym
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

first pass disconnect

Browse Source
This commit is contained in:
mfahampshire
2024-11-26 12:13:13 +01:00
parent 8e8eceb894
commit 415c30fcc9
3 changed files with 42 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sdk/rust/nym-sdk/examples/tcp_proxy_single_connection.rs
+9 -4
View File
@@ -31,7 +31,6 @@ struct ExampleMessage {
#[tokio::main]
async fn main() -> anyhow::Result<()> {
// Keep track of sent/received messages
// let counter = Arc::new(Mutex::new(0));
let counter = AtomicU8::new(0);
// Comment this out to just see println! statements from this example, as Nym client logging is very informative but quite verbose.
@@ -60,9 +59,15 @@ async fn main() -> anyhow::Result<()> {
// We'll run the instance with a long timeout since we're sending everything down the same Tcp connection, so should be using a single session.
// Within the TcpProxyClient, individual client shutdown is triggered by the timeout.
let proxy_client =
tcp_proxy::NymProxyClient::new(*proxy_nym_addr, "127.0.0.1", &client_port, 60, Some(env))
.await?;
let proxy_client = tcp_proxy::NymProxyClient::new(
*proxy_nym_addr,
"127.0.0.1",
&client_port,
20,
Some(env),
2,
)
.await?;
tokio::spawn(async move {
proxy_server.run_with_shutdown().await?;
sdk/rust/nym-sdk/src/tcp_proxy/client_pool.rs
+18 -10
View File
@@ -5,13 +5,14 @@ use std::fmt;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use tokio::sync::RwLock;
use tokio::sync::Semaphore;
use tracing::{debug, info, instrument, warn};
// Make a set # of clients (low default)
// Make sure that # of clients is always 1 above the # of incoming conn requests
// Once a client is used, kill the client & remove it from the pool
pub struct ClientPool {
clients: Arc<RwLock<Vec<Arc<MixnetClient>>>>,
semaphore: Arc<Semaphore>,
default_pool_size: usize,
conn_count: Arc<AtomicUsize>, // the actual # of connections running, denoting an incoming tcp request that is matched with a nym client
}
@@ -56,20 +57,20 @@ impl ClientPool {
pub fn new(default_pool_size: usize) -> Self {
ClientPool {
clients: Arc::new(RwLock::new(Vec::new())),
semaphore: Arc::new(Semaphore::new(default_pool_size)),
default_pool_size,
conn_count: Arc::new(AtomicUsize::new(0)),
}
}
// if clients == default, sleep
// if incoming conns > default - conn_count (aka in use clients) then make more clients
pub async fn start(&self) -> Result<()> {
loop {
// TODO double check this..
let spawned_clients = self.clients.read().await.len();
info!("Currently spawned clients: {}", spawned_clients);
debug!("Currently spawned clients: {}", spawned_clients);
if spawned_clients >= self.default_pool_size {
debug!("got enough clients already: sleeping");
debug!("Got enough clients already: sleeping");
} else {
info!("Spawning new client");
let client = loop {
@@ -93,14 +94,20 @@ impl ClientPool {
}
}
pub async fn get_mixnet_client(&self) -> Result<MixnetClient> {
todo!()
pub async fn get_mixnet_client(&self) -> Option<MixnetClient> {
let _permit = self.semaphore.acquire().await.ok()?;
let mut clients = self.clients.write().await;
clients
.pop()
.and_then(|arc_client| Arc::try_unwrap(arc_client).ok())
}
// disconnect ephemeral
// remove from vec
pub async fn disconnect_and_remove_client(&self, client: MixnetClient) -> Result<()> {
todo!()
let mut clients = self.clients.write().await;
clients.retain(|arc_client| arc_client.as_ref().nym_address() != client.nym_address());
client.disconnect().await;
self.semaphore.add_permits(1);
Ok(())
}
pub async fn get_client_count(&self) -> usize {
@@ -126,6 +133,7 @@ impl ClientPool {
pub fn clone(&self) -> Self {
Self {
clients: Arc::clone(&self.clients),
semaphore: Arc::clone(&self.semaphore),
default_pool_size: *&self.default_pool_size,
conn_count: Arc::clone(&self.conn_count),
}
sdk/rust/nym-sdk/src/tcp_proxy/tcp_proxy_client.rs
+15 -12
View File
@@ -7,7 +7,7 @@ mod client_pool;
use client_pool::ClientPool;
#[path = "utils.rs"]
mod utils;
use anyhow::Result;
use anyhow::{bail, Result};
use dashmap::DashSet;
use nym_network_defaults::setup_env;
use nym_sphinx::addressing::Recipient;
@@ -42,7 +42,7 @@ impl NymProxyClient {
env: Option<String>,
default_client_amount: usize,
) -> Result<Self> {
debug!("loading env file: {:?}", env);
debug!("Loading env file: {:?}", env);
setup_env(env); // Defaults to mainnet if empty
Ok(NymProxyClient {
server_address,
@@ -78,13 +78,16 @@ impl NymProxyClient {
let overall_counter = self.conn_pool.clone();
tokio::spawn(async move {
loop {
info!("active connections: {}", overall_counter.get_conn_count());
info!("Active connections: {}", overall_counter.get_conn_count());
tokio::time::sleep(Duration::from_secs(5)).await;
}
});
// if self.conn_pool.get_client_count().await >= DEFAULT_CLIENT_POOL_SIZE / 2 {
loop {
if self.conn_pool.get_client_count().await >= DEFAULT_CLIENT_POOL_SIZE / 2 {
if DEFAULT_CLIENT_POOL_SIZE == 1 && self.conn_pool.get_client_count().await == 1
|| self.conn_pool.get_client_count().await >= DEFAULT_CLIENT_POOL_SIZE / 2
{
let (stream, _) = listener.accept().await?;
tokio::spawn(NymProxyClient::handle_incoming(
stream,
@@ -123,14 +126,13 @@ impl NymProxyClient {
info!("Starting session: {}", session_id);
let mut client: MixnetClient = match conn_pool.get_conn_count() <= DEFAULT_CLIENT_POOL_SIZE
{
true => {
debug!("grabbing client from pool");
conn_pool.get_mixnet_client().await?
let mut client = match conn_pool.get_mixnet_client().await {
Some(client) => {
info!("Grabbing client {} from pool", client.nym_address());
client
}
false => {
debug!("not enough clients in pool, creating ephemeral client");
None => {
info!("Not enough clients in pool, creating ephemeral client");
let net = NymNetworkDetails::new_from_env();
MixnetClientBuilder::new_ephemeral()
.network_details(net)
@@ -243,7 +245,8 @@ impl NymProxyClient {
info!(" Closing write end of session: {}", session_id);
info!(" Triggering client shutdown");
// TODO change this to be a fn in the conn_pool, disconnect and remove from vec
client.disconnect().await;
// client.disconnect().await;
conn_pool.disconnect_and_remove_client(client).await?;
conn_pool.clone().decrement_conn_count()?;
info!(
"Dropped connection - current active connections: {}",