Merge pull request #4141 from ProvableHQ/feat/low_level_tweaks

vicsn · web-flow · commit 669dccd732a8 · 2026-03-26T20:15:06.000+01:00
[Feat] Low-level network tweaks
diff --git a/.circleci/config.yml b/.circleci/config.yml
@@ -819,7 +819,7 @@ workflows:
           filters:
             branches:
               only:
-                - feat/vk-migration
+                - fix/sync-race-conditions 
                 - canary
                 - testnet
                 - mainnet
diff --git a/node/router/src/heartbeat.rs b/node/router/src/heartbeat.rs
@@ -19,7 +19,6 @@ use crate::{
     NodeType,
     Outbound,
     PeerPoolHandling,
-    Router,
     bootstrap_peers,
     messages::{DisconnectReason, Message, PeerRequest},
 };
@@ -67,8 +66,6 @@ pub trait Heartbeat<N: Network>: Outbound<N> {
         self.safety_check_minimum_number_of_peers();
         self.log_connected_peers();
 
-        // Remove any stale connected peers.
-        self.remove_stale_connected_peers();
         // Remove the oldest connected peer.
         self.remove_oldest_connected_peer();
         // Keep the number of connected peers within the allowed range.
@@ -111,20 +108,6 @@ pub trait Heartbeat<N: Network>: Outbound<N> {
         }
     }
 
-    /// This function removes any connected peers that have not communicated within the predefined time.
-    fn remove_stale_connected_peers(&self) {
-        // Check if any connected peer is stale.
-        for peer in self.router().get_connected_peers() {
-            // Disconnect if the peer has not communicated back within the predefined time.
-            let elapsed = peer.last_seen.elapsed();
-            if elapsed > Router::<N>::MAX_RADIO_SILENCE {
-                warn!("Peer '{}' has not communicated in {elapsed:?}", peer.listener_addr);
-                // Disconnect from this peer.
-                self.router().disconnect(peer.listener_addr);
-            }
-        }
-    }
-
     /// Returns a sorted vector of network addresses of all removable connected peers
     /// where the first entry has the lowest priority and the last one the highest.
     ///
diff --git a/node/router/src/lib.rs b/node/router/src/lib.rs
@@ -69,7 +69,7 @@ use anyhow::Result;
 use locktick::parking_lot::{Mutex, RwLock};
 #[cfg(not(feature = "locktick"))]
 use parking_lot::{Mutex, RwLock};
-use std::{collections::HashMap, future::Future, io, net::SocketAddr, ops::Deref, sync::Arc, time::Duration};
+use std::{collections::HashMap, future::Future, io, net::SocketAddr, ops::Deref, sync::Arc};
 use tokio::task::JoinHandle;
 
 /// The default port used by the router.
@@ -147,9 +147,6 @@ impl<N: Network> Router<N> {
     /// The maximum amount of connection attempts within a 10 second threshold.
     #[cfg(not(feature = "test"))]
     const MAX_CONNECTION_ATTEMPTS: usize = 10;
-    /// The duration after which a connected peer is considered inactive or
-    /// disconnected if no message has been received in the meantime.
-    const MAX_RADIO_SILENCE: Duration = Duration::from_secs(150); // 2.5 minutes
 }
 
 impl<N: Network> Router<N> {
diff --git a/node/tcp/src/helpers/connections.rs b/node/tcp/src/helpers/connections.rs
@@ -137,3 +137,11 @@ impl Not for ConnectionSide {
         }
     }
 }
+
+impl Drop for Connection {
+    fn drop(&mut self) {
+        for task in self.tasks.iter().rev() {
+            task.abort();
+        }
+    }
+}
diff --git a/node/tcp/src/protocols/disconnect.rs b/node/tcp/src/protocols/disconnect.rs
@@ -13,9 +13,13 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::net::SocketAddr;
+use std::{net::SocketAddr, time::Duration};
 
-use tokio::sync::{mpsc, oneshot};
+use tokio::{
+    sync::{mpsc, oneshot},
+    task::JoinHandle,
+    time::timeout,
+};
 use tracing::*;
 
 #[cfg(doc)]
