Skip to main content

airio_core/
muxing.rs

1use futures::{AsyncRead, AsyncWrite};
2use std::{
3    pin::Pin,
4    task::{Context, Poll},
5};
6
7mod boxed;
8
9pub use boxed::{StreamMuxerBox, SubstreamBox};
10
11pub trait StreamMuxer {
12    type Substream: AsyncRead + AsyncWrite;
13    type Error: std::error::Error;
14
15    /// Poll 进站子流
16    fn poll_inbound(
17        self: Pin<&mut Self>,
18        cx: &mut Context<'_>,
19    ) -> Poll<Result<Self::Substream, Self::Error>>;
20
21    /// Poll 出站子流
22    fn poll_outbound(
23        self: Pin<&mut Self>,
24        cx: &mut Context<'_>,
25    ) -> Poll<Result<Self::Substream, Self::Error>>;
26
27    /// Poll 关闭多路复用器
28    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>>;
29
30    /// Poll 多路复用器事件
31    fn poll(
32        self: Pin<&mut Self>,
33        cx: &mut Context<'_>,
34    ) -> Poll<Result<StreamMuxerEvent, Self::Error>>;
35}
36
37#[derive(Debug)]
38pub enum StreamMuxerEvent {}
39
40pub trait StreamMuxerExt: StreamMuxer + Sized {
41    /// Convenience function for calling [`StreamMuxer::poll_inbound`]
42    /// for [`StreamMuxer`]s that are `Unpin`.
43    fn poll_inbound_unpin(
44        &mut self,
45        cx: &mut Context<'_>,
46    ) -> Poll<Result<Self::Substream, Self::Error>>
47    where
48        Self: Unpin,
49    {
50        Pin::new(self).poll_inbound(cx)
51    }
52
53    /// Convenience function for calling [`StreamMuxer::poll_outbound`]
54    /// for [`StreamMuxer`]s that are `Unpin`.
55    fn poll_outbound_unpin(
56        &mut self,
57        cx: &mut Context<'_>,
58    ) -> Poll<Result<Self::Substream, Self::Error>>
59    where
60        Self: Unpin,
61    {
62        Pin::new(self).poll_outbound(cx)
63    }
64
65    /// Convenience function for calling [`StreamMuxer::poll`]
66    /// for [`StreamMuxer`]s that are `Unpin`.
67    fn poll_unpin(&mut self, cx: &mut Context<'_>) -> Poll<Result<StreamMuxerEvent, Self::Error>>
68    where
69        Self: Unpin,
70    {
71        Pin::new(self).poll(cx)
72    }
73
74    /// Convenience function for calling [`StreamMuxer::poll_close`]
75    /// for [`StreamMuxer`]s that are `Unpin`.
76    fn poll_close_unpin(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>>
77    where
78        Self: Unpin,
79    {
80        Pin::new(self).poll_close(cx)
81    }
82
83    /// Returns a future for closing this [`StreamMuxer`].
84    fn close(self) -> Close<Self> {
85        Close(self)
86    }
87}
88
89impl<S> StreamMuxerExt for S where S: StreamMuxer {}
90
91pub struct Close<S>(S);
92
93impl<S> Future for Close<S>
94where
95    S: StreamMuxer + Unpin,
96{
97    type Output = Result<(), S::Error>;
98
99    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
100        self.0.poll_close_unpin(cx)
101    }
102}