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

@@ -152,6 +152,7 @@ fn spawn_connection(
 
 fn serve_socket(stream: TcpStream, timeouts: ConnectionTimeouts, handle: impl FnMut(String, Value) -> Option<Value>) {
     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.

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