[pull] master from eycorsican:master

leaf/src/option/mod.rs

@@ -177,6 +177,22 @@ lazy_static! {
         get_env_var_or("QUIC_ACCEPT_CHANNEL_SIZE", 1024)
     };
 
+    pub static ref INCOMING_ACCEPT_CONCURRENCY: usize = {
+        get_env_var_or("INCOMING_ACCEPT_CONCURRENCY", 256)
+    };
+
+    pub static ref QUIC_MAX_CONCURRENT_BIDI_STREAMS: u32 = {
+        get_env_var_or("QUIC_MAX_CONCURRENT_BIDI_STREAMS", 256)
+    };
+
+    pub static ref QUIC_MAX_IDLE_TIMEOUT_MS: u32 = {
+        get_env_var_or("QUIC_MAX_IDLE_TIMEOUT_MS", 120_000)
+    };
+
+    pub static ref QUIC_ACCEPT_QUEUE_TIMEOUT: u64 = {
+        get_env_var_or("QUIC_ACCEPT_QUEUE_TIMEOUT", 5)
+    };
+
     pub static ref AMUX_ACCEPT_CHANNEL_SIZE: usize = {
         get_env_var_or("AMUX_ACCEPT_CHANNEL_SIZE", 1024)
     };

leaf/src/proxy/chain/inbound/mod.rs

@@ -1,12 +1,12 @@
-use std::{io, pin::Pin};
+use std::{io, pin::Pin, time::Duration};
 
-use futures::stream::Stream;
+use futures::stream::{FuturesUnordered, Stream};
 use futures::{
     future::BoxFuture,
-    ready,
     task::{Context, Poll},
 };
-use tracing::{debug, warn};
+use tokio::time::timeout;
+use tracing::debug;
 
 use crate::proxy::*;
 
@@ -17,13 +17,14 @@ pub use datagram::Handler as DatagramHandler;
 pub use stream::Handler as StreamHandler;
 
 enum State {
-    WaitingIncoming,
-    Pending(usize, BoxFuture<'static, io::Result<AnyInboundTransport>>),
+    Running,
+    Closed,
 }
 
 pub struct Incoming {
     incoming: AnyIncomingTransport,
     actors: Vec<AnyInboundHandler>,
+    pending: FuturesUnordered<BoxFuture<'static, io::Result<AnyBaseInboundTransport>>>,
     state: State,
 }
 
@@ -32,81 +33,195 @@ impl Incoming {
         Incoming {
             incoming,
             actors,
-            state: State::WaitingIncoming,
+            pending: FuturesUnordered::new(),
+            state: State::Running,
         }
     }
 }
 
