Harden QUIC happy eyeballs scheduling

Signed-off-by: Andrzej Haczewski <ahaczewski@gmail.com>

Author: ahaczewski

lore-transport/src/quic/client.rs

@@ -73,6 +73,7 @@ pub struct EndpointConfig {
 const IDLE_TIMEOUT_MS: u32 = 30000;
 const KEEP_ALIVE_MS: u64 = 500;
 const HAPPY_EYEBALLS_DELAY_MS: u64 = 250;
+const HAPPY_EYEBALLS_MAX_IN_FLIGHT: usize = 10;
 pub const DEFAULT_EXPECTED_RTT_MS: u64 = 100;
 
 #[derive(Clone, Debug)]
@@ -693,6 +694,7 @@ pub async fn connect(
         .to_socket_addrs()
         .internal_with(|| format!("remote {remote_url} is invalid"))?
         .collect();
+    let remote_addrs = interleave_socket_addrs(remote_addrs);
     let server_name = config.sni_override.as_deref().unwrap_or(host.as_str());
 
     let validate_certificate = url.scheme().ends_with("s");
@@ -772,6 +774,38 @@ pub async fn connect(
     Err(ProtocolError::internal(format!("connect: {remote_url}")))
 }
 
+fn interleave_socket_addrs(remote_addrs: Vec<SocketAddr>) -> Vec<SocketAddr> {
+    let Some(first) = remote_addrs.first() else {
+        return remote_addrs;
+    };
+    let prefer_ipv6 = first.is_ipv6();
+    let (preferred, fallback): (Vec<_>, Vec<_>) = remote_addrs
+        .into_iter()
+        .partition(|addr| addr.is_ipv6() == prefer_ipv6);
+    let mut preferred = preferred.into_iter();
+    let mut fallback = fallback.into_iter();
+    let mut interleaved = Vec::with_capacity(preferred.len() + fallback.len());
+
+    loop {
+        match preferred.next() {
+            Some(addr) => interleaved.push(addr),
+            None => {
+                interleaved.extend(fallback);
+                break;
+            }
+        }
+        match fallback.next() {
+            Some(addr) => interleaved.push(addr),
+            None => {
+                interleaved.extend(preferred);
+                break;
+            }
+        }
+    }
+
+    interleaved
+}
+
 async fn connect_happy_eyeballs<T, F, Fut>(
     remote_addrs: Vec<SocketAddr>,
     attempt_delay: Duration,
@@ -802,6 +836,8 @@ where
         tokio::select! {
             result = attempts.next(), if !attempts.is_empty() => {
                 if let Some(Some(connection)) = result {
+                    // Dropping `attempts` cancels the losing Quinn handshakes because each
+                    // production future owns its `Connecting` and `Endpoint`.
                     return Some(connection);
                 }
                 if attempts.is_empty() {
@@ -813,7 +849,7 @@ where
                     delay.as_mut().reset(tokio::time::Instant::now() + attempt_delay);
                 }
             }
-            _ = &mut delay => {
+            _ = &mut delay, if attempts.len() < HAPPY_EYEBALLS_MAX_IN_FLIGHT => {
                 let Some(addr) = next_addr.take() else {
                     continue;
                 };
@@ -1244,8 +1280,12 @@ mod tests {
     use std::time::Duration;
 
     use parking_lot::Mutex;
+    use tokio::sync::Semaphore;
+    use tokio::sync::mpsc;
 
+    use super::HAPPY_EYEBALLS_MAX_IN_FLIGHT;
     use super::connect_happy_eyeballs;
+    use super::interleave_socket_addrs;
 
     fn ipv6_addr() -> SocketAddr {
         "[::1]:41337".parse().unwrap()
@@ -1255,6 +1295,41 @@ mod tests {
         "127.0.0.1:41337".parse().unwrap()
     }
 
+    #[test]
+    fn happy_eyeballs_interleaves_ipv6_first_addresses() {
+        let ipv6_second = "[::2]:41337".parse().unwrap();
+        let ipv6_third = "[::3]:41337".parse().unwrap();
+        let ipv4_second = "127.0.0.2:41337".parse().unwrap();
+
+        assert_eq!(
+            interleave_socket_addrs(vec![
+                ipv6_addr(),
+                ipv6_second,
+                ipv6_third,
+                ipv4_addr(),
+                ipv4_second,
+            ]),
+            vec![
+                ipv6_addr(),
+                ipv4_addr(),
+                ipv6_second,
+                ipv4_second,
+                ipv6_third,
+            ]
+        );
+    }
+
+    #[test]
+    fn happy_eyeballs_interleaves_ipv4_first_addresses() {
+        let ipv4_second = "127.0.0.2:41337".parse().unwrap();
+        let ipv6_second = "[::2]:41337".parse().unwrap();
+
+        assert_eq!(
+            interleave_socket_addrs(vec![ipv4_addr(), ipv4_second, ipv6_addr(), ipv6_second,]),
+            vec![ipv4_addr(), ipv6_addr(), ipv4_second, ipv6_second]
+        );
+    }
+
     #[tokio::test]
     async fn happy_eyeballs_starts_fallback_while_first_attempt_is_stalled() {
         let attempts = Arc::new(Mutex::new(Vec::new()));
@@ -1334,4 +1409,52 @@ mod tests {
 
         assert_eq!(result, None);
     }
+
+    #[tokio::test]
+    async fn happy_eyeballs_bounds_in_flight_attempts() {
+        let remote_addrs: Vec<_> = (1..=HAPPY_EYEBALLS_MAX_IN_FLIGHT + 1)
+            .map(|port| SocketAddr::new(ipv6_addr().ip(), port as u16))
+            .collect();
+        let release = Arc::new(Semaphore::new(0));
+        let attempt_release = release.clone();
+        let (started_tx, mut started_rx) = mpsc::unbounded_channel();
+
+        let task = lore_base::lore_spawn!(connect_happy_eyeballs(
+            remote_addrs.clone(),
+            Duration::from_millis(1),
+            move |addr| {
+                started_tx.send(addr).unwrap();
+                let attempt_release = attempt_release.clone();
+                async move {
+                    attempt_release.acquire().await.unwrap().forget();
+                    None::<SocketAddr>
+                }
+            },
+        ));
+
+        for expected in remote_addrs.iter().take(HAPPY_EYEBALLS_MAX_IN_FLIGHT) {
+            assert_eq!(
+                tokio::time::timeout(Duration::from_secs(2), started_rx.recv())
+                    .await
+                    .expect("attempt should start"),
+                Some(*expected)
+            );
+        }
+        assert!(
+            tokio::time::timeout(Duration::from_millis(50), started_rx.recv())
+                .await
+                .is_err(),
+            "attempts above the in-flight limit should remain queued"
+        );
+
+        release.add_permits(1);
+        assert_eq!(
+            tokio::time::timeout(Duration::from_secs(2), started_rx.recv())
+                .await
+                .expect("queued attempt should start when a slot opens"),
+            Some(remote_addrs[HAPPY_EYEBALLS_MAX_IN_FLIGHT])
+        );
+
+        task.abort();
+    }
 }