Merge pull request #142 from nymtech/feature/concurrent_connection_managers

Feature/concurrent connection managers
This commit is contained in:
Dave Hrycyszyn
2020-03-16 17:20:27 +00:00
committed by GitHub
7 changed files with 140 additions and 114 deletions
@@ -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!();
}
}
+67 -57
View File
@@ -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
View File
@@ -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) {
+9 -13
View File
@@ -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
+9 -19
View File
@@ -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) {
+6 -10
View File
@@ -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>(