Merge pull request #142 from nymtech/feature/concurrent_connection_managers
Feature/concurrent connection managers
This commit is contained in:
@@ -88,6 +88,7 @@ mod converting_mixnode_presence_into_topology_mixnode {
|
||||
|
||||
let result: Result<mix::Node, std::io::Error> = mix_presence.try_into();
|
||||
// assert!(result.is_err()) // This fails only for me. Why?
|
||||
// ¯\_(ツ)_/¯ - works on my machine (and travis)
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -1,21 +1,31 @@
|
||||
use crate::connection_manager::reconnector::ConnectionReconnector;
|
||||
use crate::connection_manager::writer::ConnectionWriter;
|
||||
use futures::channel::{mpsc, oneshot};
|
||||
use futures::task::Poll;
|
||||
use futures::AsyncWriteExt;
|
||||
use futures::{AsyncWriteExt, StreamExt};
|
||||
use log::*;
|
||||
use std::io;
|
||||
use std::net::SocketAddr;
|
||||
use std::time::Duration;
|
||||
use tokio::runtime::Handle;
|
||||
|
||||
mod reconnector;
|
||||
mod writer;
|
||||
|
||||
pub(crate) type ResponseSender = Option<oneshot::Sender<io::Result<()>>>;
|
||||
|
||||
pub(crate) type ConnectionManagerSender = mpsc::UnboundedSender<(Vec<u8>, ResponseSender)>;
|
||||
type ConnectionManagerReceiver = mpsc::UnboundedReceiver<(Vec<u8>, ResponseSender)>;
|
||||
|
||||
enum ConnectionState<'a> {
|
||||
Writing(ConnectionWriter),
|
||||
Reconnecting(ConnectionReconnector<'a>),
|
||||
}
|
||||
|
||||
pub(crate) struct ConnectionManager<'a> {
|
||||
conn_tx: ConnectionManagerSender,
|
||||
conn_rx: ConnectionManagerReceiver,
|
||||
|
||||
address: SocketAddr,
|
||||
|
||||
maximum_reconnection_backoff: Duration,
|
||||
@@ -24,12 +34,14 @@ pub(crate) struct ConnectionManager<'a> {
|
||||
state: ConnectionState<'a>,
|
||||
}
|
||||
|
||||
impl<'a> ConnectionManager<'a> {
|
||||
impl<'a> ConnectionManager<'static> {
|
||||
pub(crate) async fn new(
|
||||
address: SocketAddr,
|
||||
reconnection_backoff: Duration,
|
||||
maximum_reconnection_backoff: Duration,
|
||||
) -> ConnectionManager<'a> {
|
||||
let (conn_tx, conn_rx) = mpsc::unbounded();
|
||||
|
||||
// based on initial connection we will either have a writer or a reconnector
|
||||
let state = match tokio::net::TcpStream::connect(address).await {
|
||||
Ok(conn) => ConnectionState::Writing(ConnectionWriter::new(conn)),
|
||||
@@ -47,6 +59,8 @@ impl<'a> ConnectionManager<'a> {
|
||||
};
|
||||
|
||||
ConnectionManager {
|
||||
conn_tx,
|
||||
conn_rx,
|
||||
address,
|
||||
maximum_reconnection_backoff,
|
||||
reconnection_backoff,
|
||||
@@ -54,7 +68,27 @@ impl<'a> ConnectionManager<'a> {
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn send(&mut self, msg: &[u8]) -> io::Result<()> {
|
||||
/// consumes Self and returns channel for communication
|
||||
pub(crate) fn start(mut self, handle: &Handle) -> ConnectionManagerSender {
|
||||
let sender_clone = self.conn_tx.clone();
|
||||
handle.spawn(async move {
|
||||
while let Some(msg) = self.conn_rx.next().await {
|
||||
let (msg_content, res_ch) = msg;
|
||||
let res = self.handle_new_message(msg_content).await;
|
||||
if let Some(res_ch) = res_ch {
|
||||
if let Err(e) = res_ch.send(res) {
|
||||
error!(
|
||||
"failed to send response on the channel to the caller! - {:?}",
|
||||
e
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
sender_clone
|
||||
}
|
||||
|
||||
async fn handle_new_message(&mut self, msg: Vec<u8>) -> io::Result<()> {
|
||||
if let ConnectionState::Reconnecting(conn_reconnector) = &mut self.state {
|
||||
// do a single poll rather than await for future to completely resolve
|
||||
let new_connection = match futures::poll!(conn_reconnector) {
|
||||
@@ -74,9 +108,12 @@ impl<'a> ConnectionManager<'a> {
|
||||
// we must be in writing state if we are here, either by being here from beginning or just
|
||||
// transitioning from reconnecting
|
||||
if let ConnectionState::Writing(conn_writer) = &mut self.state {
|
||||
return match conn_writer.write_all(msg).await {
|
||||
return match conn_writer.write_all(msg.as_ref()).await {
|
||||
// if we failed to write to connection we should reconnect
|
||||
// TODO: is this true? can we fail to write to a connection while it still remains open and valid?
|
||||
|
||||
// TODO: change connection writer to somehow also poll for responses and
|
||||
// change return type of this method from io::Result<> to io::Result<Vec<u8>>
|
||||
Ok(_) => Ok(()),
|
||||
Err(e) => {
|
||||
trace!("Creating connection reconnector!");
|
||||
@@ -89,6 +126,7 @@ impl<'a> ConnectionManager<'a> {
|
||||
}
|
||||
};
|
||||
};
|
||||
unreachable!()
|
||||
|
||||
unreachable!();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,87 +1,91 @@
|
||||
use crate::connection_manager::ConnectionManager;
|
||||
use crate::connection_manager::{ConnectionManager, ConnectionManagerSender};
|
||||
use futures::channel::oneshot;
|
||||
use log::*;
|
||||
use std::collections::HashMap;
|
||||
use std::io;
|
||||
use std::net::SocketAddr;
|
||||
use std::time::Duration;
|
||||
use tokio::runtime::Handle;
|
||||
|
||||
mod connection_manager;
|
||||
|
||||
pub struct Config {
|
||||
initial_endpoints: Vec<SocketAddr>,
|
||||
initial_reconnection_backoff: Duration,
|
||||
maximum_reconnection_backoff: Duration,
|
||||
}
|
||||
|
||||
impl Config {
|
||||
pub fn new(
|
||||
initial_endpoints: Vec<SocketAddr>,
|
||||
initial_reconnection_backoff: Duration,
|
||||
maximum_reconnection_backoff: Duration,
|
||||
) -> Self {
|
||||
Config {
|
||||
initial_endpoints,
|
||||
initial_reconnection_backoff,
|
||||
maximum_reconnection_backoff,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct Client<'a> {
|
||||
connections_managers: HashMap<SocketAddr, ConnectionManager<'a>>,
|
||||
pub struct Client {
|
||||
runtime_handle: Handle,
|
||||
connections_managers: HashMap<SocketAddr, ConnectionManagerSender>,
|
||||
maximum_reconnection_backoff: Duration,
|
||||
initial_reconnection_backoff: Duration,
|
||||
}
|
||||
|
||||
impl<'a> Client<'a> {
|
||||
pub async fn new(config: Config) -> Client<'a> {
|
||||
let mut connections_managers = HashMap::new();
|
||||
for initial_endpoint in config.initial_endpoints {
|
||||
connections_managers.insert(
|
||||
initial_endpoint,
|
||||
ConnectionManager::new(
|
||||
initial_endpoint,
|
||||
config.initial_reconnection_backoff,
|
||||
config.maximum_reconnection_backoff,
|
||||
)
|
||||
.await,
|
||||
);
|
||||
}
|
||||
|
||||
impl Client {
|
||||
pub fn new(config: Config) -> Client {
|
||||
Client {
|
||||
connections_managers,
|
||||
// if the function is not called within tokio runtime context, this will panic
|
||||
// but perhaps the code should be better structured to completely avoid this call
|
||||
runtime_handle: Handle::try_current()
|
||||
.expect("The client MUST BE used within tokio runtime context"),
|
||||
connections_managers: HashMap::new(),
|
||||
initial_reconnection_backoff: config.maximum_reconnection_backoff,
|
||||
maximum_reconnection_backoff: config.initial_reconnection_backoff,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn send(&mut self, address: SocketAddr, message: &[u8]) -> io::Result<()> {
|
||||
async fn start_new_connection_manager(&self, address: SocketAddr) -> ConnectionManagerSender {
|
||||
ConnectionManager::new(
|
||||
address,
|
||||
self.initial_reconnection_backoff,
|
||||
self.maximum_reconnection_backoff,
|
||||
)
|
||||
.await
|
||||
.start(&self.runtime_handle)
|
||||
}
|
||||
|
||||
// if wait_for_response is set to true, we will get information about any possible IO errors
|
||||
// as well as (once implemented) received replies, however, this will also cause way longer
|
||||
// waiting periods
|
||||
pub async fn send(
|
||||
&mut self,
|
||||
address: SocketAddr,
|
||||
message: Vec<u8>,
|
||||
wait_for_response: bool,
|
||||
) -> io::Result<()> {
|
||||
if !self.connections_managers.contains_key(&address) {
|
||||
info!(
|
||||
"There is no existing connection to {:?} - it will be established now",
|
||||
address
|
||||
);
|
||||
|
||||
// TODO: now we're blocking to establish TCP connection this need to be changed
|
||||
// so that other connections could progress
|
||||
let new_manager = ConnectionManager::new(
|
||||
address,
|
||||
self.initial_reconnection_backoff,
|
||||
self.maximum_reconnection_backoff,
|
||||
)
|
||||
.await;
|
||||
|
||||
self.connections_managers.insert(address, new_manager);
|
||||
let new_manager_sender = self.start_new_connection_manager(address).await;
|
||||
self.connections_managers
|
||||
.insert(address, new_manager_sender);
|
||||
}
|
||||
|
||||
// to optimize later by using channels and separate tokio tasks for each connection handler
|
||||
// because right now say we want to write to addresses A and B -
|
||||
// We have to wait until we're done dealing with A before we can do anything with B
|
||||
self.connections_managers
|
||||
.get_mut(&address)
|
||||
.unwrap()
|
||||
.send(&message)
|
||||
.await
|
||||
let manager = self.connections_managers.get_mut(&address).unwrap();
|
||||
|
||||
if wait_for_response {
|
||||
let (res_tx, res_rx) = oneshot::channel();
|
||||
manager.unbounded_send((message, Some(res_tx))).unwrap();
|
||||
res_rx.await.unwrap()
|
||||
} else {
|
||||
manager.unbounded_send((message, None)).unwrap();
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -146,22 +150,22 @@ mod tests {
|
||||
let mut rt = tokio::runtime::Runtime::new().unwrap();
|
||||
let addr = "127.0.0.1:6000".parse().unwrap();
|
||||
let reconnection_backoff = Duration::from_secs(1);
|
||||
let client_config =
|
||||
Config::new(vec![addr], reconnection_backoff, 10 * reconnection_backoff);
|
||||
let client_config = Config::new(reconnection_backoff, 10 * reconnection_backoff);
|
||||
|
||||
let messages_to_send = vec![[1u8; SERVER_MSG_LEN], [2; SERVER_MSG_LEN]];
|
||||
let messages_to_send = vec![[1u8; SERVER_MSG_LEN].to_vec(), [2; SERVER_MSG_LEN].to_vec()];
|
||||
|
||||
let dummy_server = rt.block_on(DummyServer::new(addr));
|
||||
let finished_dummy_server_future = rt.spawn(dummy_server.listen_until(&CLOSE_MESSAGE));
|
||||
|
||||
let mut c = rt.block_on(Client::new(client_config));
|
||||
let mut c = rt.enter(|| Client::new(client_config));
|
||||
|
||||
for msg in &messages_to_send {
|
||||
rt.block_on(c.send(addr, msg)).unwrap();
|
||||
rt.block_on(c.send(addr, msg.clone(), true)).unwrap();
|
||||
}
|
||||
|
||||
// kill server
|
||||
rt.block_on(c.send(addr, &CLOSE_MESSAGE)).unwrap();
|
||||
rt.block_on(c.send(addr, CLOSE_MESSAGE.to_vec(), true))
|
||||
.unwrap();
|
||||
let received_messages = rt
|
||||
.block_on(finished_dummy_server_future)
|
||||
.unwrap()
|
||||
@@ -170,17 +174,22 @@ mod tests {
|
||||
assert_eq!(received_messages, messages_to_send);
|
||||
|
||||
// try to send - go into reconnection
|
||||
let post_kill_message = [3u8; SERVER_MSG_LEN];
|
||||
let post_kill_message = [3u8; SERVER_MSG_LEN].to_vec();
|
||||
|
||||
// we are trying to send to killed server
|
||||
assert!(rt.block_on(c.send(addr, &post_kill_message)).is_err());
|
||||
assert!(rt
|
||||
.block_on(c.send(addr, post_kill_message.clone(), true))
|
||||
.is_err());
|
||||
|
||||
let new_dummy_server = rt.block_on(DummyServer::new(addr));
|
||||
let new_server_future = rt.spawn(new_dummy_server.listen_until(&CLOSE_MESSAGE));
|
||||
|
||||
// keep sending after we leave reconnection backoff and reconnect
|
||||
loop {
|
||||
if rt.block_on(c.send(addr, &post_kill_message)).is_ok() {
|
||||
if rt
|
||||
.block_on(c.send(addr, post_kill_message.clone(), true))
|
||||
.is_ok()
|
||||
{
|
||||
break;
|
||||
}
|
||||
rt.block_on(
|
||||
@@ -189,7 +198,8 @@ mod tests {
|
||||
}
|
||||
|
||||
// kill the server to ensure it actually got the message
|
||||
rt.block_on(c.send(addr, &CLOSE_MESSAGE)).unwrap();
|
||||
rt.block_on(c.send(addr, CLOSE_MESSAGE.to_vec(), true))
|
||||
.unwrap();
|
||||
let new_received_messages = rt.block_on(new_server_future).unwrap().get_received();
|
||||
assert_eq!(post_kill_message.to_vec(), new_received_messages[0]);
|
||||
}
|
||||
@@ -199,21 +209,21 @@ mod tests {
|
||||
let mut rt = tokio::runtime::Runtime::new().unwrap();
|
||||
let addr = "127.0.0.1:6001".parse().unwrap();
|
||||
let reconnection_backoff = Duration::from_secs(2);
|
||||
let client_config =
|
||||
Config::new(vec![addr], reconnection_backoff, 10 * reconnection_backoff);
|
||||
let client_config = Config::new(reconnection_backoff, 10 * reconnection_backoff);
|
||||
|
||||
let messages_to_send = vec![[1u8; SERVER_MSG_LEN], [2; SERVER_MSG_LEN]];
|
||||
let messages_to_send = vec![[1u8; SERVER_MSG_LEN].to_vec(), [2; SERVER_MSG_LEN].to_vec()];
|
||||
|
||||
let dummy_server = rt.block_on(DummyServer::new(addr));
|
||||
let finished_dummy_server_future = rt.spawn(dummy_server.listen_until(&CLOSE_MESSAGE));
|
||||
|
||||
let mut c = rt.block_on(Client::new(client_config));
|
||||
let mut c = rt.enter(|| Client::new(client_config));
|
||||
|
||||
for msg in &messages_to_send {
|
||||
rt.block_on(c.send(addr, msg)).unwrap();
|
||||
rt.block_on(c.send(addr, msg.clone(), true)).unwrap();
|
||||
}
|
||||
|
||||
rt.block_on(c.send(addr, &CLOSE_MESSAGE)).unwrap();
|
||||
rt.block_on(c.send(addr, CLOSE_MESSAGE.to_vec(), true))
|
||||
.unwrap();
|
||||
|
||||
// the server future should have already been resolved
|
||||
let received_messages = rt
|
||||
|
||||
+5
-10
@@ -87,16 +87,11 @@ impl MixNode {
|
||||
|
||||
fn start_packet_forwarder(&mut self) -> mpsc::UnboundedSender<(SocketAddr, Vec<u8>)> {
|
||||
info!("Starting packet forwarder...");
|
||||
|
||||
// this can later be replaced with topology information
|
||||
let initial_addresses = vec![];
|
||||
self.runtime
|
||||
.block_on(packet_forwarding::PacketForwarder::new(
|
||||
initial_addresses,
|
||||
self.config.get_packet_forwarding_initial_backoff(),
|
||||
self.config.get_packet_forwarding_maximum_backoff(),
|
||||
))
|
||||
.start(self.runtime.handle())
|
||||
packet_forwarding::PacketForwarder::new(
|
||||
self.config.get_packet_forwarding_initial_backoff(),
|
||||
self.config.get_packet_forwarding_maximum_backoff(),
|
||||
)
|
||||
.start(self.runtime.handle())
|
||||
}
|
||||
|
||||
pub fn run(&mut self) {
|
||||
|
||||
@@ -1,24 +1,21 @@
|
||||
use futures::channel::mpsc;
|
||||
use futures::StreamExt;
|
||||
use log::*;
|
||||
use std::net::SocketAddr;
|
||||
use std::time::Duration;
|
||||
use tokio::runtime::Handle;
|
||||
|
||||
pub(crate) struct PacketForwarder<'a> {
|
||||
tcp_client: multi_tcp_client::Client<'a>,
|
||||
pub(crate) struct PacketForwarder {
|
||||
tcp_client: multi_tcp_client::Client,
|
||||
conn_tx: mpsc::UnboundedSender<(SocketAddr, Vec<u8>)>,
|
||||
conn_rx: mpsc::UnboundedReceiver<(SocketAddr, Vec<u8>)>,
|
||||
}
|
||||
|
||||
impl PacketForwarder<'static> {
|
||||
pub(crate) async fn new(
|
||||
initial_endpoints: Vec<SocketAddr>,
|
||||
impl PacketForwarder {
|
||||
pub(crate) fn new(
|
||||
initial_reconnection_backoff: Duration,
|
||||
maximum_reconnection_backoff: Duration,
|
||||
) -> PacketForwarder<'static> {
|
||||
) -> PacketForwarder {
|
||||
let tcp_client_config = multi_tcp_client::Config::new(
|
||||
initial_endpoints,
|
||||
initial_reconnection_backoff,
|
||||
maximum_reconnection_backoff,
|
||||
);
|
||||
@@ -26,7 +23,7 @@ impl PacketForwarder<'static> {
|
||||
let (conn_tx, conn_rx) = mpsc::unbounded();
|
||||
|
||||
PacketForwarder {
|
||||
tcp_client: multi_tcp_client::Client::new(tcp_client_config).await,
|
||||
tcp_client: multi_tcp_client::Client::new(tcp_client_config),
|
||||
conn_tx,
|
||||
conn_rx,
|
||||
}
|
||||
@@ -37,10 +34,9 @@ impl PacketForwarder<'static> {
|
||||
let sender_channel = self.conn_tx.clone();
|
||||
handle.spawn(async move {
|
||||
while let Some((address, packet)) = self.conn_rx.next().await {
|
||||
match self.tcp_client.send(address, &packet).await {
|
||||
Err(e) => warn!("Failed to forward packet to {:?} - {:?}", address, e),
|
||||
Ok(_) => trace!("Forwarded packet to {:?}", address),
|
||||
}
|
||||
// as a mix node we don't care about responses, we just want to fire packets
|
||||
// as quickly as possible
|
||||
self.tcp_client.send(address, packet, false).await.unwrap(); // if we're not waiting for response, we MUST get an Ok
|
||||
}
|
||||
});
|
||||
sender_channel
|
||||
|
||||
@@ -18,45 +18,35 @@ impl MixMessage {
|
||||
}
|
||||
|
||||
// TODO: put our TCP client here
|
||||
pub(crate) struct MixTrafficController<'a> {
|
||||
tcp_client: multi_tcp_client::Client<'a>,
|
||||
pub(crate) struct MixTrafficController {
|
||||
tcp_client: multi_tcp_client::Client,
|
||||
mix_rx: MixMessageReceiver,
|
||||
}
|
||||
|
||||
impl MixTrafficController<'static> {
|
||||
pub(crate) async fn new(
|
||||
initial_endpoints: Vec<SocketAddr>,
|
||||
impl MixTrafficController {
|
||||
pub(crate) fn new(
|
||||
initial_reconnection_backoff: Duration,
|
||||
maximum_reconnection_backoff: Duration,
|
||||
mix_rx: MixMessageReceiver,
|
||||
) -> Self {
|
||||
let tcp_client_config = multi_tcp_client::Config::new(
|
||||
initial_endpoints,
|
||||
initial_reconnection_backoff,
|
||||
maximum_reconnection_backoff,
|
||||
);
|
||||
|
||||
MixTrafficController {
|
||||
tcp_client: multi_tcp_client::Client::new(tcp_client_config).await,
|
||||
tcp_client: multi_tcp_client::Client::new(tcp_client_config),
|
||||
mix_rx,
|
||||
}
|
||||
}
|
||||
|
||||
async fn on_message(&mut self, mix_message: MixMessage) {
|
||||
debug!("Got a mix_message for {:?}", mix_message.0);
|
||||
match self
|
||||
.tcp_client
|
||||
.send(mix_message.0, &mix_message.1.to_bytes())
|
||||
self.tcp_client
|
||||
// TODO: possibly we might want to get an actual result here at some point
|
||||
.send(mix_message.0, mix_message.1.to_bytes(), false)
|
||||
.await
|
||||
{
|
||||
Ok(_) => trace!("sent a mix message"),
|
||||
// TODO: should there be some kind of threshold of failed messages
|
||||
// that if reached, the application blows?
|
||||
Err(e) => error!(
|
||||
"We failed to send the packet to {} - {:?}",
|
||||
mix_message.0, e
|
||||
),
|
||||
};
|
||||
.unwrap(); // if we're not waiting for response, we MUST get an Ok
|
||||
}
|
||||
|
||||
pub(crate) async fn run(&mut self) {
|
||||
|
||||
@@ -212,16 +212,12 @@ impl NymClient {
|
||||
// controller for sending sphinx packets to mixnet (either real traffic or cover traffic)
|
||||
fn start_mix_traffic_controller(&mut self, mix_rx: MixMessageReceiver) {
|
||||
info!("Starting mix trafic controller...");
|
||||
// TODO: possible optimisation: set the initial endpoints to all known mixes from layer 1
|
||||
let initial_mix_endpoints = Vec::new();
|
||||
self.runtime
|
||||
.block_on(MixTrafficController::new(
|
||||
initial_mix_endpoints,
|
||||
self.config.get_packet_forwarding_initial_backoff(),
|
||||
self.config.get_packet_forwarding_maximum_backoff(),
|
||||
mix_rx,
|
||||
))
|
||||
.start(self.runtime.handle());
|
||||
MixTrafficController::new(
|
||||
self.config.get_packet_forwarding_initial_backoff(),
|
||||
self.config.get_packet_forwarding_maximum_backoff(),
|
||||
mix_rx,
|
||||
)
|
||||
.start(self.runtime.handle());
|
||||
}
|
||||
|
||||
fn start_socket_listener<T: 'static + NymTopology>(
|
||||
|
||||
Reference in New Issue
Block a user