opt(s2n-quic-dc): skip epoll registration in happy path (#2754)

Currently, dcQUIC streams over TCP will be accepted, be registered with
epoll, attempt reading (usually fails), in <1ms the first data packet
arrives and we succeed reading, deregister the socket, and then hand off
the stream to the application for further reading.

We'd like to avoid the epoll registration as it uses extra CPU (even if
latency impact is minimal) so this patch uses the Linux-only
TCP_DEFER_ACCEPT to only accept sockets with data already available.
That's combined with lazy registration of sockets with Tokio's epoll by
only doing so if we get WouldBlock after attempting a read or write.

The net effect is a 8.8% (relative) drop in overall CPU usage in one of
our internal benchmarks which exercises short streams over loopback,
bringing CPU usage in the acceptor from 23% of the workload to 18%.

Author: Mark-Simulacrum

dc/s2n-quic-dc/src/stream/environment/tokio/tcp.rs

@@ -6,6 +6,7 @@ use crate::{
     stream::{
         environment::{tokio::Environment, Peer, SetupResult, SocketSet},
         recv::shared::RecvBuffer,
+        server::tokio::tcp::LazyBoundStream,
         TransportFeatures,
     },
 };
@@ -51,7 +52,7 @@ where
 
 /// A socket that should be reregistered with the application runtime
 pub struct Reregistered {
-    pub socket: TcpStream,
+    pub socket: LazyBoundStream,
     pub peer_addr: SocketAddress,
     pub local_port: u16,
     pub recv_buffer: RecvBuffer,

dc/s2n-quic-dc/src/stream/server/tokio.rs

@@ -468,7 +468,7 @@ impl<H: Handshake + Clone, S: event::Subscriber + Clone> Start<'_, H, S> {
             self.server.local_addr = socket.local_addr()?;
         }
 
