use crate::{
transport::{ListenerEvent, TransportError},
Multiaddr, Transport,
};
use futures::prelude::*;
use futures_timer::Delay;
use std::{error, fmt, io, pin::Pin, task::Context, task::Poll, time::Duration};
#[derive(Debug, Copy, Clone)]
pub struct TransportTimeout<InnerTrans> {
inner: InnerTrans,
outgoing_timeout: Duration,
incoming_timeout: Duration,
}
impl<InnerTrans> TransportTimeout<InnerTrans> {
pub fn new(trans: InnerTrans, timeout: Duration) -> Self {
TransportTimeout {
inner: trans,
outgoing_timeout: timeout,
incoming_timeout: timeout,
}
}
pub fn with_outgoing_timeout(trans: InnerTrans, timeout: Duration) -> Self {
TransportTimeout {
inner: trans,
outgoing_timeout: timeout,
incoming_timeout: Duration::from_secs(100 * 365 * 24 * 3600), }
}
pub fn with_ingoing_timeout(trans: InnerTrans, timeout: Duration) -> Self {
TransportTimeout {
inner: trans,
outgoing_timeout: Duration::from_secs(100 * 365 * 24 * 3600), incoming_timeout: timeout,
}
}
}
impl<InnerTrans> Transport for TransportTimeout<InnerTrans>
where
InnerTrans: Transport,
InnerTrans::Error: 'static,
{
type Output = InnerTrans::Output;
type Error = TransportTimeoutError<InnerTrans::Error>;
type Listener = TimeoutListener<InnerTrans::Listener>;
type ListenerUpgrade = Timeout<InnerTrans::ListenerUpgrade>;
type Dial = Timeout<InnerTrans::Dial>;
fn listen_on(self, addr: Multiaddr) -> Result<Self::Listener, TransportError<Self::Error>> {
let listener = self
.inner
.listen_on(addr)
.map_err(|err| err.map(TransportTimeoutError::Other))?;
let listener = TimeoutListener {
inner: listener,
timeout: self.incoming_timeout,
};
Ok(listener)
}
fn dial(self, addr: Multiaddr) -> Result<Self::Dial, TransportError<Self::Error>> {
let dial = self
.inner
.dial(addr)
.map_err(|err| err.map(TransportTimeoutError::Other))?;
Ok(Timeout {
inner: dial,
timer: Delay::new(self.outgoing_timeout),
})
}
fn dial_as_listener(self, addr: Multiaddr) -> Result<Self::Dial, TransportError<Self::Error>> {
let dial = self
.inner
.dial_as_listener(addr)
.map_err(|err| err.map(TransportTimeoutError::Other))?;
Ok(Timeout {
inner: dial,
timer: Delay::new(self.outgoing_timeout),
})
}
fn address_translation(&self, server: &Multiaddr, observed: &Multiaddr) -> Option<Multiaddr> {
self.inner.address_translation(server, observed)
}
}
#[pin_project::pin_project]
pub struct TimeoutListener<InnerStream> {
#[pin]
inner: InnerStream,
timeout: Duration,
}
impl<InnerStream, O, E> Stream for TimeoutListener<InnerStream>
where
InnerStream: TryStream<Ok = ListenerEvent<O, E>, Error = E>,
{
type Item =
Result<ListenerEvent<Timeout<O>, TransportTimeoutError<E>>, TransportTimeoutError<E>>;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let this = self.project();
let poll_out = match TryStream::try_poll_next(this.inner, cx) {
Poll::Ready(Some(Err(err))) => {
return Poll::Ready(Some(Err(TransportTimeoutError::Other(err))))
}
Poll::Ready(Some(Ok(v))) => v,
Poll::Ready(None) => return Poll::Ready(None),
Poll::Pending => return Poll::Pending,
};
let timeout = *this.timeout;
let event = poll_out
.map(move |inner_fut| Timeout {
inner: inner_fut,
timer: Delay::new(timeout),
})
.map_err(TransportTimeoutError::Other);
Poll::Ready(Some(Ok(event)))
}
}
#[pin_project::pin_project]
#[must_use = "futures do nothing unless polled"]
pub struct Timeout<InnerFut> {
#[pin]
inner: InnerFut,
timer: Delay,
}
impl<InnerFut> Future for Timeout<InnerFut>
where
InnerFut: TryFuture,
{
type Output = Result<InnerFut::Ok, TransportTimeoutError<InnerFut::Error>>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let mut this = self.project();
match TryFuture::try_poll(this.inner, cx) {
Poll::Pending => {}
Poll::Ready(Ok(v)) => return Poll::Ready(Ok(v)),
Poll::Ready(Err(err)) => return Poll::Ready(Err(TransportTimeoutError::Other(err))),
}
match Pin::new(&mut this.timer).poll(cx) {
Poll::Pending => Poll::Pending,
Poll::Ready(()) => Poll::Ready(Err(TransportTimeoutError::Timeout)),
}
}
}
#[derive(Debug)]
pub enum TransportTimeoutError<TErr> {
Timeout,
TimerError(io::Error),
Other(TErr),
}
impl<TErr> fmt::Display for TransportTimeoutError<TErr>
where
TErr: fmt::Display,
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
TransportTimeoutError::Timeout => write!(f, "Timeout has been reached"),
TransportTimeoutError::TimerError(err) => write!(f, "Error in the timer: {}", err),
TransportTimeoutError::Other(err) => write!(f, "{}", err),
}
}
}
impl<TErr> error::Error for TransportTimeoutError<TErr>
where
TErr: error::Error + 'static,
{
fn source(&self) -> Option<&(dyn error::Error + 'static)> {
match self {
TransportTimeoutError::Timeout => None,
TransportTimeoutError::TimerError(err) => Some(err),
TransportTimeoutError::Other(err) => Some(err),
}
}
}