Progress?

Author: larseggert

neqo-bin/src/client/mod.rs

@@ -34,7 +34,10 @@ use neqo_udp::RecvBuf;
 use tokio::time::Sleep;
 use url::{Host, Origin, Url};
 
-use crate::{udp::process_send, SharedArgs};
+use crate::{
+    udp::{process_send, SocketLookup},
+    SharedArgs,
+};
 
 mod http09;
 mod http3;
@@ -406,6 +409,11 @@ struct Runner<'a, H: Handler> {
     args: &'a Args,
     recv_buf: RecvBuf,
 }
+impl<'a> SocketLookup for &'a mut crate::udp::Socket {
+    fn get(&mut self, _addr: SocketAddr) -> &mut crate::udp::Socket {
+        self
+    }
+}
 
 impl<'a, H: Handler> Runner<'a, H> {
     fn new(
@@ -462,12 +470,15 @@ impl<'a, H: Handler> Runner<'a, H> {
     }
 
     async fn process_output(&mut self) -> Result<(), io::Error> {
-        process_send(
-            |_| self.socket,
+        let mut timeout = None;
+        let res = process_send(
+            self.socket,
             || self.client.process_output(Instant::now()),
-            &mut self.timeout,
+            &mut timeout,
         )
-        .await
+        .await;
+        self.timeout = timeout;
+        res
     }
 
     async fn process_multiple_input(&mut self) -> Res<()> {

neqo-bin/src/server/mod.rs

@@ -39,7 +39,10 @@ use neqo_transport::{Output, RandomConnectionIdGenerator, Version};
 use neqo_udp::RecvBuf;
 use tokio::time::Sleep;
 
-use crate::{udp::process_send, SharedArgs};
+use crate::{
+    udp::{process_send, SocketLookup},
+    SharedArgs,
+};
 
 const ANTI_REPLAY_WINDOW: Duration = Duration::from_secs(10);
 
@@ -211,6 +214,16 @@ pub struct ServerRunner {
     recv_buf: RecvBuf,
 }
 
+impl SocketLookup for &mut [(SocketAddr, crate::udp::Socket)] {
+    fn get(&mut self, addr: SocketAddr) -> &mut crate::udp::Socket {
+        let ((_host, first_socket), rest) = self.split_first_mut().unwrap();
+        rest.iter_mut()
+            .map(|(_host, socket)| socket)
+            .find(|socket| socket.local_addr().is_ok_and(|a| a == addr))
+            .unwrap_or(first_socket)
+    }
+}
+
 impl ServerRunner {
     #[must_use]
     pub fn new(
@@ -250,7 +263,7 @@ impl ServerRunner {
         mut input_dgram: Option<Datagram<&mut [u8]>>,
     ) -> Result<(), io::Error> {
         process_send(
-            |addr| Self::find_socket(sockets, addr),
+            sockets,
             || server.process(input_dgram.take(), now()),
             timeout,
         )

neqo-bin/src/udp.rs

@@ -85,6 +85,10 @@ impl Socket {
     }
 }
 
+pub trait SocketLookup {
+    fn get(&mut self, addr: SocketAddr) -> &mut Socket;
+}
+
 /// Loops around `process` while it returns `Output::Datagram`. While looping,
 /// batch together returned datagrams with the same length, destination and TOS
 /// byte into GSO-sized chunks and send them when the batch is full. If we
@@ -95,8 +99,8 @@ impl Socket {
 /// introduces a lot of unnecessary memory copies. It would be better to extend
 /// this with something like https://github.com/mozilla/neqo/pull/2138, where
 /// the stack directly writes into the GSO buffer.
-pub async fn process_send<'a>(
-    mut socket: impl FnMut(SocketAddr) -> &'a mut Socket,
+pub async fn process_send(
+    mut socket_lookup: impl SocketLookup,
     mut process: impl FnMut() -> Output,
     timeout: &mut Option<Pin<Box<Sleep>>>,
 ) -> Result<(), io::Error> {
@@ -111,7 +115,7 @@ pub async fn process_send<'a>(
         if (send || exit) && !batch_data.is_empty() {
             let meta = batch_meta.clone().unwrap();
             // Send all collected datagrams as GSO-sized chunks.
-            let socket = socket(meta.source());
+            let socket = socket_lookup.get(meta.source());
             for gso_chunk in batch_data.chunks(socket.max_gso_segments() * meta.len()) {
                 // Optimistically attempt sending datagram. In case the OS
                 // buffer is full, wait till socket is writable then try again.