1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
use crate::{
application::types::{PeerError, PeerInfo},
transport::ConnectionMetadata,
};
use diem_infallible::RwLock;
use diem_types::PeerId;
use std::{
collections::{hash_map::Entry, HashMap},
fmt::Debug,
hash::Hash,
};
pub type PeerMetadataStorage = LockingHashMap<PeerId, PeerInfo>;
pub struct LockingHashMap<Key: Clone + Debug + Eq + Hash, Value: Clone + Debug> {
map: RwLock<HashMap<Key, Value>>,
}
impl<Key, Value> LockingHashMap<Key, Value>
where
Key: Clone + Debug + Eq + Hash,
Value: Clone + Debug,
{
pub fn new() -> Self {
Self {
map: RwLock::new(HashMap::new()),
}
}
pub fn read(&self, key: &Key) -> Option<Value> {
self.map.read().get(key).cloned()
}
pub fn read_filtered<F: FnMut(&(&Key, &Value)) -> bool>(
&self,
filter: F,
) -> HashMap<Key, Value> {
self.map
.read()
.iter()
.filter(filter)
.map(|(key, value)| (key.clone(), value.clone()))
.collect()
}
pub fn keys(&self) -> Vec<Key> {
self.map.read().keys().cloned().collect()
}
pub fn read_all(&self) -> HashMap<Key, Value> {
self.map.read().clone()
}
pub fn insert(&self, key: Key, new_value: Value) {
let mut map = self.map.write();
map.entry(key)
.and_modify(|value| *value = new_value.clone())
.or_insert_with(|| new_value);
}
pub fn remove(&self, key: &Key) {
let mut map = self.map.write();
map.remove(key);
}
pub fn write<F: FnOnce(&mut Entry<Key, Value>) -> Result<(), PeerError>>(
&self,
key: Key,
modifier: F,
) -> Result<(), PeerError> {
let mut map = self.map.write();
modifier(&mut map.entry(key))
}
}
impl PeerMetadataStorage {
pub fn insert_connection(&self, connection_metadata: ConnectionMetadata) {
let peer_id = connection_metadata.remote_peer_id;
self.map
.write()
.entry(peer_id)
.and_modify(|entry| entry.active_connection = connection_metadata.clone())
.or_insert_with(|| PeerInfo::new(connection_metadata));
}
pub fn remove_connection(&self, connection_metadata: &ConnectionMetadata) {
let peer_id = connection_metadata.remote_peer_id;
let mut map = self.map.write();
if let Entry::Occupied(entry) = map.entry(peer_id) {
if entry.get().active_connection.connection_id == connection_metadata.connection_id {
entry.remove();
}
}
}
}