Repository: mintlayer/mintlayer-core
Pull request #632 (2023-01-13 13:47:00)

Fix backend event loop blocking
Make outbound connections in background.


Patch:
diff --git a/p2p/src/net/mock/backend.rs b/p2p/src/net/mock/backend.rs
index 12e564b17..a60b7ca15 100644
--- a/p2p/src/net/mock/backend.rs
+++ b/p2p/src/net/mock/backend.rs
@@ -27,11 +27,12 @@ use std::{
     time::Duration,
 };
 
-use futures::{future::join_all, TryFutureExt};
-use tokio::{
-    sync::{mpsc, oneshot},
-    time::timeout,
+use futures::{
+    future::{join_all, BoxFuture},
+    stream::FuturesUnordered,
+    FutureExt, StreamExt, TryFutureExt,
 };
+use tokio::{sync::mpsc, time::timeout};
 
 use common::chain::ChainConfig;
 use crypto::random::{make_pseudo_rng, Rng, SliceRandom};
@@ -113,6 +114,10 @@ pub struct Backend<T: TransportSocket> {
 
     /// Request manager for managing inbound/outbound requests and responses
     request_mgr: request_manager::RequestManager,
+
+    /// List of incoming commands to the backend; we put them in a queue
+    /// to make receiving commands can run concurrently with other backend operations
+    command_queue: FuturesUnordered<BackendTask<T>>,
 }
 
 impl<T> Backend<T>
@@ -143,50 +148,38 @@ where
             pending: HashMap::new(),
             peer_chan: mpsc::channel(64),
             request_mgr: request_manager::RequestManager::new(),
+            command_queue: FuturesUnordered::new(),
         }
     }
 
-    /// Try to establish connection with a remote peer
-    async fn connect(
+    /// Handle connection result to a remote peer
+    async fn handle_connect_res(
         &mut self,
         address: T::Address,
-        response: oneshot::Sender<crate::Result<()>>,
+        connection_res: crate::Result<T::Stream>,
     ) -> crate::Result<()> {
-        // For now we always respond with success
-        response.send(Ok(())).map_err(|_| P2pError::ChannelClosed)?;
-
-        match timeout(self.timeout, self.transport.connect(address.clone())).await {
-            Ok(event) => match event {
-                Ok(socket) => {
-                    let handshake_nonce = make_pseudo_rng().gen();
+        match connection_res {
+            Ok(socket) => {
+                let handshake_nonce = make_pseudo_rng().gen();
+
+                self.create_peer(
+                    socket,
+                    MockPeerId::new(),
+                    PeerRole::Outbound { handshake_nonce },
+                    address,
+                )
+            }
+            Err(err) => {
+                log::error!("Failed to establish connection: {err}");
 
-                    self.create_peer(
-                        socket,
-                        MockPeerId::new(),
-                        PeerRole::Outbound { handshake_nonce },
+                self.conn_tx
+                    .send(ConnectivityEvent::ConnectionError {
                         address,
-                    )
-                }
-                Err(err) => {
-                    log::error!("Failed to establish connection: {err}");
-
-                    self.conn_tx
-                        .send(ConnectivityEvent::ConnectionError {
-                            address,
-                            error: P2pError::DialError(DialError::ConnectionRefusedOrTimedOut),
-                        })
-                        .await
-                        .map_err(P2pError::from)
-                }
-            },
-            Err(_err) => self
-                .conn_tx
-                .send(ConnectivityEvent::ConnectionError {
-                    address,
-                    error: P2pError::DialError(DialError::ConnectionRefusedOrTimedOut),
-                })
-                .await
-                .map_err(P2pError::from),
+                        error: P2pError::DialError(DialError::ConnectionRefusedOrTimedOut),
+                    })
+                    .await
+                    .map_err(P2pError::from)
+            }
         }
     }
 
@@ -365,15 +358,11 @@ where
                 },
                 // Handle commands.
                 command = self.cmd_rx.recv() => {
-                    // TODO: commands running here will block, which means (maybe) these commands in the
-                    //       below function should be short-lived. The issue is that connect/disconnect
-                    //       functions are part of the handling and they could take a relatively long time
-                    //       we should try to figure out how to handle this problem. Maybe we need a separate
-                    //       pipeline for commands that are slow, or maybe everything here should go into a new
-                    //       task (which sounds excessive).
-                    //       Once fixed enable self_connect_channels and self_connect_noise tests.
                     self.handle_command(command.ok_or(P2pError::ChannelClosed)?).await?;
-                }
+                },
+                callback = self.command_queue.select_next_some(), if !self.command_queue.is_empty() => {
+                    callback(self).await?;
+                },
             }
         }
     }
