From 7b608b2dab589e5a403742bdfa263acf168dc683 Mon Sep 17 00:00:00 2001 From: bflatastic Date: Thu, 5 Nov 2026 06:57:04 +0700 Subject: [PATCH] Control server: close connections gracefully, so a reply is not lost to a reset A connection ended (a junk line, an idle or stalled peer, the app going away) by dropping its socket. On Windows, a reset can follow and discard replies the peer has not read yet, which made tcp_connection_is_closed_after_junk fail at random in CI with ConnectionReset. serve_socket now shuts the write side down and reads what the peer still sends until it closes, for at most a second and 1 MiB, before dropping the socket. The test goes through serve_socket, the path the app uses. --- apps/vectorcraft/src/control_server.rs | 28 ++++++++++++++++++++ apps/vectorcraft/src/control_server/tests.rs | 7 +++-- 2 files changed, 31 insertions(+), 4 deletions(-) diff --git a/apps/vectorcraft/src/control_server.rs b/apps/vectorcraft/src/control_server.rs index 508ab0daf..1a44eba28 100644 --- a/apps/vectorcraft/src/control_server.rs +++ b/apps/vectorcraft/src/control_server.rs @@ -152,6 +152,7 @@ fn spawn_connection( fn serve_socket(stream: TcpStream, timeouts: ConnectionTimeouts, handle: impl FnMut(String, Value) -> Option) { let Ok(read) = stream.try_clone() else { return }; + let closer = stream.try_clone().ok(); let read = DeadlineReader::new(read, timeouts.idle, timeouts.line); // Replies reach the writer in [`REPLY_CHUNK`] pieces ([`write_line`]), so checking the remaining // budget costs one socket option call per piece, not one per JSON token. @@ -166,6 +167,33 @@ fn serve_socket(stream: TcpStream, timeouts: ConnectionTimeouts, handle: impl Fn }, DeadlineWriter::begin_reply, ); + if let Some(s) = closer { + close_gracefully(s); + } +} + +/// How long [`close_gracefully`] waits for the peer to close its side. +const CLOSE_DRAIN: Duration = Duration::from_secs(1); + +/// The most [`close_gracefully`] reads from a peer that keeps sending. +const MAX_DRAIN: usize = 1 << 20; + +/// Ends a connection: no more replies (the peer reads them all, then end of stream), then what the +/// peer still sends is read and dropped until it closes too, for at most [`CLOSE_DRAIN`] and +/// [`MAX_DRAIN`] bytes. On Windows, dropping a socket with unread input resets the connection, and +/// the reset can discard replies the peer has not read yet. +fn close_gracefully(mut s: TcpStream) { + // Each step is best effort: a peer that is already gone needs none of them. + let _ = s.shutdown(std::net::Shutdown::Write); + let _ = s.set_read_timeout(Some(CLOSE_DRAIN)); + let deadline = std::time::Instant::now() + CLOSE_DRAIN; + let (mut buf, mut left) = ([0u8; 4096], MAX_DRAIN); + while left > 0 && std::time::Instant::now() < deadline { + match s.read(&mut buf) { + Ok(0) | Err(_) => break, + Ok(n) => left = left.saturating_sub(n), + } + } } /// One parsed line. diff --git a/apps/vectorcraft/src/control_server/tests.rs b/apps/vectorcraft/src/control_server/tests.rs index d991fd602..7aa290e01 100644 --- a/apps/vectorcraft/src/control_server/tests.rs +++ b/apps/vectorcraft/src/control_server/tests.rs @@ -334,13 +334,12 @@ fn tcp_connection_is_closed_after_junk() { let addr = listener.local_addr().unwrap(); let server = std::thread::spawn(move || { let (s, _) = listener.accept().unwrap(); - let read = s.try_clone().unwrap(); - serve_lines(read, s, |_, _| Some(json!({"ok": true}))); + serve_socket(s, ConnectionTimeouts::default(), |_, _| Some(json!({"ok": true}))); }); let mut c = TcpStream::connect(addr).unwrap(); c.set_read_timeout(Some(WAIT)).unwrap(); - // Nothing after the junk line: on Windows, closing a socket with unread input resets the - // connection, and the reset can discard the reply before the client reads it. + // The server closes gracefully ([`close_gracefully`]): the client reads both replies and then + // the end of the stream, not a reset. c.write_all(b"{\"id\":7,\"method\":\"x\"}\nhello\n").unwrap(); let mut text = String::new(); c.read_to_string(&mut text).unwrap(); // returns at EOF: the server closed the connection