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

// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0
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>;

/// A generic locking hash map with ability to read before write consistency
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()),
        }
    }

    /// Get a clone of the value
    pub fn read(&self, key: &Key) -> Option<Value> {
        self.map.read().get(key).cloned()
    }

    /// Filtered read clone based on keys or values
    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()
    }

    /// All keys of the hash map
    pub fn keys(&self) -> Vec<Key> {
        self.map.read().keys().cloned().collect()
    }

    /// Read a clone of the entire state
    pub fn read_all(&self) -> HashMap<Key, Value> {
        self.map.read().clone()
    }

    /// Insert new entry
    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);
    }

    /// Remove old entries
    pub fn remove(&self, key: &Key) {
        let mut map = self.map.write();
        map.remove(key);
    }

    /// Take in a function to modify the data, must handle concurrency control with the input function
    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();

        // Don't remove the peer if the connection doesn't match!
        if let Entry::Occupied(entry) = map.entry(peer_id) {
            // For now, remove the peer entirely, we could in the future have multiple connections for a peer
            if entry.get().active_connection.connection_id == connection_metadata.connection_id {
                entry.remove();
            }
        }
    }
}