Harden reconnect paths under UDP churn

src/udp.rs

@@ -1,5 +1,9 @@
 //! UDP error classification shared by Dosh clients, servers, and embedders.
 
+use std::net::SocketAddr;
+use std::time::{Duration, Instant};
+use tokio::net::UdpSocket;
+
 pub fn is_transient_udp_error(err: &std::io::Error) -> bool {
     matches!(
         err.kind(),
@@ -13,6 +17,54 @@ pub fn is_transient_udp_send_error(err: &std::io::Error) -> bool {
     is_transient_udp_error(err)
 }
 
+pub async fn send_udp_retrying_transient(
+    socket: &UdpSocket,
+    packet: &[u8],
+    peer: SocketAddr,
+    wait: Duration,
+) -> std::io::Result<()> {
+    let deadline = Instant::now() + wait.max(Duration::from_millis(1));
+    loop {
+        match socket.send_to(packet, peer).await {
+            Ok(_) => return Ok(()),
+            Err(err) if is_transient_udp_send_error(&err) && Instant::now() < deadline => {
+                sleep_until_retry(deadline).await;
+            }
+            Err(err) => return Err(err),
+        }
+    }
+}
+
+pub async fn recv_udp_retrying_transient(
+    socket: &UdpSocket,
+    buf: &mut [u8],
+    wait: Duration,
+) -> std::io::Result<(usize, SocketAddr)> {
+    let deadline = Instant::now() + wait.max(Duration::from_millis(1));
+    loop {
+        let now = Instant::now();
+        if now >= deadline {
+            return Err(std::io::Error::from(std::io::ErrorKind::TimedOut));
+        }
+        match tokio::time::timeout(deadline - now, socket.recv_from(buf)).await {
+            Ok(Ok(value)) => return Ok(value),
+            Ok(Err(err)) if is_transient_udp_error(&err) && Instant::now() < deadline => {
+                sleep_until_retry(deadline).await;
+            }
+            Ok(Err(err)) => return Err(err),
+            Err(_) => return Err(std::io::Error::from(std::io::ErrorKind::TimedOut)),
+        }
+    }
+}
+
+async fn sleep_until_retry(deadline: Instant) {
+    let now = Instant::now();
+    if now >= deadline {
+        return;
+    }
+    tokio::time::sleep((deadline - now).min(Duration::from_millis(20))).await;
+}
+
 #[cfg(unix)]
 fn is_transient_udp_os_error(code: i32) -> bool {
     matches!(
@@ -52,7 +104,12 @@ fn is_transient_udp_os_error(_code: i32) -> bool {
 
 #[cfg(test)]
 mod tests {
-    use super::{is_transient_udp_error, is_transient_udp_send_error};
+    use super::{
+        is_transient_udp_error, is_transient_udp_send_error, recv_udp_retrying_transient,
+        send_udp_retrying_transient,
+    };
+    use std::time::Duration;
+    use tokio::net::UdpSocket;
 
     #[test]
     fn classifies_portable_transient_udp_errors() {
@@ -70,6 +127,35 @@ mod tests {
         assert!(!is_transient_udp_send_error(&fatal));
     }
 
+    #[tokio::test]
+    async fn retrying_udp_send_and_receive_round_trip() {
+        let receiver = UdpSocket::bind("127.0.0.1:0").await.unwrap();
+        let sender = UdpSocket::bind("127.0.0.1:0").await.unwrap();
+        let addr = receiver.local_addr().unwrap();
+
+        send_udp_retrying_transient(&sender, b"hello", addr, Duration::from_millis(50))
+            .await
+            .unwrap();
+        let mut buf = [0u8; 16];
+        let (n, peer) = recv_udp_retrying_transient(&receiver, &mut buf, Duration::from_millis(50))
+            .await
+            .unwrap();
+
+        assert_eq!(&buf[..n], b"hello");
+        assert_eq!(peer, sender.local_addr().unwrap());
+    }
+
+    #[tokio::test]
+    async fn retrying_udp_receive_times_out_cleanly() {
+        let receiver = UdpSocket::bind("127.0.0.1:0").await.unwrap();
+        let mut buf = [0u8; 16];
+        let err = recv_udp_retrying_transient(&receiver, &mut buf, Duration::from_millis(1))
+            .await
+            .unwrap_err();
+
+        assert_eq!(err.kind(), std::io::ErrorKind::TimedOut);
+    }
+
     #[cfg(unix)]
     #[test]
     fn classifies_unix_network_churn_as_transient() {