Retry transport receive during UDP churn

Cargo.lock

@@ -436,7 +436,7 @@ dependencies = [
 
 [[package]]
 name = "dosh"
-version = "1.0.0-rc20"
+version = "1.0.0-rc22"
 dependencies = [
  "anyhow",
  "base64",

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "dosh"
-version = "1.0.0-rc20"
+version = "1.0.0-rc22"
 edition = "2024"
 license = "MIT"
 

src/bin/dosh-client.rs

@@ -36,7 +36,7 @@ use dosh::transport::{
     DoshTransport, SessionEvent, SessionRole, SessionTransportConfig, TransportConfig,
     TransportEvent,
 };
-use dosh::udp::is_transient_udp_send_error;
+use dosh::udp::{is_transient_udp_error, is_transient_udp_send_error};
 use serde::{Deserialize, Serialize};
 use sha2::{Digest, Sha256};
 use std::collections::BTreeMap;
@@ -4440,9 +4440,14 @@ where
             return Err(anyhow!("timed out waiting for server response"));
         }
         let remaining = wait - elapsed;
-        let (n, _) = tokio::time::timeout(remaining, socket.recv_from(&mut buf))
+        let (n, _) = match tokio::time::timeout(remaining, socket.recv_from(&mut buf))
             .await
-            .context("timed out waiting for server response")??;
+            .context("timed out waiting for server response")?
+        {
+            Ok(value) => value,
+            Err(err) if is_transient_udp_error(&err) => continue,
+            Err(err) => return Err(err.into()),
+        };
         let Ok(packet) = protocol::decode(&buf[..n]) else {
             continue;
         };
@@ -4832,7 +4837,11 @@ async fn run_terminal(
                 }
             }
             recv = socket.recv_from(&mut recv_buf) => {
-                let (n, _) = recv?;
+                let (n, _) = match recv {
+                    Ok(value) => value,
+                    Err(err) if is_transient_udp_error(&err) => continue,
+                    Err(err) => return Err(err.into()),
+                };
                 let Ok(packet) = protocol::decode(&recv_buf[..n]) else {
                     continue;
                 };
@@ -7770,8 +7779,8 @@ mod tests {
         retransmit_stream_opens, rewrite_forward_command, selected_predict_mode, selected_udp_host,
         server_version_mismatch, should_flush_terminal_input_after_contact,
         should_hold_during_startup_gate, should_hold_post_submit_input,
-        should_repaint_idle_alternate_screen, split_after_command_submit, ssh_destination_host,
-        ssh_username, ssh_with_user, ssh_command_target, startup_command, status_ssh_target,
+        should_repaint_idle_alternate_screen, split_after_command_submit, ssh_command_target,
+        ssh_destination_host, ssh_username, ssh_with_user, startup_command, status_ssh_target,
         strip_stale_mouse_reports, strip_terminal_focus_reports, terminal_private_mode_transition,
         toml_bare_key_or_quoted, unix_update_script, update_installer_url, update_version_status,
         upsert_managed_block, valid_forward_host, vscode_safe_alias, windows_update_script,

src/bin/dosh-server.rs

@@ -28,7 +28,7 @@ use dosh::protocol::{
     TicketAttachBody, TicketAttachEnvelope, TicketAttachOkEnvelope,
 };
 use dosh::pty::{PtyHandle, PtyOutput, adopt_pty_from_fd, spawn_pty_session};
-use dosh::udp::is_transient_udp_send_error;
+use dosh::udp::{is_transient_udp_error, is_transient_udp_send_error};
 use sha2::{Digest, Sha256};
 use std::collections::{BTreeMap, HashMap, HashSet, VecDeque};
 use std::fs;
@@ -262,7 +262,11 @@ async fn serve(config_path: Option<std::path::PathBuf>) -> Result<()> {
 
     let mut buf = vec![0u8; 65535];
     loop {
-        let (n, peer) = socket.recv_from(&mut buf).await?;
+        let (n, peer) = match socket.recv_from(&mut buf).await {
+            Ok(value) => value,
+            Err(err) if is_transient_udp_error(&err) => continue,
+            Err(err) => return Err(err.into()),
+        };
         if let Err(err) = handle_packet(&state, &socket, peer, &buf[..n]).await {
             eprintln!("packet from {peer}: {err:#}");
         }

src/persist.rs

@@ -23,12 +23,14 @@ use std::io::{Read, Write};
 use std::os::fd::{AsRawFd, FromRawFd, IntoRawFd, RawFd};
 use std::os::unix::net::{UnixListener, UnixStream};
 use std::path::{Path, PathBuf};
+use std::sync::Mutex;
 use std::time::Duration;
 
 /// One-byte commands a server sends to a holder over its control socket after
 /// the holder has handed back the master fd.
 const HOLDER_CMD_SHUTDOWN: u8 = b'X';
 const HOLDER_STARTUP_TIMEOUT: Duration = Duration::from_millis(750);
+static SCREEN_WRITE_LOCK: Mutex<()> = Mutex::new(());
 
 /// Magic written into a holder's `meta` file, bumped if the on-disk layout
 /// changes so a stale holder from an incompatible build is ignored.
@@ -140,13 +142,28 @@ pub fn save_screen(
         // No holder runtime for this session (non-persistent): nothing to do.
         return Ok(());
     }
+    let _guard = SCREEN_WRITE_LOCK
+        .lock()
+        .map_err(|_| anyhow!("screen persistence lock poisoned"))?;
+    let path = screen_path(&dir);
+    if existing_screen_output_seq(&path).is_some_and(|existing| existing >= output_seq) {
+        return Ok(());
+    }
     let mut buf = Vec::with_capacity(snapshot.len() + 32);
     buf.extend_from_slice(&cols.to_be_bytes());
     buf.extend_from_slice(&rows.to_be_bytes());
     buf.extend_from_slice(&output_seq.to_be_bytes());
     buf.extend_from_slice(&(snapshot.len() as u32).to_be_bytes());
     buf.extend_from_slice(snapshot);
-    atomic_write(&screen_path(&dir), &buf)
+    atomic_write(&path, &buf)
+}
+
+fn existing_screen_output_seq(path: &Path) -> Option<u64> {
+    let data = std::fs::read(path).ok()?;
+    if data.len() < 12 {
+        return None;
+    }
+    Some(u64::from_be_bytes(data[4..12].try_into().ok()?))
 }
 
 /// A restored screen for a re-adopted session.
@@ -619,6 +636,20 @@ mod tests {
         assert_eq!(loaded.snapshot, snap);
     }
 
+    #[test]
+    fn save_screen_never_overwrites_newer_snapshot_with_older_seq() {
+        let tmp = tempfile::tempdir().unwrap();
+        let sessions_dir = tmp.path();
+        ensure_runtime_dir(sessions_dir, "work").unwrap();
+
+        save_screen(sessions_dir, "work", 80, 24, 10, b"new-screen").unwrap();
+        save_screen(sessions_dir, "work", 80, 24, 9, b"old-screen").unwrap();
+
+        let loaded = load_screen(sessions_dir, "work").expect("screen restored");
+        assert_eq!(loaded.output_seq, 10);
+        assert_eq!(loaded.snapshot, b"new-screen");
+    }
+
     #[test]
     fn load_screen_absent_is_none() {
         let tmp = tempfile::tempdir().unwrap();

src/transport.rs

@@ -19,7 +19,7 @@ use crate::protocol::{
     self, CLIENT_TO_SERVER, PacketKind, ReplayWindow, SERVER_TO_CLIENT, StreamClose, StreamData,
     StreamOpen, StreamOpenOk, StreamOpenReject, StreamWindowAdjust,
 };
-use crate::udp::is_transient_udp_send_error;
+use crate::udp::{is_transient_udp_error, is_transient_udp_send_error};
 use anyhow::{Result, anyhow, bail};
 use std::collections::{BTreeMap, HashMap, HashSet, VecDeque};
 use std::net::SocketAddr;
@@ -852,7 +852,11 @@ impl DoshTransport {
     pub async fn recv(&mut self) -> Result<SessionEvent> {
         let mut buf = vec![0u8; 65535];
         loop {
-            let (n, peer) = self.socket.recv_from(&mut buf).await?;
+            let (n, peer) = match self.socket.recv_from(&mut buf).await {
+                Ok(value) => value,
+                Err(err) if transport_recv_error_is_retryable(&err) => continue,
+                Err(err) => return Err(err.into()),
+            };
             let packet = match protocol::decode(&buf[..n]) {
                 Ok(packet) => packet,
                 Err(_) => continue,
@@ -984,6 +988,10 @@ impl DoshTransport {
     }
 }
 
+fn transport_recv_error_is_retryable(err: &std::io::Error) -> bool {
+    is_transient_udp_error(err)
+}
+
 fn initial_stream_id(role: SessionRole) -> u64 {
     match role {
         SessionRole::Client => 1,
@@ -1033,6 +1041,35 @@ mod tests {
         assert!(service_target("../bad").is_err());
     }
 
+    #[test]
+    fn transport_recv_retries_network_churn_errors() {
+        for kind in [
+            std::io::ErrorKind::Interrupted,
+            std::io::ErrorKind::TimedOut,
+            std::io::ErrorKind::WouldBlock,
+        ] {
+            let err = std::io::Error::from(kind);
+            assert!(transport_recv_error_is_retryable(&err));
+        }
+        assert!(!transport_recv_error_is_retryable(&std::io::Error::from(
+            std::io::ErrorKind::PermissionDenied,
+        )));
+
+        #[cfg(unix)]
+        for code in [libc::ENETUNREACH, libc::EHOSTUNREACH, libc::ECONNRESET] {
+            assert!(transport_recv_error_is_retryable(
+                &std::io::Error::from_raw_os_error(code)
+            ));
+        }
+
+        #[cfg(windows)]
+        for code in [10051, 10054, 10065] {
+            assert!(transport_recv_error_is_retryable(
+                &std::io::Error::from_raw_os_error(code)
+            ));
+        }
+    }
+
     #[test]
     fn service_registry_matches_registered_service_opens() {
         let mut registry = ServiceRegistry::new();

src/udp.rs

@@ -1,6 +1,6 @@
 //! UDP error classification shared by Dosh clients, servers, and embedders.
 
-pub fn is_transient_udp_send_error(err: &std::io::Error) -> bool {
+pub fn is_transient_udp_error(err: &std::io::Error) -> bool {
     matches!(
         err.kind(),
         std::io::ErrorKind::Interrupted
@@ -9,6 +9,10 @@ pub fn is_transient_udp_send_error(err: &std::io::Error) -> bool {
     ) || err.raw_os_error().is_some_and(is_transient_udp_os_error)
 }
 
+pub fn is_transient_udp_send_error(err: &std::io::Error) -> bool {
+    is_transient_udp_error(err)
+}
+
 #[cfg(unix)]
 fn is_transient_udp_os_error(code: i32) -> bool {
     matches!(
@@ -48,20 +52,22 @@ fn is_transient_udp_os_error(_code: i32) -> bool {
 
 #[cfg(test)]
 mod tests {
-    use super::is_transient_udp_send_error;
+    use super::{is_transient_udp_error, is_transient_udp_send_error};
 
     #[test]
-    fn classifies_portable_transient_send_errors() {
+    fn classifies_portable_transient_udp_errors() {
         for kind in [
             std::io::ErrorKind::Interrupted,
             std::io::ErrorKind::TimedOut,
             std::io::ErrorKind::WouldBlock,
         ] {
-            assert!(is_transient_udp_send_error(&std::io::Error::from(kind)));
+            let err = std::io::Error::from(kind);
+            assert!(is_transient_udp_error(&err));
+            assert!(is_transient_udp_send_error(&err));
         }
-        assert!(!is_transient_udp_send_error(&std::io::Error::from(
-            std::io::ErrorKind::PermissionDenied,
-        )));
+        let fatal = std::io::Error::from(std::io::ErrorKind::PermissionDenied);
+        assert!(!is_transient_udp_error(&fatal));
+        assert!(!is_transient_udp_send_error(&fatal));
     }
 
     #[cfg(unix)]
@@ -78,9 +84,9 @@ mod tests {
             libc::ENETUNREACH,
             libc::ETIMEDOUT,
         ] {
-            assert!(is_transient_udp_send_error(
-                &std::io::Error::from_raw_os_error(code)
-            ));
+            let err = std::io::Error::from_raw_os_error(code);
+            assert!(is_transient_udp_error(&err));
+            assert!(is_transient_udp_send_error(&err));
         }
     }
 
@@ -90,9 +96,9 @@ mod tests {
         for code in [
             10049, 10050, 10051, 10052, 10054, 10060, 10061, 10064, 10065,
         ] {
-            assert!(is_transient_udp_send_error(
-                &std::io::Error::from_raw_os_error(code)
-            ));
+            let err = std::io::Error::from_raw_os_error(code);
+            assert!(is_transient_udp_error(&err));
+            assert!(is_transient_udp_send_error(&err));
         }
     }
 }