libp2p_perf/server/
behaviour.rs

1// Copyright 2023 Protocol Labs.
2//
3// Permission is hereby granted, free of charge, to any person obtaining a
4// copy of this software and associated documentation files (the "Software"),
5// to deal in the Software without restriction, including without limitation
6// the rights to use, copy, modify, merge, publish, distribute, sublicense,
7// and/or sell copies of the Software, and to permit persons to whom the
8// Software is furnished to do so, subject to the following conditions:
9//
10// The above copyright notice and this permission notice shall be included in
11// all copies or substantial portions of the Software.
12//
13// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
14// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
15// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
16// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
17// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
18// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
19// DEALINGS IN THE SOFTWARE.
20
21//! [`NetworkBehaviour`] of the libp2p perf server protocol.
22
23use std::{
24    collections::VecDeque,
25    task::{Context, Poll},
26};
27
28use libp2p_core::transport::PortUse;
29use libp2p_identity::PeerId;
30use libp2p_swarm::{
31    ConnectionId, FromSwarm, NetworkBehaviour, THandlerInEvent, THandlerOutEvent, ToSwarm,
32};
33
34use crate::{server::handler::Handler, Run};
35
36#[derive(Debug)]
37pub struct Event {
38    pub remote_peer_id: PeerId,
39    pub stats: Run,
40}
41
42#[derive(Default)]
43pub struct Behaviour {
44    /// Queue of actions to return when polled.
45    queued_events: VecDeque<ToSwarm<Event, THandlerInEvent<Self>>>,
46}
47
48impl Behaviour {
49    pub fn new() -> Self {
50        Self::default()
51    }
52}
53
54impl NetworkBehaviour for Behaviour {
55    type ConnectionHandler = Handler;
56    type ToSwarm = Event;
57
58    fn handle_established_inbound_connection(
59        &mut self,
60        _connection_id: ConnectionId,
61        _peer: PeerId,
62        _local_addr: &libp2p_core::Multiaddr,
63        _remote_addr: &libp2p_core::Multiaddr,
64    ) -> Result<libp2p_swarm::THandler<Self>, libp2p_swarm::ConnectionDenied> {
65        Ok(Handler::default())
66    }
67
68    fn handle_established_outbound_connection(
69        &mut self,
70        _connection_id: ConnectionId,
71        _peer: PeerId,
72        _addr: &libp2p_core::Multiaddr,
73        _role_override: libp2p_core::Endpoint,
74        _port_use: PortUse,
75    ) -> Result<libp2p_swarm::THandler<Self>, libp2p_swarm::ConnectionDenied> {
76        Ok(Handler::default())
77    }
78
79    fn on_swarm_event(&mut self, _event: FromSwarm) {}
80
81    fn on_connection_handler_event(
82        &mut self,
83        event_source: PeerId,
84        _connection_id: ConnectionId,
85        super::handler::Event { stats }: THandlerOutEvent<Self>,
86    ) {
87        self.queued_events.push_back(ToSwarm::GenerateEvent(Event {
88            remote_peer_id: event_source,
89            stats,
90        }))
91    }
92
93    #[tracing::instrument(level = "trace", name = "NetworkBehaviour::poll", skip(self))]
94    fn poll(&mut self, _: &mut Context<'_>) -> Poll<ToSwarm<Self::ToSwarm, THandlerInEvent<Self>>> {
95        if let Some(event) = self.queued_events.pop_front() {
96            return Poll::Ready(event);
97        }
98
99        Poll::Pending
100    }
101}