volans_yamux/
lib.rs

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

volans_yamux/
lib.rs
