automatically attempt to reconnect to peers if a message was received successfully

Author: Ash-L2L

lib/net/mod.rs

@@ -23,6 +23,7 @@ use crate::{
 
 mod peer;
 
+pub(crate) use peer::error::mailbox::Error as PeerConnectionMailboxError;
 use peer::{
     Connection, ConnectionContext as PeerConnectionCtxt,
     ConnectionHandle as PeerConnectionHandle,
@@ -256,6 +257,20 @@ impl Net {
         }
     }
 
+    /// Apply the provided function to the peer connection handle,
+    /// if it exists.
+    pub fn try_with_active_peer_connection<F, T>(
+        &self,
+        addr: SocketAddr,
+        f: F,
+    ) -> Option<T>
+    where
+        F: FnMut(&PeerConnectionHandle) -> T,
+    {
+        let active_peers_read = self.active_peers.read();
+        active_peers_read.get(&addr).map(f)
+    }
+
     // TODO: This should have more context.
     // Last received message, connection state, etc.
     pub fn get_active_peers(&self) -> Vec<Peer> {

lib/net/peer/channel_pool.rs

@@ -10,9 +10,12 @@ use async_lock::{Semaphore, SemaphoreGuardArc};
 use futures::{FutureExt as _, Stream, StreamExt as _, future::Either, stream};
 use tokio::task::AbortHandle;
 
-use crate::net::peer::{
-    Connection, PeerResponseItem, error, join_set,
-    message::{Heartbeat, Request},
+use crate::{
+    net::peer::{
+        Connection, PeerResponseItem, error,
+        message::{Heartbeat, Request},
+    },
+    util::join_set,
 };
 
 /// Type tags for channel limiters

lib/net/peer/error.rs

@@ -188,7 +188,7 @@ pub(in crate::net::peer) mod request_queue {
     }
 }
 
-pub(in crate::net::peer) mod mailbox {
+pub mod mailbox {
     #[derive(thiserror::Error, Debug)]
     pub enum Error {
         #[error("Heartbeat timeout")]

lib/net/peer/join_set.rs

@@ -1,5 +1,10 @@
 //! Mailbox for a peer connection task
 
+use std::sync::{
+    Arc,
+    atomic::{self, AtomicBool},
+};
+
 use futures::{
     Stream, StreamExt as _, TryFutureExt as _, channel::mpsc, stream,
 };
@@ -65,9 +70,12 @@ pub struct Receiver {
 }
 
 impl Receiver {
+    /// `received_msg_successfully` is set to `True` if a valid message is
+    /// received successfully.
     pub fn into_stream(
         self,
         connection: Connection,
+        received_msg_successfully: &Arc<AtomicBool>,
     ) -> impl Stream<Item = MailboxItem> + Unpin {
         let (peer_response_tx, peer_response_rx) = mpsc::unbounded();
         let internal_message_stream =
@@ -81,10 +89,14 @@ impl Receiver {
             .map(|err| MailboxItem::Error(err.into()));
         let peer_request_stream = stream::try_unfold((), move |()| {
             let conn = connection.clone();
+            let received_msg_successfully = received_msg_successfully.clone();
             let fut = async move {
                 let item = timeout(
                     Connection::HEARTBEAT_TIMEOUT_INTERVAL,
-                    conn.receive_request(),
+                    conn.receive_request().inspect_ok(|_| {
+                        received_msg_successfully
+                            .store(true, atomic::Ordering::SeqCst);
+                    }),
                 )
                 .map_err(|_| Error::HeartbeatTimeout)
                 .await??;
@@ -96,8 +108,12 @@ impl Receiver {
             Ok(peer_request) => MailboxItem::PeerRequest(peer_request),
             Err(err) => MailboxItem::Error(err),
         });
-        let peer_response_stream =
-            peer_response_rx.map(MailboxItem::PeerResponse);
+        let peer_response_stream = peer_response_rx.map(|resp| {
+            if resp.response.is_ok() {
+                received_msg_successfully.store(true, atomic::Ordering::SeqCst);
+            }
+            MailboxItem::PeerResponse(resp)
+        });
         stream::select_all([
             internal_message_stream.boxed(),
             heartbeat_stream.boxed(),

lib/net/peer/mailbox.rs

@@ -22,9 +22,8 @@ use crate::{
 };
 
 mod channel_pool;
-mod error;
-mod join_set;
-mod mailbox;
+pub(crate) mod error;
+pub(crate) mod mailbox;
 pub mod message;
 mod request_queue;
 mod task;
@@ -316,6 +315,8 @@ impl PeerConnectionStatus {
 /// Connection killed on drop
 pub struct ConnectionHandle {
     task: JoinHandle<()>,
+    /// Indicates that at least one message has been received successfully
+    pub(in crate::net) received_msg_successfully: Arc<AtomicBool>,
     /// Representation of [`PeerConnectionStatus`]
     pub(in crate::net) status_repr: Arc<AtomicBool>,
     /// Push messages from connection task / net task / node
@@ -328,6 +329,12 @@ impl ConnectionHandle {
             self.status_repr.load(atomic::Ordering::SeqCst),
         )
     }
+
+    /// Indicates that at least one message has been received successfully
+    pub fn received_msg_successfully(&self) -> bool {
+        self.received_msg_successfully
+            .load(atomic::Ordering::SeqCst)
+    }
 }
 
 impl Drop for ConnectionHandle {
@@ -346,15 +353,18 @@ pub fn handle(
     let (info_tx, info_rx) = mpsc::unbounded();
     let (mailbox_tx, mailbox_rx) = mailbox::new();
     let internal_message_tx = mailbox_tx.internal_message_tx.clone();
+    let received_msg_successfully = Arc::new(AtomicBool::new(false));
     let connection_task = {
         let info_tx = info_tx.clone();
+        let received_msg_successfully = received_msg_successfully.clone();
         move || async move {
             let connection_task = ConnectionTask {
                 connection,
                 ctxt,
                 info_tx,
                 mailbox_rx,
                 mailbox_tx,
+                received_msg_successfully,
             };
             connection_task.run().await
         }
@@ -373,6 +383,7 @@ pub fn handle(
     let status = PeerConnectionStatus::Connected;
     let connection_handle = ConnectionHandle {
         task,
+        received_msg_successfully,
         status_repr: Arc::new(AtomicBool::new(status.as_repr())),
         internal_message_tx,
     };
@@ -385,10 +396,12 @@ pub fn connect(
 ) -> (ConnectionHandle, mpsc::UnboundedReceiver<Info>) {
     let connection_status = PeerConnectionStatus::Connecting;
     let status_repr = Arc::new(AtomicBool::new(connection_status.as_repr()));
+    let received_msg_successfully = Arc::new(AtomicBool::new(false));
     let (info_tx, info_rx) = mpsc::unbounded();
     let (mailbox_tx, mailbox_rx) = mailbox::new();
     let internal_message_tx = mailbox_tx.internal_message_tx.clone();
     let connection_task = {
+        let received_msg_successfully = received_msg_successfully.clone();
         let status_repr = status_repr.clone();
         let info_tx = info_tx.clone();
         move || async move {
@@ -404,6 +417,7 @@ pub fn connect(
                 info_tx,
                 mailbox_rx,
                 mailbox_tx,
+                received_msg_successfully,
             };
             connection_task.run().await
         }
@@ -419,6 +433,7 @@ pub fn connect(
     });
     let connection_handle = ConnectionHandle {
         task,
+        received_msg_successfully,
         status_repr,
         internal_message_tx,
     };

lib/net/peer/mod.rs

@@ -3,6 +3,7 @@
 use std::{
     cmp::Ordering,
     collections::{HashMap, HashSet},
+    sync::{Arc, atomic::AtomicBool},
 };
 
 use fallible_iterator::FallibleIterator;
@@ -32,6 +33,8 @@ pub(in crate::net::peer) struct ConnectionTask {
     pub mailbox_rx: mailbox::Receiver,
     /// Receiver for the task's mailbox
     pub mailbox_tx: mailbox::Sender,
+    /// `True` if a valid message has been received successfully
+    pub received_msg_successfully: Arc<AtomicBool>,
 }
 
 impl ConnectionTask {
@@ -766,7 +769,9 @@ impl ConnectionTask {
         let mut peer_state = Option::<PeerStateId>::None;
         // known peer states
         let mut peer_states = HashMap::<PeerStateId, PeerState>::new();
-        let mut mailbox_stream = self.mailbox_rx.into_stream(self.connection);
+        let mut mailbox_stream = self
+            .mailbox_rx
+            .into_stream(self.connection, &self.received_msg_successfully);
         while let Some(mailbox_item) = mailbox_stream.next().await {
             match mailbox_item {
                 MailboxItem::Error(err) => return Err(err.into()),