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    /// Removes a batch of addresses all at once.
59    pub fn batch_remove<I: Iterator<Item = IpAddr>>(&self, addrs: I) {
60        let mut peers = self.0.write();
61        for addr in addrs {
62            peers.remove(&addr);
63        }
64    }
65
66    /// Returns the list of all known peers and their stats.
67    pub fn snapshot(&self) -> HashMap<IpAddr, Arc<Stats>> {
68        self.0.read().clone()
69    }
70
71    /// Registers a submission of a message to the given address.
72    pub fn register_sent_message(&self, to: IpAddr, size: usize) {
73        if let Some(stats) = self.0.read().get(&to) {
74            stats.register_sent_message(size);
75        }
76    }
77
78    /// Registers a receipt of a message to the given address.
79    pub fn register_received_message(&self, from: IpAddr, size: usize) {
80        if let Some(stats) = self.0.read().get(&from) {
81            stats.register_received_message(size);
82        }
83    }
84
85    /// Registers a failure associated with the given address.
86    pub fn register_failure(&self, addr: IpAddr) {
87        if let Some(stats) = self.0.read().get(&addr) {
88            stats.register_failure();
89        }
90    }
91}
snarkos_node_tcp/helpers/known_peers.rs

snarkos_node_tcp/helpers/
known_peers.rs