@@ -597,39 +586,101 @@ where
     }
 
     async fn handle_command(&mut self, command: Command<T>) -> crate::Result<()> {
-        match command {
+        // All handlings are separated to two parts, fast and slow:
+        // - "Slow", potentially blocking part (can't take reference to '&mut self' because they are run concurrently).
+        // - "Fast", non-blocking part (take mutable reference to self because they are run sequentially).
+        // Because the second part depends on result of the first part boxed closures are used.
+        // So the first future will return a closure that takes a reference to self and returns one more future.
+
+        let backend_task: BackendTask<T> = match command {
             Command::Connect { address, response } => {
-                self.connect(address, response).await?;
+                // For now we always respond with success
+                response.send(Ok(())).map_err(|_| P2pError::ChannelClosed)?;
+
+                let connection_fut = timeout(self.timeout, self.transport.connect(address.clone()));
+
+                async move {
+                    let connection_res = connection_fut.await.unwrap_or(Err(P2pError::DialError(
+                        DialError::ConnectionRefusedOrTimedOut,
+                    )));
+
+                    boxed_cb(move |this| {
+                        async move { this.handle_connect_res(address, connection_res).await }
+                            .boxed()
+                    })
+                }
+                .boxed()
             }
-            Command::Disconnect { peer_id, response } => {
-                let res = self.disconnect_peer(&peer_id).await;
-                response.send(res).map_err(|_| P2pError::ChannelClosed)?
+            Command::Disconnect { peer_id, response } => async move {
+                boxed_cb(move |this: &mut Self| {
+                    async move {
+                        let res = this.disconnect_peer(&peer_id).await;
+                        response.send(res).map_err(|_| P2pError::ChannelClosed)
+                    }
+                    .boxed()
+                })
             }
+            .boxed(),
             Command::SendRequest {
                 peer_id,
                 message,
                 response,
-            } => {
-                let res = self.send_request(peer_id, message).await;
-                response.send(res).map_err(|_| P2pError::ChannelClosed)?;
+            } => async move {
+                boxed_cb(move |this| {
+                    async move {
+                        let res = this.send_request(peer_id, message).await;
+                        response.send(res).map_err(|_| P2pError::ChannelClosed)
+                    }
+                    .boxed()
+                })
             }
+            .boxed(),
             Command::SendResponse {
                 request_id,
                 message,
                 response,
-            } => {
-                let res = self.send_response(request_id, message).await;
-                response.send(res).map_err(|_| P2pError::ChannelClosed)?;
+            } => async move {
+                boxed_cb(move |this| {
+                    async move {
+                        let res = this.send_response(request_id, message).await;
+                        response.send(res).map_err(|_| P2pError::ChannelClosed)
+                    }
+                    .boxed()
+                })
             }
+            .boxed(),
             Command::AnnounceData {
                 topic,
                 message,
                 response,
-            } => {
-                let res = self.announce_data(topic, message).await;
-                response.send(res).map_err(|_| P2pError::ChannelClosed)?;
+            } => async move {
+                boxed_cb(move |this| {
+                    async move {
+                        let res = this.announce_data(topic, message).await;
+                        response.send(res).map_err(|_| P2pError::ChannelClosed)
+                    }
+                    .boxed()
+                })
             }
-        }
+            .boxed(),
+        };
+
+        self.command_queue.push(backend_task);
+
         Ok(())
     }
 }
