use crate::protocol;
use futures::future::BoxFuture;
use futures::prelude::*;
use futures_timer::Delay;
use libp2p_core::{upgrade::NegotiationError, UpgradeError};
use libp2p_swarm::{
ConnectionHandler, ConnectionHandlerEvent, ConnectionHandlerUpgrErr, KeepAlive,
NegotiatedSubstream, SubstreamProtocol,
};
use std::collections::VecDeque;
use std::{
error::Error,
fmt, io,
num::NonZeroU32,
task::{Context, Poll},
time::Duration,
};
use void::Void;
#[derive(Debug, Clone)]
pub struct Config {
timeout: Duration,
interval: Duration,
max_failures: NonZeroU32,
keep_alive: bool,
}
impl Config {
pub fn new() -> Self {
Self {
timeout: Duration::from_secs(20),
interval: Duration::from_secs(15),
max_failures: NonZeroU32::new(1).expect("1 != 0"),
keep_alive: false,
}
}
pub fn with_timeout(mut self, d: Duration) -> Self {
self.timeout = d;
self
}
pub fn with_interval(mut self, d: Duration) -> Self {
self.interval = d;
self
}
pub fn with_max_failures(mut self, n: NonZeroU32) -> Self {
self.max_failures = n;
self
}
pub fn with_keep_alive(mut self, b: bool) -> Self {
self.keep_alive = b;
self
}
}
impl Default for Config {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug)]
pub enum Success {
Pong,
Ping { rtt: Duration },
}
#[derive(Debug)]
pub enum Failure {
Timeout,
Unsupported,
Other {
error: Box<dyn std::error::Error + Send + 'static>,
},
}
impl fmt::Display for Failure {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Failure::Timeout => f.write_str("Ping timeout"),
Failure::Other { error } => write!(f, "Ping error: {}", error),
Failure::Unsupported => write!(f, "Ping protocol not supported"),
}
}
}
impl Error for Failure {
fn source(&self) -> Option<&(dyn Error + 'static)> {
match self {
Failure::Timeout => None,
Failure::Other { error } => Some(&**error),
Failure::Unsupported => None,
}
}
}
pub struct Handler {
config: Config,
timer: Delay,
pending_errors: VecDeque<Failure>,
failures: u32,
outbound: Option<PingState>,
inbound: Option<PongFuture>,
state: State,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum State {
Inactive {
reported: bool,
},
Active,
}
impl Handler {
pub fn new(config: Config) -> Self {
Handler {
config,
timer: Delay::new(Duration::new(0, 0)),
pending_errors: VecDeque::with_capacity(2),
failures: 0,
outbound: None,
inbound: None,
state: State::Active,
}
}
}
impl ConnectionHandler for Handler {
type InEvent = Void;
type OutEvent = crate::Result;
type Error = Failure;
type InboundProtocol = protocol::Ping;
type OutboundProtocol = protocol::Ping;
type OutboundOpenInfo = ();
type InboundOpenInfo = ();
fn listen_protocol(&self) -> SubstreamProtocol<protocol::Ping, ()> {
SubstreamProtocol::new(protocol::Ping, ())
}
fn inject_fully_negotiated_inbound(&mut self, stream: NegotiatedSubstream, (): ()) {
self.inbound = Some(protocol::recv_ping(stream).boxed());
}
fn inject_fully_negotiated_outbound(&mut self, stream: NegotiatedSubstream, (): ()) {
self.timer.reset(self.config.timeout);
self.outbound = Some(PingState::Ping(protocol::send_ping(stream).boxed()));
}
fn inject_event(&mut self, _: Void) {}
fn inject_dial_upgrade_error(&mut self, _info: (), error: ConnectionHandlerUpgrErr<Void>) {
self.outbound = None;
let error = match error {
ConnectionHandlerUpgrErr::Upgrade(UpgradeError::Select(NegotiationError::Failed)) => {
debug_assert_eq!(self.state, State::Active);
self.state = State::Inactive { reported: false };
return;
}
ConnectionHandlerUpgrErr::Timeout => Failure::Timeout,
e => Failure::Other { error: Box::new(e) },
};
self.pending_errors.push_front(error);
}
fn connection_keep_alive(&self) -> KeepAlive {
if self.config.keep_alive {
KeepAlive::Yes
} else {
KeepAlive::No
}
}
fn poll(
&mut self,
cx: &mut Context<'_>,
) -> Poll<ConnectionHandlerEvent<protocol::Ping, (), crate::Result, Self::Error>> {
match self.state {
State::Inactive { reported: true } => {
return Poll::Pending; }
State::Inactive { reported: false } => {
self.state = State::Inactive { reported: true };
return Poll::Ready(ConnectionHandlerEvent::Custom(Err(Failure::Unsupported)));
}
State::Active => {}
}
if let Some(fut) = self.inbound.as_mut() {
match fut.poll_unpin(cx) {
Poll::Pending => {}
Poll::Ready(Err(e)) => {
log::debug!("Inbound ping error: {:?}", e);
self.inbound = None;
}
Poll::Ready(Ok(stream)) => {
self.inbound = Some(protocol::recv_ping(stream).boxed());
return Poll::Ready(ConnectionHandlerEvent::Custom(Ok(Success::Pong)));
}
}
}
loop {
if let Some(error) = self.pending_errors.pop_back() {
log::debug!("Ping failure: {:?}", error);
self.failures += 1;
if self.failures > 1 || self.config.max_failures.get() > 1 {
if self.failures >= self.config.max_failures.get() {
log::debug!("Too many failures ({}). Closing connection.", self.failures);
return Poll::Ready(ConnectionHandlerEvent::Close(error));
}
return Poll::Ready(ConnectionHandlerEvent::Custom(Err(error)));
}
}
match self.outbound.take() {
Some(PingState::Ping(mut ping)) => match ping.poll_unpin(cx) {
Poll::Pending => {
if self.timer.poll_unpin(cx).is_ready() {
self.pending_errors.push_front(Failure::Timeout);
} else {
self.outbound = Some(PingState::Ping(ping));
break;
}
}
Poll::Ready(Ok((stream, rtt))) => {
self.failures = 0;
self.timer.reset(self.config.interval);
self.outbound = Some(PingState::Idle(stream));
return Poll::Ready(ConnectionHandlerEvent::Custom(Ok(Success::Ping {
rtt,
})));
}
Poll::Ready(Err(e)) => {
self.pending_errors
.push_front(Failure::Other { error: Box::new(e) });
}
},
Some(PingState::Idle(stream)) => match self.timer.poll_unpin(cx) {
Poll::Pending => {
self.outbound = Some(PingState::Idle(stream));
break;
}
Poll::Ready(()) => {
self.timer.reset(self.config.timeout);
self.outbound = Some(PingState::Ping(protocol::send_ping(stream).boxed()));
}
},
Some(PingState::OpenStream) => {
self.outbound = Some(PingState::OpenStream);
break;
}
None => {
self.outbound = Some(PingState::OpenStream);
let protocol = SubstreamProtocol::new(protocol::Ping, ())
.with_timeout(self.config.timeout);
return Poll::Ready(ConnectionHandlerEvent::OutboundSubstreamRequest {
protocol,
});
}
}
}
Poll::Pending
}
}
type PingFuture = BoxFuture<'static, Result<(NegotiatedSubstream, Duration), io::Error>>;
type PongFuture = BoxFuture<'static, Result<NegotiatedSubstream, io::Error>>;
enum PingState {
OpenStream,
Idle(NegotiatedSubstream),
Ping(PingFuture),
}