p2p/network/identify/
p2p_network_identify_state.rs

1use crate::{network::identify::stream::P2pNetworkIdentifyStreamState, PeerId, StreamId};
2use malloc_size_of_derive::MallocSizeOf;
3use serde::{Deserialize, Serialize};
4
5#[derive(Clone, Debug, Default, Serialize, Deserialize, MallocSizeOf)]
6pub struct P2pNetworkIdentifyState {
7    pub streams: crate::network::scheduler::StreamState<P2pNetworkIdentifyStreamState>,
8}
9
10impl P2pNetworkIdentifyState {
11    pub fn find_identify_stream_state(
12        &self,
13        peer_id: &PeerId,
14        stream_id: &StreamId,
15    ) -> Option<&P2pNetworkIdentifyStreamState> {
16        self.streams.get(peer_id)?.get(stream_id)
17    }
18
19    pub fn create_identify_stream_state(
20        &mut self,
21        peer_id: &PeerId,
22        stream_id: &StreamId,
23    ) -> Result<&mut P2pNetworkIdentifyStreamState, &P2pNetworkIdentifyStreamState> {
24        match self.streams.entry(*peer_id).or_default().entry(*stream_id) {
25            std::collections::btree_map::Entry::Vacant(e) => Ok(e.insert(Default::default())),
26            std::collections::btree_map::Entry::Occupied(e) => Err(e.into_mut()),
27        }
28    }
29
30    pub fn find_identify_stream_state_mut(
31        &mut self,
32        peer_id: &PeerId,
33        stream_id: &StreamId,
34    ) -> Option<&mut P2pNetworkIdentifyStreamState> {
35        self.streams.get_mut(peer_id)?.get_mut(stream_id)
36    }
37
38    pub fn remove_identify_stream_state(&mut self, peer_id: &PeerId, stream_id: &StreamId) -> bool {
39        self.streams
40            .get_mut(peer_id)
41            .is_some_and(|m| m.remove(stream_id).is_some())
42    }
43
44    pub fn prune_peer_state(&mut self, peer_id: &PeerId) {
45        self.streams.remove(peer_id);
46    }
47}