snarkos_node_router/

heartbeat.rs

// Copyright (c) 2019-2025 Provable Inc.
// This file is part of the snarkOS library.

// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at:

// http://www.apache.org/licenses/LICENSE-2.0

// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use crate::{
    ConnectedPeer,
    NodeType,
    Outbound,
    PeerPoolHandling,
    Router,
    bootstrap_peers,
    messages::{DisconnectReason, Message, PeerRequest},
};
use snarkvm::prelude::Network;

use snarkos_node_tcp::P2P;

use colored::Colorize;
use rand::{prelude::IteratorRandom, rngs::OsRng};

/// A helper function to compute the maximum of two numbers.
/// See Rust issue 92391: https://github.com/rust-lang/rust/issues/92391.
pub const fn max(a: usize, b: usize) -> usize {
    match a > b {
        true => a,
        false => b,
    }
}

#[async_trait]
pub trait Heartbeat<N: Network>: Outbound<N> {
    /// The duration in seconds to sleep in between heartbeat executions.
    const HEARTBEAT_IN_SECS: u64 = 25; // 25 seconds
    /// The minimum number of peers required to maintain connections with.
    const MINIMUM_NUMBER_OF_PEERS: usize = 3;
    /// The median number of peers to maintain connections with.
    const MEDIAN_NUMBER_OF_PEERS: usize = max(Self::MAXIMUM_NUMBER_OF_PEERS / 2, Self::MINIMUM_NUMBER_OF_PEERS);
    /// The maximum number of peers permitted to maintain connections with.
    const MAXIMUM_NUMBER_OF_PEERS: usize = 21;
    /// The maximum number of provers to maintain connections with.
    const MAXIMUM_NUMBER_OF_PROVERS: usize = Self::MAXIMUM_NUMBER_OF_PEERS / 4;
    /// The amount of time an IP address is prohibited from connecting.
    const IP_BAN_TIME_IN_SECS: u64 = 300;

    /// Handles the heartbeat request.
    async fn heartbeat(&self) {
        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.
        self.handle_connected_peers();
        // Keep the bootstrap peers within the allowed range.
        self.handle_bootstrap_peers().await;
        // Keep the trusted peers connected.
        self.handle_trusted_peers().await;
        // Keep the puzzle request up to date.
        self.handle_puzzle_request();
        // Unban any addresses whose ban time has expired.
        self.handle_banned_ips();
    }

    /// TODO (howardwu): Consider checking minimum number of validators, to exclude clients and provers.
    /// This function performs safety checks on the setting for the minimum number of peers.
    fn safety_check_minimum_number_of_peers(&self) {
        // Perform basic sanity checks on the configuration for the number of peers.
        assert!(Self::MINIMUM_NUMBER_OF_PEERS >= 1, "The minimum number of peers must be at least 1.");
        assert!(Self::MINIMUM_NUMBER_OF_PEERS <= Self::MAXIMUM_NUMBER_OF_PEERS);
        assert!(Self::MINIMUM_NUMBER_OF_PEERS <= Self::MEDIAN_NUMBER_OF_PEERS);
        assert!(Self::MEDIAN_NUMBER_OF_PEERS <= Self::MAXIMUM_NUMBER_OF_PEERS);
        assert!(Self::MAXIMUM_NUMBER_OF_PROVERS <= Self::MAXIMUM_NUMBER_OF_PEERS);
    }

    /// This function logs the connected peers.
    fn log_connected_peers(&self) {
        // Log the connected peers.
        let connected_peers = self.router().connected_peers();
        let connected_peers_fmt = format!("{connected_peers:?}").dimmed();
        match connected_peers.len() {
            0 => warn!("No connected peers"),
            1 => debug!("Connected to 1 peer: {connected_peers_fmt}"),
            num_connected => debug!("Connected to {num_connected} peers {connected_peers_fmt}"),
        }
    }

