hopr_transport_network/
track.rs

1use std::sync::Arc;
2
3use hopr_api::PeerId;
4
5use super::observation::Observations;
6
7/// Tracker of [`Observations`] for network peers.
8///
9/// This structure maintains a mapping between [`PeerId`] and their associated
10/// [`Observations`], allowing for efficient tracking and updating of peer telemetry data.
11///
12/// It can be combined with other objects to offer a complete view of the network state in regards
13/// to immediate peer probing.
14#[derive(Debug, Default, Clone)]
15pub struct NetworkPeerTracker {
16    peers: Arc<dashmap::DashMap<PeerId, Observations>>,
17}
18
19impl NetworkPeerTracker {
20    pub fn new() -> Self {
21        Self {
22            peers: Arc::new(dashmap::DashMap::new()),
23        }
24    }
25
26    #[inline]
27    pub fn add(&self, peer: PeerId) {
28        self.peers.entry(peer).or_default();
29    }
30
31    #[inline]
32    pub fn alter<F>(&self, peer: &PeerId, f: F)
33    where
34        F: FnOnce(&PeerId, Observations) -> Observations,
35    {
36        self.peers.alter(peer, f);
37    }
38
39    #[inline]
40    pub fn get(&self, peer: &PeerId) -> Option<Observations> {
41        self.peers.get(peer).map(|o| *o.value())
42    }
43
44    #[inline]
45    pub fn remove(&self, peer: &PeerId) {
46        self.peers.remove(peer);
47    }
48
49    /// The number of currently tracked peers.
50    #[inline]
51    pub fn len(&self) -> usize {
52        self.peers.len()
53    }
54
55    /// Check whether there are no tracked peers.
56    #[inline]
57    pub fn is_empty(&self) -> bool {
58        self.peers.len() == 0
59    }
60
61    #[inline]
62    pub fn iter_keys(&self) -> impl Iterator<Item = PeerId> + '_ {
63        self.peers.iter().map(|entry| *entry.key())
64    }
65}
66
67#[cfg(test)]
68mod tests {
69    use anyhow::Context;
70
71    use super::*;
72
73    #[test]
74    fn peer_tracker_adding_a_peer_adds_a_default_observation() -> anyhow::Result<()> {
75        let tracker = NetworkPeerTracker::new();
76
77        let peer = PeerId::random();
78
79        tracker.add(peer.clone());
80
81        assert_eq!(
82            tracker.get(&peer).context("should contain a value")?,
83            Observations::default()
84        );
85
86        Ok(())
87    }
88
89    #[test]
90    fn peer_tracker_adding_multiple_different_peers_results_in_higher_count() -> anyhow::Result<()> {
91        let tracker = NetworkPeerTracker::new();
92
93        const NUM_PEERS: usize = 10;
94
95        for _ in 0..NUM_PEERS {
96            tracker.add(PeerId::random());
97        }
98
99        assert_eq!(tracker.len(), NUM_PEERS);
100
101        Ok(())
102    }
103
104    #[test]
105    fn peer_tracker_should_reflect_the_alteration_changes() -> anyhow::Result<()> {
106        let tracker = NetworkPeerTracker::new();
107
108        let peer = PeerId::random();
109
110        tracker.add(peer.clone());
111        tracker.alter(&peer, |_, mut o| {
112            o.msg_sent += 1;
113            o
114        });
115
116        let obs = tracker.get(&peer).context("should contain a value")?;
117        assert_eq!(obs.msg_sent, 1);
118        assert_eq!(obs.ack_received, 0);
119
120        Ok(())
121    }
122
123    #[test]
124    fn peer_tracker_should_not_reflect_alterations_on_non_existent_peers() -> anyhow::Result<()> {
125        let tracker = NetworkPeerTracker::new();
126
127        let peer = PeerId::random();
128
129        tracker.alter(&peer, |_, mut o| {
130            o.msg_sent += 1;
131            o
132        });
133
134        assert!(tracker.get(&peer).is_none());
135
136        Ok(())
137    }
138}