volans_muxing/
lib.rs

1use volans_core::{
2    StreamMuxer, UpgradeInfo,
3    upgrade::{InboundConnectionUpgrade, OutboundConnectionUpgrade},
4};
5use futures::{AsyncRead, AsyncWrite, future, ready};
6pub use muxing::{Connection, ConnectionError, Endpoint, Stream};
7use std::{
8    collections::VecDeque,
9    io, iter,
10    pin::Pin,
11    task::{Context, Poll, Waker},
12};
13
14#[derive(Debug)]
15pub struct Muxer<C> {
16    connection: Connection<C>,
17    inbound_stream_buffer: VecDeque<Stream>,
18    inbound_stream_waker: Option<Waker>,
19}
20
21impl<C> Muxer<C>
22where
23    C: AsyncRead + AsyncWrite + Unpin + 'static,
24{
25    pub fn new(connection: Connection<C>) -> Self {
26        Muxer {
27            connection,
28            inbound_stream_buffer: VecDeque::with_capacity(MAX_BUFFERED_INBOUND_STREAMS),
29            inbound_stream_waker: None,
30        }
31    }
32
33    fn poll_inner(&mut self, cx: &mut Context<'_>) -> Poll<Result<Stream, ConnectionError>> {
34        let stream =
35            ready!(self.connection.poll_next_inbound(cx)?).ok_or(ConnectionError::Closed)?;
36        Poll::Ready(Ok(stream))
37    }
38}
39
40const MAX_BUFFERED_INBOUND_STREAMS: usize = 256;
41
42impl<C> StreamMuxer for Muxer<C>
43where
44    C: AsyncRead + AsyncWrite + Unpin + 'static,
45{
46    type Substream = Stream;
47    type Error = ConnectionError;
48
49    fn poll_inbound(
50        mut self: Pin<&mut Self>,
51        cx: &mut Context<'_>,
52    ) -> Poll<Result<Self::Substream, Self::Error>> {
53        if let Some(stream) = self.inbound_stream_buffer.pop_front() {
54            return Poll::Ready(Ok(stream));
55        }
56        if let Poll::Ready(res) = self.poll_inner(cx) {
57            return Poll::Ready(res);
58        }
59        self.inbound_stream_waker = Some(cx.waker().clone());
60        Poll::Pending
61    }
62
63    fn poll_outbound(
64        mut self: Pin<&mut Self>,
65        cx: &mut Context<'_>,
66    ) -> Poll<Result<Self::Substream, Self::Error>> {
67        self.as_mut().connection.poll_new_outbound(cx)
68    }
69
70    #[tracing::instrument(level = "trace", name = "StreamMuxer::poll", skip(self, cx))]
71    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
72        let mut this = self.as_mut();
73        let inbound_stream = ready!(this.poll_inner(cx))?;
74
75        if this.inbound_stream_buffer.len() >= MAX_BUFFERED_INBOUND_STREAMS {
76            tracing::warn!(
77                "{}: Inbound stream buffer is full, dropping stream:",
78                inbound_stream
79            );
80            drop(inbound_stream);
81        } else {
82            this.inbound_stream_buffer.push_back(inbound_stream);
83            if let Some(waker) = this.inbound_stream_waker.take() {
84                waker.wake();
85            }
86        }
87        // 马上唤醒任务
88        cx.waker().wake_by_ref();
89        Poll::Pending
90    }
91
92    #[tracing::instrument(level = "trace", name = "StreamMuxer::poll_close", skip(self, cx))]
93    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
94        self.as_mut().connection.poll_close(cx)
95    }
96}
97
98#[derive(Debug, Clone)]
99pub struct Config(muxing::Config);
100
101impl Config {
102    pub fn new() -> Self {
103        Config(muxing::Config::default())
104    }
105
106    pub fn set_max_active_streams(&mut self, max_active_streams: usize) -> &mut Self {
107        self.0.set_max_active_streams(max_active_streams);
108        self
109    }
110
111    pub fn set_read_after_close(&mut self, read_after_close: bool) -> &mut Self {
112        self.0.set_read_after_close(read_after_close);
113        self
114    }
115}
116
117impl Default for Config {
118    fn default() -> Self {
119        Config(muxing::Config::default())
120    }
121}
122
123impl UpgradeInfo for Config {
124    type Info = &'static str;
125    type InfoIter = iter::Once<Self::Info>;
126
127    fn protocol_info(&self) -> Self::InfoIter {
128        iter::once("/v1/muxing")
129    }
130}
131
132impl<C> InboundConnectionUpgrade<C> for Config
133where
134    C: AsyncRead + AsyncWrite + Unpin + 'static,
135{
136    type Output = Muxer<C>;
137    type Error = io::Error;
138    type Future = future::Ready<Result<Self::Output, Self::Error>>;
139
140    fn upgrade_inbound(self, socket: C, _info: Self::Info) -> Self::Future {
141        let connection = Connection::new(socket, self.0, Endpoint::Server);
142        future::ready(Ok(Muxer::new(connection)))
143    }
144}
145
146impl<C> OutboundConnectionUpgrade<C> for Config
147where
148    C: AsyncRead + AsyncWrite + Unpin + 'static,
149{
150    type Output = Muxer<C>;
151    type Error = io::Error;
152    type Future = future::Ready<Result<Self::Output, Self::Error>>;
153
154    fn upgrade_outbound(self, socket: C, _info: Self::Info) -> Self::Future {
155        let connection = Connection::new(socket, self.0, Endpoint::Client);
156        future::ready(Ok(Muxer::new(connection)))
157    }
158}
volans_muxing/lib.rs

volans_muxing/
lib.rs