    /// 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.
    ///
    /// Rules:
    ///     - Trusted peers and bootstrap nodes are not removable.
    ///     - Peers that we are currently syncing with are not removable.
    ///     - Connections that have not been seen in a while are considered lower priority.
    fn get_removable_peers(&self) -> Vec<ConnectedPeer<N>> {
        // Are we synced already? (cache this here, so it does not need to be recomputed)
        let is_block_synced = self.is_block_synced();

        // Sort by priority, where lowest priority will be at the beginning
        // of the vector.
        // Note, that this gives equal priority to clients and provers, which
        // we might want to change in the future.
        let mut peers = self.router().filter_connected_peers(|peer| {
            !peer.trusted
                && peer.node_type != NodeType::BootstrapClient
                && !self.router().cache.contains_inbound_block_request(&peer.listener_addr) // This peer is currently syncing from us.
                && (is_block_synced || self.router().cache.num_outbound_block_requests(&peer.listener_addr) == 0) // We are currently syncing from this peer.
        });
        peers.sort_by_key(|peer| peer.last_seen);

        peers
    }

    /// This function removes the peer that we have not heard from the longest,
    /// to keep the connections fresh.
    /// It only triggers if the router is above the minimum number of connected peers.
    fn remove_oldest_connected_peer(&self) {
        // Skip if the node is not requesting peers.
        if self.router().trusted_peers_only() {
            return;
        }

        // Skip if the router is at or below the minimum number of connected peers.
        if self.router().number_of_connected_peers() <= Self::MINIMUM_NUMBER_OF_PEERS {
            return;
        }

        // Disconnect from the oldest connected peer, which is the first entry in the list
        // of removable peers.
        // Do nothing, if the list is empty.
        if let Some(oldest) = self.get_removable_peers().first().map(|peer| peer.listener_addr) {
            info!("Disconnecting from '{oldest}' (periodic refresh of peers)");
            let _ = self.router().send(oldest, Message::Disconnect(DisconnectReason::PeerRefresh.into()));
            self.router().disconnect(oldest);
        }
    }

    /// This function keeps the number of connected peers within the allowed range.
    fn handle_connected_peers(&self) {
        // Initialize an RNG.
        let rng = &mut OsRng;

        // Obtain the number of connected peers.
        let num_connected = self.router().number_of_connected_peers();
        // Obtain the number of connected provers.
        let num_connected_provers = self.router().filter_connected_peers(|peer| peer.node_type.is_prover()).len();

        // Determine the maximum number of peers and provers to keep.
        let (max_peers, max_provers) = (Self::MAXIMUM_NUMBER_OF_PEERS, Self::MAXIMUM_NUMBER_OF_PROVERS);

        // Compute the number of surplus peers.
        let num_surplus_peers = num_connected.saturating_sub(max_peers);
        // Compute the number of surplus provers.
        let num_surplus_provers = num_connected_provers.saturating_sub(max_provers);
        // Compute the number of provers remaining connected.
        let num_remaining_provers = num_connected_provers.saturating_sub(num_surplus_provers);
        // Compute the number of surplus clients and validators.
        let num_surplus_clients_validators = num_surplus_peers.saturating_sub(num_remaining_provers);

        if num_surplus_provers > 0 || num_surplus_clients_validators > 0 {
            debug!(
                "Exceeded maximum number of connected peers, disconnecting from ({num_surplus_provers} + {num_surplus_clients_validators}) peers"
            );

            // Determine the provers to disconnect from.
            let provers_to_disconnect = self
                .router()
                .filter_connected_peers(|peer| peer.node_type.is_prover() && !peer.trusted)
                .into_iter()
                .choose_multiple(rng, num_surplus_provers);

            // Determine the clients and validators to disconnect from.
            let peers_to_disconnect = self
                .get_removable_peers()
                .into_iter()
                .filter(|peer| !peer.node_type.is_prover()) // remove provers as those are handled separately
                .take(num_surplus_clients_validators);

            // Proceed to send disconnect requests to these peers.
            for peer in peers_to_disconnect.chain(provers_to_disconnect) {
                // TODO (howardwu): Remove this after specializing this function.
                if self.router().node_type().is_prover() {
                    continue;
                }

                let peer_addr = peer.listener_addr;
                info!("Disconnecting from '{peer_addr}' (exceeded maximum connections)");
                self.router().send(peer_addr, Message::Disconnect(DisconnectReason::TooManyPeers.into()));
                // Disconnect from this peer.
                self.router().disconnect(peer_addr);
            }
        }

        // Obtain the number of connected peers.
        let num_connected = self.router().number_of_connected_peers();
        // Compute the number of deficit peers.
        let num_deficient = Self::MEDIAN_NUMBER_OF_PEERS.saturating_sub(num_connected);

        if num_deficient > 0 {
            // Initialize an RNG.
            let rng = &mut OsRng;

            // Attempt to connect to more peers, separately choosing from those at a greater block
            // height, and those whose height is lower or unknown to us.
            let own_height = self.router().ledger.latest_block_height();
            let (higher_peers, other_peers): (Vec<_>, Vec<_>) = self
                .router()
                .get_candidate_peers()
                .into_iter()
                .partition(|peer| peer.last_height_seen.map(|h| h > own_height).unwrap_or(false));
            // We may not know of half of `num_deficient` candidates; account for it using `min`.
            let num_higher_peers = num_deficient.div_ceil(2).min(higher_peers.len());
            for peer in higher_peers.into_iter().choose_multiple(rng, num_higher_peers) {
                self.router().connect(peer.listener_addr);
            }
            for peer in other_peers.into_iter().choose_multiple(rng, num_deficient - num_higher_peers) {
                self.router().connect(peer.listener_addr);
            }

            if !self.router().trusted_peers_only() {
                // Request more peers from the connected peers.
                for peer_ip in self.router().connected_peers().into_iter().choose_multiple(rng, 3) {
                    self.router().send(peer_ip, Message::PeerRequest(PeerRequest));
                }
            }
        }
    }

