pub(crate) mod protocol;
pub use protocol::ProtocolSupport;
use crate::codec::Codec;
use crate::handler::protocol::{RequestProtocol, ResponseProtocol};
use crate::{RequestId, EMPTY_QUEUE_SHRINK_THRESHOLD};
use futures::{channel::oneshot, future::BoxFuture, prelude::*, stream::FuturesUnordered};
use instant::Instant;
use libp2p_swarm::handler::{
ConnectionEvent, DialUpgradeError, FullyNegotiatedInbound, FullyNegotiatedOutbound,
ListenUpgradeError,
};
use libp2p_swarm::{
handler::{ConnectionHandler, ConnectionHandlerEvent, KeepAlive, StreamUpgradeError},
SubstreamProtocol,
};
use smallvec::SmallVec;
use std::{
collections::VecDeque,
fmt,
sync::{
atomic::{AtomicU64, Ordering},
Arc,
},
task::{Context, Poll},
time::Duration,
};
pub struct Handler<TCodec>
where
TCodec: Codec,
{
inbound_protocols: SmallVec<[TCodec::Protocol; 2]>,
codec: TCodec,
keep_alive_timeout: Duration,
substream_timeout: Duration,
keep_alive: KeepAlive,
pending_events: VecDeque<Event<TCodec>>,
outbound: VecDeque<RequestProtocol<TCodec>>,
inbound: FuturesUnordered<
BoxFuture<
'static,
Result<
(
(RequestId, TCodec::Request),
oneshot::Sender<TCodec::Response>,
),
oneshot::Canceled,
>,
>,
>,
inbound_request_id: Arc<AtomicU64>,
}
impl<TCodec> Handler<TCodec>
where
TCodec: Codec + Send + Clone + 'static,
{
pub(super) fn new(
inbound_protocols: SmallVec<[TCodec::Protocol; 2]>,
codec: TCodec,
keep_alive_timeout: Duration,
substream_timeout: Duration,
inbound_request_id: Arc<AtomicU64>,
) -> Self {
Self {
inbound_protocols,
codec,
keep_alive: KeepAlive::Yes,
keep_alive_timeout,
substream_timeout,
outbound: VecDeque::new(),
inbound: FuturesUnordered::new(),
pending_events: VecDeque::new(),
inbound_request_id,
}
}
fn on_fully_negotiated_inbound(
&mut self,
FullyNegotiatedInbound {
protocol: sent,
info: request_id,
}: FullyNegotiatedInbound<
<Self as ConnectionHandler>::InboundProtocol,
<Self as ConnectionHandler>::InboundOpenInfo,
>,
) {
if sent {
self.pending_events
.push_back(Event::ResponseSent(request_id))
} else {
self.pending_events
.push_back(Event::ResponseOmission(request_id))
}
}
fn on_dial_upgrade_error(
&mut self,
DialUpgradeError { info, error }: DialUpgradeError<
<Self as ConnectionHandler>::OutboundOpenInfo,
<Self as ConnectionHandler>::OutboundProtocol,
>,
) {
match error {
StreamUpgradeError::Timeout => {
self.pending_events.push_back(Event::OutboundTimeout(info));
}
StreamUpgradeError::NegotiationFailed => {
self.pending_events
.push_back(Event::OutboundUnsupportedProtocols(info));
}
StreamUpgradeError::Apply(e) => {
log::debug!("outbound stream {info} failed: {e}");
}
StreamUpgradeError::Io(e) => {
log::debug!("outbound stream {info} failed: {e}");
}
}
}
fn on_listen_upgrade_error(
&mut self,
ListenUpgradeError { error, info }: ListenUpgradeError<
<Self as ConnectionHandler>::InboundOpenInfo,
<Self as ConnectionHandler>::InboundProtocol,
>,
) {
log::debug!("inbound stream {info} failed: {error}");
}
}
pub enum Event<TCodec>
where
TCodec: Codec,
{
Request {
request_id: RequestId,
request: TCodec::Request,
sender: oneshot::Sender<TCodec::Response>,
},
Response {
request_id: RequestId,
response: TCodec::Response,
},
ResponseSent(RequestId),
ResponseOmission(RequestId),
OutboundTimeout(RequestId),
OutboundUnsupportedProtocols(RequestId),
}
impl<TCodec: Codec> fmt::Debug for Event<TCodec> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Event::Request {
request_id,
request: _,
sender: _,
} => f
.debug_struct("Event::Request")
.field("request_id", request_id)
.finish(),
Event::Response {
request_id,
response: _,
} => f
.debug_struct("Event::Response")
.field("request_id", request_id)
.finish(),
Event::ResponseSent(request_id) => f
.debug_tuple("Event::ResponseSent")
.field(request_id)
.finish(),
Event::ResponseOmission(request_id) => f
.debug_tuple("Event::ResponseOmission")
.field(request_id)
.finish(),
Event::OutboundTimeout(request_id) => f
.debug_tuple("Event::OutboundTimeout")
.field(request_id)
.finish(),
Event::OutboundUnsupportedProtocols(request_id) => f
.debug_tuple("Event::OutboundUnsupportedProtocols")
.field(request_id)
.finish(),
}
}
}
impl<TCodec> ConnectionHandler for Handler<TCodec>
where
TCodec: Codec + Send + Clone + 'static,
{
type FromBehaviour = RequestProtocol<TCodec>;
type ToBehaviour = Event<TCodec>;
type Error = void::Void;
type InboundProtocol = ResponseProtocol<TCodec>;
type OutboundProtocol = RequestProtocol<TCodec>;
type OutboundOpenInfo = RequestId;
type InboundOpenInfo = RequestId;
fn listen_protocol(&self) -> SubstreamProtocol<Self::InboundProtocol, Self::InboundOpenInfo> {
let (rq_send, rq_recv) = oneshot::channel();
let (rs_send, rs_recv) = oneshot::channel();
let request_id = RequestId(self.inbound_request_id.fetch_add(1, Ordering::Relaxed));
let proto = ResponseProtocol {
protocols: self.inbound_protocols.clone(),
codec: self.codec.clone(),
request_sender: rq_send,
response_receiver: rs_recv,
request_id,
};
self.inbound
.push(rq_recv.map_ok(move |rq| (rq, rs_send)).boxed());
SubstreamProtocol::new(proto, request_id).with_timeout(self.substream_timeout)
}
fn on_behaviour_event(&mut self, request: Self::FromBehaviour) {
self.keep_alive = KeepAlive::Yes;
self.outbound.push_back(request);
}
fn connection_keep_alive(&self) -> KeepAlive {
self.keep_alive
}
#[allow(deprecated)]
fn poll(
&mut self,
cx: &mut Context<'_>,
) -> Poll<
ConnectionHandlerEvent<RequestProtocol<TCodec>, RequestId, Self::ToBehaviour, Self::Error>,
> {
if let Some(event) = self.pending_events.pop_front() {
return Poll::Ready(ConnectionHandlerEvent::NotifyBehaviour(event));
} else if self.pending_events.capacity() > EMPTY_QUEUE_SHRINK_THRESHOLD {
self.pending_events.shrink_to_fit();
}
while let Poll::Ready(Some(result)) = self.inbound.poll_next_unpin(cx) {
match result {
Ok(((id, rq), rs_sender)) => {
self.keep_alive = KeepAlive::Yes;
return Poll::Ready(ConnectionHandlerEvent::NotifyBehaviour(Event::Request {
request_id: id,
request: rq,
sender: rs_sender,
}));
}
Err(oneshot::Canceled) => {
}
}
}
if let Some(request) = self.outbound.pop_front() {
let info = request.request_id;
return Poll::Ready(ConnectionHandlerEvent::OutboundSubstreamRequest {
protocol: SubstreamProtocol::new(request, info)
.with_timeout(self.substream_timeout),
});
}
debug_assert!(self.outbound.is_empty());
if self.outbound.capacity() > EMPTY_QUEUE_SHRINK_THRESHOLD {
self.outbound.shrink_to_fit();
}
#[allow(deprecated)]
if self.inbound.is_empty() && self.keep_alive.is_yes() {
let until = Instant::now() + self.substream_timeout + self.keep_alive_timeout;
self.keep_alive = KeepAlive::Until(until);
}
Poll::Pending
}
fn on_connection_event(
&mut self,
event: ConnectionEvent<
Self::InboundProtocol,
Self::OutboundProtocol,
Self::InboundOpenInfo,
Self::OutboundOpenInfo,
>,
) {
match event {
ConnectionEvent::FullyNegotiatedInbound(fully_negotiated_inbound) => {
self.on_fully_negotiated_inbound(fully_negotiated_inbound)
}
ConnectionEvent::FullyNegotiatedOutbound(FullyNegotiatedOutbound {
protocol: response,
info: request_id,
}) => {
self.pending_events.push_back(Event::Response {
request_id,
response,
});
}
ConnectionEvent::DialUpgradeError(dial_upgrade_error) => {
self.on_dial_upgrade_error(dial_upgrade_error)
}
ConnectionEvent::ListenUpgradeError(listen_upgrade_error) => {
self.on_listen_upgrade_error(listen_upgrade_error)
}
ConnectionEvent::AddressChange(_)
| ConnectionEvent::LocalProtocolsChange(_)
| ConnectionEvent::RemoteProtocolsChange(_) => {}
}
}
}