From da73df4b2eba41f2882a0a2eb283e6f770c6eb51 Mon Sep 17 00:00:00 2001 From: DWAA1660 Date: Tue, 13 Oct 2026 16:03:53 -0600 Subject: [PATCH] Script sockets: per-socket locks, a write deadline, and close interrupts (#1248) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every socket operation locked one global table and held it through the blocking write_all/flush, and streams had no write timeout (accepted ones were explicitly blocking). A peer that stopped reading kept the script blocked for good, and every other socket call, close included, waited behind it. - The table maps handles to Arc, each with its own Mutex. The table is locked only to look a socket up, add or remove one, never during network I/O, so one stalled socket doesn't hold up the others. - `write` sets a write timeout from the socket's `timeout` (10 s by default, as reads) and reports "write timed out … the other side isn't reading" instead of blocking. Failures to set the read or write timeout, or to make an accepted stream blocking, are reported instead of ignored. - `close` removes the socket and shuts its stream down through a second handle, so a call blocked on it fails at once. Tests (socket.rs, native only): a 64 MB write to a loopback peer that never reads fails with "timed out" after a 1 s timeout; while one socket is blocked writing, another answers at once, and closing the blocked one returns at once and ends its write. Both fail on main (the write never gives up; the second socket waits on the table lock). Co-Authored-By: Claude Opus 5.5 --- crates/script/src/socket.rs | 168 +++++++++++++++++++++++++++++++----- 1 file changed, 145 insertions(+), 23 deletions(-) diff --git a/crates/script/src/socket.rs b/crates/script/src/socket.rs index 6ac0d41d..4030332d 100644 --- a/crates/script/src/socket.rs +++ b/crates/script/src/socket.rs @@ -50,8 +50,8 @@ mod imp { mod imp { use std::collections::BTreeMap; use std::io::{ErrorKind, Read, Write}; - use std::net::{TcpListener, TcpStream, ToSocketAddrs}; - use std::sync::Mutex; + use std::net::{Shutdown, TcpListener, TcpStream, ToSocketAddrs}; + use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; use super::*; @@ -61,9 +61,18 @@ mod imp { Listener(TcpListener), } + /// One open socket. Reading, writing and polling lock only it, never the whole table, so a + /// socket blocked on a slow peer doesn't hold up the others (#1248). + struct Entry { + sock: Mutex, + /// A second handle on a stream, so `close` can shut it down while another call is + /// blocked reading or writing it. + wake: Option, + } + struct Table { next: u32, - socks: BTreeMap, + socks: BTreeMap>, } fn table() -> &'static Mutex { @@ -72,13 +81,31 @@ mod imp { } fn add(s: Sock) -> u32 { + let wake = match &s { + Sock::Stream { s, .. } => s.try_clone().ok(), + Sock::Listener(_) => None, + }; + let entry = Arc::new(Entry { sock: Mutex::new(s), wake }); let mut t = table().lock().unwrap_or_else(|e| e.into_inner()); let id = t.next; t.next += 1; - t.socks.insert(id, s); + t.socks.insert(id, entry); id } + /// The socket `id`, if open. The table is locked only for the lookup. + fn get(id: u32) -> Option> { + table().lock().unwrap_or_else(|e| e.into_inner()).socks.get(&id).cloned() + } + + /// 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()) { + let _ = w.shutdown(Shutdown::Both); + } + } + 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)) } @@ -98,8 +125,9 @@ mod imp { if utf8 { String::from_utf8_lossy(b).into_owned() } else { b.iter().map(|&c| c as char).collect() } } - /// Fill a stream's buffer until `done(buf)` or the timeout / end of stream. - fn fill(s: &mut TcpStream, buf: &mut Vec, eof: &mut bool, limit: Duration, done: impl Fn(&[u8]) -> bool) { + /// Fill a stream's buffer until `done(buf)` or the timeout / end of stream. Fails only when + /// the read deadline can't be set. + fn fill(s: &mut TcpStream, buf: &mut Vec, eof: &mut bool, limit: Duration, done: impl Fn(&[u8]) -> bool) -> std::io::Result<()> { let end = Instant::now() + limit; let mut chunk = [0u8; 4096]; while !done(buf) && !*eof { @@ -107,7 +135,7 @@ mod imp { if left.is_zero() { break; } - let _ = s.set_read_timeout(Some(left.max(Duration::from_millis(1)))); + s.set_read_timeout(Some(left.max(Duration::from_millis(1))))?; match s.read(&mut chunk) { Ok(0) => *eof = true, Ok(n) => buf.extend_from_slice(&chunk[..n]), @@ -119,6 +147,7 @@ mod imp { Err(_) => *eof = true, } } + Ok(()) } pub(super) fn run(op: &str, id: u32, p: &J) -> J { @@ -156,29 +185,27 @@ mod imp { Err(e) => err(format!("cannot listen on port {port}: {e}")), } } + "close" => { + close(id); + json!({}) + } _ => { - let mut t = table().lock().unwrap_or_else(|e| e.into_inner()); - if op == "close" { - t.socks.remove(&id); - return json!({}); - } - let Some(sock) = t.socks.get_mut(&id) else { + let Some(entry) = get(id) else { return match op { "connected" => json!({"value": false}), "eof" => json!({"value": true}), _ => err("the socket is not open"), }; }; - match (op, sock) { + let mut sock = entry.sock.lock().unwrap_or_else(|e| e.into_inner()); + match (op, &mut *sock) { ("poll", Sock::Listener(l)) => match l.accept() { Ok((s, peer)) => { - let _ = s.set_nonblocking(false); + if let Err(e) = s.set_nonblocking(false) { + return err(format!("cannot use the connection from {peer}: {e}")); + } let host = peer.to_string(); - let new = Sock::Stream { s, buf: vec![], eof: false }; - let nid = t.next; - t.next += 1; - t.socks.insert(nid, new); - json!({"id": nid, "host": host}) + json!({"id": add(Sock::Stream { s, buf: vec![], eof: false }), "host": host}) } Err(_) => json!({"id": null}), }, @@ -202,19 +229,25 @@ mod imp { let count = p.get("count").and_then(J::as_i64).unwrap_or(-1); let limit = timeout(p); if count < 0 { - fill(s, buf, eof, limit, |_| false); + if let Err(e) = fill(s, buf, eof, limit, |_| false) { + return err(format!("cannot set the read timeout: {e}")); + } let all = std::mem::take(buf); json!({"data": decode(&all, utf8)}) } else { let n = count as usize; - fill(s, buf, eof, limit, |b| b.len() >= n); + if let Err(e) = fill(s, buf, eof, limit, |b| b.len() >= n) { + return err(format!("cannot set the read timeout: {e}")); + } let take = n.min(buf.len()); let out: Vec = buf.drain(..take).collect(); json!({"data": decode(&out, utf8)}) } } ("readln", Sock::Stream { s, buf, eof }) => { - fill(s, buf, eof, timeout(p), |b| b.contains(&b'\n')); + if let Err(e) = fill(s, buf, eof, timeout(p), |b| b.contains(&b'\n')) { + return err(format!("cannot set the read timeout: {e}")); + } let line: Vec = match buf.iter().position(|&c| c == b'\n') { Some(i) => { let mut l: Vec = buf.drain(..=i).collect(); @@ -229,9 +262,19 @@ mod imp { json!({"data": decode(&line, utf8)}) } ("write", Sock::Stream { s, .. }) => { + // 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()) { 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)", + limit.as_secs_f64() + )), Err(e) => err(e), } } @@ -242,3 +285,82 @@ mod imp { } } } + +#[cfg(test)] +#[cfg(not(target_arch = "wasm32"))] +mod tests { + use std::net::{TcpListener, TcpStream}; + use std::sync::mpsc; + use std::time::{Duration, Instant}; + + use serde_json::{Value as J, json}; + + use super::imp::run; + + /// A loopback peer that accepts one connection and never reads from it; returns its + /// `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(); + let host = l.local_addr().unwrap().to_string(); + let (stop, wait) = mpsc::channel::<()>(); + std::thread::spawn(move || { + let (_s, _): (TcpStream, _) = l.accept().unwrap(); + let _ = wait.recv_timeout(Duration::from_secs(120)); + }); + (host, stop) + } + + 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. + fn big() -> String { + "x".repeat(64 * 1024 * 1024) + } + + /// Run `f` on a thread; its result, or `None` if it takes longer than `limit`. + fn within(limit: Duration, f: impl FnOnce() -> T + Send + 'static) -> Option { + let (tx, rx) = mpsc::channel(); + std::thread::spawn(move || { + let _ = tx.send(f()); + }); + rx.recv_timeout(limit).ok() + } + + #[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(); + 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()); + run("close", id, &J::Null); + } + + #[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(); + 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)); + // 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}"); + // Closing A returns at once and ends A's write. + within(Duration::from_secs(3), move || run("close", a, &J::Null)).expect("close returned while A was blocked"); + let r = blocked.recv_timeout(Duration::from_secs(5)).expect("A's write ended when A was closed"); + assert!(r.get("error").is_some(), "{r}"); + run("close", b, &J::Null); + } +}