-        let socket = tokio::net::TcpListener::from_std(socket)?;
+        let socket = tokio::io::unix::AsyncFd::new(socket)?;
         let id = self.id();
 
         let acceptor = tcp::Acceptor::new(
@@ -480,7 +480,7 @@ impl<H: Handshake + Clone, S: event::Subscriber + Clone> Start<'_, H, S> {
             self.backlog,
             self.accept_flavor,
             self.linger,
-        )
+        )?
         .run();
 
         if self.span.is_disabled() {

dc/s2n-quic-dc/src/stream/server/tokio/tcp.rs

@@ -11,20 +11,23 @@ use crate::{
 };
 use core::{future::poll_fn, task::Poll};
 use s2n_quic_core::{inet::SocketAddress, time::Clock};
-use std::time::Duration;
-use tokio::net::TcpListener;
+use std::{net::TcpListener, os::fd::AsRawFd, time::Duration};
+use tokio::io::unix::AsyncFd;
 use tracing::debug;
 
 mod fresh;
+mod lazy;
 mod manager;
 mod worker;
 
+pub(crate) use lazy::LazyBoundStream;
+
 pub struct Acceptor<Sub>
 where
     Sub: Subscriber + Clone,
 {
     sender: accept::Sender<Sub>,
-    socket: TcpListener,
+    socket: AsyncFd<TcpListener>,
     env: Environment<Sub>,
     secrets: secret::Map,
     backlog: usize,
@@ -39,14 +42,14 @@ where
     #[inline]
     pub fn new(
         id: usize,
-        socket: TcpListener,
+        socket: AsyncFd<TcpListener>,
         sender: &accept::Sender<Sub>,
         env: &Environment<Sub>,
         secrets: &secret::Map,
         backlog: usize,
         accept_flavor: accept::Flavor,
         linger: Option<Duration>,
-    ) -> Self {
+    ) -> std::io::Result<Self> {
         let acceptor = Self {
             sender: sender.clone(),
             socket,
@@ -57,7 +60,27 @@ where
             linger,
         };
 
-        if let Ok(addr) = acceptor.socket.local_addr() {
+        #[cfg(target_os = "linux")]
+        {
+            let res = unsafe {
+                libc::setsockopt(
+                    acceptor.socket.get_ref().as_raw_fd(),
+                    libc::SOL_TCP,
+                    libc::TCP_DEFER_ACCEPT,
+                    // This is how many seconds elapse before the kernel will accept a stream
+                    // without any data and return it to userspace. Any number of seconds is
+                    // arguably too many in our domain (we'd expect data in milliseconds) but in
+                    // practice this value shouldn't matter much.
+                    &1u32 as *const _ as *const _,
+                    std::mem::size_of::<u32>() as libc::socklen_t,
+                )
+            };
+            if res != 0 {
+                return Err(std::io::Error::last_os_error());
+            }
+        }
+
+        if let Ok(addr) = acceptor.socket.get_ref().local_addr() {
             let local_address: SocketAddress = addr.into();
             acceptor.env.endpoint_publisher().on_acceptor_tcp_started(
                 event::builder::AcceptorTcpStarted {
@@ -68,7 +91,7 @@ where
             );
         }
 
-        acceptor
+        Ok(acceptor)
     }
 
     pub async fn run(mut self) {
@@ -103,7 +126,7 @@ where
 
                 workers.insert(
                     remote_address,
-                    socket,
+                    LazyBoundStream::Std(socket),
                     self.linger,
                     &mut context,
                     subscriber_ctx,

dc/s2n-quic-dc/src/stream/server/tokio/tcp/fresh.rs

@@ -4,7 +4,7 @@
 use crate::event::{self, EndpointPublisher};
 use core::task::{Context, Poll};
 use s2n_quic_core::inet::SocketAddress;
-use std::{collections::VecDeque, io};
+use std::{collections::VecDeque, io, task::ready};
 
 /// Converts the kernel's TCP FIFO accept queue to LIFO
 ///
@@ -107,13 +107,20 @@ pub trait Listener {
     fn poll_accept(&mut self, cx: &mut Context) -> Poll<io::Result<(Self::Stream, SocketAddress)>>;
 }
 
-impl Listener for tokio::net::TcpListener {
-    type Stream = tokio::net::TcpStream;
+impl Listener for tokio::io::unix::AsyncFd<std::net::TcpListener> {
+    type Stream = std::net::TcpStream;
 
     #[inline]
     fn poll_accept(&mut self, cx: &mut Context) -> Poll<io::Result<(Self::Stream, SocketAddress)>> {
-        (*self)
-            .poll_accept(cx)
-            .map_ok(|(socket, remote_address)| (socket, remote_address.into()))
+        loop {
+            let mut guard = ready!(self.poll_read_ready(cx))?;
+            let (socket, remote_addr) = match guard.try_io(|listener| listener.get_ref().accept()) {
+                Ok(v) => v?,
+                // arm the waker via poll_read_ready if WouldBlock returned.
+                Err(_) => continue,
+            };
+            socket.set_nonblocking(true)?;
+            return Poll::Ready(Ok((socket, remote_addr.into())));
+        }
     }
 }

dc/s2n-quic-dc/src/stream/server/tokio/tcp/lazy.rs

@@ -0,0 +1,154 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+use crate::{
+    msg,
+    stream::socket::{fd::tcp, Flags, Socket},
+};
+use s2n_quic_core::inet::ExplicitCongestionNotification;
+use std::{
+    io::{self, ErrorKind, Write},
+    net::TcpStream as StdTcpStream,
+    os::fd::AsRawFd,
+    pin::Pin,
+    task::Poll,
+    time::Duration,
+};
+use tokio::{io::AsyncWrite as _, net::TcpStream as TokioTcpStream};
+
+pub enum LazyBoundStream {
+    Tokio(TokioTcpStream),
+    Std(StdTcpStream),
+    // needed for moving between the previous two while only having &mut access.
+    TempEmpty,
+}
+
+impl LazyBoundStream {
+    pub fn set_nodelay(&self, nodelay: bool) -> io::Result<()> {
+        match self {
+            LazyBoundStream::Tokio(s) => s.set_nodelay(nodelay),
+            LazyBoundStream::Std(s) => s.set_nodelay(nodelay),
+            LazyBoundStream::TempEmpty => unreachable!(),
+        }
+    }
+
+    pub fn set_linger(&self, linger: Option<Duration>) -> io::Result<()> {
+        match self {
+            LazyBoundStream::Tokio(s) => s.set_linger(linger),
+            LazyBoundStream::Std(s) => {
+                // Once it stabilizes we can switch to the std function
+                // https://github.com/rust-lang/rust/issues/88494
+                let res = unsafe {
+                    libc::setsockopt(
+                        s.as_raw_fd(),
+                        libc::SOL_SOCKET,
+                        libc::SO_LINGER,
+                        &libc::linger {
+                            l_onoff: linger.is_some() as libc::c_int,
+                            l_linger: linger.unwrap_or_default().as_secs() as libc::c_int,
+                        } as *const _ as *const _,
+                        std::mem::size_of::<libc::linger>() as libc::socklen_t,
+                    )
+                };
+                if res != 0 {
+                    return Err(std::io::Error::last_os_error());
+                }
+
+                Ok(())
+            }
+            LazyBoundStream::TempEmpty => unreachable!(),
+        }
+    }
+
+    pub fn into_std(self) -> io::Result<StdTcpStream> {
+        match self {
+            LazyBoundStream::Tokio(s) => s.into_std(),
+            LazyBoundStream::Std(s) => Ok(s),
+            LazyBoundStream::TempEmpty => unreachable!(),
+        }
+    }
+
+    pub fn poll_write(
+        &mut self,
+        cx: &mut std::task::Context,
+        buffer: &[u8],
+    ) -> std::task::Poll<io::Result<usize>> {
+        loop {
+            match self {
+                LazyBoundStream::Tokio(stream) => return Pin::new(stream).poll_write(cx, buffer),
+                LazyBoundStream::Std(stream) => match stream.write(buffer) {
+                    Ok(v) => return Poll::Ready(Ok(v)),
+                    Err(e) => {
+                        if e.kind() == ErrorKind::WouldBlock {
+                            let LazyBoundStream::Std(stream) =
+                                std::mem::replace(self, LazyBoundStream::TempEmpty)
+                            else {
+                                unreachable!();
+                            };
+                            *self = LazyBoundStream::Tokio(TokioTcpStream::from_std(stream)?);
+                        } else {
+                            return Poll::Ready(Err(e));
+                        }
+                    }
+                },
+                LazyBoundStream::TempEmpty => unreachable!(),
+            }
+        }
+    }
+
+    pub fn poll_recv_buffer(
+        &mut self,
+        cx: &mut std::task::Context,
+        buffer: &mut msg::recv::Message,
+    ) -> std::task::Poll<io::Result<usize>> {
+        loop {
+            match self {
+                LazyBoundStream::Tokio(stream) => {
+                    return Pin::new(stream).poll_recv_buffer(cx, buffer)
+                }
+                LazyBoundStream::Std(stream) => {
+                    let res = buffer.recv_with(|_addr, cmsg, buffer| {
+                        loop {
+                            let flags = Flags::default();
+                            let res = tcp::recv(&*stream, buffer, flags);
+
+                            match res {
+                                Ok(len) => {
+                                    // we don't need ECN markings from TCP since it handles that logic for us
+                                    cmsg.set_ecn(ExplicitCongestionNotification::NotEct);
+
+                                    // TCP doesn't have segments so just set it to 0 (which will indicate a single
+                                    // stream of bytes)
+                                    cmsg.set_segment_len(0);
+
+                                    return Ok(len);
+                                }
+                                Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {
+                                    // try the operation again if we were interrupted
+                                    continue;
+                                }
+                                Err(err) => return Err(err),
+                            }
+                        }
+                    });
+                    match res {
+                        Ok(v) => return Poll::Ready(Ok(v)),
+                        Err(e) => {
+                            if e.kind() == ErrorKind::WouldBlock {
+                                let LazyBoundStream::Std(stream) =
+                                    std::mem::replace(self, LazyBoundStream::TempEmpty)
+                                else {
+                                    unreachable!();
+                                };
+                                *self = LazyBoundStream::Tokio(TokioTcpStream::from_std(stream)?);
+                            } else {
+                                return Poll::Ready(Err(e));
+                            }
+                        }
+                    }
+                }
+                LazyBoundStream::TempEmpty => unreachable!(),
+            }
+        }
+    }
+}

dc/s2n-quic-dc/src/stream/server/tokio/tcp/manager.rs

@@ -370,7 +370,7 @@ where
     }
 }
 
-pub trait Worker {
+pub(crate) trait Worker {
     type Context;
     type ConnectionContext;
     type Stream;