ckb-testkit 0.0.1

ckb testkit
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 crossbeam::channel::{unbounded, Receiver, Sender};
use p2p::{bytes::Bytes, context::SessionContext, multiaddr::Multiaddr, ProtocolId, SessionId};
use std::collections::HashMap;

/// Shared state between protocol handlers and service handler. As it is used across multiple
/// protocols, it should be wrapped into `Arc<RwLock<SharedState>>`.
pub struct SharedState {
    /// Session manager, #{ session.id => ( session, #{ protocol.id => mailbox } ) }
    session_manager: HashMap<
        SessionId,
        (
            SessionContext,
            HashMap<ProtocolId, (Sender<Bytes>, Receiver<Bytes>)>,
        ),
    >,
}

impl SharedState {
    /// Create a shared state.
    pub fn new() -> Self {
        Self {
            session_manager: HashMap::new(),
        }
    }

    pub fn add_session(&mut self, session: SessionContext) {
        self.session_manager
            .entry(session.id)
            .or_insert_with(|| (session.clone(), HashMap::new()));
    }

    pub fn remove_session(&mut self, session_id: &SessionId) -> Option<SessionContext> {
        self.session_manager
            .remove(session_id)
            .map(|(session, _mailbox)| session)
    }

    /// Get session by node address
    pub fn get_session(&self, address: &Multiaddr) -> Option<SessionContext> {
        for (session, _) in self.session_manager.values() {
            if &session.address == address {
                return Some(session.clone());
            }
        }
        None
    }

    pub fn add_protocol(&mut self, session: &SessionContext, protocol_id: ProtocolId) {
        let (protocol_mailbox_sender, protocol_mailbox_receiver) = unbounded::<Bytes>();
        self.session_manager
            .entry(session.id)
            .or_insert_with(|| (session.clone(), HashMap::new()))
            .1
            .insert(
                protocol_id,
                (protocol_mailbox_sender, protocol_mailbox_receiver),
            );
    }

    pub fn remove_protocol(&mut self, session_id: &SessionId, protocol_id: &ProtocolId) {
        let _ = self
            .session_manager
            .get_mut(session_id)
            .map(|(_session, mailbox)| mailbox.remove(protocol_id));
    }

    pub fn get_protocol_sender(
        &self,
        session_id: &SessionId,
        protocol_id: &ProtocolId,
    ) -> Option<Sender<Bytes>> {
        self.session_manager
            .get(session_id)
            .and_then(|(_, mailbox)| {
                mailbox
                    .get(protocol_id)
                    .map(|(sender, _receiver)| sender.clone())
            })
    }

    pub fn get_protocol_receiver(
        &self,
        session_id: &SessionId,
        protocol_id: &ProtocolId,
    ) -> Option<Receiver<Bytes>> {
        self.session_manager
            .get(session_id)
            .and_then(|(_, mailbox)| {
                mailbox
                    .get(protocol_id)
                    .map(|(_sender, receiver)| receiver.clone())
            })
    }

    pub fn get_opened_protocol_ids(&self, session_id: &SessionId) -> Option<Vec<ProtocolId>> {
        self.session_manager
            .get(session_id)
            .map(|(_session, mailbox)| mailbox.keys().cloned().collect())
    }

    /// Return all sessions
    pub fn get_sessions(&self) -> Vec<&SessionContext> {
        self.session_manager
            .values()
            .map(|(session, _)| session)
            .collect()
    }
}