Fix polling in handled_node.rs (#531)

This commit is contained in:
Pierre Krieger
2018-10-01 14:24:01 +02:00
committed by GitHub
parent a9eefe1015
commit 4ff7d686a4

View File

@ -195,46 +195,44 @@ where
type Error = IoError; type Error = IoError;
fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> { fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
// We extract the value from `self.node` and put it back in place if `NotReady`.
if let Some(mut node) = self.node.take() {
loop { loop {
match node.poll() { let mut node_not_ready = false;
Ok(Async::NotReady) => {
self.node = Some(node); match self.node.as_mut().map(|n| n.poll()) {
break; Some(Ok(Async::NotReady)) | None => {},
}, Some(Ok(Async::Ready(Some(NodeEvent::InboundSubstream { substream })))) => {
Ok(Async::Ready(Some(NodeEvent::InboundSubstream { substream }))) => {
self.handler.inject_substream(substream, NodeHandlerEndpoint::Listener); self.handler.inject_substream(substream, NodeHandlerEndpoint::Listener);
}, },
Ok(Async::Ready(Some(NodeEvent::OutboundSubstream { user_data, substream }))) => { Some(Ok(Async::Ready(Some(NodeEvent::OutboundSubstream { user_data, substream })))) => {
let endpoint = NodeHandlerEndpoint::Dialer(user_data); let endpoint = NodeHandlerEndpoint::Dialer(user_data);
self.handler.inject_substream(substream, endpoint); self.handler.inject_substream(substream, endpoint);
}, },
Ok(Async::Ready(None)) => { Some(Ok(Async::Ready(None))) => {
// Breaking from the loop without putting back the node. node_not_ready = true;
self.node = None;
self.handler.shutdown(); self.handler.shutdown();
break;
}, },
Ok(Async::Ready(Some(NodeEvent::Multiaddr(result)))) => { Some(Ok(Async::Ready(Some(NodeEvent::Multiaddr(result))))) => {
self.handler.inject_multiaddr(result); self.handler.inject_multiaddr(result);
}, },
Ok(Async::Ready(Some(NodeEvent::OutboundClosed { user_data }))) => { Some(Ok(Async::Ready(Some(NodeEvent::OutboundClosed { user_data })))) => {
self.handler.inject_outbound_closed(user_data); self.handler.inject_outbound_closed(user_data);
}, },
Ok(Async::Ready(Some(NodeEvent::InboundClosed))) => { Some(Ok(Async::Ready(Some(NodeEvent::InboundClosed)))) => {
self.handler.inject_inbound_closed(); self.handler.inject_inbound_closed();
}, },
Err(err) => { Some(Err(err)) => {
// Breaking from the loop without putting back the node. self.node = None;
return Err(err); return Err(err);
}, },
} }
}
}
loop {
match self.handler.poll() { match self.handler.poll() {
Ok(Async::NotReady) => break, Ok(Async::NotReady) => {
if node_not_ready {
break;
}
},
Ok(Async::Ready(Some(NodeHandlerEvent::OutboundSubstreamRequest(user_data)))) => { Ok(Async::Ready(Some(NodeHandlerEvent::OutboundSubstreamRequest(user_data)))) => {
if let Some(node) = self.node.as_mut() { if let Some(node) = self.node.as_mut() {
match node.open_substream(user_data) { match node.open_substream(user_data) {
@ -264,7 +262,7 @@ where
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
use futures::{future, task}; use futures::future;
use muxing::StreamMuxer; use muxing::StreamMuxer;
use tokio::runtime::current_thread; use tokio::runtime::current_thread;