@@ -31,10 +35,18 @@ pub trait Disconnect: P2P
 where
     Self: Clone + Send + Sync + 'static,
 {
+    /// The maximum time allowed for the on_disconnect hook to execute.
+    /// If the hook exceeds this time, it will be aborted to ensure the node cleans up
+    /// resources promptly.
+    const TIMEOUT: Duration = Duration::from_secs(3);
+
     /// Attaches the behavior specified in [`Disconnect::handle_disconnect`] to every occurrence of the
     /// node disconnecting from a peer.
     async fn enable_disconnect(&self) {
-        let (from_node_sender, mut from_node_receiver) = mpsc::unbounded_channel::<(SocketAddr, oneshot::Sender<()>)>();
+        let (from_node_sender, mut from_node_receiver) = mpsc::channel::<(
+            SocketAddr,
+            oneshot::Sender<(JoinHandle<()>, oneshot::Receiver<()>)>,
+        )>(self.tcp().config().max_connections as usize);
 
         // use a channel to know when the disconnect task is ready
         let (tx, rx) = oneshot::channel::<()>();
@@ -47,13 +59,20 @@ where
 
             while let Some((peer_addr, notifier)) = from_node_receiver.recv().await {
                 let self_clone2 = self_clone.clone();
-                tokio::spawn(async move {
+                // create a channel for waiting on completion
+                let (done_tx, done_rx) = oneshot::channel();
+                let handle = tokio::spawn(async move {
                     // perform the specified extra actions
-                    self_clone2.handle_disconnect(peer_addr).await;
+                    if timeout(Self::TIMEOUT, self_clone2.handle_disconnect(peer_addr)).await.is_err() {
+                        warn!(parent: self_clone2.tcp().span(), "Disconnect logic timed out for {peer_addr}");
+                    }
                     // notify the node that the extra actions have concluded
                     // and that the related connection can be dropped
-                    let _ = notifier.send(()); // can't really fail
+                    let _ = done_tx.send(());
                 });
+                // provide the node with a handle to the scheduled task,
+                // and a receiver that will notify it of its completion
+                let _ = notifier.send((handle, done_rx)); // can't really fail
             }
         });
         let _ = rx.await;
diff --git a/node/tcp/src/protocols/handshake.rs b/node/tcp/src/protocols/handshake.rs
@@ -45,7 +45,8 @@ where
 
     /// Prepares the node to perform specified network handshakes.
     async fn enable_handshake(&self) {
-        let (from_node_sender, mut from_node_receiver) = mpsc::unbounded_channel::<ReturnableConnection>();
+        let (from_node_sender, mut from_node_receiver) =
+            mpsc::channel::<ReturnableConnection>(self.tcp().config().max_connections as usize);
 
         // use a channel to know when the handshake task is ready
         let (tx, rx) = oneshot::channel();
diff --git a/node/tcp/src/protocols/mod.rs b/node/tcp/src/protocols/mod.rs
@@ -20,7 +20,10 @@
 use std::{io, net::SocketAddr};
 
 use once_cell::race::OnceBox;
-use tokio::sync::{mpsc, oneshot};
+use tokio::{
+    sync::{mpsc, oneshot},
+    task::JoinHandle,
+};
 
 use crate::connections::Connection;
 
@@ -36,13 +39,17 @@ pub use on_connect::OnConnect;
 pub use reading::Reading;
 pub use writing::Writing;
 
+// The value returned to the node by the OnDisconnect protocol is a bit complex,
+// so use an alias to break it down.
+type OnDisconnectBundle = (JoinHandle<()>, oneshot::Receiver<()>);
+
 #[derive(Default)]
 pub(crate) struct Protocols {
     pub(crate) handshake: OnceBox<ProtocolHandler<Connection, io::Result<Connection>>>,
     pub(crate) reading: OnceBox<ProtocolHandler<Connection, io::Result<Connection>>>,
     pub(crate) writing: OnceBox<writing::WritingHandler>,
-    pub(crate) on_connect: OnceBox<ProtocolHandler<SocketAddr, ()>>,
-    pub(crate) disconnect: OnceBox<ProtocolHandler<SocketAddr, ()>>,
+    pub(crate) on_connect: OnceBox<ProtocolHandler<SocketAddr, JoinHandle<()>>>,
+    pub(crate) disconnect: OnceBox<ProtocolHandler<SocketAddr, OnDisconnectBundle>>,
 }
 
 /// An object sent to a protocol handler task; the task assumes control of a protocol-relevant item `T`,
@@ -52,15 +59,15 @@ pub(crate) type ReturnableItem<T, U> = (T, oneshot::Sender<U>);
 
 pub(crate) type ReturnableConnection = ReturnableItem<Connection, io::Result<Connection>>;
 
-pub(crate) struct ProtocolHandler<T, U>(mpsc::UnboundedSender<ReturnableItem<T, U>>);
+pub(crate) struct ProtocolHandler<T, U>(mpsc::Sender<ReturnableItem<T, U>>);
 
 pub(crate) trait Protocol<T, U> {
-    fn trigger(&self, item: ReturnableItem<T, U>);
+    async fn trigger(&self, item: ReturnableItem<T, U>);
 }
 
 impl<T, U> Protocol<T, U> for ProtocolHandler<T, U> {
-    fn trigger(&self, item: ReturnableItem<T, U>) {
+    async fn trigger(&self, item: ReturnableItem<T, U>) {
         // ignore errors; they can only happen if a disconnect interrupts the protocol setup process
-        let _ = self.0.send(item);
+        let _ = self.0.send(item).await;
     }
 }
diff --git a/node/tcp/src/protocols/on_connect.rs b/node/tcp/src/protocols/on_connect.rs
@@ -35,7 +35,10 @@ where
     /// Attaches the behavior specified in [`OnConnect::on_connect`] right after every successful
     /// handshake.
     async fn enable_on_connect(&self) {
-        let (from_node_sender, mut from_node_receiver) = mpsc::unbounded_channel::<(SocketAddr, oneshot::Sender<()>)>();
+        let (from_node_sender, mut from_node_receiver) = mpsc::channel::<(
+            SocketAddr,
+            oneshot::Sender<tokio::task::JoinHandle<()>>,
+        )>(self.tcp().config().max_connections as usize);
 
         // use a channel to know when the on_connect task is ready
         let (tx, rx) = oneshot::channel::<()>();
@@ -52,12 +55,12 @@ where
 
             while let Some((addr, notifier)) = from_node_receiver.recv().await {
                 let self_clone2 = self_clone.clone();
-                tokio::spawn(async move {
+                let handle = tokio::spawn(async move {
                     // perform the specified initial actions
                     self_clone2.on_connect(addr).await;
-                    // notify the node that the initial actions have concluded
-                    let _ = notifier.send(()); // can't really fail
                 });
+                // notify the node that the initial actions have concluded
+                let _ = notifier.send(handle); // can't really fail
             }
         });
         let _ = rx.await;
diff --git a/node/tcp/src/protocols/reading.rs b/node/tcp/src/protocols/reading.rs
@@ -25,10 +25,15 @@ use crate::{
 use async_trait::async_trait;
 use bytes::BytesMut;
 use futures_util::StreamExt;
-use std::{io, net::SocketAddr};
+use std::{
+    io,
+    net::SocketAddr,
+    time::{Duration, Instant},
+};
 use tokio::{
     io::AsyncRead,
     sync::{mpsc, oneshot},
+    time::timeout,
 };
 use tokio_util::codec::{Decoder, FramedRead};
 use tracing::*;
@@ -60,6 +65,9 @@ where
     /// The default value is 1024KiB.
     const INITIAL_BUFFER_SIZE: usize = 1024 * 1024;
 
+    /// The maximum time the node will wait for a new message before considering the connection dead.
+    const IDLE_TIMEOUT: Duration = Duration::from_secs(150);
+
     /// The final (deserialized) type of inbound messages.
     type Message: Send;
 
@@ -68,7 +76,7 @@ where
 
     /// Prepares the node to receive messages.
     async fn enable_reading(&self) {
-        let (conn_sender, mut conn_receiver) = mpsc::unbounded_channel();
+        let (conn_sender, mut conn_receiver) = mpsc::channel(self.tcp().config().max_connections as usize);
 
         // use a channel to know when the reading task is ready
         let (tx_reading, rx_reading) = oneshot::channel();
@@ -168,27 +176,57 @@ impl<R: Reading> ReadingInternal for R {
             // this task gets aborted, so there is no need for a dedicated timeout
             let _ = rx_conn_ready.await;
 
-            while let Some(bytes) = framed.next().await {
-                match bytes {
-                    Ok(msg) => {
+            // dropped message log suppression helpers
+            let mut dropped_count: usize = 0;
+            let mut last_drop_log = Instant::now();
+
+            loop {
+                let next_frame_future = framed.next();
+                let read_result = match timeout(Self::IDLE_TIMEOUT, next_frame_future).await {
+                    Ok(res) => res, // IO completed (success or error)
+                    Err(_) => {
+                        debug!(parent: node.span(), "connection with {addr} timed out due to inactivity");
+                        break;
+                    }
+                };
+                match read_result {
+                    Some(Ok(msg)) => {
                         // send the message for further processing
                         if let Err(e) = inbound_message_sender.try_send(msg) {
-                            error!(parent: node.span(), "can't process a message from {addr}: {e}");
                             node.stats().register_failure();
-                            if matches!(e, mpsc::error::TrySendError::Closed(_)) {
-                                break;
+                            match e {
+                                mpsc::error::TrySendError::Full(_) => {
+                                    // avoid log flooding
+                                    dropped_count += 1;
+                                    if last_drop_log.elapsed() >= Duration::from_secs(1) {
+                                        warn_about_dropped_messages(
+                                            &node,
+                                            addr,
+                                            &mut dropped_count,
+                                            &mut last_drop_log,
+                                        );
+                                    }
+                                }
+                                mpsc::error::TrySendError::Closed(_) => {
+                                    error!(parent: node.span(), "inbound channel closed for {addr}");
+                                    break;
+                                }
                             }
+                        } else if dropped_count != 0 {
+                            warn_about_dropped_messages(&node, addr, &mut dropped_count, &mut last_drop_log);
+                            debug!(parent: node.span(), "the inbound queue for {addr} is no longer saturated");
                         }
                         #[cfg(feature = "metrics")]
                         metrics::increment_gauge(metrics::tcp::TCP_TASKS, 1f64);
                     }
-                    Err(e) => {
+                    Some(Err(e)) => {
                         error!(parent: node.span(), "can't read from {addr}: {e}");
                         node.known_peers().register_failure(addr.ip());
                         if node.config().fatal_io_errors.contains(&e.kind()) {
                             break;
                         }
                     }
+                    None => break, // end of stream
                 }
             }
 
@@ -245,3 +283,15 @@ impl<D: Decoder> Decoder for CountingCodec<D> {
         Ok(ret)
     }
 }
+
+/// Warns that some messages were dropped and resets the related counters.
+fn warn_about_dropped_messages(node: &Tcp, addr: SocketAddr, dropped_count: &mut usize, last_drop_log: &mut Instant) {
+    warn!(
+        parent: node.span(),
+        "dropped {dropped_count} messages from {addr} due\
+        to inbound queue saturation",
+    );
+    // reset counters
+    *dropped_count = 0;
+    *last_drop_log = Instant::now();
+}
diff --git a/node/tcp/src/protocols/writing.rs b/node/tcp/src/protocols/writing.rs
diff --git a/node/tcp/src/tcp.rs b/node/tcp/src/tcp.rs

Original file line number	Diff line number	Diff line change
`@@ -137,3 +137,11 @@ impl Not for ConnectionSide {`
`137`	`137`	`}`
`138`	`138`	`}`
`139`	`139`	`}`
	`140`	`+`
	`141`	`+impl Drop for Connection {`
	`142`	`+ fn drop(&mut self) {`
	`143`	`+ for task in self.tasks.iter().rev() {`
	`144`	`+ task.abort();`
	`145`	`+ }`
	`146`	`+ }`
	`147`	`+}`