    /// This function keeps the number of bootstrap peers within the allowed range.
    async fn handle_bootstrap_peers(&self) {
        // Return early if we are in trusted peers only mode.
        if self.router().trusted_peers_only() {
            return;
        }
        // Split the bootstrap peers into connected and candidate lists.
        let mut candidate_bootstrap = Vec::new();
        let connected_bootstrap =
            self.router().filter_connected_peers(|peer| peer.node_type == NodeType::BootstrapClient);
        for bootstrap_ip in bootstrap_peers::<N>(self.router().is_dev()) {
            if !connected_bootstrap.iter().any(|peer| peer.listener_addr == bootstrap_ip) {
                candidate_bootstrap.push(bootstrap_ip);
            }
        }
        // If there are not enough connected bootstrap peers, connect to more.
        if connected_bootstrap.is_empty() {
            // Initialize an RNG.
            let rng = &mut OsRng;
            // Attempt to connect to a bootstrap peer.
            if let Some(peer_ip) = candidate_bootstrap.into_iter().choose(rng) {
                match self.router().connect(peer_ip) {
                    Some(hdl) => {
                        let result = hdl.await;
                        if let Err(err) = result {
                            warn!("Failed to connect to bootstrap peer at {peer_ip}: {err}");
                        }
                    }
                    None => warn!("Could not initiate connect to bootstrap peer at {peer_ip}"),
                }
            }
        }
        // Determine if the node is connected to more bootstrap peers than allowed.
        let num_surplus = connected_bootstrap.len().saturating_sub(1);
        if num_surplus > 0 {
            // Initialize an RNG.
            let rng = &mut OsRng;
            // Proceed to send disconnect requests to these bootstrap peers.
            for peer in connected_bootstrap.into_iter().choose_multiple(rng, num_surplus) {
                info!("Disconnecting from '{}' (exceeded maximum bootstrap)", peer.listener_addr);
                self.router().send(peer.listener_addr, Message::Disconnect(DisconnectReason::TooManyPeers.into()));
                // Disconnect from this peer.
                self.router().disconnect(peer.listener_addr);
            }
        }
    }

    /// This function attempts to connect to any disconnected trusted peers.
    async fn handle_trusted_peers(&self) {
        // Ensure that the trusted nodes are connected.
        let handles: Vec<_> = self
            .router()
            .unconnected_trusted_peers()
            .iter()
            .filter_map(|listener_addr| {
                debug!("Attempting to (re-)connect to trusted peer `{listener_addr}`");
                let hdl = self.router().connect(*listener_addr);
                if hdl.is_none() {
                    warn!("Could not initiate connection to trusted peer at `{listener_addr}`");
                }
                hdl
            })
            .collect();

        for result in futures::future::join_all(handles).await {
            if let Err(err) = result {
                warn!("Could not connect to trusted peer: {err}");
            }
        }
    }

    /// This function updates the puzzle if network has updated.
    fn handle_puzzle_request(&self) {
        // No-op
    }

    // Remove addresses whose ban time has expired.
    fn handle_banned_ips(&self) {
        self.router().tcp().banned_peers().remove_old_bans(Self::IP_BAN_TIME_IN_SECS);
    }
}