libp2p_perf/client/
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 client protocol.
22
23use std::{
24    collections::{HashSet, VecDeque},
25    task::{Context, Poll},
26};
27
28use libp2p_core::{transport::PortUse, Multiaddr};
29use libp2p_identity::PeerId;
30use libp2p_swarm::{
31    derive_prelude::ConnectionEstablished, ConnectionClosed, ConnectionId, FromSwarm,
32    NetworkBehaviour, NotifyHandler, THandlerInEvent, THandlerOutEvent, ToSwarm,
33};
34
35use super::{RunError, RunId};
36use crate::{client::handler::Handler, RunParams, RunUpdate};
37
38#[derive(Debug)]
39pub struct Event {
40    pub id: RunId,
41    pub result: Result<RunUpdate, RunError>,
42}
43
44#[derive(Default)]
45pub struct Behaviour {
46    /// Queue of actions to return when polled.
47    queued_events: VecDeque<ToSwarm<Event, THandlerInEvent<Self>>>,
48    /// Set of connected peers.
49    connected: HashSet<PeerId>,
50}
51
52impl Behaviour {
53    pub fn new() -> Self {
54        Self::default()
55    }
56
57    pub fn perf(&mut self, server: PeerId, params: RunParams) -> Result<RunId, NotConnected> {
58        if !self.connected.contains(&server) {
59            return Err(NotConnected {});
60        }
61
62        let id = RunId::next();
63
64        self.queued_events.push_back(ToSwarm::NotifyHandler {
65            peer_id: server,
66            handler: NotifyHandler::Any,
67            event: crate::client::handler::Command { id, params },
68        });
69
70        Ok(id)
71    }
72}
73
74#[derive(thiserror::Error, Debug)]
75pub struct NotConnected();
76
77impl std::fmt::Display for NotConnected {
78    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
79        write!(f, "not connected to peer")
80    }
81}
82
83impl NetworkBehaviour for Behaviour {
84    type ConnectionHandler = Handler;
85    type ToSwarm = Event;
86
87    fn handle_established_outbound_connection(
88        &mut self,
89        _connection_id: ConnectionId,
90        _peer: PeerId,
91        _addr: &Multiaddr,
92        _role_override: libp2p_core::Endpoint,
93        _port_use: PortUse,
94    ) -> Result<libp2p_swarm::THandler<Self>, libp2p_swarm::ConnectionDenied> {
95        Ok(Handler::default())
96    }
97
98    fn handle_established_inbound_connection(
99        &mut self,
100        _connection_id: ConnectionId,
101        _peer: PeerId,
102        _local_addr: &Multiaddr,
103        _remote_addr: &Multiaddr,
104    ) -> Result<libp2p_swarm::THandler<Self>, libp2p_swarm::ConnectionDenied> {
105        Ok(Handler::default())
106    }
107
108    fn on_swarm_event(&mut self, event: FromSwarm) {
109        match event {
110            FromSwarm::ConnectionEstablished(ConnectionEstablished { peer_id, .. }) => {
111                self.connected.insert(peer_id);
112            }
113            FromSwarm::ConnectionClosed(ConnectionClosed {
114                peer_id,
115                connection_id: _,
116                endpoint: _,
117                remaining_established,
118                ..
119            }) => {
120                if remaining_established == 0 {
121                    assert!(self.connected.remove(&peer_id));
122                }
123            }
124            _ => {}
125        }
126    }
127
128    fn on_connection_handler_event(
129        &mut self,
130        _event_source: PeerId,
131        _connection_id: ConnectionId,
132        super::handler::Event { id, result }: THandlerOutEvent<Self>,
133    ) {
134        self.queued_events
135            .push_back(ToSwarm::GenerateEvent(Event { id, result }));
136    }
137
138    #[tracing::instrument(level = "trace", name = "NetworkBehaviour::poll", skip(self))]
139    fn poll(&mut self, _: &mut Context<'_>) -> Poll<ToSwarm<Self::ToSwarm, THandlerInEvent<Self>>> {
140        if let Some(event) = self.queued_events.pop_front() {
141            return Poll::Ready(event);
142        }
143
144        Poll::Pending
145    }
146}
libp2p_perf/client/behaviour.rs

libp2p_perf/client/
behaviour.rs
