use mio for Tcp management in bootstrap system

massa-bootstrap/Cargo.toml

@@ -20,6 +20,7 @@ tracing = "0.1"
 substruct = { git = "https://github.com/sydhds/substruct" }
 socket2 = "0.4.7"
 crossbeam = "0.8.2"
+mio =  { version = "0.8", features = ["net", "os-poll"] }
 
 # custom modules
 massa_async_pool = { path = "../massa-async-pool" }

massa-bootstrap/src/establisher.rs

@@ -5,10 +5,12 @@ use std::{
     net::{SocketAddr, TcpListener, TcpStream},
 };
 
+use crate::{error::BootstrapError, listener::PollEvent};
+
 /// Specifies a common interface that can be used by standard, or mockers
 #[cfg_attr(test, mockall::automock)]
-pub trait BSListener {
-    fn accept(&mut self) -> io::Result<(TcpStream, SocketAddr)>;
+pub trait BSEventPoller {
+    fn poll(&mut self) -> Result<PollEvent, BootstrapError>;
 }
 
 /// Specifies a common interface that can be used by standard, or mockers
@@ -49,14 +51,14 @@ impl DefaultListener {
     }
 }
 
-impl BSListener for DefaultListener {
+impl BSEventPoller for DefaultListener {
     /// Accepts a new incoming connection from this listener.
-    fn accept(&mut self) -> io::Result<(TcpStream, SocketAddr)> {
+    fn poll(&mut self) -> Result<PollEvent, BootstrapError> {
         // accept
         let (sock, mut remote_addr) = self.0.accept()?;
         // normalize address
         remote_addr.set_ip(remote_addr.ip().to_canonical());
-        Ok((sock, remote_addr))
+        Ok(PollEvent::NewConnection((sock, remote_addr)))
     }
 }
 /// Initiates a connection with given timeout in milliseconds

massa-bootstrap/src/lib.rs

@@ -23,6 +23,8 @@ mod client;
 mod client_binder;
 mod error;
 mod establisher;
+mod listener;
+pub use listener::BootstrapTcpListener;
 mod messages;
 mod server;
 mod server_binder;

massa-bootstrap/src/listener.rs

@@ -0,0 +1,84 @@
+use std::net::{SocketAddr, TcpListener, TcpStream};
+
+use mio::net::TcpListener as MioTcpListener;
+
+use mio::{Events, Interest, Poll, Token, Waker};
+use tracing::info;
+
+use crate::error::BootstrapError;
+use crate::establisher::BSEventPoller;
+
+const NEW_CONNECTION: Token = Token(0);
+const STOP_LISTENER: Token = Token(10);
+
+/// TODO: this should be crate-private. currently needed for models testing
+pub struct BootstrapTcpListener {
+    poll: Poll,
+    events: Events,
+    server: TcpListener,
+    // HACK : create variable to move ownership of mio_server to the thread
+    // if mio_server is not moved, poll does not receive any event from listener
+    _mio_server: MioTcpListener,
+}
+
+pub struct BootstrapListenerStopHandle(Waker);
+
+pub enum PollEvent {
+    NewConnection((TcpStream, SocketAddr)),
+    Stop,
+}
+impl BootstrapTcpListener {
+    /// Setup a mio-listener that functions as a `select!` on a connection, or a waker
+    ///
+    /// * `addr` - the address to listen on
+    pub fn new(addr: SocketAddr) -> Result<(BootstrapListenerStopHandle, Self), BootstrapError> {
+        info!("Starting bootstrap listener on {}", &addr);
+        let server = TcpListener::bind(addr)?;
+        let mut mio_server =
+            MioTcpListener::from_std(server.try_clone().expect("Unable to clone server socket"));
+
+        let poll = Poll::new()?;
+
+        // wake up the poll when we want to stop the listener
+        let waker = BootstrapListenerStopHandle(Waker::new(poll.registry(), STOP_LISTENER)?);
+
+        poll.registry()
+            .register(&mut mio_server, NEW_CONNECTION, Interest::READABLE)?;
+
+        // TODO use config for capacity ?
+        let events = Events::with_capacity(32);
+        Ok((
+            waker,
+            BootstrapTcpListener {
+                poll,
+                server,
+                events,
+                _mio_server: mio_server,
+            },
+        ))
+    }
+}
+
+impl BSEventPoller for BootstrapTcpListener {
+    fn poll(&mut self) -> Result<PollEvent, BootstrapError> {
+        self.poll.poll(&mut self.events, None).unwrap();
+
+        // Confirm that we are not being signalled to shut down
+        if self.events.iter().any(|ev| ev.token() == STOP_LISTENER) {
+            return Ok(PollEvent::Stop);
+        }
+
+        // Ther could be more than one connection ready, but we want to re-check for the stop
+        // signal after processing each connection.
+        Ok(PollEvent::NewConnection(
+            self.server.accept().map_err(BootstrapError::from)?,
+        ))
+    }
+}
+
+impl BootstrapListenerStopHandle {
+    /// Stop the bootstrap listener.
+    pub fn stop(&self) -> Result<(), BootstrapError> {
+        self.0.wake().map_err(BootstrapError::from)
+    }
+}