Close the client side when the upstream closes, in inferproxy

Implemented-By: OpenCode session (model recorded in docs/implementer-log.md)

crates/inferproxy/src/lib.rs

@@ -108,28 +108,34 @@ pub fn serve(listener: UnixListener, upstream: String, limits: Limits) -> io::Re
 }
 
 fn handle(client: UnixStream, upstream: String, open: Arc<AtomicUsize>) {
-    let _guard = OpenGuard::new(open);
     let server = match TcpStream::connect(&upstream) {
         Ok(s) => s,
         Err(_) => return,
     };
-    let _ = forward(&client, &server);
-    drop(_guard);
+    let c2s = {
+        let _guard = OpenGuard::new(open);
+        match forward(&client, &server) {
+            Ok(c2s) => c2s,
+            Err(_) => return,
+        }
+    };
+    let _ = client.shutdown(Shutdown::Both);
+    let _ = server.shutdown(Shutdown::Both);
+    let _ = c2s.join();
 }
 
-fn forward(client: &UnixStream, server: &TcpStream) -> io::Result<()> {
+fn forward(client: &UnixStream, server: &TcpStream) -> io::Result<thread::JoinHandle<()>> {
     let mut c2s_in = client.try_clone()?;
     let mut s2c_out = client.try_clone()?;
     let mut s2c_in = server.try_clone()?;
     let mut c2s_out = server.try_clone()?;
-    thread::scope(|s| {
-        s.spawn(move || {
-            let _ = copy(&mut c2s_in, &mut c2s_out);
-            let _ = c2s_out.shutdown(Shutdown::Write);
-        });
-        s.spawn(move || {
-            let _ = copy(&mut s2c_in, &mut s2c_out);
-        });
+    let c2s = thread::spawn(move || {
+        let _ = copy(&mut c2s_in, &mut c2s_out);
+        let _ = c2s_out.shutdown(Shutdown::Write);
     });
-    Ok(())
+    let _ = thread::spawn(move || {
+        let _ = copy(&mut s2c_in, &mut s2c_out);
+    })
+    .join();
+    Ok(c2s)
 }

crates/inferproxy/tests/forward.rs

@@ -154,3 +154,44 @@ fn an_unreachable_upstream_closes_the_client() {
     let path = start("127.0.0.1:1".to_string(), Limits::default());
     assert!(was_refused(&path));
 }
+
+/// An upstream that answers at once and closes, whether or not the client has finished sending.
+/// This is what `llama-server` does, and what a server that dies mid-answer looks like.
+fn answer_and_close_upstream() -> String {
+    let listener = TcpListener::bind("127.0.0.1:0").unwrap();
+    let addr = listener.local_addr().unwrap().to_string();
+    thread::spawn(move || {
+        for stream in listener.incoming() {
+            let mut stream = stream.unwrap();
+            let mut first = [0u8; 1];
+            let _ = stream.read(&mut first);
+            stream.write_all(b"answer").unwrap();
+            // dropping `stream` closes it
+        }
+    });
+    addr
+}
+
+#[test]
+fn the_upstreams_close_reaches_a_client_that_is_still_sending() {
+    // The HTTP client in loopd never half-closes: it keeps its sending side open until it drops
+    // the connection. When the server closes, the proxy must close towards the client at once,
+    // or the client only learns of a dead server from its own timeout.
+    let path = start(answer_and_close_upstream(), Limits::default());
+    let mut s = UnixStream::connect(&path).unwrap();
+    s.set_read_timeout(Some(Duration::from_secs(2))).unwrap();
+    s.write_all(b"request").unwrap();
+    let started = std::time::Instant::now();
+    let mut got = Vec::new();
+    let result = s.read_to_end(&mut got);
+    assert!(
+        result.is_ok(),
+        "the read must end with EOF, not a timeout: {result:?}"
+    );
+    assert_eq!(got, b"answer");
+    assert!(
+        started.elapsed() < Duration::from_millis(1000),
+        "EOF took {:?}",
+        started.elapsed()
+    );
+}