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); + } +}