Merge #1891: fix(script): interrupt blocked socket writes on close (@TimotheusG)

Cargo.lock

@@ -1602,6 +1602,7 @@ dependencies = [
  "effectcraft-render",
  "effectcraft-time",
  "serde_json",
+ "socket2",
  "web-time",
 ]
 
@@ -4827,6 +4828,16 @@ dependencies = [
  "serde",
 ]
 
+[[package]]
+name = "socket2"
+version = "0.6.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c3d1e2c7f27f8d4cb10542a02c49005dbd6e93095799d6f3be745fae9f8fedd4"
+dependencies = [
+ "libc",
+ "windows-sys 0.61.2",
+]
+
 [[package]]
 name = "spin"
 version = "0.9.9"

crates/script/Cargo.toml

@@ -27,3 +27,6 @@ boa_engine = { workspace = true, features = ["js"] }
 [dev-dependencies]
 effectcraft-expr = { workspace = true }
 effectcraft-psd = { workspace = true }
+
+[target.'cfg(not(target_arch = "wasm32"))'.dev-dependencies]
+socket2 = "0.6"

crates/script/src/socket.rs

@@ -51,6 +51,7 @@ mod imp {
     use std::collections::BTreeMap;
     use std::io::{ErrorKind, Read, Write};
     use std::net::{Shutdown, TcpListener, TcpStream, ToSocketAddrs};
+    use std::sync::atomic::{AtomicBool, Ordering};
     use std::sync::{Arc, Mutex};
     use std::time::{Duration, Instant};
 
@@ -65,9 +66,11 @@ mod imp {
     /// socket blocked on a slow peer doesn't hold up the others (#1248).
     struct Entry {
         sock: Mutex<Sock>,
-        /// A second handle on a stream, so `close` can shut it down while another call is
+        /// A second handle on a stream, so `close` can wake a blocked read.
-        /// blocked reading or writing it.
         wake: Option<TcpStream>,
+        /// Closed or unusable after a mode error. Writes poll it because Windows sends may not
+        /// be interrupted by shutdown alone; queued operations recheck it after locking.
+        closed: AtomicBool,
     }
 
     struct Table {
@@ -85,7 +88,7 @@ mod imp {
             Sock::Stream { s, .. } => s.try_clone().ok(),
             Sock::Listener(_) => None,
         };
-        let entry = Arc::new(Entry { sock: Mutex::new(s), wake });
+        let entry = Arc::new(Entry { sock: Mutex::new(s), wake, closed: AtomicBool::new(false) });
         let mut t = table().lock().unwrap_or_else(|e| e.into_inner());
         let id = t.next;
         t.next += 1;
@@ -101,11 +104,23 @@ mod imp {
     /// Forget socket `id`; a call blocked on it fails at once.
     fn close(id: u32) {
         let entry = table().lock().unwrap_or_else(|e| e.into_inner()).socks.remove(&id);
-        if let Some(w) = entry.as_ref().and_then(|e| e.wake.as_ref()) {
+        if let Some(entry) = entry {
-            let _ = w.shutdown(Shutdown::Both);
+            entry.closed.store(true, Ordering::Release);
+            if let Some(w) = &entry.wake {
+                let _ = w.shutdown(Shutdown::Both);
+            }
         }
     }
 
+    /// A test-only handle for establishing backpressure before exercising the public operations.
+    #[cfg(test)]
+    pub(super) fn test_stream(id: u32) -> TcpStream {
+        let entry = get(id).unwrap();
+        let sock = entry.sock.lock().unwrap();
+        let Sock::Stream { s, .. } = &*sock else { panic!("expected a stream") };
+        s.try_clone().unwrap()
+    }
+
     fn timeout(p: &J) -> Duration {
         Duration::from_secs_f64(p.get("timeout").and_then(J::as_f64).unwrap_or(10.0).clamp(0.0, 3600.0))
     }
@@ -150,7 +165,56 @@ mod imp {
         Ok(())
     }
 
+    /// A failed mode transition leaves the stream unusable; reject queued operations and wake I/O.
+    fn change_mode(s: &TcpStream, entry: &Entry, nonblocking: bool, mode: &impl Fn(&TcpStream, bool) -> std::io::Result<()>) -> std::io::Result<()> {
+        mode(s, nonblocking).map_err(|e| {
+            entry.closed.store(true, Ordering::Release);
+            let _ = s.shutdown(Shutdown::Both);
+            let action = if nonblocking { "enable nonblocking mode" } else { "restore blocking mode" };
+            std::io::Error::other(format!("socket is unusable: cannot {action}: {e}"))
+        })
+    }
+
+    /// Write with one deadline and an explicit close signal. On Windows, shutdown alone does
+    /// not reliably interrupt a blocking send, so never leave a send blocked inside the OS.
+    /// The caller holds this socket's lock; restore blocking mode or retire the connection.
+    fn write(s: &mut TcpStream, data: &[u8], limit: Duration, entry: &Entry, mode: &impl Fn(&TcpStream, bool) -> std::io::Result<()>) -> std::io::Result<()> {
+        change_mode(s, entry, true, mode)?;
+        let end = Instant::now() + limit;
+        let result = (|| {
+            let mut pending = data;
+            while !pending.is_empty() {
+                if entry.closed.load(Ordering::Acquire) {
+                    return Err(std::io::Error::new(ErrorKind::ConnectionAborted, "the socket was closed"));
+                }
+                let left = end.saturating_duration_since(Instant::now());
+                if left.is_zero() {
+                    return Err(std::io::Error::from(ErrorKind::TimedOut));
+                }
+                match s.write(pending) {
+                    Ok(0) => return Err(std::io::Error::from(ErrorKind::WriteZero)),
+                    Ok(n) => pending = pending.get(n..).ok_or_else(|| std::io::Error::from(ErrorKind::InvalidData))?,
+                    Err(e) if e.kind() == ErrorKind::WouldBlock => std::thread::sleep(left.min(Duration::from_millis(5))),
+                    Err(e) if e.kind() == ErrorKind::Interrupted => continue,
+                    Err(e) => return Err(e),
+                }
+            }
+            Ok(())
+        })();
+        match (result, change_mode(s, entry, false, mode)) {
+            (Err(write), Err(restore)) => Err(std::io::Error::other(format!("{restore}; write also failed: {write}"))),
+            (result, Ok(())) => result,
+            (Ok(()), Err(restore)) => Err(restore),
+        }
+    }
+
+    /// Dispatch an authorized operation without holding the socket table during I/O.
     pub(super) fn run(op: &str, id: u32, p: &J) -> J {
+        run_with_mode(op, id, p, TcpStream::set_nonblocking)
+    }
+
+    /// Dispatch with a mode setter so exceptional OS failures can be injected per operation.
+    pub(super) fn run_with_mode(op: &str, id: u32, p: &J, mode: impl Fn(&TcpStream, bool) -> std::io::Result<()>) -> J {
         let utf8 = p.get("encoding").and_then(J::as_str).is_some_and(|e| e.eq_ignore_ascii_case("utf-8") || e.eq_ignore_ascii_case("utf8"));
         match op {
             "open" => {
@@ -198,6 +262,15 @@ mod imp {
                     };
                 };
                 let mut sock = entry.sock.lock().unwrap_or_else(|e| e.into_inner());
+                // Recheck after locking: close or a failed mode transition may have retired it
+                // while this caller waited with an Arc obtained from the table.
+                if entry.closed.load(Ordering::Acquire) {
+                    return match op {
+                        "connected" => json!({"value": false}),
+                        "eof" => json!({"value": true}),
+                        _ => err("the socket is closed or unusable"),
+                    };
+                }
                 match (op, &mut *sock) {
                     ("poll", Sock::Listener(l)) => match l.accept() {
                         Ok((s, peer)) => {
@@ -212,14 +285,18 @@ mod imp {
                     ("connected", Sock::Stream { s, buf, eof }) => {
                         if !*eof && buf.is_empty() {
                             // Peek without blocking to notice a closed peer.
-                            let _ = s.set_nonblocking(true);
+                            if let Err(e) = change_mode(s, &entry, true, &mode) {
+                                return err(e);
+                            }
                             let mut b = [0u8; 1];
                             match s.peek(&mut b) {
                                 Ok(0) => *eof = true,
                                 Err(e) if e.kind() != ErrorKind::WouldBlock => *eof = true,
                                 _ => {}
                             }
-                            let _ = s.set_nonblocking(false);
+                            if let Err(e) = change_mode(s, &entry, false, &mode) {
+                                return err(e);
+                            }
                         }
                         json!({"value": !(*eof && buf.is_empty())})
                     }
@@ -265,11 +342,8 @@ mod imp {
                         // A peer that stops reading fills the buffers: give up after `timeout`
                         // instead of blocking the script for good (#1248).
                         let limit = timeout(p).max(Duration::from_millis(1));
-                        if let Err(e) = s.set_write_timeout(Some(limit)) {
-                            return err(format!("cannot set the write timeout: {e}"));
-                        }
                         let data = p.get("data").and_then(J::as_str).unwrap_or("");
-                        match s.write_all(&encode(data, utf8)).and_then(|_| s.flush()) {
+                        match write(s, &encode(data, utf8), limit, &entry, &mode) {
                             Ok(()) => json!({"ok": true}),
                             Err(e) if matches!(e.kind(), ErrorKind::WouldBlock | ErrorKind::TimedOut) => err(format!(
                                 "write timed out after {} s: the other side isn't reading (some of the data may have been sent)",
@@ -289,6 +363,7 @@ mod imp {
 #[cfg(test)]
 #[cfg(not(target_arch = "wasm32"))]
 mod tests {
+    use std::io::{ErrorKind, Read, Write};
     use std::net::{TcpListener, TcpStream};
     use std::sync::mpsc;
     use std::time::{Duration, Instant};
@@ -301,6 +376,7 @@ mod tests {
     /// `host:port` and the accepted stream's keeper (drop it to let the connection go).
     fn stalled_peer() -> (String, mpsc::Sender<()>) {
         let l = TcpListener::bind("127.0.0.1:0").unwrap();
+        socket2::SockRef::from(&l).set_recv_buffer_size(4096).unwrap();
         let host = l.local_addr().unwrap().to_string();
         let (stop, wait) = mpsc::channel::<()>();
         std::thread::spawn(move || {
@@ -310,14 +386,32 @@ mod tests {
         (host, stop)
     }
 
+    /// Open a test connection through the production socket dispatcher.
     fn open(host: &str) -> u32 {
         let r = run("open", 0, &json!({"host": host, "timeout": 5}));
         r["id"].as_u64().unwrap_or_else(|| panic!("open {host}: {r}")) as u32
     }
 
-    /// More than the loopback buffers hold, so a peer that doesn't read stops the write.
+    /// Fill bounded socket buffers nonblocking, rather than assuming a 64 MiB write stalls.
-    fn big() -> String {
+    /// In particular, Windows may accept that entire write with its default send buffering.
-        "x".repeat(64 * 1024 * 1024)
+    fn backpressure(id: u32) {
+        let mut stream = super::imp::test_stream(id);
+        socket2::SockRef::from(&stream).set_send_buffer_size(4096).unwrap();
+        stream.set_nonblocking(true).unwrap();
+        let start = Instant::now();
+        let chunk = [b'x'; 4096];
+        let mut sent = 0;
+        loop {
+            assert!(sent < 16 * 1024 * 1024 && start.elapsed() < Duration::from_secs(5), "could not establish backpressure");
+            match stream.write(&chunk) {
+                Ok(0) => panic!("peer closed while establishing backpressure"),
+                Ok(n) => sent += n,
+                Err(e) if e.kind() == ErrorKind::WouldBlock => break,
+                Err(e) if e.kind() == ErrorKind::Interrupted => continue,
+                Err(e) => panic!("establishing backpressure: {e}"),
+            }
+        }
+        stream.set_nonblocking(false).unwrap();
     }
 
     /// Run `f` on a thread; its result, or `None` if it takes longer than `limit`.
@@ -329,31 +423,36 @@ mod tests {
         rx.recv_timeout(limit).ok()
     }
 
+    /// A bounded write times out even when the peer never drains its receive buffer.
     #[test]
     fn a_write_to_a_peer_that_stops_reading_fails_after_the_timeout() {
         // #1248: it used to block for good.
         let (host, _keep) = stalled_peer();
         let id = open(&host);
-        let data = big();
+        backpressure(id);
+        let data = "x".repeat(64 * 1024);
         let t0 = Instant::now();
         let r = within(Duration::from_secs(20), move || run("write", id, &json!({"data": data, "timeout": 1}))).expect("the write gave up");
         assert!(r["error"].as_str().is_some_and(|e| e.contains("timed out")), "{r}");
-        assert!(t0.elapsed() < Duration::from_secs(15), "{:?}", t0.elapsed());
+        assert!(t0.elapsed() >= Duration::from_millis(500) && t0.elapsed() < Duration::from_secs(15), "{:?}", t0.elapsed());
         run("close", id, &J::Null);
     }
 
+    /// Closing a stalled stream interrupts its writer without blocking another socket.
     #[test]
     fn a_blocked_socket_doesnt_hold_up_the_others_and_close_interrupts_it() {
         let (host_a, _keep_a) = stalled_peer();
         let (host_b, _keep_b) = stalled_peer();
         let (a, b) = (open(&host_a), open(&host_b));
         // A write on A that would wait a minute.
-        let data = big();
+        backpressure(a);
+        let data = "x".repeat(64 * 1024);
         let (tx, blocked) = mpsc::channel();
         std::thread::spawn(move || {
             let _ = tx.send(run("write", a, &json!({"data": data, "timeout": 60})));
         });
         std::thread::sleep(Duration::from_millis(500));
+        assert!(matches!(blocked.try_recv(), Err(mpsc::TryRecvError::Empty)), "A must still be writing before close");
         // Socket B answers meanwhile (the table isn't locked by A's write).
         let r = within(Duration::from_secs(3), move || run("connected", b, &json!({}))).expect("socket B answered while A was blocked");
         assert_eq!(r["value"], true, "{r}");
@@ -363,4 +462,86 @@ mod tests {
         assert!(r.get("error").is_some(), "{r}");
         run("close", b, &J::Null);
     }
+
+    /// Failed mode changes retire the stream; subsequent I/O fails instead of using it.
+    #[test]
+    fn mode_restoration_failure_retires_streams() {
+        for (op, fail_nonblocking) in [("write", false), ("connected", false), ("write", true), ("connected", true)] {
+            let (host, _keep) = stalled_peer();
+            let id = open(&host);
+            let result = super::imp::run_with_mode(op, id, &json!({"data": "hello"}), |s, nonblocking| {
+                if nonblocking == fail_nonblocking { Err(std::io::Error::other("injected mode failure")) } else { s.set_nonblocking(nonblocking) }
+            });
+            assert!(result["error"].as_str().is_some_and(|e| e.contains("injected mode failure")), "{op}: {result}");
+            assert_eq!(run("connected", id, &J::Null)["value"], false, "{op}: broken socket remains connected");
+            assert_eq!(run("eof", id, &J::Null)["value"], true);
+            for operation in ["read", "readln", "write"] {
+                let r = run(operation, id, &json!({"count": 1, "timeout": 0.01, "data": "x"}));
+                assert!(r.get("error").is_some(), "{operation} must reject the unusable stream: {r}");
+            }
+            run("close", id, &J::Null);
+            run("close", id, &J::Null);
+        }
+    }
+
+    /// A preceding write error cannot hide the failure that makes the stream unusable.
+    #[test]
+    fn restoration_failure_preserves_the_write_error_too() {
+        let (host, _keep) = stalled_peer();
+        let id = open(&host);
+        let result = super::imp::run_with_mode("write", id, &json!({"data": "hello"}), |s, nonblocking| {
+            if nonblocking {
+                s.set_nonblocking(true)?;
+                s.shutdown(std::net::Shutdown::Write)
+            } else {
+                Err(std::io::Error::other("injected restore failure"))
+            }
+        });
+        let error = result["error"].as_str().unwrap();
+        assert!(error.contains("injected restore failure") && error.contains("write also failed"), "{error}");
+        assert_eq!(run("connected", id, &J::Null)["value"], false);
+        run("close", id, &J::Null);
+    }
+
+    /// Both outgoing and accepted sockets deliver complete data, then return to blocking reads.
+    #[test]
+    fn writes_deliver_all_data_and_restore_blocking_reads() {
+        for accepted in [false, true] {
+            let (id, mut peer) = if accepted {
+                let listener = run("listen", 0, &json!({"port": 0}));
+                let listener_id = listener["id"].as_u64().unwrap() as u32;
+                let peer = TcpStream::connect(("127.0.0.1", listener["port"].as_u64().unwrap() as u16)).unwrap();
+                let start = Instant::now();
+                let id = loop {
+                    if let Some(id) = run("poll", listener_id, &J::Null)["id"].as_u64() {
+                        break id as u32;
+                    }
+                    assert!(start.elapsed() < Duration::from_secs(5));
+                    std::thread::sleep(Duration::from_millis(1));
+                };
+                run("close", listener_id, &J::Null);
+                (id, peer)
+            } else {
+                let listener = TcpListener::bind("127.0.0.1:0").unwrap();
+                let id = open(&listener.local_addr().unwrap().to_string());
+                (id, listener.accept().unwrap().0)
+            };
+            socket2::SockRef::from(&super::imp::test_stream(id)).set_send_buffer_size(4096).unwrap();
+            peer.set_read_timeout(Some(Duration::from_secs(5))).unwrap();
+            let reader = std::thread::spawn(move || {
+                let mut data = vec![0; 256 * 1024];
+                peer.read_exact(&mut data).unwrap();
+                assert!(data.iter().all(|b| *b == b'x'));
+                // A read on a socket accidentally left nonblocking would return before this.
+                std::thread::sleep(Duration::from_millis(100));
+                peer.write_all(b"ack").unwrap();
+            });
+            let r = within(Duration::from_secs(10), move || run("write", id, &json!({"data": "x".repeat(256 * 1024), "timeout": 5}))).unwrap();
+            assert_eq!(r["ok"], true, "{r}");
+            let r = run("read", id, &json!({"count": 3, "timeout": 2}));
+            assert_eq!(r["data"], "ack", "{r}");
+            run("close", id, &J::Null);
+            reader.join().unwrap();
+        }
+    }
 }