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

// Copyright 2019 Conflux Foundation. All rights reserved.
// Conflux is free software and distributed under GNU General Public License.
// See http://www.gnu.org/licenses/

use crate::{message::MsgId, sync::message::msgid};
use cfx_types::H256;
use malloc_size_of_derive::MallocSizeOf as DeriveMallocSizeOf;
use parking_lot::{RwLock, RwLockReadGuard, RwLockWriteGuard};
use std::collections::HashSet;

#[derive(Hash, Eq, PartialEq, Debug, DeriveMallocSizeOf)]
pub enum Key {
    Hash(H256),
    Num(u64),
    Id(u32),
}

/// Common key container for all inflight requests. The supported message types
/// are all registered in the Default constructor.
#[derive(DeriveMallocSizeOf)]
pub struct KeyContainer {
    keys: Vec<Option<RwLock<HashSet<Key>>>>,
}

impl Default for KeyContainer {
    fn default() -> Self {
        let mut keys: Vec<Option<RwLock<HashSet<Key>>>> = Default::default();
        for _ in 0..256 {
            keys.push(None);
        }
        keys[msgid::GET_BLOCK_HASHES_BY_EPOCH as usize] =
            Some(Default::default());
        keys[msgid::GET_BLOCK_HEADERS as usize] = Some(Default::default());
        keys[msgid::GET_BLOCKS as usize] = Some(Default::default());
        keys[msgid::GET_TRANSACTIONS as usize] = Some(Default::default());
        keys[msgid::GET_TRANSACTIONS_FROM_TX_HASHES as usize] =
            Some(Default::default());
        keys[msgid::NET_INFLIGHT_BLOCKS as usize] = Some(Default::default());

        KeyContainer { keys }
    }
}

impl KeyContainer {
    pub fn read(&self, msg_type: MsgId) -> RwLockReadGuard<HashSet<Key>> {
        self.keys[msg_type as usize]
            .as_ref()
            .expect("msg not supported")
            .read()
    }

    pub fn write(&self, msg_type: MsgId) -> RwLockWriteGuard<HashSet<Key>> {
        self.keys[msg_type as usize]
            .as_ref()
            .expect("msg not supported")
            .write()
    }

    pub fn add(&mut self, msg_type: MsgId, key: Key) -> bool {
        self.write(msg_type).insert(key)
    }

    pub fn remove(&mut self, msg_type: MsgId, key: Key) -> bool {
        self.write(msg_type).remove(&key)
    }
}