use futures::prelude::*;
use crate::muxing;
use smallvec::SmallVec;
use std::fmt;
use std::io::Error as IoError;
use std::sync::Arc;
pub struct NodeStream<TMuxer, TUserData>
where
TMuxer: muxing::StreamMuxer,
{
muxer: Arc<TMuxer>,
inbound_state: StreamState,
outbound_state: StreamState,
outbound_substreams: SmallVec<[(TUserData, TMuxer::OutboundSubstream); 8]>,
}
pub type Substream<TMuxer> = muxing::SubstreamRef<Arc<TMuxer>>;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum StreamState {
Open,
Shutdown,
Flush,
Closed
}
pub enum NodeEvent<TMuxer, TUserData>
where
TMuxer: muxing::StreamMuxer,
{
InboundSubstream {
substream: Substream<TMuxer>,
},
OutboundSubstream {
user_data: TUserData,
substream: Substream<TMuxer>,
},
OutboundClosed {
user_data: TUserData,
},
InboundClosed,
}
#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq, PartialOrd, Ord)]
pub struct OutboundSubstreamId(usize);
impl<TMuxer, TUserData> NodeStream<TMuxer, TUserData>
where
TMuxer: muxing::StreamMuxer,
{
#[inline]
pub fn new(muxer: TMuxer) -> Self {
NodeStream {
muxer: Arc::new(muxer),
inbound_state: StreamState::Open,
outbound_state: StreamState::Open,
outbound_substreams: SmallVec::new(),
}
}
pub fn open_substream(&mut self, user_data: TUserData) -> Result<(), TUserData> {
if self.outbound_state != StreamState::Open {
return Err(user_data);
}
let raw = self.muxer.open_outbound();
self.outbound_substreams.push((user_data, raw));
Ok(())
}
#[inline]
pub fn is_inbound_open(&self) -> bool {
self.inbound_state == StreamState::Open
}
#[inline]
pub fn is_outbound_open(&self) -> bool {
self.outbound_state == StreamState::Open
}
pub fn close(mut self) -> Vec<TUserData> {
self.cancel_outgoing()
}
pub fn cancel_outgoing(&mut self) -> Vec<TUserData> {
let mut out = Vec::with_capacity(self.outbound_substreams.len());
for (user_data, outbound) in self.outbound_substreams.drain() {
out.push(user_data);
self.muxer.destroy_outbound(outbound);
}
out
}
pub fn shutdown_all(&mut self) {
if self.inbound_state == StreamState::Open {
self.inbound_state = StreamState::Shutdown
}
if self.outbound_state == StreamState::Open {
self.outbound_state = StreamState::Shutdown
}
}
fn poll_shutdown(&mut self) -> Poll<(), IoError> {
use self::StreamState::*;
loop {
match (self.inbound_state, self.outbound_state) {
(Open, Open) | (Open, Closed) | (Closed, Open) | (Closed, Closed) => {
return Ok(Async::Ready(()))
}
(Shutdown, Shutdown) => {
if let Async::Ready(()) = self.muxer.shutdown(muxing::Shutdown::All)? {
self.inbound_state = StreamState::Flush;
self.outbound_state = StreamState::Flush;
continue
}
return Ok(Async::NotReady)
}
(Shutdown, _) => {
if let Async::Ready(()) = self.muxer.shutdown(muxing::Shutdown::Inbound)? {
self.inbound_state = StreamState::Flush;
continue
}
return Ok(Async::NotReady)
}
(_, Shutdown) => {
if let Async::Ready(()) = self.muxer.shutdown(muxing::Shutdown::Outbound)? {
self.outbound_state = StreamState::Flush;
continue
}
return Ok(Async::NotReady)
}
(Flush, Open) => {
if let Async::Ready(()) = self.muxer.flush_all()? {
self.inbound_state = StreamState::Closed;
continue
}
return Ok(Async::NotReady)
}
(Open, Flush) => {
if let Async::Ready(()) = self.muxer.flush_all()? {
self.outbound_state = StreamState::Closed;
continue
}
return Ok(Async::NotReady)
}
(Flush, Flush) | (Flush, Closed) | (Closed, Flush) => {
if let Async::Ready(()) = self.muxer.flush_all()? {
self.inbound_state = StreamState::Closed;
self.outbound_state = StreamState::Closed;
continue
}
return Ok(Async::NotReady)
}
}
}
}
}
impl<TMuxer, TUserData> Stream for NodeStream<TMuxer, TUserData>
where
TMuxer: muxing::StreamMuxer,
{
type Item = NodeEvent<TMuxer, TUserData>;
type Error = IoError;
fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
if self.poll_shutdown()?.is_not_ready() {
return Ok(Async::NotReady)
}
if self.inbound_state == StreamState::Open {
match self.muxer.poll_inbound()? {
Async::Ready(Some(substream)) => {
let substream = muxing::substream_from_ref(self.muxer.clone(), substream);
return Ok(Async::Ready(Some(NodeEvent::InboundSubstream {
substream,
})));
}
Async::Ready(None) => {
self.inbound_state = StreamState::Closed;
return Ok(Async::Ready(Some(NodeEvent::InboundClosed)));
}
Async::NotReady => {}
}
}
for n in (0..self.outbound_substreams.len()).rev() {
let (user_data, mut outbound) = self.outbound_substreams.swap_remove(n);
match self.muxer.poll_outbound(&mut outbound) {
Ok(Async::Ready(Some(substream))) => {
let substream = muxing::substream_from_ref(self.muxer.clone(), substream);
self.muxer.destroy_outbound(outbound);
return Ok(Async::Ready(Some(NodeEvent::OutboundSubstream {
user_data,
substream,
})));
}
Ok(Async::Ready(None)) => {
self.outbound_state = StreamState::Closed;
self.muxer.destroy_outbound(outbound);
return Ok(Async::Ready(Some(NodeEvent::OutboundClosed { user_data })));
}
Ok(Async::NotReady) => {
self.outbound_substreams.push((user_data, outbound));
}
Err(err) => {
self.muxer.destroy_outbound(outbound);
return Err(err);
}
}
}
if self.inbound_state == StreamState::Closed
&& self.outbound_state == StreamState::Closed
&& self.outbound_substreams.is_empty()
{
return Ok(Async::Ready(None))
}
Ok(Async::NotReady)
}
}
impl<TMuxer, TUserData> fmt::Debug for NodeStream<TMuxer, TUserData>
where
TMuxer: muxing::StreamMuxer,
{
fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
f.debug_struct("NodeStream")
.field("inbound_state", &self.inbound_state)
.field("outbound_state", &self.outbound_state)
.field("outbound_substreams", &self.outbound_substreams.len())
.finish()
}
}
impl<TMuxer, TUserData> Drop for NodeStream<TMuxer, TUserData>
where
TMuxer: muxing::StreamMuxer,
{
fn drop(&mut self) {
for (_, outbound) in self.outbound_substreams.drain() {
self.muxer.destroy_outbound(outbound);
}
}
}
impl<TMuxer, TUserData> fmt::Debug for NodeEvent<TMuxer, TUserData>
where
TMuxer: muxing::StreamMuxer,
TMuxer::Substream: fmt::Debug,
TUserData: fmt::Debug,
{
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match self {
NodeEvent::InboundSubstream { substream } => {
f.debug_struct("NodeEvent::OutboundClosed")
.field("substream", substream)
.finish()
},
NodeEvent::OutboundSubstream { user_data, substream } => {
f.debug_struct("NodeEvent::OutboundSubstream")
.field("user_data", user_data)
.field("substream", substream)
.finish()
},
NodeEvent::OutboundClosed { user_data } => {
f.debug_struct("NodeEvent::OutboundClosed")
.field("user_data", user_data)
.finish()
},
NodeEvent::InboundClosed => {
f.debug_struct("NodeEvent::InboundClosed")
.finish()
},
}
}
}
#[cfg(test)]
mod node_stream {
use super::{NodeEvent, NodeStream};
use crate::tests::dummy_muxer::{DummyMuxer, DummyConnectionState};
use assert_matches::assert_matches;
use futures::prelude::*;
use tokio_mock_task::MockTask;
fn build_node_stream() -> NodeStream<DummyMuxer, Vec<u8>> {
let muxer = DummyMuxer::new();
NodeStream::<_, Vec<u8>>::new(muxer)
}
#[test]
fn can_open_outbound_substreams_until_an_outbound_channel_is_closed() {
let mut muxer = DummyMuxer::new();
muxer.set_outbound_connection_state(DummyConnectionState::Closed);
let mut ns = NodeStream::<_, Vec<u8>>::new(muxer);
assert!(ns.open_substream(vec![1,2,3]).is_ok());
let out = ns.poll();
assert_matches!(out, Ok(Async::Ready(Some(node_event))) => {
assert_matches!(node_event, NodeEvent::OutboundClosed{user_data} => {
assert_eq!(user_data, vec![1,2,3])
})
});
assert_matches!(ns.open_substream(vec![22]), Err(user_data) => {
assert_eq!(user_data, vec![22]);
});
}
#[test]
fn query_inbound_outbound_state() {
let ns = build_node_stream();
assert!(ns.is_inbound_open());
assert!(ns.is_outbound_open());
}
#[test]
fn query_inbound_state() {
let mut muxer = DummyMuxer::new();
muxer.set_inbound_connection_state(DummyConnectionState::Closed);
let mut ns = NodeStream::<_, Vec<u8>>::new(muxer);
assert_matches!(ns.poll(), Ok(Async::Ready(Some(node_event))) => {
assert_matches!(node_event, NodeEvent::InboundClosed)
});
assert!(!ns.is_inbound_open());
}
#[test]
fn query_outbound_state() {
let mut muxer = DummyMuxer::new();
muxer.set_outbound_connection_state(DummyConnectionState::Closed);
let mut ns = NodeStream::<_, Vec<u8>>::new(muxer);
assert!(ns.is_outbound_open());
ns.open_substream(vec![1]).unwrap();
let poll_result = ns.poll();
assert_matches!(poll_result, Ok(Async::Ready(Some(node_event))) => {
assert_matches!(node_event, NodeEvent::OutboundClosed{user_data} => {
assert_eq!(user_data, vec![1])
})
});
assert!(!ns.is_outbound_open(), "outbound connection should be closed after polling");
}
#[test]
fn closing_a_node_stream_destroys_substreams_and_returns_submitted_user_data() {
let mut ns = build_node_stream();
ns.open_substream(vec![2]).unwrap();
ns.open_substream(vec![3]).unwrap();
ns.open_substream(vec![5]).unwrap();
let user_data_submitted = ns.close();
assert_eq!(user_data_submitted, vec![
vec![2], vec![3], vec![5]
]);
}
#[test]
fn poll_returns_not_ready_when_there_is_nothing_to_do() {
let mut task = MockTask::new();
task.enter(|| {
let mut muxer = DummyMuxer::new();
muxer.set_inbound_connection_state(DummyConnectionState::Pending);
muxer.set_outbound_connection_state(DummyConnectionState::Pending);
let mut ns = NodeStream::<_, Vec<u8>>::new(muxer);
assert_matches!(ns.poll(), Ok(Async::NotReady));
});
}
#[test]
fn poll_closes_the_node_stream_when_no_more_work_can_be_done() {
let mut muxer = DummyMuxer::new();
muxer.set_inbound_connection_state(DummyConnectionState::Closed);
muxer.set_outbound_connection_state(DummyConnectionState::Closed);
let mut ns = NodeStream::<_, Vec<u8>>::new(muxer);
ns.open_substream(vec![]).unwrap();
ns.poll().unwrap();
ns.poll().unwrap();
ns.poll().unwrap();
assert_matches!(ns.poll(), Ok(Async::Ready(None)));
}
#[test]
fn poll_sets_up_substreams_yielding_them_in_reverse_order() {
let mut muxer = DummyMuxer::new();
muxer.set_inbound_connection_state(DummyConnectionState::Closed);
muxer.set_outbound_connection_state(DummyConnectionState::Opened);
let mut ns = NodeStream::<_, Vec<u8>>::new(muxer);
ns.open_substream(vec![1]).unwrap();
ns.open_substream(vec![2]).unwrap();
ns.poll().unwrap();
assert_matches!(ns.poll(), Ok(Async::Ready(Some(node_event))) => {
assert_matches!(node_event, NodeEvent::OutboundSubstream{ user_data, substream:_ } => {
assert_eq!(user_data, vec![2]);
})
});
assert_matches!(ns.poll(), Ok(Async::Ready(Some(node_event))) => {
assert_matches!(node_event, NodeEvent::OutboundSubstream{ user_data, substream: _ } => {
assert_eq!(user_data, vec![1]);
})
});
}
#[test]
fn poll_keeps_outbound_substreams_when_the_outgoing_connection_is_not_ready() {
let mut muxer = DummyMuxer::new();
muxer.set_inbound_connection_state(DummyConnectionState::Closed);
muxer.set_outbound_connection_state(DummyConnectionState::Pending);
let mut ns = NodeStream::<_, Vec<u8>>::new(muxer);
ns.open_substream(vec![1]).unwrap();
ns.poll().unwrap();
ns.poll().unwrap();
assert!(ns.is_outbound_open());
assert!(format!("{:?}", ns).contains("outbound_substreams: 1"));
}
#[test]
fn poll_returns_incoming_substream() {
let mut muxer = DummyMuxer::new();
muxer.set_inbound_connection_state(DummyConnectionState::Opened);
let mut ns = NodeStream::<_, Vec<u8>>::new(muxer);
assert_matches!(ns.poll(), Ok(Async::Ready(Some(node_event))) => {
assert_matches!(node_event, NodeEvent::InboundSubstream{ substream: _ });
});
}
}