daemon: treat heartbeat ack timeout like the send timeout

libshpool/src/daemon/server.rs

@@ -1111,13 +1111,18 @@ impl Server {
         let (tty_size_change_ack_tx, tty_size_change_ack_rx) = crossbeam_channel::bounded(0);
 
         let (heartbeat_tx, heartbeat_rx) = crossbeam_channel::bounded(0);
-        let (heartbeat_ack_tx, heartbeat_ack_rx) = crossbeam_channel::bounded(0);
+        // One slot rather than a rendezvous, so the shell->client thread can
+        // always deposit an ack and get back to its select loop. If the
+        // heartbeat thread gave up waiting, the ack it abandoned sits here and
+        // is discarded by request id on the next pass.
+        let (heartbeat_ack_tx, heartbeat_ack_rx) = crossbeam_channel::bounded(1);
 
         // We make this buffered to avoid blocking during a broadcast. There is
         // no ack chan so we can afford to buffer a bit.
         let (maybe_switch_tx, maybe_switch_rx) = crossbeam_channel::bounded(10);
 
         let shell_to_client_ctl = Arc::new(Mutex::new(shell::ShellToClientCtl {
+            next_heartbeat_id: std::sync::atomic::AtomicU64::new(0),
             client_connection: client_connection_tx,
             client_connection_ack: client_connection_ack_rx,
             tty_size_change: tty_size_change_tx,

libshpool/src/daemon/shell.rs

@@ -245,10 +245,12 @@ pub struct ShellToClientArgs {
     pub client_connection_ack: crossbeam_channel::Sender<ClientConnectionStatus>,
     pub tty_size_change: crossbeam_channel::Receiver<TtySize>,
     pub tty_size_change_ack: crossbeam_channel::Sender<()>,
-    pub heartbeat: crossbeam_channel::Receiver<()>,
+    // Carries the request id that the ack must echo back.
+    pub heartbeat: crossbeam_channel::Receiver<u64>,
     pub maybe_switch: crossbeam_channel::Receiver<MaybeSwitch>,
-    // true if the client is still live, false if it has hung up on us
-    pub heartbeat_ack: crossbeam_channel::Sender<bool>,
+    // The request id of the heartbeat this ack answers, and a flag that is
+    // true if the client is still live, false if it has hung up on us.
+    pub heartbeat_ack: crossbeam_channel::Sender<(u64, bool)>,
     pub child_exit_notifier: Arc<ExitNotifier>,
 }
 
@@ -410,7 +412,7 @@ impl SessionInner {
                             }
                         }
                     }
-                    recv(args.heartbeat) -> _ => {
+                    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()) {
@@ -431,7 +433,9 @@ impl SessionInner {
                             false
                         };
 
-                        args.heartbeat_ack.send(client_present)
+                        test_hooks::emit("daemon-wrote-heartbeat");
+
+                        args.heartbeat_ack.send((request_id.unwrap_or(0), client_present))
                             .context("sending heartbeat ack")?;
                     }
                     recv(args.maybe_switch) -> maybe_switch => {
@@ -956,7 +960,7 @@ impl SessionInner {
             .spawn_scoped(scope, move || -> anyhow::Result<()> {
                 let _s1 = span!(Level::INFO, "heartbeat", s = self.name, cid = conn_id).entered();
 
-                loop {
+                'heartbeat: loop {
                     trace!("checking stop_rx");
                     let stop_early = common::sleep_unless(
                         consts::HEARTBEAT_DURATION,
@@ -969,9 +973,14 @@ impl SessionInner {
                     }
                     {
                         let shell_to_client_ctl = self.shell_to_client_ctl.lock();
+
+                        let request_id = shell_to_client_ctl
+                            .next_heartbeat_id
+                            .fetch_add(1, Ordering::Relaxed);
+
                         match shell_to_client_ctl
                             .heartbeat
-                            .send_timeout((), SHELL_TO_CLIENT_CTL_TIMEOUT)
+                            .send_timeout(request_id, SHELL_TO_CLIENT_CTL_TIMEOUT)
                         {
                             // If the channel is disconnected, it means that the shell exited and
                             // the shell->client process exited cleanly. We should not raise a
@@ -989,16 +998,28 @@ impl SessionInner {
                             }
                             _ => {}
                         }
-                        let client_present = match shell_to_client_ctl
-                            .heartbeat_ack
-                            .recv_timeout(SHELL_TO_CLIENT_CTL_TIMEOUT)
-                        {
-                            // If the channel is disconnected, it means that the shell exited and
-                            // the shell->client process exited cleanly. We should not raise a
-                            // ruckus.
-                            Err(crossbeam_channel::RecvTimeoutError::Disconnected) => return Ok(()),
-                            Err(e) => return Err(e).context("waiting for heartbeat ack"),
-                            Ok(client_present) => client_present,
+                        let client_present = loop {
+                            match shell_to_client_ctl
+                                .heartbeat_ack
+                                .recv_timeout(SHELL_TO_CLIENT_CTL_TIMEOUT)
+                            {
+                                // If the channel is disconnected, it means that the shell exited
+                                // and the shell->client process exited cleanly. We should not
+                                // raise a ruckus.
+                                Err(crossbeam_channel::RecvTimeoutError::Disconnected) => {
+                                    return Ok(())
+                                }
+                                // Like the send timeout above, a slow ack just means
+                                // the shell->client thread is busy, not that the
+                                // client is gone. A dead client still gets noticed
+                                // when the write to it fails.
+                                Err(crossbeam_channel::RecvTimeoutError::Timeout) => {
+                                    continue 'heartbeat
+                                }
+                                // An ack for a request we already gave up on.
+                                Ok((ack_id, _)) if ack_id != request_id => continue,
+                                Ok((_, client_present)) => break client_present,
+                            }
                         };
                         if !client_present {
                             // Bail from the thread to get the rest of the
@@ -1100,11 +1121,19 @@ pub struct ShellToClientCtl {
     pub tty_size_change_ack: crossbeam_channel::Receiver<()>,
 
     // A control channel telling the shell->client thread to issue
-    // a heartbeat to check if the client is still listening.
-    pub heartbeat: crossbeam_channel::Sender<()>,
-    // True if the client is still listening, false if it has hung up
-    // on us.
-    pub heartbeat_ack: crossbeam_channel::Receiver<bool>,
+    // a heartbeat to check if the client is still listening. The payload
+    // is a request id, which the ack echoes back.
+    pub heartbeat: crossbeam_channel::Sender<u64>,
+    // The request id of the heartbeat this ack answers, and a flag that is
+    // true if the client is still listening, false if it has hung up on us.
+    //
+    // The heartbeat thread gives up waiting for an ack that takes too long,
+    // so a late ack can still turn up afterwards. The id is what lets the
+    // next read tell that ack apart from its own and discard it.
+    pub heartbeat_ack: crossbeam_channel::Receiver<(u64, bool)>,
+
+    /// The id to give the next heartbeat request.
+    pub next_heartbeat_id: std::sync::atomic::AtomicU64,
 
     /// A control channel telling the shell->client thread to
     /// broadcast the given MaybeSwitch. There is no ack channel

shpool/tests/regression.rs

@@ -301,6 +301,43 @@ fn concurrent_attach_to_existing_session_race() -> anyhow::Result<()> {
     Ok(())
 }
 
+/// Regression test for a slow heartbeat ack killing the session. The
+/// shell->client thread can be busy when the heartbeat thread asks it for an
+/// ack (generating a large session restore buffer is the usual cause). Treating
+/// that as fatal tore down the thread scope along with the session's
+/// shell->client thread, leaving a session that could never be attached again.
+///
+/// We stall the shell->client thread right after it writes a heartbeat but
+/// before it acks, hold it past the ack timeout, then check the session still
+/// works.
+#[test]
+#[timeout(30000)]
+fn slow_heartbeat_ack_does_not_wedge_session() -> anyhow::Result<()> {
+    let mut daemon_proc = support::daemon::Proc::new("norc.toml", DaemonArgs::default())
+        .context("starting daemon proc")?;
+
+    let mut attach_proc =
+        daemon_proc.attach("sh1", Default::default()).context("starting attach proc")?;
+    daemon_proc.await_event("daemon-bidi-stream-enter")?;
+
+    daemon_proc.send_event_command("pause-at daemon-wrote-heartbeat")?;
+    daemon_proc.await_event("paused-at daemon-wrote-heartbeat")?;
+
+    // SHELL_TO_CLIENT_CTL_TIMEOUT is 300ms, so this guarantees the ack recv
+    // times out at least once.
+    std::thread::sleep(Duration::from_millis(600));
+
+    daemon_proc.send_event_command("release daemon-wrote-heartbeat")?;
+
+    // On buggy code the heartbeat thread has already returned an error, the
+    // scope has unwound, and this command never produces output.
+    let mut line_matcher = attach_proc.line_matcher()?;
+    attach_proc.run_cmd("echo still-alive")?;
+    line_matcher.scan_until_re("still-alive$")?;
+
+    Ok(())
+}
+
 /// Regression test for a bug where shpool would abort the attach process
 /// if the MOTD pager exited normally (e.g. less EOF). It should transition
 /// to the shell instead.