fix: preserve sessions after client write resets

libshpool/src/daemon/shell.rs

@@ -423,24 +423,32 @@ impl SessionInner {
                     recv(args.heartbeat) -> request_id => {
                         let client_present = if let ClientConnectionMsg::New(conn) = &mut client_conn {
                             let chunk = Chunk { kind: ChunkKind::Heartbeat, buf: &[] };
-                            match chunk.write_to(&mut conn.sink).and_then(|_| conn.sink.flush()) {
+                            let write_result = test_hooks::client_write_error("heartbeat")
+                                .map_or_else(|| chunk.write_to(&mut conn.sink).and_then(|_| conn.sink.flush()), Err);
+                            match write_result {
                                 Ok(_) => {
                                     trace!("wrote heartbeat");
                                     true
                                 }
-                                Err(e) if e.kind() == io::ErrorKind::BrokenPipe => {
-                                    trace!("client hangup: {:?}", e);
-                                    false
-                                }
                                 Err(e) => {
-                                    error!("unexpected IO error while writing heartbeat: {}", e);
-                                    return Err(e).context("writing heartbeat")?;
+                                    info!("heartbeat write err, assuming hangup: {:?}", e);
+                                    // A failed client write must not terminate the
+                                    // session's long-lived output worker. Shut down
+                                    // every clone so client->shell also wakes up.
+                                    if let Err(err) = shutdown_socket(&conn.stream, net::Shutdown::Both) {
+                                        warn!("shutting down client after heartbeat error: {:?}", err);
+                                    }
+                                    false
                                 }
                             }
                         } else {
                             false
                         };
 
+                        if !client_present {
+                            client_conn = ClientConnectionMsg::Disconnect;
+                        }
+
                         test_hooks::emit("daemon-wrote-heartbeat");
 
                         args.heartbeat_ack.send((request_id.unwrap_or(0), client_present))
@@ -469,16 +477,18 @@ impl SessionInner {
                         }
 
                         let chunk = Chunk { kind: ChunkKind::MaybeSwitch, buf: &encoded[..] };
-                        match chunk.write_to(&mut conn.sink).and_then(|_| conn.sink.flush()) {
+                        let write_result = test_hooks::client_write_error("maybe-switch")
+                            .map_or_else(|| chunk.write_to(&mut conn.sink).and_then(|_| conn.sink.flush()), Err);
+                        match write_result {
                             Ok(_) => {
                                 trace!("wrote MaybeSwitch");
                             }
-                            Err(e) if e.kind() == io::ErrorKind::BrokenPipe => {
-                                trace!("writing MaybeSwitch: client hangup: {:?}", e);
-                            }
                             Err(e) => {
-                                error!("unexpected IO error while writing heartbeat: {}", e);
-                                return Err(e).context("writing MaybeSwitch")?;
+                                info!("MaybeSwitch write err, assuming hangup: {:?}", e);
+                                if let Err(err) = shutdown_socket(&conn.stream, net::Shutdown::Both) {
+                                    warn!("shutting down client after MaybeSwitch error: {:?}", err);
+                                }
+                                client_conn = ClientConnectionMsg::Disconnect;
                             }
                         }
                     }

libshpool/src/test_hooks.rs

@@ -45,6 +45,29 @@ pub fn emit(_event: &str) {
     // a no-op normally
 }
 
+/// Inject one client write reset per daemon for the regression tests. This
+/// leaves the peer open so the test also checks that shutdown wakes the input
+/// thread, rather than relying on the client to close its own socket.
+#[cfg(feature = "test_hooks")]
+pub fn client_write_error(kind: &str) -> Option<std::io::Error> {
+    use std::sync::atomic::{AtomicBool, Ordering};
+
+    static INJECTED: AtomicBool = AtomicBool::new(false);
+    if std::env::var("SHPOOL_TEST__CLIENT_WRITE_RESET").as_deref() == Ok(kind)
+        && !INJECTED.swap(true, Ordering::Relaxed)
+    {
+        emit("daemon-injected-client-write-error");
+        Some(std::io::Error::from_raw_os_error(libc::ECONNRESET))
+    } else {
+        None
+    }
+}
+
+#[cfg(not(feature = "test_hooks"))]
+pub fn client_write_error(_kind: &str) -> Option<std::io::Error> {
+    None
+}
+
 #[cfg(feature = "test_hooks")]
 pub fn scoped(event: &str) -> ScopedEvent {
     ScopedEvent::new(event)

shpool/tests/regression.rs

@@ -11,6 +11,68 @@ mod support;
 
 use crate::support::{daemon::DaemonArgs, tmpdir};
 
+fn client_write_reset_preserves_session(kind: &str) -> anyhow::Result<()> {
+    let mut daemon_proc = support::daemon::Proc::new(
+        "norc.toml",
+        DaemonArgs {
+            extra_env: vec![(String::from("SHPOOL_TEST__CLIENT_WRITE_RESET"), kind.to_owned())],
+            ..DaemonArgs::default()
+        },
+    )?;
+    let mut original_client = daemon_proc.attach("main", Default::default())?;
+    daemon_proc.await_event("daemon-bidi-stream-enter")?;
+
+    let sessions: serde_json::Value = serde_json::from_slice(&daemon_proc.list_json()?.stdout)?;
+    let shell_pid = sessions["sessions"][0]["shell_pid"].as_i64().unwrap();
+
+    if kind == "maybe-switch" {
+        assert!(daemon_proc.var_set("repro", "changed")?.status.success());
+    }
+    daemon_proc.await_event("daemon-injected-client-write-error")?;
+    daemon_proc.wait_until_list_matches(|out| out.contains("disconnected"))?;
+    // The peer did not close voluntarily. It must be released by the daemon.
+    original_client.proc.wait()?;
+    signal::kill(Pid::from_raw(shell_pid as i32), None)?;
+    assert!(
+        !fs::read_to_string(&daemon_proc.log_file)?.contains("error in shell->client:"),
+        "a client write reset must not terminate the session's output worker"
+    );
+
+    for _ in 0..3 {
+        let mut client = daemon_proc.attach("main", Default::default())?;
+        let mut output = client.line_matcher()?;
+        client.run_cmd("printf 'reset-%s\\n' recovered")?;
+        output.scan_until_re("reset-recovered$")?;
+
+        let sessions: serde_json::Value = serde_json::from_slice(&daemon_proc.list_json()?.stdout)?;
+        assert_eq!(sessions["sessions"][0]["shell_pid"].as_i64(), Some(shell_pid));
+        assert_eq!(sessions["sessions"][0]["status"], "Attached");
+
+        assert!(daemon_proc.detach(vec![String::from("main")])?.status.success());
+        client.proc.wait()?;
+        daemon_proc.wait_until_list_matches(|out| out.contains("disconnected"))?;
+    }
+    let sessions: serde_json::Value = serde_json::from_slice(&daemon_proc.list_json()?.stdout)?;
+    assert_eq!(sessions["sessions"][0]["attachments"], serde_json::json!([]));
+    assert!(daemon_proc.kill(vec![String::from("main")])?.status.success());
+    Ok(())
+}
+
+/// ECONNRESET on a heartbeat used to terminate the output worker while the
+/// shell stayed alive, permanently preventing subsequent attaches.
+#[test]
+#[timeout(15000)]
+fn heartbeat_write_reset_preserves_session() -> anyhow::Result<()> {
+    client_write_reset_preserves_session("heartbeat")
+}
+
+/// MaybeSwitch writes share the heartbeat's client-disconnect failure mode.
+#[test]
+#[timeout(15000)]
+fn maybe_switch_write_reset_preserves_session() -> anyhow::Result<()> {
+    client_write_reset_preserves_session("maybe-switch")
+}
+
 /// Regression test for a deadlock where the global `shells` mutex is held
 /// during `spawn_subshell` -> `wait_for_startup`. If `wait_for_startup`
 /// blocks (e.g. the shell never produces the startup sentinel), the mutex