+
+// Some boilerplate types and a function for blocking tasks handling
+
+type BackendTask<T> = BoxFuture<'static, BackendTaskCallback<T>>;
+
+type BackendTaskCallback<T> = Box<dyn FnOnce(&mut Backend<T>) -> BackendTaskFut + Send>;
+
+type BackendTaskFut<'a> = BoxFuture<'a, crate::Result<()>>;
+
+fn boxed_cb<T: TransportSocket, F: FnOnce(&mut Backend<T>) -> BackendTaskFut + Send + 'static>(
+    f: F,
+) -> BackendTaskCallback<T> {
+    Box::new(f)
+}
diff --git a/p2p/src/net/mock/mod.rs b/p2p/src/net/mock/mod.rs
index b1793acbc..47a343344 100644
--- a/p2p/src/net/mock/mod.rs
+++ b/p2p/src/net/mock/mod.rs
@@ -591,15 +591,11 @@ mod tests {
         self_connect::<TestTransportTcp, TcpTransportSocket>().await;
     }
 
-    // Test fails because of event loop blocking in Backend::connect
-    #[ignore]
     #[tokio::test]
     async fn self_connect_channels() {
         self_connect::<TestTransportChannel, MockChannelTransport>().await;
     }
 
-    // Test fails because of event loop blocking in Backend::connect
-    #[ignore]
     #[tokio::test]
     async fn self_connect_noise() {
         self_connect::<TestTransportNoise, NoiseTcpTransport>().await;
diff --git a/p2p/src/net/mock/transport/impls/channel.rs b/p2p/src/net/mock/transport/impls/channel.rs
index 7690ddff4..ee6c78f9f 100644
--- a/p2p/src/net/mock/transport/impls/channel.rs
+++ b/p2p/src/net/mock/transport/impls/channel.rs
@@ -22,6 +22,7 @@ use std::{
 };
 
 use async_trait::async_trait;
+use futures::future::BoxFuture;
 use once_cell::sync::Lazy;
 use tokio::{
     io::DuplexStream,
@@ -125,25 +126,29 @@ impl TransportSocket for MockChannelTransport {
         })
     }
 
-    async fn connect(&self, address: Self::Address) -> Result<Self::Stream> {
+    fn connect(&self, address: Self::Address) -> BoxFuture<'static, crate::Result<Self::Stream>> {
         // A connection can only be established to a known address.
         assert_ne!(ZERO_ADDRESS, address);
 
-        let server_sender = CONNECTIONS
-            .lock()
-            .expect("Connections mutex is poisoned")
-            .get(&address)
-            .ok_or(P2pError::DialError(DialError::NoAddresses))?
-            .clone();
+        let local_address = self.local_address;
 
-        let (connect_sender, connect_receiver) = oneshot::channel();
-        server_sender
-            .send((self.local_address, connect_sender))
-            .map_err(|_| P2pError::DialError(DialError::NoAddresses))?;
+        Box::pin(async move {
+            let server_sender = CONNECTIONS
+                .lock()
+                .expect("Connections mutex is poisoned")
+                .get(&address)
+                .ok_or(P2pError::DialError(DialError::NoAddresses))?
+                .clone();
 
-        let channel = connect_receiver.await.map_err(|_| P2pError::ChannelClosed)?;
+            let (connect_sender, connect_receiver) = oneshot::channel();
+            server_sender
+                .send((local_address, connect_sender))
+                .map_err(|_| P2pError::DialError(DialError::NoAddresses))?;
 
-        Ok(channel)
+            let channel = connect_receiver.await.map_err(|_| P2pError::ChannelClosed)?;
+
+            Ok(channel)
+        })
     }
 }
 
diff --git a/p2p/src/net/mock/transport/impls/stream_adapter/wrapped_transport/wrapped_socket.rs b/p2p/src/net/mock/transport/impls/stream_adapter/wrapped_transport/wrapped_socket.rs
index 7d536b7b0..e0b0e5010 100644
--- a/p2p/src/net/mock/transport/impls/stream_adapter/wrapped_transport/wrapped_socket.rs
+++ b/p2p/src/net/mock/transport/impls/stream_adapter/wrapped_transport/wrapped_socket.rs
@@ -14,6 +14,7 @@
 // limitations under the License.
 
 use async_trait::async_trait;
+use futures::future::BoxFuture;
 
 use crate::net::{
     mock::transport::{impls::stream_adapter::traits::StreamAdapter, TransportSocket},
@@ -58,9 +59,13 @@ impl<S: StreamAdapter<T::Stream>, T: TransportSocket> TransportSocket
         Ok(AdaptedListener::new(stream_adapter, listener))
     }
 
