volans_swarm/connection/
inbound.rs

1use std::{
2    pin::Pin,
3    task::{Context, Poll},
4    time::Duration,
5};
6
7use futures::{
8    Stream, StreamExt,
9    stream::{self, FuturesUnordered},
10};
11use volans_core::muxing::{Closing, StreamMuxerBox, StreamMuxerExt};
12
13use crate::{
14    ConnectionHandler, ConnectionHandlerEvent, InboundStreamHandler, InboundUpgradeSend,
15    StreamUpgradeError,
16    connection::{ConnectionController, Shutdown, StreamUpgrade, compute_new_shutdown},
17    error::ConnectionError,
18    substream::ActiveStreamCounter,
19};
20
21pub struct InboundConnection<THandler>
22where
23    THandler: InboundStreamHandler,
24{
25    muxer: StreamMuxerBox,
26    handler: THandler,
27    negotiating_in: FuturesUnordered<
28        StreamUpgrade<
29            THandler::InboundUserData,
30            <THandler::InboundUpgrade as InboundUpgradeSend>::Output,
31            <THandler::InboundUpgrade as InboundUpgradeSend>::Error,
32        >,
33    >,
34    max_negotiating_inbound_streams: usize,
35    stream_counter: ActiveStreamCounter,
36    closing: bool,
37    idle_timeout: Duration,
38    shutdown: Shutdown,
39}
40
41impl<THandler> Unpin for InboundConnection<THandler> where THandler: InboundStreamHandler {}
42
43impl<THandler> InboundConnection<THandler>
44where
45    THandler: InboundStreamHandler,
46{
47    pub fn new(
48        muxer: StreamMuxerBox,
49        handler: THandler,
50        max_negotiating_inbound_streams: usize,
51        idle_timeout: Duration,
52    ) -> Self {
53        Self {
54            muxer,
55            handler,
56            negotiating_in: FuturesUnordered::new(),
57            max_negotiating_inbound_streams,
58            stream_counter: ActiveStreamCounter::new(),
59            closing: false,
60            idle_timeout,
61            shutdown: Shutdown::None,
62        }
63    }
64
65    pub fn close(
66        self,
67    ) -> (
68        Pin<Box<dyn Stream<Item = <THandler as ConnectionHandler>::Event> + Send>>,
69        Closing<StreamMuxerBox>,
70    ) {
71        let Self {
72            muxer, mut handler, ..
73        } = self;
74
75        let stream = stream::poll_fn(move |cx| handler.poll_close(cx)).boxed();
76
77        (stream, muxer.close())
78    }
79
80    pub fn handle_action(&mut self, action: THandler::Action) {
81        self.handler.handle_action(action);
82    }
83
84    #[tracing::instrument(level = "debug", name = "Connection::poll", skip(self, cx))]
85    pub fn poll(&mut self, cx: &mut Context<'_>) -> Poll<Result<THandler::Event, ConnectionError>> {
86        let Self {
87            muxer,
88            handler,
89            negotiating_in,
90            max_negotiating_inbound_streams,
91            stream_counter,
92            closing,
93            idle_timeout,
94            shutdown,
95            ..
96        } = self;
97        loop {
98            if *closing {
99                // 如果连接已关闭，直接返回关闭事件
100                return Poll::Ready(Err(ConnectionError::Closing));
101            }
102
103            match handler.poll(cx) {
104                Poll::Pending => {}
105                // 处理器发生事件
106                Poll::Ready(ConnectionHandlerEvent::Notify(event)) => {
107                    return Poll::Ready(Ok(event));
108                }
109                // 关闭连接
110                Poll::Ready(ConnectionHandlerEvent::CloseConnection) => {
111                    *closing = true;
112                    continue;
113                }
114            }
115
116            match negotiating_in.poll_next_unpin(cx) {
117                Poll::Pending | Poll::Ready(None) => {}
118                Poll::Ready(Some((info, Ok(protocol)))) => {
119                    handler.on_fully_negotiated(info, protocol);
120                    continue;
121                }
122                Poll::Ready(Some((info, Err(StreamUpgradeError::Apply(error))))) => {
123                    handler.on_upgrade_error(info, error);
124                    continue;
125                }
126                Poll::Ready(Some((_, Err(StreamUpgradeError::Timeout)))) => {
127                    tracing::debug!("inbound stream upgrade timed out");
128                    continue;
129                }
130                Poll::Ready(Some((_, Err(StreamUpgradeError::NegotiationFailed)))) => {
131                    tracing::debug!("inbound stream upgrade negotiation failed");
132                    continue;
133                }
134                Poll::Ready(Some((_, Err(StreamUpgradeError::Io(error))))) => {
135                    tracing::debug!("inbound stream upgrade IO error: {:?}", error);
136                    continue;
137                }
138            }
139
140            if negotiating_in.is_empty() && stream_counter.no_active_streams() {
141                if let Some(new_timeout) =
142                    compute_new_shutdown(handler.connection_keep_alive(), shutdown, *idle_timeout)
143                {
144                    *shutdown = new_timeout;
145                }
146                match shutdown {
147                    Shutdown::None => {}
148                    Shutdown::Asap => return Poll::Ready(Err(ConnectionError::KeepAliveTimeout)),
149                    Shutdown::Later(delay) => match Future::poll(Pin::new(delay), cx) {
150                        Poll::Ready(_) => {
151                            return Poll::Ready(Err(ConnectionError::KeepAliveTimeout));
152                        }
153                        Poll::Pending => {}
154                    },
155                }
156            } else {
157                *shutdown = Shutdown::None;
158            }
159
160            // 处理多路复用器的事件
161            match muxer.poll_unpin(cx)? {
162                Poll::Pending => {}
163                Poll::Ready(()) => {}
164            }
165
166            if negotiating_in.len() < *max_negotiating_inbound_streams {
167                match muxer.poll_inbound_unpin(cx)? {
168                    Poll::Pending => {}
169                    Poll::Ready(substream) => {
170                        let protocol = handler.listen_protocol();
171                        negotiating_in.push(StreamUpgrade::new_inbound(
172                            substream,
173                            protocol,
174                            stream_counter.clone(),
175                        ));
176                        continue;
177                    }
178                }
179            }
180
181            return Poll::Pending;
182        }
183    }
184}
185
186impl<THandler> ConnectionController<THandler> for InboundConnection<THandler>
187where
188    THandler: InboundStreamHandler,
189{
190    fn close(
191        self,
192    ) -> (
193        Pin<Box<dyn Stream<Item = <THandler as ConnectionHandler>::Event> + Send>>,
194        Closing<StreamMuxerBox>,
195    ) {
196        self.close()
197    }
198
199    fn handle_action(&mut self, action: THandler::Action) {
200        self.handle_action(action)
201    }
202
203    fn poll(&mut self, cx: &mut Context<'_>) -> Poll<Result<THandler::Event, ConnectionError>> {
204        self.poll(cx)
205    }
206}
volans_swarm/connection/inbound.rs

volans_swarm/connection/
inbound.rs
