From 9b81c2eb685cbd81c1d4e021c8f0e1156d29f507 Mon Sep 17 00:00:00 2001 From: Jedrzej Stuczynski Date: Tue, 25 Feb 2020 15:54:50 +0000 Subject: [PATCH] Initial version of the reconnecting tcp client --- .../src/connection_manager/mod.rs | 95 +++++++ .../src/connection_manager/reconnector.rs | 76 ++++++ .../src/connection_manager/writer.rs | 46 ++++ common/clients/multi-tcp-client/src/lib.rs | 247 ++++++------------ common/clients/multi-tcp-client/src/main.rs | 20 ++ 5 files changed, 323 insertions(+), 161 deletions(-) create mode 100644 common/clients/multi-tcp-client/src/connection_manager/mod.rs create mode 100644 common/clients/multi-tcp-client/src/connection_manager/reconnector.rs create mode 100644 common/clients/multi-tcp-client/src/connection_manager/writer.rs create mode 100644 common/clients/multi-tcp-client/src/main.rs diff --git a/common/clients/multi-tcp-client/src/connection_manager/mod.rs b/common/clients/multi-tcp-client/src/connection_manager/mod.rs new file mode 100644 index 0000000000..a3ee712392 --- /dev/null +++ b/common/clients/multi-tcp-client/src/connection_manager/mod.rs @@ -0,0 +1,95 @@ +use crate::connection_manager::reconnector::ConnectionReconnector; +use crate::connection_manager::writer::ConnectionWriter; +use futures::task::Poll; +use futures::AsyncWriteExt; +use log::*; +use std::io; +use std::net::SocketAddr; +use std::time::Duration; + +mod reconnector; +mod writer; + +enum ConnectionState { + Writing(ConnectionWriter), + Reconnecting(ConnectionReconnector), +} + +pub(crate) struct ConnectionManager { + address: SocketAddr, + + maximum_reconnection_backoff: Duration, + reconnection_backoff: Duration, + + state: ConnectionState, +} + +impl ConnectionManager { + pub(crate) async fn new( + address: SocketAddr, + reconnection_backoff: Duration, + maximum_reconnection_backoff: Duration, + ) -> Self { + // 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)), + Err(e) => { + warn!( + "failed to establish initial connection to {} ({}). Going into reconnection mode", + address, e + ); + ConnectionState::Reconnecting(ConnectionReconnector::new( + address, + reconnection_backoff, + maximum_reconnection_backoff, + )) + } + }; + + ConnectionManager { + address, + maximum_reconnection_backoff, + reconnection_backoff, + state, + } + } + + pub(crate) async fn send(&mut self, msg: &[u8]) -> io::Result<()> { + if let ConnectionState::Reconnecting(conn_reconnector) = &mut self.state { + // do a single poll rather than await for future to completely resolve + // TODO: if we call poll ourselves here, will the Waker still call it itself later on? + 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", + )) + } + Poll::Ready(conn) => conn, + }; + + println!("We managed to reconnect!"); + 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(res) => Ok(res), + Err(e) => { + println!("Creating connection reconnector!"); + self.state = ConnectionState::Reconnecting(ConnectionReconnector::new( + self.address, + self.reconnection_backoff, + self.maximum_reconnection_backoff, + )); + Err(e) + } + }; + }; + unreachable!() + } +} diff --git a/common/clients/multi-tcp-client/src/connection_manager/reconnector.rs b/common/clients/multi-tcp-client/src/connection_manager/reconnector.rs new file mode 100644 index 0000000000..df8d6c0703 --- /dev/null +++ b/common/clients/multi-tcp-client/src/connection_manager/reconnector.rs @@ -0,0 +1,76 @@ +use log::*; +use std::future::Future; +use std::io; +use std::net::SocketAddr; +use std::pin::Pin; +use std::task::{Context, Poll}; +use std::time::Duration; + +pub(crate) struct ConnectionReconnector { + address: SocketAddr, + connection: Pin>>>, + + current_retry_attempt: u32, + + current_backoff_delay: tokio::time::Delay, + maximum_reconnection_backoff: Duration, + + reconnection_backoff: Duration, +} + +impl ConnectionReconnector { + pub(crate) fn new( + address: SocketAddr, + reconnection_backoff: Duration, + maximum_reconnection_backoff: Duration, + ) -> Self { + ConnectionReconnector { + address, + connection: Box::pin(tokio::net::TcpStream::connect(address)), + current_backoff_delay: tokio::time::delay_for(Duration::new(0, 0)), // if we can re-establish connection on first try without any backoff that's perfect + current_retry_attempt: 0, + maximum_reconnection_backoff, + reconnection_backoff, + } + } +} + +impl Future for ConnectionReconnector { + type Output = tokio::net::TcpStream; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + // see if we are still in exponential backoff + if Pin::new(&mut self.current_backoff_delay) + .poll(cx) + .is_pending() + { + return Poll::Pending; + }; + + // see if we managed to resolve the connection yet + match Pin::new(&mut self.connection).poll(cx) { + Poll::Pending => Poll::Pending, + Poll::Ready(Err(e)) => { + warn!( + "we failed to re-establish connection to {} - {:?} (attempt {})", + self.address, e, self.current_retry_attempt + ); + self.current_retry_attempt += 1; + + // we failed to re-establish connection - continue exponential backoff + let next_delay = std::cmp::min( + self.maximum_reconnection_backoff, + 2_u32.pow(self.current_retry_attempt) * self.reconnection_backoff, + ); + + self.current_backoff_delay + .reset(tokio::time::Instant::now() + next_delay); + + self.connection = Box::pin(tokio::net::TcpStream::connect(self.address)); + + Poll::Pending + } + Poll::Ready(Ok(conn)) => Poll::Ready(conn), + } + } +} diff --git a/common/clients/multi-tcp-client/src/connection_manager/writer.rs b/common/clients/multi-tcp-client/src/connection_manager/writer.rs new file mode 100644 index 0000000000..32d76694f3 --- /dev/null +++ b/common/clients/multi-tcp-client/src/connection_manager/writer.rs @@ -0,0 +1,46 @@ +use futures::task::{Context, Poll}; +use futures::AsyncWrite; +use std::io; +use std::pin::Pin; +use tokio::prelude::*; + +pub(crate) struct ConnectionWriter { + connection: tokio::net::TcpStream, +} + +impl ConnectionWriter { + pub(crate) fn new(connection: tokio::net::TcpStream) -> Self { + ConnectionWriter { connection } + } +} + +impl AsyncWrite for ConnectionWriter { + fn poll_write( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + use tokio::io::AsyncWrite; + + let mut read_buf = [0; 1]; + match Pin::new(&mut self.connection).poll_read(cx, &mut read_buf) { + // at least try the obvious check for if connection is definitely down + // TODO: can we do anything else? + Poll::Ready(Ok(n)) if n == 0 => Poll::Ready(Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + "trying to write to closed connection", + ))), + _ => Pin::new(&mut self.connection).poll_write(cx, buf), + } + } + + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + use tokio::io::AsyncWrite; + Pin::new(&mut self.connection).poll_flush(cx) + } + + fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + use tokio::io::AsyncWrite; + Pin::new(&mut self.connection).poll_shutdown(cx) + } +} diff --git a/common/clients/multi-tcp-client/src/lib.rs b/common/clients/multi-tcp-client/src/lib.rs index 697acc93fa..c5c46bf359 100644 --- a/common/clients/multi-tcp-client/src/lib.rs +++ b/common/clients/multi-tcp-client/src/lib.rs @@ -1,222 +1,83 @@ -use futures::task::{Context, Poll}; -use futures::{AsyncWrite, AsyncWriteExt}; +use crate::connection_manager::ConnectionManager; use std::collections::HashMap; use std::io; use std::net::SocketAddr; -use std::pin::Pin; use std::str; use std::time::Duration; -use tokio::prelude::*; +use tokio::prelude::*; // this import is actually required for socket.read() call -struct ConnectionWriter { - connection: tokio::net::TcpStream, - - reconnection_backoff: Duration, - maximum_reconnection_backoff: Duration, - current_reconnection_backoff: Duration, -} - -impl ConnectionWriter { - fn new( - connection: tokio::net::TcpStream, - initial_reconnection_backoff: Duration, - maximum_reconnection_backoff: Duration, - ) -> Self { - ConnectionWriter { - connection, - reconnection_backoff: initial_reconnection_backoff, - -struct ConnectionReconnector { - address: SocketAddr, - connection: Pin>>>, - - current_retry_attempt: u32, - maximum_retry_attempts: u32, - - current_backoff_delay: tokio::time::Delay, - maximum_reconnection_backoff: Duration, - - reconnection_backoff: Duration, -} - -impl ConnectionReconnector { - fn new( - address: SocketAddr, - maximum_retry_attempts: u32, - reconnection_backoff: Duration, - maximum_reconnection_backoff: Duration, - ) -> Self { - ConnectionReconnector { - address, - connection: Box::pin(tokio::net::TcpStream::connect(address)), - current_backoff_delay: tokio::time::delay_for(Duration::new(0, 0)), // if we can re-establish connection on first try without any backoff that's perfect - current_retry_attempt: 0, - maximum_reconnection_backoff, - maximum_retry_attempts, - reconnection_backoff, - } - } -} - -impl Drop for ConnectionWriter { - fn drop(&mut self) { - // try to cleanly shutdown connection on going out of scope - if let Err(e) = self.connection.shutdown(std::net::Shutdown::Both) { - eprintln!("Failed to cleanly shutdown the connection - {:?}", e); -impl Future for ConnectionReconnector { - type Output = io::Result; - - fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - // see if we are still in exponential backoff - if Pin::new(&mut self.current_backoff_delay) - .poll(cx) - .is_pending() - { - return Poll::Pending; - }; - - // see if we managed to resolve the connection yet - match Pin::new(&mut self.connection).poll(cx) { - Poll::Pending => Poll::Pending, - Poll::Ready(Err(e)) => { - warn!( - "we failed to re-establish connection to {} - {:?}", - self.address, e - ); - self.current_retry_attempt += 1; - - // check if we reached our limit - if self.current_retry_attempt == self.maximum_retry_attempts { - return Poll::Ready(Err(io::Error::new( - io::ErrorKind::Other, - "Reached maximum number of retry attempts", - ))); - } - - // we failed to re-establish connection - continue exponential backoff - let next_delay = std::cmp::min( - self.maximum_reconnection_backoff, - 2_u32.pow(self.current_retry_attempt) * self.reconnection_backoff, - ); - - self.current_backoff_delay - .reset(tokio::time::Instant::now() + next_delay); - - Poll::Pending - } - Poll::Ready(Ok(conn)) => Poll::Ready(Ok(conn)), - } - } -} - -impl AsyncWrite for ConnectionWriter { - fn poll_write( - mut self: Pin<&mut Self>, - cx: &mut Context<'_>, - buf: &[u8], - ) -> Poll> { - use tokio::io::AsyncWrite; - - let mut read_buf = [0; 1]; - match Pin::new(&mut self.connection).poll_read(cx, &mut read_buf) { - // at least try the obvious check if connection is definitely down - // can't do more than that - Poll::Ready(Ok(n)) if n == 0 => Poll::Ready(Err(io::Error::new( - io::ErrorKind::UnexpectedEof, - "trying to write to closed connection", - ))), - _ => Pin::new(&mut self.connection).poll_write(cx, buf), - } - } - - fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - use tokio::io::AsyncWrite; - Pin::new(&mut self.connection).poll_flush(cx) - } - - fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - use tokio::io::AsyncWrite; - Pin::new(&mut self.connection).poll_shutdown(cx) - } -} +mod connection_manager; pub struct Config { initial_endpoints: Vec, - initial_reconnection_backoff: Duration, + reconnection_backoff: Duration, maximum_reconnection_backoff: Duration, } impl Config { pub fn new( initial_endpoints: Vec, - initial_reconnection_backoff: Duration, + reconnection_backoff: Duration, maximum_reconnection_backoff: Duration, ) -> Self { Config { initial_endpoints, - initial_reconnection_backoff, + reconnection_backoff, maximum_reconnection_backoff, } } } pub struct Client { - connections_writers: HashMap, + connections_managers: HashMap, } impl Client { pub async fn new(config: Config) -> Client { - let mut connections_writers = HashMap::new(); + let mut connections_managers = HashMap::new(); for endpoint in config.initial_endpoints { - connections_writers.insert( + connections_managers.insert( endpoint, - ConnectionWriter::new( - tokio::net::TcpStream::connect(endpoint).await.unwrap(), - config.initial_reconnection_backoff, + ConnectionManager::new( + endpoint, + config.reconnection_backoff, config.maximum_reconnection_backoff, - ), + ) + .await, ); } Client { - connections_writers, + connections_managers, } } pub async fn send(&mut self, address: SocketAddr, message: &[u8]) -> io::Result<()> { println!("sending {:?}", str::from_utf8(message)); - if !self.connections_writers.contains_key(&address) { + if !self.connections_managers.contains_key(&address) { return Err(io::Error::new( io::ErrorKind::AddrNotAvailable, - "address not in the list", + "address not in the list - dynamic connections not yet supported", )); } + // let (tx, rx) = oneshot::channel(); + // 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 - if let Err(e) = self - .connections_writers + self.connections_managers .get_mut(&address) .unwrap() - .write_all(&message) + .send(&message) .await - { - println!( - "Failed to write to socket - {:?}. Presumably we need to reconnect!", - e - ); - // TODO: reconnection - } - - Ok(()) } } #[cfg(test)] mod tests { use super::*; - use std::time; + use std::{env, time}; const CLOSE_MESSAGE: [u8; 3] = [0, 0, 0]; @@ -267,8 +128,72 @@ mod tests { } } + #[test] + fn client_reconnects_to_server_after_it_went_down() { + let num_test_threads = env::var("RUST_TEST_THREADS").expect("Number of test threads must be set to 1 to prevent tests interacting with themselves! (RUST_TEST_THREADS=1)"); + assert_eq!(num_test_threads, "1", "Number of test threads must be set to 1 to prevent tests interacting with themselves! (RUST_TEST_THREADS=1)"); + + let mut rt = tokio::runtime::Runtime::new().unwrap(); + let addr = "127.0.0.1:5000".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![b"foomp1", b"foomp2"]; + let finished_dummy_server_future = + rt.spawn(DummyServer::new().listen_until(addr, CLOSE_MESSAGE.as_ref())); + + 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( + async move { tokio::time::delay_for(time::Duration::from_millis(50)).await }, + ); + } + + // kill server + rt.block_on(c.send(addr, CLOSE_MESSAGE.as_ref())).unwrap(); + rt.block_on(async move { tokio::time::delay_for(time::Duration::from_millis(50)).await }); + + // try to send - go into reconnection + let post_kill_message = b"new foomp"; + + // we are trying to send to killed server + assert!(rt.block_on(c.send(addr, post_kill_message)).is_err()); + // 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); + + let new_server_future = + rt.spawn(DummyServer::new().listen_until(addr, CLOSE_MESSAGE.as_ref())); + + // 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 }, + ); + } + 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.as_ref())).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 num_test_threads = env::var("RUST_TEST_THREADS").expect("Number of test threads must be set to 1 to prevent tests interacting with themselves! (RUST_TEST_THREADS=1)"); + assert_eq!(num_test_threads, "1", "Number of test threads must be set to 1 to prevent tests interacting with themselves! (RUST_TEST_THREADS=1)"); + let mut rt = tokio::runtime::Runtime::new().unwrap(); let addr = "127.0.0.1:5000".parse().unwrap(); let reconnection_backoff = Duration::from_secs(2); diff --git a/common/clients/multi-tcp-client/src/main.rs b/common/clients/multi-tcp-client/src/main.rs new file mode 100644 index 0000000000..669972eae7 --- /dev/null +++ b/common/clients/multi-tcp-client/src/main.rs @@ -0,0 +1,20 @@ +use multi_tcp_client::{Client, Config}; +use std::time; +use std::time::Duration; +use tokio::prelude::*; +use tokio::runtime::Runtime; + +fn main() { + let mut rt = Runtime::new().unwrap(); + let addr = "127.0.0.1:5000".parse().unwrap(); + let reconnection_backoff = Duration::from_secs(1); + + let client_config = Config::new(vec![addr], reconnection_backoff, 10 * reconnection_backoff); + + let mut c = rt.block_on(Client::new(client_config)); + + for _ in 0..50 { + rt.block_on(c.send(addr, b"foomp\n")); + rt.block_on(async move { tokio::time::delay_for(time::Duration::from_millis(250)).await }); + } +}