fluencelabs/rust-libp2p
1
0
Fork 0
mirror of https://github.com/fluencelabs/rust-libp2p synced 2025-06-27 00:31:35 +00:00
Code Issues Projects Releases Wiki Activity

protocols/kademlia: Refactor stream state machines (#3074)

Browse Source 
* Refactor `InboundSubstreamState` to implement `Stream`

This allows us to use `stream::SelectAll` instead of iterating
through the states ourselves. It also allows us to fix an existing
TODO where we don't properly close a stream.

* Refactor `OutboundSubstreamState` to implement `Stream`

* Remove unnecessary keep_alive changes

The same thing is happening further down.

* Remove code duplication in answering requests

* Debug assert that we can answer every request
This commit is contained in:
Thomas Eizinger
2022-11-03 06:20:55 +11:00
committed by GitHub
parent f4ce1fe9ae
commit 53d445e61c
1 changed files with 362 additions and 438 deletions
Download Patch File Download Diff File Expand all files Collapse all files

800
protocols/kad/src/handler.rs
View File

@ -24,6 +24,7 @@ use crate::protocol::{
}; };
use crate::record::{self, Record}; use crate::record::{self, Record};
use futures::prelude::*; use futures::prelude::*;
use futures::stream::SelectAll;
use instant::Instant; use instant::Instant;
use libp2p_core::{ use libp2p_core::{
either::EitherOutput, either::EitherOutput,
@ -35,6 +36,7 @@ use libp2p_swarm::{
KeepAlive, NegotiatedSubstream, SubstreamProtocol, KeepAlive, NegotiatedSubstream, SubstreamProtocol,
}; };
use log::trace; use log::trace;
use std::task::Waker;
use std::{ use std::{
error, fmt, io, marker::PhantomData, pin::Pin, task::Context, task::Poll, time::Duration, error, fmt, io, marker::PhantomData, pin::Pin, task::Context, task::Poll, time::Duration,
}; };
@ -56,7 +58,9 @@ impl<T> KademliaHandlerProto<T> {
} }
} }
impl<T: Clone + fmt::Debug + Send + 'static> IntoConnectionHandler for KademliaHandlerProto<T> { impl<T: Clone + fmt::Debug + Send + 'static + Unpin> IntoConnectionHandler
for KademliaHandlerProto<T>
{
type Handler = KademliaHandler<T>; type Handler = KademliaHandler<T>;
fn into_handler(self, remote_peer_id: &PeerId, endpoint: &ConnectedPoint) -> Self::Handler { fn into_handler(self, remote_peer_id: &PeerId, endpoint: &ConnectedPoint) -> Self::Handler {
@ -87,10 +91,10 @@ pub struct KademliaHandler<TUserData> {
next_connec_unique_id: UniqueConnecId, next_connec_unique_id: UniqueConnecId,
/// List of active outbound substreams with the state they are in. /// List of active outbound substreams with the state they are in.
outbound_substreams: Vec<OutboundSubstreamState<TUserData>>, outbound_substreams: SelectAll<OutboundSubstreamState<TUserData>>,
/// List of active inbound substreams with the state they are in. /// List of active inbound substreams with the state they are in.
inbound_substreams: Vec<InboundSubstreamState>, inbound_substreams: SelectAll<InboundSubstreamState<TUserData>>,
/// Until when to keep the connection alive. /// Until when to keep the connection alive.
keep_alive: KeepAlive, keep_alive: KeepAlive,
@ -137,7 +141,7 @@ pub struct KademliaHandlerConfig {
enum OutboundSubstreamState<TUserData> { enum OutboundSubstreamState<TUserData> {
/// We haven't started opening the outgoing substream yet. /// We haven't started opening the outgoing substream yet.
/// Contains the request we want to send, and the user data if we expect an answer. /// Contains the request we want to send, and the user data if we expect an answer.
PendingOpen(KadRequestMsg, Option<TUserData>), PendingOpen(SubstreamProtocol<KademliaProtocolConfig, (KadRequestMsg, Option<TUserData>)>),
/// Waiting to send a message to the remote. /// Waiting to send a message to the remote.
PendingSend( PendingSend(
KadOutStreamSink<NegotiatedSubstream>, KadOutStreamSink<NegotiatedSubstream>,
@ -153,10 +157,13 @@ enum OutboundSubstreamState<TUserData> {
ReportError(KademliaHandlerQueryErr, TUserData), ReportError(KademliaHandlerQueryErr, TUserData),
/// The substream is being closed. /// The substream is being closed.
Closing(KadOutStreamSink<NegotiatedSubstream>), Closing(KadOutStreamSink<NegotiatedSubstream>),
/// The substream is complete and will not perform any more work.
Done,
Poisoned,
} }
/// State of an active inbound substream. /// State of an active inbound substream.
enum InboundSubstreamState { enum InboundSubstreamState<TUserData> {
/// Waiting for a request from the remote. /// Waiting for a request from the remote.
WaitingMessage { WaitingMessage {
/// Whether it is the first message to be awaited on this stream. /// Whether it is the first message to be awaited on this stream.
@ -165,7 +172,11 @@ enum InboundSubstreamState {
substream: KadInStreamSink<NegotiatedSubstream>, substream: KadInStreamSink<NegotiatedSubstream>,
}, },
/// Waiting for the user to send a [`KademliaHandlerIn`] event containing the response. /// Waiting for the user to send a [`KademliaHandlerIn`] event containing the response.
WaitingUser(UniqueConnecId, KadInStreamSink<NegotiatedSubstream>), WaitingUser(
UniqueConnecId,
KadInStreamSink<NegotiatedSubstream>,
Option<Waker>,
),
/// Waiting to send an answer back to the remote. /// Waiting to send an answer back to the remote.
PendingSend( PendingSend(
UniqueConnecId, UniqueConnecId,
@ -176,48 +187,63 @@ enum InboundSubstreamState {
PendingFlush(UniqueConnecId, KadInStreamSink<NegotiatedSubstream>), PendingFlush(UniqueConnecId, KadInStreamSink<NegotiatedSubstream>),
/// The substream is being closed. /// The substream is being closed.
Closing(KadInStreamSink<NegotiatedSubstream>), Closing(KadInStreamSink<NegotiatedSubstream>),
/// The substream was cancelled in favor of a new one.
Cancelled,
Poisoned {
phantom: PhantomData<TUserData>,
},
} }
impl<TUserData> OutboundSubstreamState<TUserData> { impl<TUserData> InboundSubstreamState<TUserData> {
/// Tries to close the substream. fn try_answer_with(
/// &mut self,
/// If the substream is not ready to be closed, returns it back. id: KademliaRequestId,
fn try_close(&mut self, cx: &mut Context<'_>) -> Poll<()> { msg: KadResponseMsg,
match self { ) -> Result<(), KadResponseMsg> {
OutboundSubstreamState::PendingOpen(_, _) match std::mem::replace(
| OutboundSubstreamState::ReportError(_, _) => Poll::Ready(()), self,
OutboundSubstreamState::PendingSend(ref mut stream, _, _) InboundSubstreamState::Poisoned {
| OutboundSubstreamState::PendingFlush(ref mut stream, _) phantom: PhantomData,
| OutboundSubstreamState::WaitingAnswer(ref mut stream, _) },
| OutboundSubstreamState::Closing(ref mut stream) => { ) {
match Sink::poll_close(Pin::new(stream), cx) { InboundSubstreamState::WaitingUser(conn_id, substream, mut waker)
Poll::Ready(_) => Poll::Ready(()), if conn_id == id.connec_unique_id =>
Poll::Pending => Poll::Pending, {
*self = InboundSubstreamState::PendingSend(conn_id, substream, msg);
if let Some(waker) = waker.take() {
waker.wake();
} }
Ok(())
}
other => {
*self = other;
Err(msg)
} }
} }
} }
}
impl InboundSubstreamState { fn close(&mut self) {
/// Tries to close the substream. match std::mem::replace(
/// self,
/// If the substream is not ready to be closed, returns it back. InboundSubstreamState::Poisoned {
fn try_close(&mut self, cx: &mut Context<'_>) -> Poll<()> { phantom: PhantomData,
match self { },
InboundSubstreamState::WaitingMessage { ) {
substream: ref mut stream, InboundSubstreamState::WaitingMessage { substream, .. }
.. | InboundSubstreamState::WaitingUser(_, substream, _)
| InboundSubstreamState::PendingSend(_, substream, _)
| InboundSubstreamState::PendingFlush(_, substream)
| InboundSubstreamState::Closing(substream) => {
*self = InboundSubstreamState::Closing(substream);
} }
| InboundSubstreamState::WaitingUser(_, ref mut stream) InboundSubstreamState::Cancelled => {
| InboundSubstreamState::PendingSend(_, ref mut stream, _) *self = InboundSubstreamState::Cancelled;
| InboundSubstreamState::PendingFlush(_, ref mut stream)
| InboundSubstreamState::Closing(ref mut stream) => {
match Sink::poll_close(Pin::new(stream), cx) {
Poll::Ready(_) => Poll::Ready(()),
Poll::Pending => Poll::Pending,
}
} }
InboundSubstreamState::Poisoned { .. } => unreachable!(),
} }
} }
} }
@ -469,7 +495,7 @@ pub enum KademliaHandlerIn<TUserData> {
/// Unique identifier for a request. Must be passed back in order to answer a request from /// Unique identifier for a request. Must be passed back in order to answer a request from
/// the remote. /// the remote.
#[derive(Debug, PartialEq, Eq)] #[derive(Debug, PartialEq, Eq, Copy, Clone)]
pub struct KademliaRequestId { pub struct KademliaRequestId {
/// Unique identifier for an incoming connection. /// Unique identifier for an incoming connection.
connec_unique_id: UniqueConnecId, connec_unique_id: UniqueConnecId,
@ -479,7 +505,10 @@ pub struct KademliaRequestId {
#[derive(Debug, Copy, Clone, PartialEq, Eq)] #[derive(Debug, Copy, Clone, PartialEq, Eq)]
struct UniqueConnecId(u64); struct UniqueConnecId(u64);
impl<TUserData> KademliaHandler<TUserData> { impl<TUserData> KademliaHandler<TUserData>
where
TUserData: Unpin,
{
/// Create a [`KademliaHandler`] using the given configuration. /// Create a [`KademliaHandler`] using the given configuration.
pub fn new( pub fn new(
config: KademliaHandlerConfig, config: KademliaHandlerConfig,
@ -503,7 +532,7 @@ impl<TUserData> KademliaHandler<TUserData> {
impl<TUserData> ConnectionHandler for KademliaHandler<TUserData> impl<TUserData> ConnectionHandler for KademliaHandler<TUserData>
where where
TUserData: Clone + fmt::Debug + Send + 'static, TUserData: Clone + fmt::Debug + Send + 'static + Unpin,
{ {
type InEvent = KademliaHandlerIn<TUserData>; type InEvent = KademliaHandlerIn<TUserData>;
type OutEvent = KademliaHandlerEvent<TUserData>; type OutEvent = KademliaHandlerEvent<TUserData>;
@ -560,14 +589,14 @@ where
} }
if self.inbound_substreams.len() == MAX_NUM_INBOUND_SUBSTREAMS { if self.inbound_substreams.len() == MAX_NUM_INBOUND_SUBSTREAMS {
if let Some(position) = self.inbound_substreams.iter().position(|s| { if let Some(s) = self.inbound_substreams.iter_mut().find(|s| {
matches!( matches!(
s, s,
// An inbound substream waiting to be reused. // An inbound substream waiting to be reused.
InboundSubstreamState::WaitingMessage { first: false, .. } InboundSubstreamState::WaitingMessage { first: false, .. }
) )
}) { }) {
self.inbound_substreams.remove(position); *s = InboundSubstreamState::Cancelled;
log::warn!( log::warn!(
"New inbound substream to {:?} exceeds inbound substream limit. \ "New inbound substream to {:?} exceeds inbound substream limit. \
Removed older substream waiting to be reused.", Removed older substream waiting to be reused.",
@ -597,153 +626,93 @@ where
fn inject_event(&mut self, message: KademliaHandlerIn<TUserData>) { fn inject_event(&mut self, message: KademliaHandlerIn<TUserData>) {
match message { match message {
KademliaHandlerIn::Reset(request_id) => { KademliaHandlerIn::Reset(request_id) => {
let pos = self if let Some(state) = self
.inbound_substreams .inbound_substreams
.iter() .iter_mut()
.position(|state| match state { .find(|state| match state {
InboundSubstreamState::WaitingUser(conn_id, _) => { InboundSubstreamState::WaitingUser(conn_id, _, _) => {
conn_id == &request_id.connec_unique_id conn_id == &request_id.connec_unique_id
} }
_ => false, _ => false,
}); })
if let Some(pos) = pos { {
// TODO: we don't properly close down the substream state.close();
let waker = futures::task::noop_waker();
let mut cx = Context::from_waker(&waker);
let _ = self.inbound_substreams.remove(pos).try_close(&mut cx);
} }
} }
KademliaHandlerIn::FindNodeReq { key, user_data } => { KademliaHandlerIn::FindNodeReq { key, user_data } => {
let msg = KadRequestMsg::FindNode { key }; let msg = KadRequestMsg::FindNode { key };
self.outbound_substreams self.outbound_substreams
.push(OutboundSubstreamState::PendingOpen(msg, Some(user_data))); .push(OutboundSubstreamState::PendingOpen(SubstreamProtocol::new(
self.config.protocol_config.clone(),
(msg, Some(user_data)),
)));
} }
KademliaHandlerIn::FindNodeRes { KademliaHandlerIn::FindNodeRes {
closer_peers, closer_peers,
request_id, request_id,
} => { } => self.answer_pending_request(request_id, KadResponseMsg::FindNode { closer_peers }),
let pos = self
.inbound_substreams
.iter()
.position(|state| match state {
InboundSubstreamState::WaitingUser(ref conn_id, _) => {
conn_id == &request_id.connec_unique_id
}
_ => false,
});
if let Some(pos) = pos {
let (conn_id, substream) = match self.inbound_substreams.remove(pos) {
InboundSubstreamState::WaitingUser(conn_id, substream) => {
(conn_id, substream)
}
_ => unreachable!(),
};
let msg = KadResponseMsg::FindNode { closer_peers };
self.inbound_substreams
.push(InboundSubstreamState::PendingSend(conn_id, substream, msg));
}
}
KademliaHandlerIn::GetProvidersReq { key, user_data } => { KademliaHandlerIn::GetProvidersReq { key, user_data } => {
let msg = KadRequestMsg::GetProviders { key }; let msg = KadRequestMsg::GetProviders { key };
self.outbound_substreams self.outbound_substreams
.push(OutboundSubstreamState::PendingOpen(msg, Some(user_data))); .push(OutboundSubstreamState::PendingOpen(SubstreamProtocol::new(
self.config.protocol_config.clone(),
(msg, Some(user_data)),
)));
} }
KademliaHandlerIn::GetProvidersRes { KademliaHandlerIn::GetProvidersRes {
closer_peers, closer_peers,
provider_peers, provider_peers,
request_id, request_id,
} => { } => self.answer_pending_request(
let pos = self request_id,
.inbound_substreams KadResponseMsg::GetProviders {
.iter() closer_peers,
.position(|state| matches!(state, InboundSubstreamState::WaitingUser(ref conn_id, _) if conn_id == &request_id.connec_unique_id)); provider_peers,
},
if let Some(pos) = pos { ),
let (conn_id, substream) = match self.inbound_substreams.remove(pos) {
InboundSubstreamState::WaitingUser(conn_id, substream) => {
(conn_id, substream)
}
_ => unreachable!(),
};
let msg = KadResponseMsg::GetProviders {
closer_peers,
provider_peers,
};
self.inbound_substreams
.push(InboundSubstreamState::PendingSend(conn_id, substream, msg));
}
}
KademliaHandlerIn::AddProvider { key, provider } => { KademliaHandlerIn::AddProvider { key, provider } => {
let msg = KadRequestMsg::AddProvider { key, provider }; let msg = KadRequestMsg::AddProvider { key, provider };
self.outbound_substreams self.outbound_substreams
.push(OutboundSubstreamState::PendingOpen(msg, None)); .push(OutboundSubstreamState::PendingOpen(SubstreamProtocol::new(
self.config.protocol_config.clone(),
(msg, None),
)));
} }
KademliaHandlerIn::GetRecord { key, user_data } => { KademliaHandlerIn::GetRecord { key, user_data } => {
let msg = KadRequestMsg::GetValue { key }; let msg = KadRequestMsg::GetValue { key };
self.outbound_substreams self.outbound_substreams
.push(OutboundSubstreamState::PendingOpen(msg, Some(user_data))); .push(OutboundSubstreamState::PendingOpen(SubstreamProtocol::new(
self.config.protocol_config.clone(),
(msg, Some(user_data)),
)));
} }
KademliaHandlerIn::PutRecord { record, user_data } => { KademliaHandlerIn::PutRecord { record, user_data } => {
let msg = KadRequestMsg::PutValue { record }; let msg = KadRequestMsg::PutValue { record };
self.outbound_substreams self.outbound_substreams
.push(OutboundSubstreamState::PendingOpen(msg, Some(user_data))); .push(OutboundSubstreamState::PendingOpen(SubstreamProtocol::new(
self.config.protocol_config.clone(),
(msg, Some(user_data)),
)));
} }
KademliaHandlerIn::GetRecordRes { KademliaHandlerIn::GetRecordRes {
record, record,
closer_peers, closer_peers,
request_id, request_id,
} => { } => {
let pos = self self.answer_pending_request(
.inbound_substreams request_id,
.iter() KadResponseMsg::GetValue {
.position(|state| match state {
InboundSubstreamState::WaitingUser(ref conn_id, _) => {
conn_id == &request_id.connec_unique_id
}
_ => false,
});
if let Some(pos) = pos {
let (conn_id, substream) = match self.inbound_substreams.remove(pos) {
InboundSubstreamState::WaitingUser(conn_id, substream) => {
(conn_id, substream)
}
_ => unreachable!(),
};
let msg = KadResponseMsg::GetValue {
record, record,
closer_peers, closer_peers,
}; },
self.inbound_substreams );
.push(InboundSubstreamState::PendingSend(conn_id, substream, msg));
}
} }
KademliaHandlerIn::PutRecordRes { KademliaHandlerIn::PutRecordRes {
key, key,
request_id, request_id,
value, value,
} => { } => {
let pos = self self.answer_pending_request(request_id, KadResponseMsg::PutValue { key, value });
.inbound_substreams
.iter()
.position(|state| matches!(state, InboundSubstreamState::WaitingUser(ref conn_id, _) if conn_id == &request_id.connec_unique_id));
if let Some(pos) = pos {
let (conn_id, substream) = match self.inbound_substreams.remove(pos) {
InboundSubstreamState::WaitingUser(conn_id, substream) => {
(conn_id, substream)
}
_ => unreachable!(),
};
let msg = KadResponseMsg::PutValue { key, value };
self.inbound_substreams
.push(InboundSubstreamState::PendingSend(conn_id, substream, msg));
}
} }
} }
} }
@ -776,10 +745,6 @@ where
Self::Error, Self::Error,
>, >,
> { > {
if self.outbound_substreams.is_empty() && self.inbound_substreams.is_empty() {
return Poll::Pending;
}
if let ProtocolStatus::Confirmed = self.protocol_status { if let ProtocolStatus::Confirmed = self.protocol_status {
self.protocol_status = ProtocolStatus::Reported; self.protocol_status = ProtocolStatus::Reported;
return Poll::Ready(ConnectionHandlerEvent::Custom( return Poll::Ready(ConnectionHandlerEvent::Custom(
@ -789,69 +754,12 @@ where
)); ));
} }
// We remove each element from `outbound_substreams` one by one and add them back. if let Poll::Ready(Some(event)) = self.outbound_substreams.poll_next_unpin(cx) {
for n in (0..self.outbound_substreams.len()).rev() { return Poll::Ready(event);
let mut substream = self.outbound_substreams.swap_remove(n);
loop {
match advance_outbound_substream(substream, self.config.protocol_config.clone(), cx)
{
(Some(new_state), Some(event), _) => {
self.outbound_substreams.push(new_state);
return Poll::Ready(event);
}
(None, Some(event), _) => {
if self.outbound_substreams.is_empty() {
self.keep_alive =
KeepAlive::Until(Instant::now() + self.config.idle_timeout);
}
return Poll::Ready(event);
}
(Some(new_state), None, false) => {
self.outbound_substreams.push(new_state);
break;
}
(Some(new_state), None, true) => {
substream = new_state;
continue;
}
(None, None, _) => {
break;
}
}
}
} }
// We remove each element from `inbound_substreams` one by one and add them back. if let Poll::Ready(Some(event)) = self.inbound_substreams.poll_next_unpin(cx) {
for n in (0..self.inbound_substreams.len()).rev() { return Poll::Ready(event);
let mut substream = self.inbound_substreams.swap_remove(n);
loop {
match advance_inbound_substream(substream, cx) {
(Some(new_state), Some(event), _) => {
self.inbound_substreams.push(new_state);
return Poll::Ready(event);
}
(None, Some(event), _) => {
if self.inbound_substreams.is_empty() {
self.keep_alive =
KeepAlive::Until(Instant::now() + self.config.idle_timeout);
}
return Poll::Ready(event);
}
(Some(new_state), None, false) => {
self.inbound_substreams.push(new_state);
break;
}
(Some(new_state), None, true) => {
substream = new_state;
continue;
}
(None, None, _) => {
break;
}
}
}
} }
if self.outbound_substreams.is_empty() && self.inbound_substreams.is_empty() { if self.outbound_substreams.is_empty() && self.inbound_substreams.is_empty() {
@ -865,6 +773,24 @@ where
} }
} }
impl<TUserData> KademliaHandler<TUserData>
where
TUserData: 'static + Clone + Send + Unpin + fmt::Debug,
{
fn answer_pending_request(&mut self, request_id: KademliaRequestId, mut msg: KadResponseMsg) {
for state in self.inbound_substreams.iter_mut() {
match state.try_answer_with(request_id, msg) {
Ok(()) => return,
Err(m) => {
msg = m;
}
}
}
debug_assert!(false, "Cannot find inbound substream for {request_id:?}")
}
}
impl Default for KademliaHandlerConfig { impl Default for KademliaHandlerConfig {
fn default() -> Self { fn default() -> Self {
KademliaHandlerConfig { KademliaHandlerConfig {
@ -875,247 +801,245 @@ impl Default for KademliaHandlerConfig {
} }
} }
/// Advances one outbound substream. impl<TUserData> Stream for OutboundSubstreamState<TUserData>
/// where
/// Returns the new state for that substream, an event to generate, and whether the substream TUserData: Unpin,
/// should be polled again. {
fn advance_outbound_substream<TUserData>( type Item = ConnectionHandlerEvent<
state: OutboundSubstreamState<TUserData>, KademliaProtocolConfig,
upgrade: KademliaProtocolConfig, (KadRequestMsg, Option<TUserData>),
cx: &mut Context<'_>, KademliaHandlerEvent<TUserData>,
) -> ( io::Error,
Option<OutboundSubstreamState<TUserData>>, >;
Option<
ConnectionHandlerEvent< fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
KademliaProtocolConfig, let this = self.get_mut();
(KadRequestMsg, Option<TUserData>),
KademliaHandlerEvent<TUserData>, loop {
io::Error, match std::mem::replace(this, OutboundSubstreamState::Poisoned) {
>, OutboundSubstreamState::PendingOpen(protocol) => {
>, *this = OutboundSubstreamState::Done;
bool, return Poll::Ready(Some(ConnectionHandlerEvent::OutboundSubstreamRequest {
) { protocol,
match state { }));
OutboundSubstreamState::PendingOpen(msg, user_data) => { }
let ev = ConnectionHandlerEvent::OutboundSubstreamRequest { OutboundSubstreamState::PendingSend(mut substream, msg, user_data) => {
protocol: SubstreamProtocol::new(upgrade, (msg, user_data)), match substream.poll_ready_unpin(cx) {
}; Poll::Ready(Ok(())) => match substream.start_send_unpin(msg) {
(None, Some(ev), false) Ok(()) => {
} *this = OutboundSubstreamState::PendingFlush(substream, user_data);
OutboundSubstreamState::PendingSend(mut substream, msg, user_data) => { }
match Sink::poll_ready(Pin::new(&mut substream), cx) { Err(error) => {
Poll::Ready(Ok(())) => match Sink::start_send(Pin::new(&mut substream), msg) { *this = OutboundSubstreamState::Done;
Ok(()) => ( let event = user_data.map(|user_data| {
Some(OutboundSubstreamState::PendingFlush(substream, user_data)), ConnectionHandlerEvent::Custom(
None, KademliaHandlerEvent::QueryError {
true, error: KademliaHandlerQueryErr::Io(error),
), user_data,
Err(error) => { },
let event = user_data.map(|user_data| { )
ConnectionHandlerEvent::Custom(KademliaHandlerEvent::QueryError { });
return Poll::Ready(event);
}
},
Poll::Pending => {
*this = OutboundSubstreamState::PendingSend(substream, msg, user_data);
return Poll::Pending;
}
Poll::Ready(Err(error)) => {
*this = OutboundSubstreamState::Done;
let event = user_data.map(|user_data| {
ConnectionHandlerEvent::Custom(KademliaHandlerEvent::QueryError {
error: KademliaHandlerQueryErr::Io(error),
user_data,
})
});
return Poll::Ready(event);
}
}
}
OutboundSubstreamState::PendingFlush(mut substream, user_data) => {
match substream.poll_flush_unpin(cx) {
Poll::Ready(Ok(())) => {
if let Some(user_data) = user_data {
*this = OutboundSubstreamState::WaitingAnswer(substream, user_data);
} else {
*this = OutboundSubstreamState::Closing(substream);
}
}
Poll::Pending => {
*this = OutboundSubstreamState::PendingFlush(substream, user_data);
return Poll::Pending;
}
Poll::Ready(Err(error)) => {
*this = OutboundSubstreamState::Done;
let event = user_data.map(|user_data| {
ConnectionHandlerEvent::Custom(KademliaHandlerEvent::QueryError {
error: KademliaHandlerQueryErr::Io(error),
user_data,
})
});
return Poll::Ready(event);
}
}
}
OutboundSubstreamState::WaitingAnswer(mut substream, user_data) => {
match substream.poll_next_unpin(cx) {
Poll::Ready(Some(Ok(msg))) => {
*this = OutboundSubstreamState::Closing(substream);
let event = process_kad_response(msg, user_data);
return Poll::Ready(Some(ConnectionHandlerEvent::Custom(event)));
}
Poll::Pending => {
*this = OutboundSubstreamState::WaitingAnswer(substream, user_data);
return Poll::Pending;
}
Poll::Ready(Some(Err(error))) => {
*this = OutboundSubstreamState::Done;
let event = KademliaHandlerEvent::QueryError {
error: KademliaHandlerQueryErr::Io(error), error: KademliaHandlerQueryErr::Io(error),
user_data, user_data,
}) };
});
(None, event, false) return Poll::Ready(Some(ConnectionHandlerEvent::Custom(event)));
}
Poll::Ready(None) => {
*this = OutboundSubstreamState::Done;
let event = KademliaHandlerEvent::QueryError {
error: KademliaHandlerQueryErr::Io(
io::ErrorKind::UnexpectedEof.into(),
),
user_data,
};
return Poll::Ready(Some(ConnectionHandlerEvent::Custom(event)));
}
}
}
OutboundSubstreamState::ReportError(error, user_data) => {
*this = OutboundSubstreamState::Done;
let event = KademliaHandlerEvent::QueryError { error, user_data };
return Poll::Ready(Some(ConnectionHandlerEvent::Custom(event)));
}
OutboundSubstreamState::Closing(mut stream) => match stream.poll_close_unpin(cx) {
Poll::Ready(Ok(())) | Poll::Ready(Err(_)) => return Poll::Ready(None),
Poll::Pending => {
*this = OutboundSubstreamState::Closing(stream);
return Poll::Pending;
} }
}, },
Poll::Pending => ( OutboundSubstreamState::Done => {
Some(OutboundSubstreamState::PendingSend( *this = OutboundSubstreamState::Done;
substream, msg, user_data, return Poll::Ready(None);
)),
None,
false,
),
Poll::Ready(Err(error)) => {
let event = user_data.map(|user_data| {
ConnectionHandlerEvent::Custom(KademliaHandlerEvent::QueryError {
error: KademliaHandlerQueryErr::Io(error),
user_data,
})
});
(None, event, false)
} }
} OutboundSubstreamState::Poisoned => unreachable!(),
}
OutboundSubstreamState::PendingFlush(mut substream, user_data) => {
match Sink::poll_flush(Pin::new(&mut substream), cx) {
Poll::Ready(Ok(())) => {
if let Some(user_data) = user_data {
(
Some(OutboundSubstreamState::WaitingAnswer(substream, user_data)),
None,
true,
)
} else {
(Some(OutboundSubstreamState::Closing(substream)), None, true)
}
}
Poll::Pending => (
Some(OutboundSubstreamState::PendingFlush(substream, user_data)),
None,
false,
),
Poll::Ready(Err(error)) => {
let event = user_data.map(|user_data| {
ConnectionHandlerEvent::Custom(KademliaHandlerEvent::QueryError {
error: KademliaHandlerQueryErr::Io(error),
user_data,
})
});
(None, event, false)
}
}
}
OutboundSubstreamState::WaitingAnswer(mut substream, user_data) => {
match Stream::poll_next(Pin::new(&mut substream), cx) {
Poll::Ready(Some(Ok(msg))) => {
let new_state = OutboundSubstreamState::Closing(substream);
let event = process_kad_response(msg, user_data);
(
Some(new_state),
Some(ConnectionHandlerEvent::Custom(event)),
true,
)
}
Poll::Pending => (
Some(OutboundSubstreamState::WaitingAnswer(substream, user_data)),
None,
false,
),
Poll::Ready(Some(Err(error))) => {
let event = KademliaHandlerEvent::QueryError {
error: KademliaHandlerQueryErr::Io(error),
user_data,
};
(None, Some(ConnectionHandlerEvent::Custom(event)), false)
}
Poll::Ready(None) => {
let event = KademliaHandlerEvent::QueryError {
error: KademliaHandlerQueryErr::Io(io::ErrorKind::UnexpectedEof.into()),
user_data,
};
(None, Some(ConnectionHandlerEvent::Custom(event)), false)
}
}
}
OutboundSubstreamState::ReportError(error, user_data) => {
let event = KademliaHandlerEvent::QueryError { error, user_data };
(None, Some(ConnectionHandlerEvent::Custom(event)), false)
}
OutboundSubstreamState::Closing(mut stream) => {
match Sink::poll_close(Pin::new(&mut stream), cx) {
Poll::Ready(Ok(())) => (None, None, false),
Poll::Pending => (Some(OutboundSubstreamState::Closing(stream)), None, false),
Poll::Ready(Err(_)) => (None, None, false),
} }
} }
} }
} }
/// Advances one inbound substream.
/// impl<TUserData> Stream for InboundSubstreamState<TUserData>
/// Returns the new state for that substream, an event to generate, and whether the substream where
/// should be polled again. TUserData: Unpin,
fn advance_inbound_substream<TUserData>( {
state: InboundSubstreamState, type Item = ConnectionHandlerEvent<
cx: &mut Context<'_>, KademliaProtocolConfig,
) -> ( (KadRequestMsg, Option<TUserData>),
Option<InboundSubstreamState>, KademliaHandlerEvent<TUserData>,
Option< io::Error,
ConnectionHandlerEvent< >;
KademliaProtocolConfig,
(KadRequestMsg, Option<TUserData>), fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
KademliaHandlerEvent<TUserData>, let this = self.get_mut();
io::Error,
>, loop {
>, match std::mem::replace(
bool, this,
) { Self::Poisoned {
match state { phantom: PhantomData,
InboundSubstreamState::WaitingMessage { },
first, ) {
connection_id, InboundSubstreamState::WaitingMessage {
mut substream,
} => match Stream::poll_next(Pin::new(&mut substream), cx) {
Poll::Ready(Some(Ok(msg))) => {
if let Ok(ev) = process_kad_request(msg, connection_id) {
(
Some(InboundSubstreamState::WaitingUser(connection_id, substream)),
Some(ConnectionHandlerEvent::Custom(ev)),
false,
)
} else {
(Some(InboundSubstreamState::Closing(substream)), None, true)
}
}
Poll::Pending => (
Some(InboundSubstreamState::WaitingMessage {
first, first,
connection_id, connection_id,
substream, mut substream,
}), } => match substream.poll_next_unpin(cx) {
None, Poll::Ready(Some(Ok(msg))) => {
false, if let Ok(ev) = process_kad_request(msg, connection_id) {
), *this =
Poll::Ready(None) => { InboundSubstreamState::WaitingUser(connection_id, substream, None);
trace!("Inbound substream: EOF"); return Poll::Ready(Some(ConnectionHandlerEvent::Custom(ev)));
(None, None, false) } else {
} *this = InboundSubstreamState::Closing(substream);
Poll::Ready(Some(Err(e))) => { }
trace!("Inbound substream error: {:?}", e); }
(None, None, false) Poll::Pending => {
} *this = InboundSubstreamState::WaitingMessage {
}, first,
InboundSubstreamState::WaitingUser(id, substream) => ( connection_id,
Some(InboundSubstreamState::WaitingUser(id, substream)), substream,
None, };
false, return Poll::Pending;
), }
InboundSubstreamState::PendingSend(id, mut substream, msg) => { Poll::Ready(None) => {
match Sink::poll_ready(Pin::new(&mut substream), cx) { return Poll::Ready(None);
Poll::Ready(Ok(())) => match Sink::start_send(Pin::new(&mut substream), msg) { }
Ok(()) => ( Poll::Ready(Some(Err(e))) => {
Some(InboundSubstreamState::PendingFlush(id, substream)), trace!("Inbound substream error: {:?}", e);
None, return Poll::Ready(None);
true, }
),
Err(_) => (None, None, false),
}, },
Poll::Pending => ( InboundSubstreamState::WaitingUser(id, substream, _) => {
Some(InboundSubstreamState::PendingSend(id, substream, msg)), *this =
None, InboundSubstreamState::WaitingUser(id, substream, Some(cx.waker().clone()));
false,
), return Poll::Pending;
Poll::Ready(Err(_)) => (None, None, false), }
} InboundSubstreamState::PendingSend(id, mut substream, msg) => {
} match substream.poll_ready_unpin(cx) {
InboundSubstreamState::PendingFlush(id, mut substream) => { Poll::Ready(Ok(())) => match substream.start_send_unpin(msg) {
match Sink::poll_flush(Pin::new(&mut substream), cx) { Ok(()) => {
Poll::Ready(Ok(())) => ( *this = InboundSubstreamState::PendingFlush(id, substream);
Some(InboundSubstreamState::WaitingMessage { }
first: false, Err(_) => return Poll::Ready(None),
connection_id: id, },
substream, Poll::Pending => {
}), *this = InboundSubstreamState::PendingSend(id, substream, msg);
None, return Poll::Pending;
true, }
), Poll::Ready(Err(_)) => return Poll::Ready(None),
Poll::Pending => ( }
Some(InboundSubstreamState::PendingFlush(id, substream)), }
None, InboundSubstreamState::PendingFlush(id, mut substream) => {
false, match substream.poll_flush_unpin(cx) {
), Poll::Ready(Ok(())) => {
Poll::Ready(Err(_)) => (None, None, false), *this = InboundSubstreamState::WaitingMessage {
} first: false,
} connection_id: id,
InboundSubstreamState::Closing(mut stream) => { substream,
match Sink::poll_close(Pin::new(&mut stream), cx) { };
Poll::Ready(Ok(())) => (None, None, false), }
Poll::Pending => (Some(InboundSubstreamState::Closing(stream)), None, false), Poll::Pending => {
Poll::Ready(Err(_)) => (None, None, false), *this = InboundSubstreamState::PendingFlush(id, substream);
return Poll::Pending;
}
Poll::Ready(Err(_)) => return Poll::Ready(None),
}
}
InboundSubstreamState::Closing(mut stream) => match stream.poll_close_unpin(cx) {
Poll::Ready(Ok(())) | Poll::Ready(Err(_)) => return Poll::Ready(None),
Poll::Pending => {
*this = InboundSubstreamState::Closing(stream);
return Poll::Pending;
}
},
InboundSubstreamState::Poisoned { .. } => unreachable!(),
InboundSubstreamState::Cancelled => return Poll::Ready(None),
} }
} }
} }