hopr_transport_network/
track.rs1use std::sync::Arc;
2
3use hopr_api::PeerId;
4
5use super::observation::Observations;
6
7#[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 #[inline]
51 pub fn len(&self) -> usize {
52 self.peers.len()
53 }
54
55 #[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}