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

use std::{
    collections::{HashMap, HashSet},
    sync::Arc,
};

use massa_models::endorsement::EndorsementId;
use massa_protocol_exports::PeerId;
use parking_lot::RwLock;
use schnellru::{ByLength, LruMap};

/// Cache of endorsements
pub struct EndorsementCache {
    /// List of endorsements we checked recently
    pub checked_endorsements: LruMap<EndorsementId, ()>,
    /// List of endorsements known by peers
    pub endorsements_known_by_peer: HashMap<PeerId, LruMap<EndorsementId, ()>>,
    /// Maximum number of endorsements known by a peer
    pub max_known_endorsements_by_peer: u32,
}

impl EndorsementCache {
    /// Create a new EndorsementCache
    pub fn new(max_known_endorsements: u32, max_known_endorsements_by_peer: u32) -> Self {
        Self {
            checked_endorsements: LruMap::new(ByLength::new(max_known_endorsements)),
            endorsements_known_by_peer: HashMap::new(),
            max_known_endorsements_by_peer,
        }
    }

    /// Mark a list of endorsement IDs prefixes as known by a peer
    pub fn insert_peer_known_endorsements(
        &mut self,
        peer_id: &PeerId,
        endorsements: &[EndorsementId],
    ) {
        let known_endorsements = self
            .endorsements_known_by_peer
            .entry(*peer_id)
            .or_insert_with(|| LruMap::new(ByLength::new(self.max_known_endorsements_by_peer)));
        for endorsement in endorsements {
            known_endorsements.insert(*endorsement, ());
        }
    }

    /// Mark an endorsement ID as checked by us
    pub fn insert_checked_endorsement(&mut self, enrodsement_id: EndorsementId) {
        self.checked_endorsements.insert(enrodsement_id, ());
    }

    /// Update caches to remove all data from disconnected peers
    pub fn update_cache(&mut self, peers_connected: &HashSet<PeerId>) {
        // Remove disconnected peers from cache
        self.endorsements_known_by_peer
            .retain(|peer_id, _| peers_connected.contains(peer_id));

        // Add new connected peers to cache
        for peer_id in peers_connected {
            match self.endorsements_known_by_peer.entry(*peer_id) {
                std::collections::hash_map::Entry::Occupied(_) => {}
                std::collections::hash_map::Entry::Vacant(entry) => {
                    entry.insert(LruMap::new(ByLength::new(
                        self.max_known_endorsements_by_peer,
                    )));
                }
            }
        }
    }
}

pub type SharedEndorsementCache = Arc<RwLock<EndorsementCache>>;