fix: preserve sessions after client write resets

libshpool/src/daemon/shell.rs

@@ -476,24 +476,35 @@ 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 = if test_hooks::maybe_fail("heartbeat-write") {
+                                Err(io::Error::from_raw_os_error(libc::ECONNRESET))
+                            } else {
+                                chunk.write_to(&mut conn.sink).and_then(|_| conn.sink.flush())
+                            };
+                            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))
@@ -522,16 +533,21 @@ 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 = if test_hooks::maybe_fail("maybe-switch-write") {
+                            Err(io::Error::from_raw_os_error(libc::ECONNRESET))
+                        } else {
+                            chunk.write_to(&mut conn.sink).and_then(|_| conn.sink.flush())
+                        };
+                        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,65 @@ pub fn emit(_event: &str) {
     // a no-op normally
 }
 
+/// Inject failures on selected calls to a named failure point. Configure
+/// SHPOOL_TEST__FAILURE_POINTS with comma-separated `name:n_1:n_2:...` entries.
+/// Calls are counted from 1, independently for each point, across the process.
+/// For example, `heartbeat-write:2:4` fails on the second and fourth calls.
+#[cfg(feature = "test_hooks")]
+pub fn maybe_fail(point: &str) -> bool {
+    use std::{collections::HashMap, sync::OnceLock};
+
+    #[derive(Default)]
+    struct FailurePoint {
+        calls: usize,
+        fail_on: HashSet<usize>,
+    }
+
+    static FAILURE_POINTS: OnceLock<Mutex<HashMap<String, FailurePoint>>> = OnceLock::new();
+    let points = FAILURE_POINTS.get_or_init(|| {
+        let mut points = HashMap::new();
+        for spec in std::env::var("SHPOOL_TEST__FAILURE_POINTS")
+            .unwrap_or_default()
+            .split(',')
+            .map(str::trim)
+            .filter(|spec| !spec.is_empty())
+        {
+            let (name, calls) =
+                spec.split_once(':').expect("failure points must use name:call[:call...]");
+            let name = name.trim();
+            assert!(!name.is_empty(), "failure point names must not be empty");
+            let failure = points.entry(name.to_owned()).or_insert_with(FailurePoint::default);
+            for call in calls.split(':') {
+                let call = call
+                    .trim()
+                    .parse::<usize>()
+                    .expect("failure point calls must be positive integers");
+                assert!(call > 0, "failure point calls are counted from 1");
+                failure.fail_on.insert(call);
+            }
+        }
+        Mutex::new(points)
+    });
+    let injected = {
+        let mut points = points.lock();
+        if let Some(failure) = points.get_mut(point) {
+            failure.calls += 1;
+            failure.fail_on.remove(&failure.calls)
+        } else {
+            false
+        }
+    };
+    if injected {
+        emit(&format!("injected-failure-{point}"));
+    }
+    injected
+}
+
+#[cfg(not(feature = "test_hooks"))]
+pub fn maybe_fail(_point: &str) -> bool {
+    false
+}
+
 #[cfg(feature = "test_hooks")]
 pub fn scoped(event: &str) -> ScopedEvent {
     ScopedEvent::new(event)
@@ -223,3 +282,68 @@ impl TestHookServer {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    #[test]
+    fn failure_points_from_environment() {
+        // Use a subprocess to isolate the environment and process-wide state
+        // from other tests, without changing the environment of running
+        // threads.
+        if std::env::var_os("SHPOOL_TEST__FAILURE_POINTS_CHILD").is_none() {
+            let status = std::process::Command::new(std::env::current_exe().unwrap())
+                .args(["--exact", "test_hooks::tests::failure_points_from_environment"])
+                .env("SHPOOL_TEST__FAILURE_POINTS_CHILD", "1")
+                .env(
+                    "SHPOOL_TEST__FAILURE_POINTS",
+                    "first:4:2:4, second:1:3,,concurrent:2:4:7,other-concurrent:1:8",
+                )
+                .status()
+                .unwrap();
+            assert!(status.success());
+            return;
+        }
+
+        assert!(!super::maybe_fail("unknown"));
+        assert!(!super::maybe_fail(""));
+        let enabled = cfg!(feature = "test_hooks");
+        for call in 1..=6 {
+            assert_eq!(super::maybe_fail("first"), enabled && [2, 4].contains(&call));
+            assert_eq!(super::maybe_fail("second"), enabled && [1, 3].contains(&call));
+        }
+
+        let workers: Vec<_> = (0..8)
+            .map(|_| {
+                std::thread::spawn(|| {
+                    (super::maybe_fail("concurrent"), super::maybe_fail("other-concurrent"))
+                })
+            })
+            .collect();
+        let results: Vec<_> = workers.into_iter().map(|worker| worker.join().unwrap()).collect();
+        assert_eq!(results.iter().filter(|(first, _)| *first).count(), if enabled { 3 } else { 0 });
+        assert_eq!(
+            results.iter().filter(|(_, second)| *second).count(),
+            if enabled { 2 } else { 0 }
+        );
+        assert!(!super::maybe_fail("concurrent"));
+        assert!(!super::maybe_fail("other-concurrent"));
+    }
+
+    #[test]
+    fn failure_points_disabled_by_default() {
+        if std::env::var_os("SHPOOL_TEST__FAILURE_POINTS_CHILD").is_none() {
+            let status = std::process::Command::new(std::env::current_exe().unwrap())
+                .args(["--exact", "test_hooks::tests::failure_points_disabled_by_default"])
+                .env("SHPOOL_TEST__FAILURE_POINTS_CHILD", "1")
+                .env_remove("SHPOOL_TEST__FAILURE_POINTS")
+                .status()
+                .unwrap();
+            assert!(status.success());
+            return;
+        }
+
+        for _ in 0..3 {
+            assert!(!super::maybe_fail("heartbeat-write"));
+        }
+    }
+}

shpool/tests/regression.rs

@@ -11,6 +11,58 @@ mod support;
 
 use crate::support::{daemon::DaemonArgs, tmpdir};
 
+/// A MaybeSwitch write reset must disconnect the client while preserving the
+/// output worker and original shell for subsequent attaches.
+#[test]
+#[timeout(15000)]
+fn client_write_reset_preserves_session() -> anyhow::Result<()> {
+    let mut daemon_proc = support::daemon::Proc::new(
+        "norc.toml",
+        DaemonArgs {
+            extra_env: vec![(
+                String::from("SHPOOL_TEST__FAILURE_POINTS"),
+                String::from("maybe-switch-write:1"),
+            )],
+            ..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();
+
+    assert!(daemon_proc.var_set("repro", "changed")?.status.success());
+    daemon_proc.await_event("injected-failure-maybe-switch-write")?;
+    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(())
+}
+
 /// 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