perf: don't allocate in UDP send & recv path

neqo-bin/src/client/http09.rs

@@ -177,15 +177,13 @@ impl TryFrom<&State> for CloseState {
 }
 
 impl super::Client for Connection {
-    fn process_output(&mut self, now: Instant) -> Output {
-        self.process_output(now)
-    }
-
-    fn process_multiple_input<'a, I>(&mut self, dgrams: I, now: Instant)
-    where
-        I: IntoIterator<Item = &'a Datagram>,
-    {
-        self.process_multiple_input(dgrams, now);
+    fn process<'a>(
+        &mut self,
+        input: Option<Datagram<&[u8]>>,
+        now: Instant,
+        write_buffer: &'a mut Vec<u8>,
+    ) -> Output<&'a [u8]> {
+        self.process_into(input, now, write_buffer)
     }
 
     fn close<S>(&mut self, now: Instant, app_error: neqo_transport::AppError, msg: S)

neqo-bin/src/client/http3.rs

@@ -124,19 +124,17 @@ impl TryFrom<Http3State> for CloseState {
 }
 
 impl super::Client for Http3Client {
-    fn is_closed(&self) -> Result<CloseState, CloseReason> {
-        self.state().try_into()
-    }
-
-    fn process_output(&mut self, now: Instant) -> Output {
-        self.process_output(now)
+    fn process<'a>(
+        &mut self,
+        input: Option<Datagram<&[u8]>>,
+        now: Instant,
+        write_buffer: &'a mut Vec<u8>,
+    ) -> Output<&'a [u8]> {
+        self.process_into(input, now, write_buffer)
     }
 
