cfxcore/sync/message/
keys.rs

1// Copyright 2019 Conflux Foundation. All rights reserved.
2// Conflux is free software and distributed under GNU General Public License.
3// See http://www.gnu.org/licenses/
4
5use crate::{message::MsgId, sync::message::msgid};
6use cfx_types::H256;
7use malloc_size_of_derive::MallocSizeOf as DeriveMallocSizeOf;
8use parking_lot::{RwLock, RwLockReadGuard, RwLockWriteGuard};
9use std::collections::HashSet;
10
11#[derive(Hash, Eq, PartialEq, Debug, DeriveMallocSizeOf)]
12pub enum Key {
13    Hash(H256),
14    Num(u64),
15    Id(u32),
16}
17
18/// Common key container for all inflight requests. The supported message types
19/// are all registered in the Default constructor.
20#[derive(DeriveMallocSizeOf)]
21pub struct KeyContainer {
22    keys: Vec<Option<RwLock<HashSet<Key>>>>,
23}
24
25impl Default for KeyContainer {
26    fn default() -> Self {
27        let mut keys: Vec<Option<RwLock<HashSet<Key>>>> = Default::default();
28        for _ in 0..256 {
29            keys.push(None);
30        }
31        keys[msgid::GET_BLOCK_HASHES_BY_EPOCH as usize] =
32            Some(Default::default());
33        keys[msgid::GET_BLOCK_HEADERS as usize] = Some(Default::default());
34        keys[msgid::GET_BLOCKS as usize] = Some(Default::default());
35        keys[msgid::GET_TRANSACTIONS as usize] = Some(Default::default());
36        keys[msgid::GET_TRANSACTIONS_FROM_TX_HASHES as usize] =
37            Some(Default::default());
38        keys[msgid::NET_INFLIGHT_BLOCKS as usize] = Some(Default::default());
39
40        KeyContainer { keys }
41    }
42}
43
44impl KeyContainer {
45    pub fn read(&self, msg_type: MsgId) -> RwLockReadGuard<'_, HashSet<Key>> {
46        self.keys[msg_type as usize]
47            .as_ref()
48            .expect("msg not supported")
49            .read()
50    }
51
52    pub fn write(&self, msg_type: MsgId) -> RwLockWriteGuard<'_, HashSet<Key>> {
53        self.keys[msg_type as usize]
54            .as_ref()
55            .expect("msg not supported")
56            .write()
57    }
58
59    pub fn add(&mut self, msg_type: MsgId, key: Key) -> bool {
60        self.write(msg_type).insert(key)
61    }
62
63    pub fn remove(&mut self, msg_type: MsgId, key: Key) -> bool {
64        self.write(msg_type).remove(&key)
65    }
66}
cfxcore/sync/message/keys.rs

cfxcore/sync/message/
keys.rs
