p2p/network/identify/
p2p_network_identify_state.rs1use 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}