+async fn run_stream_actors(
+    mut stream: AnyStream,
+    mut sess: Session,
+    actors: Vec<AnyInboundHandler>,
+    handshake_timeout: Duration,
+) -> io::Result<AnyBaseInboundTransport> {
+    for actor in actors {
+        let transport = timeout(handshake_timeout, actor.stream()?.handle(sess.clone(), stream))
+            .await
+            .map_err(|_| io::Error::new(io::ErrorKind::TimedOut, "incoming stream handle timed out"))??;
+        match transport {
+            InboundTransport::Stream(new_stream, new_sess) => {
+                stream = new_stream;
+                sess = new_sess;
+            }
+            InboundTransport::Datagram(socket, sess) => {
+                return Ok(AnyBaseInboundTransport::Datagram(socket, sess));
+            }
+            _ => {
+                return Err(io::Error::other("invalid chain inbound incoming stream transport"));
+            }
+        }
+    }
+    Ok(AnyBaseInboundTransport::Stream(stream, sess))
+}
+
+async fn run_datagram_actors(
+    mut socket: AnyInboundDatagram,
+    mut sess: Option<Session>,
+    actors: Vec<AnyInboundHandler>,
+    handshake_timeout: Duration,
+) -> io::Result<AnyBaseInboundTransport> {
+    for actor in actors {
+        let transport = timeout(handshake_timeout, actor.datagram()?.handle(socket))
+            .await
+            .map_err(|_| io::Error::new(io::ErrorKind::TimedOut, "incoming datagram handle timed out"))??;
+        match transport {
+            InboundTransport::Stream(stream, sess) => {
+                return Ok(AnyBaseInboundTransport::Stream(stream, sess));
+            }
+            InboundTransport::Datagram(new_socket, new_sess) => {
+                socket = new_socket;
+                sess = new_sess;
+            }
+            _ => {
+                return Err(io::Error::other(
+                    "invalid chain inbound incoming datagram transport",
+                ));
+            }
+        }
+    }
+    Ok(AnyBaseInboundTransport::Datagram(socket, sess))
+}
+
+async fn run_actors(
+    transport: AnyBaseInboundTransport,
+    actors: Vec<AnyInboundHandler>,
+    handshake_timeout: Duration,
+) -> io::Result<AnyBaseInboundTransport> {
+    match transport {
+        AnyBaseInboundTransport::Stream(stream, sess) => {
+            run_stream_actors(stream, sess, actors, handshake_timeout).await
+        }
+        AnyBaseInboundTransport::Datagram(socket, sess) => {
+            run_datagram_actors(socket, sess, actors, handshake_timeout).await
+        }
+        AnyBaseInboundTransport::Empty => Err(io::Error::other("empty chain inbound transport")),
+    }
+}
+
 impl Stream for Incoming {
     // TODO io::Result<(...)>
     type Item = AnyBaseInboundTransport;
 
     fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let max_concurrency = (*crate::option::INCOMING_ACCEPT_CONCURRENCY).max(1);
+        let handshake_timeout = Duration::from_secs(*crate::option::INBOUND_ACCEPT_TIMEOUT);
+
         loop {
-            match self.state {
-                // Polling for the underlying transport.
-                State::WaitingIncoming => {
-                    let transport = ready!(Stream::poll_next(Pin::new(&mut self.incoming), cx));
-                    match transport {
-                        // Only reliable transports are eligible to handle TCP requests.
-                        Some(AnyBaseInboundTransport::Stream(stream, sess)) => {
-                            assert!(!self.actors.is_empty()); // FIXME
-                            let sess = sess.clone();
-                            let a = self.actors[0].clone();
-                            // Create the task for handling the first actor.
-                            let t = Box::pin(async move { a.stream()?.handle(sess, stream).await });
-                            self.state = State::Pending(0, t);
-                        }
-                        Some(_) => {
-                            warn!("unexpected non-stream chain inbound incoming transport");
-                            continue;
-                        }
-                        None => {
-                            return Poll::Ready(None);
-                        }
+            while matches!(self.state, State::Running) && self.pending.len() < max_concurrency {
+                match Stream::poll_next(Pin::new(&mut self.incoming), cx) {
+                    Poll::Ready(Some(transport)) => {
+                        let actors = self.actors.clone();
+                        self.pending
+                            .push(Box::pin(run_actors(transport, actors, handshake_timeout)));
+                    }
+                    Poll::Ready(None) => {
+                        self.state = State::Closed;
+                        break;
                     }
+                    Poll::Pending => break,
                 }
-                // Polling for output transport from actors[idx].
-                State::Pending(idx, ref mut task) => {
-                    match ready!(task.as_mut().poll(cx)) {
-                        Ok(InboundTransport::Stream(new_stream, new_sess)) => {
-                            // If this is the output transport from the task of the last actor,
-                            // return it.
-                            if idx + 1 >= self.actors.len() {
-                                self.state = State::WaitingIncoming;
-                                return Poll::Ready(Some(AnyBaseInboundTransport::Stream(
-                                    new_stream, new_sess,
-                                )));
-                            }
-                            // Otherwise proceed with a new task for the next actor.
-                            let new_sess = new_sess.clone();
-                            let a = self.actors[idx + 1].clone();
-                            let t =
-                                Box::pin(
-                                    async move { a.stream()?.handle(new_sess, new_stream).await },
-                                );
-                            self.state = State::Pending(idx + 1, t);
-                        }
-                        Ok(InboundTransport::Datagram(socket, sess)) => {
-                            // FIXME Assume the last one, but not necessary the last one?
-                            self.state = State::WaitingIncoming;
-                            return Poll::Ready(Some(AnyBaseInboundTransport::Datagram(
-                                socket, sess,
-                            )));
-                        }
-                        Err(e) => {
-                            debug!("chain inbound incoming error: {}", e);
-                            self.state = State::WaitingIncoming;
-                            continue;
-                        }
-                        _ => {
-                            warn!("unexpected chain inbound incoming transport");
-                            self.state = State::WaitingIncoming;
-                            return Poll::Ready(None);
-                        }
+            }
+
+            match Stream::poll_next(Pin::new(&mut self.pending), cx) {
+                Poll::Ready(Some(Ok(transport))) => {
+                    return Poll::Ready(Some(transport));
+                }
+                Poll::Ready(Some(Err(e))) => {
+                    debug!("chain inbound incoming error: {}", e);
+                    continue;
+                }
+                Poll::Ready(None) if matches!(self.state, State::Closed) => {
+                    return Poll::Ready(None);
+                }
+                Poll::Ready(None) | Poll::Pending => {
+                    if matches!(self.state, State::Closed) && self.pending.is_empty() {
+                        return Poll::Ready(None);
                     }
+                    return Poll::Pending;
                 }
             }
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+    use std::time::Duration;
+
+    use async_trait::async_trait;
+    use futures::StreamExt;
+    use tokio::io::duplex;
+    use tokio::time::timeout;
+
+    use crate::proxy::inbound::Handler as InboundHandlerImpl;
+    use crate::session::StreamId;
+
+    use super::*;
+
+    struct DelayByStreamIdHandler;
+
+    #[async_trait]
+    impl InboundStreamHandler for DelayByStreamIdHandler {
+        async fn handle<'a>(
+            &'a self,
+            sess: Session,
+            stream: AnyStream,
+        ) -> io::Result<AnyInboundTransport> {
+            let delay = match sess.stream_id {
+                Some(StreamId::U64(1)) => Duration::from_millis(200),
+                Some(StreamId::U64(2)) => Duration::from_millis(10),
+                _ => Duration::from_millis(0),
+            };
+            tokio::time::sleep(delay).await;
+            Ok(InboundTransport::Stream(stream, sess))
+        }
+    }
+
+    #[tokio::test]
+    async fn incoming_does_not_block_fast_streams_behind_slow_ones() {
+        let actor: AnyInboundHandler = Arc::new(InboundHandlerImpl::new(
+            "mock".to_string(),
+            Some(Arc::new(DelayByStreamIdHandler)),
+            None,
+        ));
+        let actors = vec![actor];
+
+        let (stream1, _) = duplex(64);
+        let (stream2, _) = duplex(64);
+
+        let mut sess1 = Session::default();
+        sess1.stream_id = Some(StreamId::U64(1));
+        let mut sess2 = Session::default();
+        sess2.stream_id = Some(StreamId::U64(2));
+
+        let incoming = futures::stream::iter(vec![
+            AnyBaseInboundTransport::Stream(Box::new(stream1), sess1),
+            AnyBaseInboundTransport::Stream(Box::new(stream2), sess2),
+        ]);
+        let mut incoming = Incoming::new(Box::new(incoming), actors);
+
+        let first = timeout(Duration::from_millis(100), incoming.next())
+            .await
+            .expect("fast stream should not be blocked")
+            .expect("stream should be yielded");
+
+        match first {
+            AnyBaseInboundTransport::Stream(_, sess) => {
+                assert_eq!(sess.stream_id, Some(StreamId::U64(2)));
+            }
+            _ => panic!("expected stream transport"),
+        }
+    }
+}

leaf/src/proxy/quic/inbound/datagram.rs

@@ -11,6 +11,7 @@ use quinn::{RecvStream, SendStream};
 use rustls::pki_types::{CertificateDer, PrivateKeyDer};
 use rustls_pemfile::{certs, pkcs8_private_keys, rsa_private_keys};
 use tokio::sync::mpsc::{channel, Receiver, Sender};
+use tokio::time::{timeout, Duration};
 use tracing::{debug, trace, warn};
 
 use crate::{proxy::*, session::Session, session::StreamId};
@@ -119,10 +120,12 @@ impl Handler {
             quinn::crypto::rustls::QuicServerConfig::try_from(crypto).unwrap(),
         ));
         let mut transport_config = quinn::TransportConfig::default();
-        transport_config.max_concurrent_bidi_streams(quinn::VarInt::from_u32(64));
-        transport_config.max_idle_timeout(Some(quinn::IdleTimeout::from(quinn::VarInt::from_u32(
-            300_000,
-        ))));
+        transport_config.max_concurrent_bidi_streams(quinn::VarInt::from_u32(
+            *crate::option::QUIC_MAX_CONCURRENT_BIDI_STREAMS,
+        ));
+        transport_config.max_idle_timeout(Some(quinn::IdleTimeout::from(
+            quinn::VarInt::from_u32(*crate::option::QUIC_MAX_IDLE_TIMEOUT_MS),
+        )));
         transport_config
             .congestion_controller_factory(Arc::new(quinn::congestion::BbrConfig::default()));
         server_config.transport_config(Arc::new(transport_config));
@@ -133,20 +136,30 @@ impl Handler {
 
 async fn handle_conn(
     stream_tx: Sender<(SocketAddr, (SendStream, RecvStream))>,
-    remote_addr: &SocketAddr,
+    remote_addr: SocketAddr,
     conn: quinn::Connecting,
 ) -> Result<()> {
     let (conn, _) = conn
         .into_0rtt()
         .map_err(|_| anyhow!("convert 0rtt failed"))?;
+    let send_timeout = Duration::from_secs(*crate::option::QUIC_ACCEPT_QUEUE_TIMEOUT);
     trace!("quic handling connection from {}", remote_addr);
     loop {
         let s = conn.accept_bi().await?;
         trace!("quic accepted stream from {}", remote_addr);
         if stream_tx.capacity() == 0 {
             warn!("quic accept channel full");
         }
-        let _ = stream_tx.send((*remote_addr, s)).await;
+        match timeout(send_timeout, stream_tx.send((remote_addr, s))).await {
+            Ok(Ok(())) => {}
+            Ok(Err(_)) => return Ok(()),
+            Err(_) => {
+                return Err(anyhow!(
+                    "quic accept queue remained full for {:?}, dropping connection",
+                    send_timeout
+                ));
+            }
+        }
     }
 }
 
@@ -169,8 +182,7 @@ impl InboundDatagramHandler for Handler {
                     let remote_addr = incoming.remote_address();
                     match incoming.accept() {
                         Ok(connecting) => {
-                            if let Err(e) = handle_conn(stream_tx_c, &remote_addr, connecting).await
-                            {
+                            if let Err(e) = handle_conn(stream_tx_c, remote_addr, connecting).await {
                                 debug!(
                                     "handle quic connection from {} failed: {}",
                                     &remote_addr, e