snarkos_node_tcp/helpers/
known_peers.rs

1// Copyright (c) 2019-2025 Provable Inc.
2// This file is part of the snarkOS library.
3
4// Licensed under the Apache License, Version 2.0 (the "License");
5// you may not use this file except in compliance with the License.
6// You may obtain a copy of the License at:
7
8// http://www.apache.org/licenses/LICENSE-2.0
9
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16use std::{
17    collections::{HashMap, hash_map::Entry},
18    net::IpAddr,
19    sync::Arc,
20    time::Instant,
21};
22
23#[cfg(feature = "locktick")]
24use locktick::parking_lot::RwLock;
25#[cfg(not(feature = "locktick"))]
26use parking_lot::RwLock;
27
28use crate::Stats;
29
30/// Contains statistics related to Tcp's peers, currently connected or not.
31#[derive(Default)]
32pub struct KnownPeers(RwLock<HashMap<IpAddr, Arc<Stats>>>);
33
34impl KnownPeers {
35    /// Adds an address to the list of known peers.
36    pub fn add(&self, addr: IpAddr) {
37        let timestamp = Instant::now();
38        match self.0.write().entry(addr) {
39            Entry::Vacant(entry) => {
40                entry.insert(Arc::new(Stats::new(timestamp)));
41            }
42            Entry::Occupied(entry) => {
43                *entry.get().timestamp.write() = timestamp;
44            }
45        }
46    }
47
48    /// Returns the stats for the given peer.
49    pub fn get(&self, addr: IpAddr) -> Option<Arc<Stats>> {
50        self.0.read().get(&addr).map(Arc::clone)
51    }
52
53    /// Removes an address from the list of known peers.
54    pub fn remove(&self, addr: IpAddr) -> Option<Arc<Stats>> {
55        self.0.write().remove(&addr)
56    }
57
58    /// Returns the list of all known peers and their stats.
59    pub fn snapshot(&self) -> HashMap<IpAddr, Arc<Stats>> {
60        self.0.read().clone()
61    }
62
63    /// Registers a submission of a message to the given address.
64    pub fn register_sent_message(&self, to: IpAddr, size: usize) {
65        if let Some(stats) = self.0.read().get(&to) {
66            stats.register_sent_message(size);
67        }
68    }
69
70    /// Registers a receipt of a message to the given address.
71    pub fn register_received_message(&self, from: IpAddr, size: usize) {
72        if let Some(stats) = self.0.read().get(&from) {
73            stats.register_received_message(size);
74        }
75    }
76
77    /// Registers a failure associated with the given address.
78    pub fn register_failure(&self, addr: IpAddr) {
79        if let Some(stats) = self.0.read().get(&addr) {
80            stats.register_failure();
81        }
82    }
83}
snarkos_node_tcp/helpers/known_peers.rs

snarkos_node_tcp/helpers/
known_peers.rs