-    fn process_multiple_input<'a, I>(&mut self, dgrams: I, now: Instant)
-    where
-        I: IntoIterator<Item = &'a Datagram>,
-    {
-        self.process_multiple_input(dgrams, now);
+    fn is_closed(&self) -> Result<CloseState, CloseReason> {
+        self.state().try_into()
     }
 
     fn close<S>(&mut self, now: Instant, app_error: AppError, msg: S)

neqo-bin/src/client/mod.rs

@@ -361,10 +361,12 @@ enum CloseState {
 
 /// Network client, e.g. [`neqo_transport::Connection`] or [`neqo_http3::Http3Client`].
 trait Client {
-    fn process_output(&mut self, now: Instant) -> Output;
-    fn process_multiple_input<'a, I>(&mut self, dgrams: I, now: Instant)
-    where
-        I: IntoIterator<Item = &'a Datagram>;
+    fn process<'a>(
+        &mut self,
+        input: Option<Datagram<&[u8]>>,
+        now: Instant,
+        write_buffer: &'a mut Vec<u8>,
+    ) -> Output<&'a [u8]>;
     fn has_events(&self) -> bool;
     fn close<S>(&mut self, now: Instant, app_error: AppError, msg: S)
     where
@@ -380,13 +382,15 @@ struct Runner<'a, H: Handler> {
     handler: H,
     timeout: Option<Pin<Box<Sleep>>>,
     args: &'a Args,
+    recv_buf: Vec<u8>,
+    send_buf: Vec<u8>,
 }
 
 impl<'a, H: Handler> Runner<'a, H> {
     async fn run(mut self) -> Res<Option<ResumptionToken>> {
         loop {
             let handler_done = self.handler.handle(&mut self.client)?;
-            self.process_output().await?;
+            self.process().await?;
             if self.client.has_events() {
                 continue;
             }
@@ -407,7 +411,7 @@ impl<'a, H: Handler> Runner<'a, H> {
             }
 
             match ready(self.socket, self.timeout.as_mut()).await? {
-                Ready::Socket => self.process_multiple_input().await?,
+                Ready::Socket => {}
                 Ready::Timeout => {
                     self.timeout = None;
                 }
@@ -421,37 +425,39 @@ impl<'a, H: Handler> Runner<'a, H> {
         Ok(self.handler.take_token())
     }
 
-    async fn process_output(&mut self) -> Result<(), io::Error> {
+    async fn process(&mut self) -> Result<(), io::Error> {
+        let mut should_read = true;
         loop {
-            match self.client.process_output(Instant::now()) {
+            // TODO: Cleanup?
+            let dgram = should_read
+                .then(|| self.socket.recv(&self.local_addr, &mut self.recv_buf))
+                .transpose()?
+                .flatten();
+            should_read = dgram.is_some();
+
+            match self
+                .client
+                .process(dgram, Instant::now(), &mut self.send_buf)
+            {
                 Output::Datagram(dgram) => {
                     self.socket.writable().await?;
-                    self.socket.send(&dgram)?;
+                    self.socket.send(dgram)?;
+                    // TODO: Or should we do this right after using it?
+                    self.send_buf.clear();
+                    continue;
                 }
                 Output::Callback(new_timeout) => {
                     qdebug!("Setting timeout of {:?}", new_timeout);
                     self.timeout = Some(Box::pin(tokio::time::sleep(new_timeout)));
-                    break;
                 }
                 Output::None => {
                     qdebug!("Output::None");
-                    break;
                 }
             }
-        }
 
-        Ok(())
-    }
-
-    async fn process_multiple_input(&mut self) -> Res<()> {
-        loop {
-            let dgrams = self.socket.recv(&self.local_addr)?;
-            if dgrams.is_empty() {
+            if !should_read {
                 break;
             }
-            self.client
-                .process_multiple_input(dgrams.iter(), Instant::now());
-            self.process_output().await?;
         }
 
         Ok(())
@@ -568,6 +574,8 @@ pub async fn client(mut args: Args) -> Res<()> {
                     local_addr: real_local,
                     socket: &mut socket,
                     timeout: None,
+                    recv_buf: Vec::with_capacity(neqo_udp::RECV_BUF_SIZE),
+                    send_buf: Vec::new(),
                 }
                 .run()
                 .await?
@@ -584,6 +592,8 @@ pub async fn client(mut args: Args) -> Res<()> {
                     local_addr: real_local,
                     socket: &mut socket,
                     timeout: None,
+                    recv_buf: Vec::with_capacity(neqo_udp::RECV_BUF_SIZE),
+                    send_buf: Vec::new(),
                 }
                 .run()
                 .await?

neqo-bin/src/server/http09.rs

@@ -185,8 +185,13 @@ impl HttpServer {
 }
 
 impl super::HttpServer for HttpServer {
-    fn process(&mut self, dgram: Option<&Datagram>, now: Instant) -> Output {
-        self.server.process(dgram, now)
+    fn process_into<'a>(
+        &mut self,
+        dgram: Option<Datagram<&[u8]>>,
+        now: Instant,
+        write_buffer: &'a mut Vec<u8>,
+    ) -> Output<&'a [u8]> {
+        self.server.process_2(dgram, now, write_buffer)
     }
 
     fn process_events(&mut self, now: Instant) {

neqo-bin/src/server/http3.rs

@@ -79,8 +79,13 @@ impl Display for HttpServer {
 }
 
 impl super::HttpServer for HttpServer {
-    fn process(&mut self, dgram: Option<&Datagram>, now: Instant) -> neqo_http3::Output {
-        self.server.process(dgram, now)
+    fn process_into<'a>(
+        &mut self,
+        dgram: Option<Datagram<&[u8]>>,
+        now: Instant,
+        write_buffer: &'a mut Vec<u8>,
+    ) -> neqo_http3::Output<&'a [u8]> {
+        self.server.process_into(dgram, now, write_buffer)
     }
 
     fn process_events(&mut self, _now: Instant) {

neqo-bin/src/server/mod.rs

@@ -194,7 +194,12 @@ fn qns_read_response(filename: &str) -> Result<Vec<u8>, io::Error> {
 
 #[allow(clippy::module_name_repetitions)]
 pub trait HttpServer: Display {
-    fn process(&mut self, dgram: Option<&Datagram>, now: Instant) -> Output;
+    fn process_into<'a>(
+        &mut self,
+        dgram: Option<Datagram<&[u8]>>,
+        now: Instant,
+        write_buffer: &'a mut Vec<u8>,
+    ) -> Output<&'a [u8]>;
     fn process_events(&mut self, now: Instant);
     fn has_events(&self) -> bool;
 }
@@ -205,6 +210,8 @@ pub struct ServerRunner {
     server: Box<dyn HttpServer>,
     timeout: Option<Pin<Box<Sleep>>>,
     sockets: Vec<(SocketAddr, crate::udp::Socket)>,
+    recv_buf: Vec<u8>,
+    send_buf: Vec<u8>,
 }
 
 impl ServerRunner {
@@ -219,41 +226,62 @@ impl ServerRunner {
             server,
             timeout: None,
             sockets,
+            recv_buf: Vec::with_capacity(neqo_udp::RECV_BUF_SIZE),
+            send_buf: vec![],
         }
     }
 
-    /// Tries to find a socket, but then just falls back to sending from the first.
-    fn find_socket(&mut self, addr: SocketAddr) -> &mut crate::udp::Socket {
-        let ((_host, first_socket), rest) = self.sockets.split_first_mut().unwrap();
-        rest.iter_mut()
-            .map(|(_host, socket)| socket)
-            .find(|socket| {
-                socket
-                    .local_addr()
-                    .ok()
-                    .map_or(false, |socket_addr| socket_addr == addr)
-            })
-            .unwrap_or(first_socket)
-    }
-
-    async fn process(&mut self, mut dgram: Option<&Datagram>) -> Result<(), io::Error> {
+    // TODO: Could as well call it UDP IO now, given that it does both in and output.
+    async fn process(&mut self, mut inx: Option<usize>) -> Result<(), io::Error> {
         loop {
-            match self.server.process(dgram.take(), (self.now)()) {
+            let mut dgram = if let Some(i) = inx {
+                let (host, socket) = self.sockets.get_mut(i).unwrap();
+                let dgram = socket.recv(host, &mut self.recv_buf)?;
+                if dgram.is_none() {
+                    // TODO: Better way to not try reading again?
+                    inx.take();
+                }
+                dgram
+            } else {
+                None
+            };
+
+            match self
+                .server
+                .process_into(dgram.take(), (self.now)(), &mut self.send_buf)
+            {
                 Output::Datagram(dgram) => {
-                    let socket = self.find_socket(dgram.source());
+                    let socket = {
+                        let addr = dgram.source();
+                        let ((_host, first_socket), rest) = self.sockets.split_first_mut().unwrap();
+                        rest.iter_mut()
+                            .map(|(_host, socket)| socket)
+                            .find(|socket| {
+                                socket
+                                    .local_addr()
+                                    .ok()
+                                    .map_or(false, |socket_addr| socket_addr == addr)
+                            })
+                            .unwrap_or(first_socket)
+                    };
                     socket.writable().await?;
-                    socket.send(&dgram)?;
+                    socket.send(dgram)?;
+                    // TODO: Or should we do this right after using it?
+                    self.send_buf.clear();
+                    continue;
                 }
                 Output::Callback(new_timeout) => {
                     qdebug!("Setting timeout of {:?}", new_timeout);
                     self.timeout = Some(Box::pin(tokio::time::sleep(new_timeout)));
-                    break;
-                }
-                Output::None => {
-                    break;
                 }
+                Output::None => {}
+            }
+
+            if inx.is_none() {
+                break;
             }
         }
+
         Ok(())
     }
 
@@ -287,16 +315,10 @@ impl ServerRunner {
             }
 
             match self.ready().await? {
-                Ready::Socket(inx) => loop {
-                    let (host, socket) = self.sockets.get_mut(inx).unwrap();
-                    let dgrams = socket.recv(host)?;
-                    if dgrams.is_empty() {
-                        break;
-                    }
-                    for dgram in dgrams {
-                        self.process(Some(&dgram)).await?;
-                    }
-                },
+                Ready::Socket(inx) => {
+                    // TODO: Passing the index here to only borrow &mut self in process. Better way?
+                    self.process(Some(inx)).await?
+                }
                 Ready::Timeout => {
                     self.timeout = None;
                     self.process(None).await?;

neqo-bin/src/udp.rs

@@ -47,22 +47,27 @@ impl Socket {
     }
 
     /// Send a [`Datagram`] on the given [`Socket`].
-    pub fn send(&self, d: &Datagram) -> io::Result<()> {
+    pub fn send(&self, d: Datagram<&[u8]>) -> io::Result<()> {
         self.inner.try_io(tokio::io::Interest::WRITABLE, || {
             neqo_udp::send_inner(&self.state, (&self.inner).into(), d)
         })
     }
 
     /// Receive a batch of [`Datagram`]s on the given [`Socket`], each set with
     /// the provided local address.
-    pub fn recv(&self, local_address: &SocketAddr) -> Result<Vec<Datagram>, io::Error> {
+    pub fn recv<'a>(
+        &self,
+        local_address: &SocketAddr,
+        recv_buf: &'a mut Vec<u8>,
+    ) -> Result<Option<Datagram<&'a [u8]>>, io::Error> {
         self.inner
             .try_io(tokio::io::Interest::READABLE, || {
-                neqo_udp::recv_inner(local_address, &self.state, &self.inner)
+                neqo_udp::recv_inner_2(local_address, &self.state, &self.inner, recv_buf)
             })
+            .map(Some)
             .or_else(|e| {
                 if e.kind() == io::ErrorKind::WouldBlock {
-                    Ok(vec![])
+                    Ok(None)
                 } else {
                     Err(e)
                 }