-    async fn connect(&self, address: Self::Address) -> Result<Self::Stream> {
-        let base = self.base_transport.connect(address).await?;
-        let stream = self.stream_adapter.handshake(base, Role::Outbound).await?;
-        Ok(stream)
+    fn connect(&self, address: Self::Address) -> BoxFuture<'static, crate::Result<Self::Stream>> {
+        let base = self.base_transport.connect(address);
+        let stream_adapter = self.stream_adapter.clone();
+        Box::pin(async move {
+            let base = base.await?;
+            let stream = stream_adapter.handshake(base, Role::Outbound).await?;
+            Ok(stream)
+        })
     }
 }
diff --git a/p2p/src/net/mock/transport/impls/tcp.rs b/p2p/src/net/mock/transport/impls/tcp.rs
index 372b3976e..3c37ca9f1 100644
--- a/p2p/src/net/mock/transport/impls/tcp.rs
+++ b/p2p/src/net/mock/transport/impls/tcp.rs
@@ -16,7 +16,7 @@
 use std::net::{IpAddr, SocketAddr};
 
 use async_trait::async_trait;
-use futures::{stream::FuturesUnordered, StreamExt};
+use futures::{future::BoxFuture, stream::FuturesUnordered, StreamExt};
 use tokio::net::{TcpListener, TcpStream};
 
 use crate::{
@@ -61,9 +61,11 @@ impl TransportSocket for TcpTransportSocket {
         TcpTransportListener::new(addresses)
     }
 
-    async fn connect(&self, address: Self::Address) -> Result<Self::Stream> {
-        let stream = TcpStream::connect(address).await?;
-        Ok(stream)
+    fn connect(&self, address: Self::Address) -> BoxFuture<'static, crate::Result<Self::Stream>> {
+        Box::pin(async move {
+            let stream = TcpStream::connect(address).await?;
+            Ok(stream)
+        })
     }
 }
 
diff --git a/p2p/src/net/mock/transport/traits/socket.rs b/p2p/src/net/mock/transport/traits/socket.rs
index f4f70d4ae..0c2e7b4e4 100644
--- a/p2p/src/net/mock/transport/traits/socket.rs
+++ b/p2p/src/net/mock/transport/traits/socket.rs
@@ -16,6 +16,7 @@
 use std::{fmt::Debug, hash::Hash, str::FromStr};
 
 use async_trait::async_trait;
+use futures::future::BoxFuture;
 
 use crate::{net::AsBannableAddress, Result};
 
@@ -52,5 +53,5 @@ pub trait TransportSocket: Send + Sync + 'static {
     async fn bind(&self, address: Vec<Self::Address>) -> Result<Self::Listener>;
 
     /// Open a connection to the given address.
-    async fn connect(&self, address: Self::Address) -> Result<Self::Stream>;
+    fn connect(&self, address: Self::Address) -> BoxFuture<'static, crate::Result<Self::Stream>>;
 }


Test patch:
diff --git a/p2p/src/net/mock/transport/impls/stream_adapter/wrapped_transport/tests.rs b/p2p/src/net/mock/transport/impls/stream_adapter/wrapped_transport/tests.rs
index ab3e17770..56b728468 100644
--- a/p2p/src/net/mock/transport/impls/stream_adapter/wrapped_transport/tests.rs
+++ b/p2p/src/net/mock/transport/impls/stream_adapter/wrapped_transport/tests.rs
@@ -20,7 +20,7 @@ use std::{
 
 use crate::testing_utils::{TestTransportChannel, TestTransportMaker, TestTransportTcp};
 use async_trait::async_trait;
-use futures::StreamExt;
+use futures::{future::BoxFuture, StreamExt};
 use tokio::{
     io::{AsyncReadExt, AsyncWriteExt},
     time::timeout,
@@ -145,8 +145,8 @@ impl TransportSocket for TestMockTransport {
         })
     }
 
-    async fn connect(&self, address: Self::Address) -> crate::Result<Self::Stream> {
-        self.transport.connect(address).await
+    fn connect(&self, address: Self::Address) -> BoxFuture<'static, crate::Result<Self::Stream>> {
+        Box::pin(self.transport.connect(address))
     }
 }