Changed connection manager to be accessed via a channel
This commit is contained in:
Jedrzej Stuczynski
@@ -1,21 +1,31 @@
|
||||
use crate::connection_manager::reconnector::ConnectionReconnector;
|
||||
use crate::connection_manager::writer::ConnectionWriter;
|
||||
use crate::error_reader::ConnectionErrorSender;
|
||||
use futures::channel::mpsc;
|
||||
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 ConnectionManagerSender = mpsc::UnboundedSender<Vec<u8>>;
|
||||
type ConnectionManagerReceiver = mpsc::UnboundedReceiver<Vec<u8>>;
|
||||
|
||||
enum ConnectionState<'a> {
|
||||
Writing(ConnectionWriter),
|
||||
Reconnecting(ConnectionReconnector<'a>),
|
||||
}
|
||||
|
||||
pub(crate) struct ConnectionManager<'a> {
|
||||
conn_tx: ConnectionManagerSender,
|
||||
conn_rx: ConnectionManagerReceiver,
|
||||
|
||||
errors_tx: ConnectionErrorSender,
|
||||
address: SocketAddr,
|
||||
|
||||
maximum_reconnection_backoff: Duration,
|
||||
@@ -24,12 +34,15 @@ 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,
|
||||
errors_tx: ConnectionErrorSender,
|
||||
) -> 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 +60,9 @@ impl<'a> ConnectionManager<'a> {
|
||||
};
|
||||
|
||||
ConnectionManager {
|
||||
conn_tx,
|
||||
conn_rx,
|
||||
errors_tx,
|
||||
address,
|
||||
maximum_reconnection_backoff,
|
||||
reconnection_backoff,
|
||||
@@ -54,41 +70,55 @@ 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 {
|
||||
self.handle_new_message(msg).await;
|
||||
}
|
||||
});
|
||||
sender_clone
|
||||
}
|
||||
|
||||
async fn handle_new_message(&mut self, msg: Vec<u8>) {
|
||||
info!("sending to {:?}", self.address);
|
||||
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) {
|
||||
Poll::Pending => {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::BrokenPipe,
|
||||
"connection is broken - reconnection is in progress",
|
||||
))
|
||||
self.errors_tx
|
||||
.unbounded_send((
|
||||
self.address,
|
||||
Err(io::Error::new(
|
||||
io::ErrorKind::BrokenPipe,
|
||||
"connection is broken - reconnection is in progress",
|
||||
)),
|
||||
))
|
||||
.unwrap();
|
||||
return;
|
||||
}
|
||||
Poll::Ready(conn) => conn,
|
||||
};
|
||||
|
||||
debug!("Managed to reconnect to {}!", self.address);
|
||||
info!("Managed to reconnect to {}!", self.address);
|
||||
self.state = ConnectionState::Writing(ConnectionWriter::new(new_connection));
|
||||
}
|
||||
|
||||
// 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 {
|
||||
// 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?
|
||||
Ok(_) => Ok(()),
|
||||
Err(e) => {
|
||||
trace!("Creating connection reconnector!");
|
||||
self.state = ConnectionState::Reconnecting(ConnectionReconnector::new(
|
||||
self.address,
|
||||
self.reconnection_backoff,
|
||||
self.maximum_reconnection_backoff,
|
||||
));
|
||||
Err(e)
|
||||
}
|
||||
};
|
||||
if let Err(e) = conn_writer.write_all(msg.as_ref()).await {
|
||||
info!("Creating connection reconnector!");
|
||||
self.state = ConnectionState::Reconnecting(ConnectionReconnector::new(
|
||||
self.address,
|
||||
self.reconnection_backoff,
|
||||
self.maximum_reconnection_backoff,
|
||||
));
|
||||
self.errors_tx
|
||||
.unbounded_send((self.address, Err(e)))
|
||||
.unwrap();
|
||||
}
|
||||
};
|
||||
unreachable!()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,226 +1,227 @@
|
||||
use crate::connection_manager::ConnectionManager;
|
||||
use crate::connection_manager::{ConnectionManager, ConnectionManagerSender};
|
||||
use crate::error_reader::{ConnectionErrorReader, ConnectionErrorSender};
|
||||
use futures::channel::mpsc;
|
||||
use log::*;
|
||||
use std::collections::HashMap;
|
||||
use std::io;
|
||||
use std::net::SocketAddr;
|
||||
use std::time::Duration;
|
||||
use tokio::runtime::Handle;
|
||||
|
||||
mod connection_manager;
|
||||
mod error_reader;
|
||||
|
||||
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,
|
||||
errors_tx: ConnectionErrorSender,
|
||||
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 async fn start_new(config: Config) -> Client {
|
||||
let (errors_tx, errors_rx) = mpsc::unbounded();
|
||||
let errors_reader = ConnectionErrorReader::new(errors_rx);
|
||||
|
||||
Client {
|
||||
connections_managers,
|
||||
let client = Client {
|
||||
// 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"),
|
||||
errors_tx,
|
||||
connections_managers: HashMap::new(),
|
||||
initial_reconnection_backoff: config.maximum_reconnection_backoff,
|
||||
maximum_reconnection_backoff: config.initial_reconnection_backoff,
|
||||
}
|
||||
};
|
||||
|
||||
errors_reader.start(&client.runtime_handle);
|
||||
client
|
||||
}
|
||||
|
||||
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,
|
||||
self.errors_tx.clone(),
|
||||
)
|
||||
.await
|
||||
.start(&self.runtime_handle)
|
||||
}
|
||||
|
||||
pub async fn send(&mut self, address: SocketAddr, message: Vec<u8>) {
|
||||
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
|
||||
.unbounded_send(message)
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::str;
|
||||
use std::time;
|
||||
use tokio::prelude::*;
|
||||
|
||||
const SERVER_MSG_LEN: usize = 16;
|
||||
const CLOSE_MESSAGE: [u8; SERVER_MSG_LEN] = [0; SERVER_MSG_LEN];
|
||||
|
||||
struct DummyServer {
|
||||
received_buf: Vec<Vec<u8>>,
|
||||
listener: tokio::net::TcpListener,
|
||||
}
|
||||
|
||||
impl DummyServer {
|
||||
async fn new(address: SocketAddr) -> Self {
|
||||
DummyServer {
|
||||
received_buf: Vec::new(),
|
||||
listener: tokio::net::TcpListener::bind(address).await.unwrap(),
|
||||
}
|
||||
}
|
||||
|
||||
fn get_received(&self) -> Vec<Vec<u8>> {
|
||||
self.received_buf.clone()
|
||||
}
|
||||
|
||||
// this is only used in tests so slightly higher logging levels are fine
|
||||
async fn listen_until(mut self, close_message: &[u8]) -> Self {
|
||||
let (mut socket, _) = self.listener.accept().await.unwrap();
|
||||
loop {
|
||||
let mut buf = [0u8; SERVER_MSG_LEN];
|
||||
match socket.read(&mut buf).await {
|
||||
Ok(n) if n == 0 => {
|
||||
info!("Remote connection closed");
|
||||
return self;
|
||||
}
|
||||
Ok(n) => {
|
||||
info!("received ({}) - {:?}", n, str::from_utf8(buf[..n].as_ref()));
|
||||
|
||||
if buf[..n].as_ref() == close_message {
|
||||
info!("closing...");
|
||||
socket.shutdown(std::net::Shutdown::Both).unwrap();
|
||||
return self;
|
||||
} else {
|
||||
self.received_buf.push(buf[..n].to_vec());
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
panic!("failed to read from socket; err = {:?}", e);
|
||||
}
|
||||
};
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn client_reconnects_to_server_after_it_went_down() {
|
||||
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 messages_to_send = vec![[1u8; SERVER_MSG_LEN], [2; SERVER_MSG_LEN]];
|
||||
|
||||
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));
|
||||
|
||||
for msg in &messages_to_send {
|
||||
rt.block_on(c.send(addr, msg)).unwrap();
|
||||
}
|
||||
|
||||
// kill server
|
||||
rt.block_on(c.send(addr, &CLOSE_MESSAGE)).unwrap();
|
||||
let received_messages = rt
|
||||
.block_on(finished_dummy_server_future)
|
||||
.unwrap()
|
||||
.get_received();
|
||||
|
||||
assert_eq!(received_messages, messages_to_send);
|
||||
|
||||
// try to send - go into reconnection
|
||||
let post_kill_message = [3u8; SERVER_MSG_LEN];
|
||||
|
||||
// we are trying to send to killed server
|
||||
assert!(rt.block_on(c.send(addr, &post_kill_message)).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() {
|
||||
break;
|
||||
}
|
||||
rt.block_on(
|
||||
async move { tokio::time::delay_for(time::Duration::from_millis(50)).await },
|
||||
);
|
||||
}
|
||||
|
||||
// kill the server to ensure it actually got the message
|
||||
rt.block_on(c.send(addr, &CLOSE_MESSAGE)).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]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn server_receives_all_sent_messages_when_up() {
|
||||
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 messages_to_send = vec![[1u8; SERVER_MSG_LEN], [2; SERVER_MSG_LEN]];
|
||||
|
||||
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));
|
||||
|
||||
for msg in &messages_to_send {
|
||||
rt.block_on(c.send(addr, msg)).unwrap();
|
||||
}
|
||||
|
||||
rt.block_on(c.send(addr, &CLOSE_MESSAGE)).unwrap();
|
||||
|
||||
// the server future should have already been resolved
|
||||
let received_messages = rt
|
||||
.block_on(finished_dummy_server_future)
|
||||
.unwrap()
|
||||
.get_received();
|
||||
|
||||
assert_eq!(received_messages, messages_to_send);
|
||||
}
|
||||
}
|
||||
// #[cfg(test)]
|
||||
// mod tests {
|
||||
// use super::*;
|
||||
// use std::str;
|
||||
// use std::time;
|
||||
// use tokio::prelude::*;
|
||||
//
|
||||
// const SERVER_MSG_LEN: usize = 16;
|
||||
// const CLOSE_MESSAGE: [u8; SERVER_MSG_LEN] = [0; SERVER_MSG_LEN];
|
||||
//
|
||||
// struct DummyServer {
|
||||
// received_buf: Vec<Vec<u8>>,
|
||||
// listener: tokio::net::TcpListener,
|
||||
// }
|
||||
//
|
||||
// impl DummyServer {
|
||||
// async fn new(address: SocketAddr) -> Self {
|
||||
// DummyServer {
|
||||
// received_buf: Vec::new(),
|
||||
// listener: tokio::net::TcpListener::bind(address).await.unwrap(),
|
||||
// }
|
||||
// }
|
||||
//
|
||||
// fn get_received(&self) -> Vec<Vec<u8>> {
|
||||
// self.received_buf.clone()
|
||||
// }
|
||||
//
|
||||
// // this is only used in tests so slightly higher logging levels are fine
|
||||
// async fn listen_until(mut self, close_message: &[u8]) -> Self {
|
||||
// let (mut socket, _) = self.listener.accept().await.unwrap();
|
||||
// loop {
|
||||
// let mut buf = [0u8; SERVER_MSG_LEN];
|
||||
// match socket.read(&mut buf).await {
|
||||
// Ok(n) if n == 0 => {
|
||||
// info!("Remote connection closed");
|
||||
// return self;
|
||||
// }
|
||||
// Ok(n) => {
|
||||
// info!("received ({}) - {:?}", n, str::from_utf8(buf[..n].as_ref()));
|
||||
//
|
||||
// if buf[..n].as_ref() == close_message {
|
||||
// info!("closing...");
|
||||
// socket.shutdown(std::net::Shutdown::Both).unwrap();
|
||||
// return self;
|
||||
// } else {
|
||||
// self.received_buf.push(buf[..n].to_vec());
|
||||
// }
|
||||
// }
|
||||
// Err(e) => {
|
||||
// panic!("failed to read from socket; err = {:?}", e);
|
||||
// }
|
||||
// };
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
//
|
||||
// #[test]
|
||||
// fn client_reconnects_to_server_after_it_went_down() {
|
||||
// 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 messages_to_send = vec![[1u8; SERVER_MSG_LEN], [2; SERVER_MSG_LEN]];
|
||||
//
|
||||
// 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));
|
||||
//
|
||||
// for msg in &messages_to_send {
|
||||
// rt.block_on(c.send(addr, msg)).unwrap();
|
||||
// }
|
||||
//
|
||||
// // kill server
|
||||
// rt.block_on(c.send(addr, &CLOSE_MESSAGE)).unwrap();
|
||||
// let received_messages = rt
|
||||
// .block_on(finished_dummy_server_future)
|
||||
// .unwrap()
|
||||
// .get_received();
|
||||
//
|
||||
// assert_eq!(received_messages, messages_to_send);
|
||||
//
|
||||
// // try to send - go into reconnection
|
||||
// let post_kill_message = [3u8; SERVER_MSG_LEN];
|
||||
//
|
||||
// // we are trying to send to killed server
|
||||
// assert!(rt.block_on(c.send(addr, &post_kill_message)).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() {
|
||||
// break;
|
||||
// }
|
||||
// rt.block_on(
|
||||
// async move { tokio::time::delay_for(time::Duration::from_millis(50)).await },
|
||||
// );
|
||||
// }
|
||||
//
|
||||
// // kill the server to ensure it actually got the message
|
||||
// rt.block_on(c.send(addr, &CLOSE_MESSAGE)).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]);
|
||||
// }
|
||||
//
|
||||
// #[test]
|
||||
// fn server_receives_all_sent_messages_when_up() {
|
||||
// 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 messages_to_send = vec![[1u8; SERVER_MSG_LEN], [2; SERVER_MSG_LEN]];
|
||||
//
|
||||
// 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));
|
||||
//
|
||||
// for msg in &messages_to_send {
|
||||
// rt.block_on(c.send(addr, msg)).unwrap();
|
||||
// }
|
||||
//
|
||||
// rt.block_on(c.send(addr, &CLOSE_MESSAGE)).unwrap();
|
||||
//
|
||||
// // the server future should have already been resolved
|
||||
// let received_messages = rt
|
||||
// .block_on(finished_dummy_server_future)
|
||||
// .unwrap()
|
||||
// .get_received();
|
||||
//
|
||||
// assert_eq!(received_messages, messages_to_send);
|
||||
// }
|
||||
// }
|
||||
|
||||
Reference in New Issue
Block a user