use futures::prelude::*;
use crate::muxing;
use smallvec::SmallVec;
use std::{fmt, io::Error as IoError, pin::Pin, sync::Arc, task::Context, task::Poll};
pub struct NodeStream<TMuxer, TUserData>
where
TMuxer: muxing::StreamMuxer,
{
muxer: Arc<TMuxer>,
outbound_substreams: SmallVec<[(TUserData, TMuxer::OutboundSubstream); 8]>,
}
pub struct Close<TMuxer> {
muxer: Arc<TMuxer>,
}
pub type Substream<TMuxer> = muxing::SubstreamRef<Arc<TMuxer>>;
pub enum NodeEvent<TMuxer, TUserData>
where
TMuxer: muxing::StreamMuxer,
{
InboundSubstream {
substream: Substream<TMuxer>,
},
OutboundSubstream {
user_data: TUserData,
substream: Substream<TMuxer>,
},
}
#[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),
outbound_substreams: SmallVec::new(),
}
}
pub fn open_substream(&mut self, user_data: TUserData) {
let raw = self.muxer.open_outbound();
self.outbound_substreams.push((user_data, raw));
}
pub fn is_remote_acknowledged(&self) -> bool {
self.muxer.is_remote_acknowledged()
}
#[must_use]
pub fn close(mut self) -> (Close<TMuxer>, Vec<TUserData>) {
let substreams = self.cancel_outgoing();
let close = Close { muxer: self.muxer.clone() };
(close, substreams)
}
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 poll(&mut self, cx: &mut Context) -> Poll<Result<NodeEvent<TMuxer, TUserData>, IoError>> {
match self.muxer.poll_inbound(cx) {
Poll::Ready(Ok(substream)) => {
let substream = muxing::substream_from_ref(self.muxer.clone(), substream);
return Poll::Ready(Ok(NodeEvent::InboundSubstream {
substream,
}));
}
Poll::Ready(Err(err)) => return Poll::Ready(Err(err.into())),
Poll::Pending => {}
}
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(cx, &mut outbound) {
Poll::Ready(Ok(substream)) => {
let substream = muxing::substream_from_ref(self.muxer.clone(), substream);
self.muxer.destroy_outbound(outbound);
return Poll::Ready(Ok(NodeEvent::OutboundSubstream {
user_data,
substream,
}));
}
Poll::Pending => {
self.outbound_substreams.push((user_data, outbound));
}
Poll::Ready(Err(err)) => {
self.muxer.destroy_outbound(outbound);
return Poll::Ready(Err(err.into()));
}
}
}
Poll::Pending
}
}
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("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> Future for Close<TMuxer>
where
TMuxer: muxing::StreamMuxer,
{
type Output = Result<(), IoError>;
fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
match self.muxer.close(cx) {
Poll::Pending => Poll::Pending,
Poll::Ready(Ok(())) => Poll::Ready(Ok(())),
Poll::Ready(Err(err)) => Poll::Ready(Err(err.into())),
}
}
}
impl<TMuxer> fmt::Debug for Close<TMuxer>
where
TMuxer: muxing::StreamMuxer,
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
f.debug_struct("Close")
.finish()
}
}
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()
},
}
}
}