fluencelabs/rust-libp2p
1
0
Fork 0
mirror of https://github.com/fluencelabs/rust-libp2p synced 2025-06-26 16:21:39 +00:00
Code Issues Projects Releases Wiki Activity

Simplify the implementation of Ping (#884)

Browse Source
This commit is contained in:
Pierre Krieger
2019-01-28 15:06:07 +01:00
committed by GitHub
parent df923526ca
commit 073df709dd
4 changed files with 111 additions and 756 deletions
Download Patch File Download Diff File Expand all files Collapse all files

323
protocols/ping/src/dial_handler.rs
View File

@ -1,323 +0,0 @@
// Copyright 2018 Parity Technologies (UK) Ltd.
//
// Permission is hereby granted, free of charge, to any person obtaining a
// copy of this software and associated documentation files (the "Software"),
// to deal in the Software without restriction, including without limitation
// the rights to use, copy, modify, merge, publish, distribute, sublicense,
// and/or sell copies of the Software, and to permit persons to whom the
// Software is furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
// DEALINGS IN THE SOFTWARE.
use crate::protocol::{Ping, PingDialer};
use futures::prelude::*;
use libp2p_core::{
OutboundUpgrade,
ProtocolsHandler,
ProtocolsHandlerEvent,
protocols_handler::ProtocolsHandlerUpgrErr,
upgrade::DeniedUpgrade
};
use log::warn;
use std::{
io, mem,
time::{Duration, Instant},
};
use tokio_io::{AsyncRead, AsyncWrite};
use tokio_timer::{self, Delay};
use void::{Void, unreachable};
/// Protocol handler that handles pinging the remote at a regular period.
///
/// If the remote doesn't respond, produces an error that closes the connection.
pub struct PeriodicPingHandler<TSubstream> {
/// Configuration for the ping protocol.
ping_config: Ping<Instant>,
/// State of the outgoing ping.
out_state: OutState<TSubstream>,
/// Duration after which we consider that a ping failed.
ping_timeout: Duration,
/// After a ping succeeded, wait this long before the next ping.
delay_to_next_ping: Duration,
/// If true, we switch to the `Disabled` state if the remote doesn't support the ping protocol.
/// If false, we close the connection.
tolerate_unsupported: bool,
}
/// State of the outgoing ping substream.
enum OutState<TSubstream> {
/// We need to open a new substream.
NeedToOpen {
/// Timeout after which we decide that it's not going to work out.
///
/// Theoretically the handler should be polled immediately after we set the state to
/// `NeedToOpen` and then we immediately transition away from it. However if the local node
/// is for some reason busy, creating the `Delay` here avoids being overly generous with
/// the ping timeout.
expires: Delay,
},
/// Upgrading a substream to use ping.
///
/// We produced a substream open request, and are waiting for it to be upgraded to a full
/// ping-powered substream.
Upgrading {
/// Timeout after which we decide that it's not going to work out.
///
/// The user of the `ProtocolsHandler` should ensure that there's a timeout when upgrading,
/// but by storing a timeout here as well we ensure that we keep track of how long the
/// ping has lasted.
expires: Delay,
},
/// We sent a ping and we are waiting for the pong.
WaitingForPong {
/// Substream where we should receive the pong.
substream: PingDialer<TSubstream, Instant>,
/// Timeout after which we decide that we're not going to receive the pong.
expires: Delay,
},
/// We received a pong and now we have nothing to do except wait a bit before sending the
/// next ping.
Idle {
/// The substream to use to send pings.
substream: PingDialer<TSubstream, Instant>,
/// When to send the ping next.
next_ping: Delay,
},
/// The ping dialer is disabled. Don't do anything.
Disabled,
/// The dialer has been closed.
Shutdown,
/// Something bad happened during the previous polling.
Poisoned,
}
/// Event produced by the periodic pinger.
#[derive(Debug, Copy, Clone)]
pub enum OutEvent {
/// Started pinging the remote. This can be used to print a diagnostic message in the logs.
PingStart,
/// The node has successfully responded to a ping.
PingSuccess(Duration),
}
impl<TSubstream> PeriodicPingHandler<TSubstream> {
/// Builds a new `PeriodicPingHandler`.
pub fn new() -> PeriodicPingHandler<TSubstream> {
let ping_timeout = Duration::from_secs(30);
PeriodicPingHandler {
ping_config: Default::default(),
out_state: OutState::NeedToOpen {
expires: Delay::new(Instant::now() + ping_timeout),
},
ping_timeout,
delay_to_next_ping: Duration::from_secs(15),
tolerate_unsupported: false,
}
}
}
impl<TSubstream> Default for PeriodicPingHandler<TSubstream> {
#[inline]
fn default() -> Self {
PeriodicPingHandler::new()
}
}
impl<TSubstream> ProtocolsHandler for PeriodicPingHandler<TSubstream>
where
TSubstream: AsyncRead + AsyncWrite,
{
type InEvent = Void;
type OutEvent = OutEvent;
type Error = io::Error; // TODO: more precise error type
type Substream = TSubstream;
type InboundProtocol = DeniedUpgrade;
type OutboundProtocol = Ping<Instant>;
type OutboundOpenInfo = ();
#[inline]
fn listen_protocol(&self) -> Self::InboundProtocol {
DeniedUpgrade
}
fn inject_fully_negotiated_inbound(&mut self, protocol: Void) {
unreachable(protocol)
}
fn inject_fully_negotiated_outbound(
&mut self,
mut substream: <Self::OutboundProtocol as OutboundUpgrade<TSubstream>>::Output,
_info: Self::OutboundOpenInfo
) {
if let OutState::Upgrading { expires } = mem::replace(&mut self.out_state, OutState::Poisoned) {
// We always upgrade with the intent of immediately pinging.
substream.ping(Instant::now());
self.out_state = OutState::WaitingForPong { substream, expires }
}
}
fn inject_event(&mut self, _: Self::InEvent) {}
fn inject_inbound_closed(&mut self) {}
#[inline]
fn inject_dial_upgrade_error(&mut self, _: Self::OutboundOpenInfo, _: ProtocolsHandlerUpgrErr<<Self::OutboundProtocol as OutboundUpgrade<Self::Substream>>::Error>) {
// In case of error while upgrading, there's not much we can do except shut down.
// TODO: we assume that the error is about ping not being supported, which is not
// necessarily the case
if self.tolerate_unsupported {
self.out_state = OutState::Disabled;
} else {
self.out_state = OutState::Shutdown;
}
}
#[inline]
fn connection_keep_alive(&self) -> bool {
false
}
fn shutdown(&mut self) {
// Put `Shutdown` in `self.out_state` if we don't have any substream open.
// Otherwise, keep the state as it is but call `shutdown()` on the substream. This
// guarantees that the dialer will return `None` at some point.
match self.out_state {
OutState::WaitingForPong {
ref mut substream, ..
} => substream.shutdown(),
OutState::Idle {
ref mut substream, ..
} => substream.shutdown(),
ref mut s => *s = OutState::Shutdown,
}
}
fn poll(
&mut self,
) -> Poll<
ProtocolsHandlerEvent<Self::OutboundProtocol, Self::OutboundOpenInfo, Self::OutEvent>,
io::Error,
> {
// Shortcut for polling a `tokio_timer::Delay`
macro_rules! poll_delay {
($delay:expr => { NotReady => $notready:expr, Ready => $ready:expr, }) => (
match $delay.poll() {
Ok(Async::NotReady) => $notready,
Ok(Async::Ready(())) => $ready,
Err(err) => {
warn!(target: "sub-libp2p", "Ping timer errored: {:?}", err);
return Err(io::Error::new(io::ErrorKind::Other, err));
}
}
)
}
match mem::replace(&mut self.out_state, OutState::Poisoned) {
OutState::Shutdown | OutState::Poisoned => {
// This shuts down the whole connection with the remote.
Ok(Async::Ready(ProtocolsHandlerEvent::Shutdown))
},
OutState::Disabled => {
Ok(Async::NotReady)
}
// Need to open an outgoing substream.
OutState::NeedToOpen { expires } => {
// Note that we ignore the expiration here, as it's pretty unlikely to happen.
// The expiration is only here to be transmitted to the `Upgrading`.
self.out_state = OutState::Upgrading { expires };
Ok(Async::Ready(
ProtocolsHandlerEvent::OutboundSubstreamRequest {
upgrade: self.ping_config,
info: (),
},
))
}
// Waiting for the upgrade to be negotiated.
OutState::Upgrading { mut expires } => poll_delay!(expires => {
NotReady => {
self.out_state = OutState::Upgrading { expires };
Ok(Async::NotReady)
},
Ready => {
self.out_state = OutState::Shutdown;
Err(io::Error::new(io::ErrorKind::Other, "unresponsive node"))
},
}),
// Waiting for the pong.
OutState::WaitingForPong { mut substream, mut expires } => {
// We start by dialing the substream, leaving one last chance for it to
// produce the pong even if the expiration happened.
match substream.poll()? {
Async::Ready(Some(started)) => {
self.out_state = OutState::Idle {
substream,
next_ping: Delay::new(Instant::now() + self.delay_to_next_ping),
};
let ev = OutEvent::PingSuccess(started.elapsed());
return Ok(Async::Ready(ProtocolsHandlerEvent::Custom(ev)));
}
Async::NotReady => {}
Async::Ready(None) => {
self.out_state = OutState::Shutdown;
return Ok(Async::Ready(ProtocolsHandlerEvent::Shutdown));
}
}
// Check the expiration.
poll_delay!(expires => {
NotReady => {
self.out_state = OutState::WaitingForPong { substream, expires };
// Both `substream` and `expires` and not ready, so it's fine to return
// not ready.
Ok(Async::NotReady)
},
Ready => {
self.out_state = OutState::Shutdown;
Err(io::Error::new(io::ErrorKind::Other, "unresponsive node"))
},
})
}
OutState::Idle { mut substream, mut next_ping } => {
// Poll the future that fires when we need to ping the node again.
poll_delay!(next_ping => {
NotReady => {
self.out_state = OutState::Idle { substream, next_ping };
Ok(Async::NotReady)
},
Ready => {
let expires = Delay::new(Instant::now() + self.ping_timeout);
substream.ping(Instant::now());
self.out_state = OutState::WaitingForPong { substream, expires };
Ok(Async::Ready(ProtocolsHandlerEvent::Custom(OutEvent::PingStart)))
},
})
}
}
}
}

15
protocols/ping/src/lib.rs
View File

@ -31,14 +31,12 @@
//!
//! When a ping succeeds, a `PingSuccess` event is generated, indicating the time the ping took.
pub mod dial_handler;
pub mod listen_handler;
pub mod protocol;
use futures::prelude::*;
use libp2p_core::either::EitherOutput;
use libp2p_core::swarm::{ConnectedPoint, NetworkBehaviour, NetworkBehaviourAction, PollParameters};
use libp2p_core::{protocols_handler::ProtocolsHandler, protocols_handler::ProtocolsHandlerSelect, Multiaddr, PeerId};
use libp2p_core::protocols_handler::{OneShotHandler, ProtocolsHandler};
use libp2p_core::{Multiaddr, PeerId};
use std::{marker::PhantomData, time::Duration};
use tokio_io::{AsyncRead, AsyncWrite};
@ -85,15 +83,14 @@ impl<TSubstream> NetworkBehaviour for Ping<TSubstream>
where
TSubstream: AsyncRead + AsyncWrite,
{
type ProtocolsHandler = ProtocolsHandlerSelect<listen_handler::PingListenHandler<TSubstream>, dial_handler::PeriodicPingHandler<TSubstream>>;
type ProtocolsHandler = OneShotHandler<TSubstream, protocol::Ping, protocol::Ping, protocol::PingOutput>;
type OutEvent = PingEvent;
fn new_handler(&mut self) -> Self::ProtocolsHandler {
listen_handler::PingListenHandler::new()
.select(dial_handler::PeriodicPingHandler::new())
OneShotHandler::default()
}
fn addresses_of_peer(&self, peer_id: &PeerId) -> Vec<Multiaddr> {
fn addresses_of_peer(&self, _peer_id: &PeerId) -> Vec<Multiaddr> {
Vec::new()
}
@ -106,7 +103,7 @@ where
source: PeerId,
event: <Self::ProtocolsHandler as ProtocolsHandler>::OutEvent,
) {
if let EitherOutput::Second(dial_handler::OutEvent::PingSuccess(time)) = event {
if let protocol::PingOutput::Ping(time) = event {
self.events.push(PingEvent::PingSuccess {
peer: source,
time,

146
protocols/ping/src/listen_handler.rs
View File

@ -1,146 +0,0 @@
// Copyright 2018 Parity Technologies (UK) Ltd.
//
// Permission is hereby granted, free of charge, to any person obtaining a
// copy of this software and associated documentation files (the "Software"),
// to deal in the Software without restriction, including without limitation
// the rights to use, copy, modify, merge, publish, distribute, sublicense,
// and/or sell copies of the Software, and to permit persons to whom the
// Software is furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
// DEALINGS IN THE SOFTWARE.
use crate::protocol::{Ping, PingListener};
use arrayvec::ArrayVec;
use futures::prelude::*;
use libp2p_core::{
InboundUpgrade,
OutboundUpgrade,
ProtocolsHandler,
ProtocolsHandlerEvent,
protocols_handler::ProtocolsHandlerUpgrErr,
upgrade::DeniedUpgrade
};
use log::warn;
use tokio_io::{AsyncRead, AsyncWrite};
use void::{Void, unreachable};
/// Handler for handling pings received from a remote.
pub struct PingListenHandler<TSubstream> {
/// Configuration for the ping protocol.
ping_config: Ping<()>,
/// The ping substreams that were opened by the remote.
/// Note that we only accept a certain number of substreams, after which we refuse new ones
/// to avoid being DDoSed.
ping_in_substreams: ArrayVec<[PingListener<TSubstream>; 8]>,
/// If true, we're in the shutdown process and we shouldn't accept new substreams.
shutdown: bool,
}
impl<TSubstream> PingListenHandler<TSubstream> {
/// Builds a new `PingListenHandler`.
pub fn new() -> PingListenHandler<TSubstream> {
PingListenHandler {
ping_config: Default::default(),
shutdown: false,
ping_in_substreams: ArrayVec::new(),
}
}
}
impl<TSubstream> Default for PingListenHandler<TSubstream> {
#[inline]
fn default() -> Self {
PingListenHandler::new()
}
}
impl<TSubstream> ProtocolsHandler for PingListenHandler<TSubstream>
where
TSubstream: AsyncRead + AsyncWrite,
{
type InEvent = Void;
type OutEvent = Void;
type Error = Void;
type Substream = TSubstream;
type InboundProtocol = Ping<()>;
type OutboundProtocol = DeniedUpgrade;
type OutboundOpenInfo = ();
#[inline]
fn listen_protocol(&self) -> Self::InboundProtocol {
self.ping_config
}
fn inject_fully_negotiated_inbound(
&mut self,
protocol: <Self::InboundProtocol as InboundUpgrade<TSubstream>>::Output
) {
if self.shutdown {
return;
}
let _ = self.ping_in_substreams.try_push(protocol);
}
fn inject_fully_negotiated_outbound(&mut self, protocol: Void, _info: Self::OutboundOpenInfo) {
unreachable(protocol)
}
#[inline]
fn inject_event(&mut self, _: Self::InEvent) {}
#[inline]
fn inject_inbound_closed(&mut self) {}
#[inline]
fn inject_dial_upgrade_error(&mut self, _: Self::OutboundOpenInfo, _: ProtocolsHandlerUpgrErr<<Self::OutboundProtocol as OutboundUpgrade<Self::Substream>>::Error>) {}
#[inline]
fn connection_keep_alive(&self) -> bool {
false
}
#[inline]
fn shutdown(&mut self) {
for ping in self.ping_in_substreams.iter_mut() {
ping.shutdown();
}
self.shutdown = true;
}
fn poll(
&mut self,
) -> Poll<
ProtocolsHandlerEvent<Self::OutboundProtocol, Self::OutboundOpenInfo, Self::OutEvent>,
Self::Error,
> {
// Removes each substream one by one, and pushes them back if they're not ready (which
// should be the case 99% of the time).
for n in (0..self.ping_in_substreams.len()).rev() {
let mut ping = self.ping_in_substreams.swap_remove(n);
match ping.poll() {
Ok(Async::Ready(())) => {}
Ok(Async::NotReady) => self.ping_in_substreams.push(ping),
Err(err) => warn!(target: "sub-libp2p", "Remote ping substream errored: {:?}", err),
}
}
// Special case if shutting down.
if self.shutdown && self.ping_in_substreams.is_empty() {
return Ok(Async::Ready(ProtocolsHandlerEvent::Shutdown));
}
Ok(Async::NotReady)
}
}

383
protocols/ping/src/protocol.rs
View File

@ -18,32 +18,27 @@
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
// DEALINGS IN THE SOFTWARE.
use bytes::{BufMut, Bytes, BytesMut};
use futures::{prelude::*, future::{self, FutureResult}, try_ready};
use futures::{prelude::*, future, try_ready};
use libp2p_core::{InboundUpgrade, OutboundUpgrade, UpgradeInfo};
use log::debug;
use rand::{distributions::Standard, prelude::*, rngs::EntropyRng};
use std::collections::VecDeque;
use std::io::Error as IoError;
use std::{iter, marker::PhantomData, mem};
use tokio_codec::{Decoder, Encoder, Framed};
use std::{io, iter, time::Duration, time::Instant};
use tokio_io::{AsyncRead, AsyncWrite};
/// Represents a prototype for an upgrade to handle the ping protocol.
///
/// According to the design of libp2p, this struct would normally contain the configuration options
/// for the protocol, but in the case of `Ping` no configuration is required.
#[derive(Debug, Copy, Clone)]
pub struct Ping<TUserData = ()>(PhantomData<TUserData>);
/// The protocol works the following way:
///
/// - Dialer sends 32 bytes of random data.
/// - Listener receives the data and sends it back.
/// - Dialer receives the data and verifies that it matches what it sent.
///
/// The dialer produces a `Duration`, which corresponds to the time between when we flushed the
/// substream and when we received back the payload.
#[derive(Default, Debug, Copy, Clone)]
pub struct Ping;
impl<TUserData> Default for Ping<TUserData> {
#[inline]
fn default() -> Self {
Ping(PhantomData)
}
}
impl<TUserData> UpgradeInfo for Ping<TUserData> {
impl UpgradeInfo for Ping {
type Info = &'static [u8];
type InfoIter = iter::Once<Self::Info>;
@ -52,269 +47,139 @@ impl<TUserData> UpgradeInfo for Ping<TUserData> {
}
}
impl<TSocket, TUserData> InboundUpgrade<TSocket> for Ping<TUserData>
impl<TSocket> InboundUpgrade<TSocket> for Ping
where
TSocket: AsyncRead + AsyncWrite,
{
type Output = PingListener<TSocket>;
type Error = IoError;
type Future = FutureResult<Self::Output, Self::Error>;
type Output = ();
type Error = io::Error;
type Future = future::Map<future::AndThen<future::AndThen<future::AndThen<tokio_io::io::ReadExact<TSocket, [u8; 32]>, tokio_io::io::WriteAll<TSocket, [u8; 32]>, fn((TSocket, [u8; 32])) -> tokio_io::io::WriteAll<TSocket, [u8; 32]>>, tokio_io::io::Flush<TSocket>, fn((TSocket, [u8; 32])) -> tokio_io::io::Flush<TSocket>>, tokio_io::io::Shutdown<TSocket>, fn(TSocket) -> tokio_io::io::Shutdown<TSocket>>, fn(TSocket) -> ()>;
#[inline]
fn upgrade_inbound(self, socket: TSocket, _: Self::Info) -> Self::Future {
let listener = PingListener {
inner: Framed::new(socket, Codec),
state: PingListenerState::Listening,
};
future::ok(listener)
tokio_io::io::read_exact(socket, [0; 32])
.and_then::<fn(_) -> _, _>(|(socket, buffer)| tokio_io::io::write_all(socket, buffer))
.and_then::<fn(_) -> _, _>(|(socket, _)| tokio_io::io::flush(socket))
.and_then::<fn(_) -> _, _>(|socket| tokio_io::io::shutdown(socket))
.map(|_| ())
}
}
impl<TSocket, TUserData> OutboundUpgrade<TSocket> for Ping<TUserData>
impl<TSocket> OutboundUpgrade<TSocket> for Ping
where
TSocket: AsyncRead + AsyncWrite,
{
type Output = PingDialer<TSocket, TUserData>;
type Error = IoError;
type Future = FutureResult<Self::Output, Self::Error>;
type Output = Duration;
type Error = io::Error;
type Future = PingDialer<TSocket>;
#[inline]
fn upgrade_outbound(self, socket: TSocket, _: Self::Info) -> Self::Future {
let dialer = PingDialer {
inner: Framed::new(socket, Codec),
need_writer_flush: false,
needs_close: false,
sent_pings: VecDeque::with_capacity(4),
rng: EntropyRng::default(),
pings_to_send: VecDeque::with_capacity(4),
};
future::ok(dialer)
}
}
/// Sends pings and receives the pongs.
///
/// Implements `Stream`. The stream indicates when we receive a pong.
pub struct PingDialer<TSocket, TUserData> {
/// The underlying socket.
inner: Framed<TSocket, Codec>,
/// If true, need to flush the sink.
need_writer_flush: bool,
/// If true, need to close the sink.
needs_close: bool,
/// List of pings that have been sent to the remote and that are waiting for an answer.
sent_pings: VecDeque<(Bytes, TUserData)>,
/// Random number generator for the ping payload.
rng: EntropyRng,
/// List of pings to send to the remote.
pings_to_send: VecDeque<(Bytes, TUserData)>,
}
impl<TSocket, TUserData> PingDialer<TSocket, TUserData> {
/// Sends a ping to the remote.
///
/// The stream will produce an event containing the user data when we receive the pong.
pub fn ping(&mut self, user_data: TUserData) {
let payload: [u8; 32] = self.rng.sample(Standard);
let payload: [u8; 32] = EntropyRng::default().sample(Standard);
debug!("Preparing for ping with payload {:?}", payload);
self.pings_to_send.push_back((Bytes::from(payload.to_vec()), user_data));
PingDialer {
inner: PingDialerInner::Write {
inner: tokio_io::io::write_all(socket, payload),
},
}
}
}
impl<TSocket, TUserData> PingDialer<TSocket, TUserData>
where TSocket: AsyncRead + AsyncWrite,
{
/// Call this when the ping dialer needs to shut down. After this, the `Stream` is guaranteed
/// to finish soon-ish.
#[inline]
pub fn shutdown(&mut self) {
self.needs_close = true;
}
}
impl<TSocket, TUserData> Stream for PingDialer<TSocket, TUserData>
where TSocket: AsyncRead + AsyncWrite,
{
type Item = TUserData;
type Error = IoError;
fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
if self.needs_close {
try_ready!(self.inner.close());
return Ok(Async::Ready(None));
}
while let Some((ping, user_data)) = self.pings_to_send.pop_front() {
match self.inner.start_send(ping.clone()) {
Ok(AsyncSink::Ready) => self.need_writer_flush = true,
Ok(AsyncSink::NotReady(_)) => {
self.pings_to_send.push_front((ping, user_data));
break;
},
Err(err) => return Err(err),
}
self.sent_pings.push_back((ping, user_data));
}
if self.need_writer_flush {
match self.inner.poll_complete() {
Ok(Async::Ready(())) => self.need_writer_flush = false,
Ok(Async::NotReady) => (),
Err(err) => return Err(err),
}
}
loop {
match self.inner.poll() {
Ok(Async::Ready(Some(pong))) => {
if let Some(pos) = self.sent_pings.iter().position(|&(ref p, _)| p == &pong) {
let (_, user_data) = self.sent_pings.remove(pos)
.expect("Grabbed a valid position just above");
return Ok(Async::Ready(Some(user_data)));
} else {
debug!("Received pong that doesn't match what we sent: {:?}", pong);
}
},
Ok(Async::NotReady) => break,
Ok(Async::Ready(None)) => {
// Notify the current task so that we poll again.
self.needs_close = true;
try_ready!(self.inner.close());
return Ok(Async::Ready(None));
}
Err(err) => return Err(err),
}
}
Ok(Async::NotReady)
}
}
/// Listens to incoming pings and answers them.
/// Sends a ping and receives a pong.
///
/// Implements `Future`. The future terminates when the underlying socket closes.
pub struct PingListener<TSocket> {
/// The underlying socket.
inner: Framed<TSocket, Codec>,
/// State of the listener.
state: PingListenerState,
/// Implements `Future`. Finishes when the pong has arrived and has been verified.
pub struct PingDialer<TSocket> {
inner: PingDialerInner<TSocket>,
}
#[derive(Debug)]
enum PingListenerState {
/// We are waiting for the next ping on the socket.
Listening,
/// We are trying to send a pong.
Sending(Bytes),
/// We are flushing the underlying sink.
Flushing,
/// We are shutting down everything.
Closing,
/// A panic happened during the processing.
Poisoned,
enum PingDialerInner<TSocket> {
Write {
inner: tokio_io::io::WriteAll<TSocket, [u8; 32]>,
},
Flush {
inner: tokio_io::io::Flush<TSocket>,
ping_payload: [u8; 32],
},
Read {
inner: tokio_io::io::ReadExact<TSocket, [u8; 32]>,
ping_payload: [u8; 32],
started: Instant,
},
Shutdown {
inner: tokio_io::io::Shutdown<TSocket>,
ping_time: Duration,
},
}
impl<TSocket> PingListener<TSocket>
where TSocket: AsyncRead + AsyncWrite
impl<TSocket> Future for PingDialer<TSocket>
where TSocket: AsyncRead + AsyncWrite,
{
/// Call this when the ping listener needs to shut down. After this, the `Future` is guaranteed
/// to finish soon-ish.
#[inline]
pub fn shutdown(&mut self) {
self.state = PingListenerState::Closing;
}
}
impl<TSocket> Future for PingListener<TSocket>
where TSocket: AsyncRead + AsyncWrite
{
type Item = ();
type Error = IoError;
type Item = Duration;
type Error = io::Error;
fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
loop {
match mem::replace(&mut self.state, PingListenerState::Poisoned) {
PingListenerState::Listening => {
match self.inner.poll() {
Ok(Async::Ready(Some(payload))) => {
debug!("Received ping (payload={:?}); sending back", payload);
self.state = PingListenerState::Sending(payload.freeze())
},
Ok(Async::Ready(None)) => self.state = PingListenerState::Closing,
Ok(Async::NotReady) => {
self.state = PingListenerState::Listening;
return Ok(Async::NotReady);
},
Err(err) => return Err(err),
let new_state = match self.inner {
PingDialerInner::Write { ref mut inner } => {
let (socket, ping_payload) = try_ready!(inner.poll());
PingDialerInner::Flush {
inner: tokio_io::io::flush(socket),
ping_payload,
}
},
PingListenerState::Sending(data) => {
match self.inner.start_send(data) {
Ok(AsyncSink::Ready) => self.state = PingListenerState::Flushing,
Ok(AsyncSink::NotReady(data)) => {
self.state = PingListenerState::Sending(data);
return Ok(Async::NotReady);
},
Err(err) => return Err(err),
PingDialerInner::Flush { ref mut inner, ping_payload } => {
let socket = try_ready!(inner.poll());
let started = Instant::now();
PingDialerInner::Read {
inner: tokio_io::io::read_exact(socket, [0; 32]),
ping_payload,
started,
}
},
PingListenerState::Flushing => {
match self.inner.poll_complete() {
Ok(Async::Ready(())) => self.state = PingListenerState::Listening,
Ok(Async::NotReady) => {
self.state = PingListenerState::Flushing;
return Ok(Async::NotReady);
},
Err(err) => return Err(err),
PingDialerInner::Read { ref mut inner, ping_payload, started } => {
let (socket, obtained) = try_ready!(inner.poll());
let ping_time = started.elapsed();
if obtained != ping_payload {
return Err(io::Error::new(io::ErrorKind::InvalidData,
"Received ping payload doesn't match expected"));
}
PingDialerInner::Shutdown {
inner: tokio_io::io::shutdown(socket),
ping_time,
}
},
PingListenerState::Closing => {
match self.inner.close() {
Ok(Async::Ready(())) => return Ok(Async::Ready(())),
Ok(Async::NotReady) => {
self.state = PingListenerState::Closing;
return Ok(Async::NotReady);
},
Err(err) => return Err(err),
}
PingDialerInner::Shutdown { ref mut inner, ping_time } => {
let _ = try_ready!(inner.poll());
return Ok(Async::Ready(ping_time));
},
PingListenerState::Poisoned => panic!("Poisoned or errored PingListener"),
}
};
self.inner = new_state;
}
}
}
// Implementation of the `Codec` trait of tokio-io. Splits frames into groups of 32 bytes.
#[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
struct Codec;
impl Decoder for Codec {
type Item = BytesMut;
type Error = IoError;
/// Enum to merge the output of `Ping` for the dialer and listener.
#[derive(Debug, Copy, Clone)]
pub enum PingOutput {
/// Received a ping and sent back a pong.
Pong,
/// Sent a ping and received back a pong. Contains the ping time.
Ping(Duration),
}
impl From<Duration> for PingOutput {
#[inline]
fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<BytesMut>, IoError> {
if buf.len() >= 32 {
Ok(Some(buf.split_to(32)))
} else {
Ok(None)
}
fn from(duration: Duration) -> PingOutput {
PingOutput::Ping(duration)
}
}
impl Encoder for Codec {
type Item = Bytes;
type Error = IoError;
impl From<()> for PingOutput {
#[inline]
fn encode(&mut self, mut data: Bytes, buf: &mut BytesMut) -> Result<(), IoError> {
if !data.is_empty() {
let split = 32 * (1 + ((data.len() - 1) / 32));
buf.reserve(split);
buf.put(data.split_to(split));
}
Ok(())
fn from(_: ()) -> PingOutput {
PingOutput::Pong
}
}
@ -323,7 +188,7 @@ mod tests {
use tokio_tcp::{TcpListener, TcpStream};
use super::Ping;
use futures::{Future, Stream};
use libp2p_core::upgrade::{InboundUpgrade, OutboundUpgrade};
use libp2p_core::upgrade;
// TODO: rewrite tests with the MemoryTransport
@ -337,54 +202,16 @@ mod tests {
.into_future()
.map_err(|(e, _)| e)
.and_then(|(c, _)| {
Ping::<()>::default().upgrade_inbound(c.unwrap(), b"/ipfs/ping/1.0.0")
})
.flatten();
upgrade::apply_inbound(c.unwrap(), Ping::default()).map_err(|_| panic!())
});
let client = TcpStream::connect(&listener_addr)
.and_then(|c| {
Ping::<()>::default().upgrade_outbound(c, b"/ipfs/ping/1.0.0")
})
.and_then(|mut pinger| {
pinger.ping(());
pinger.into_future().map(|_| ()).map_err(|_| panic!())
upgrade::apply_outbound(c, Ping::default()).map_err(|_| panic!())
})
.map(|_| ());
let mut runtime = tokio::runtime::Runtime::new().unwrap();
runtime.block_on(server.select(client).map_err(|_| panic!())).unwrap();
}
#[test]
fn multipings() {
// Check that we can send multiple pings in a row and it will still work.
let listener = TcpListener::bind(&"127.0.0.1:0".parse().unwrap()).unwrap();
let listener_addr = listener.local_addr().unwrap();
let server = listener
.incoming()
.into_future()
.map_err(|(e, _)| e)
.and_then(|(c, _)| {
Ping::<u32>::default().upgrade_inbound(c.unwrap(), b"/ipfs/ping/1.0.0")
})
.flatten();
let client = TcpStream::connect(&listener_addr)
.and_then(|c| {
Ping::<u32>::default().upgrade_outbound(c, b"/ipfs/ping/1.0.0")
})
.and_then(|mut pinger| {
for n in 0..20 {
pinger.ping(n);
}
pinger.take(20)
.collect()
.map(|val| { assert_eq!(val, (0..20).collect::<Vec<_>>()); })
.map_err(|_| panic!())
});
let mut runtime = tokio::runtime::Runtime::new().unwrap();
runtime.block_on(server.select(client)).unwrap_or_else(|_| panic!());
}
}