diff --git a/libshpool/src/daemon/server.rs b/libshpool/src/daemon/server.rs index 20814ea0..a5e1bfdf 100644 --- a/libshpool/src/daemon/server.rs +++ b/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, diff --git a/libshpool/src/daemon/shell.rs b/libshpool/src/daemon/shell.rs index ad65b0fb..dd0e2e2f 100644 --- a/libshpool/src/daemon/shell.rs +++ b/libshpool/src/daemon/shell.rs @@ -245,10 +245,12 @@ pub struct ShellToClientArgs { pub client_connection_ack: crossbeam_channel::Sender, pub tty_size_change: crossbeam_channel::Receiver, 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, pub maybe_switch: crossbeam_channel::Receiver, - // true if the client is still live, false if it has hung up on us - pub heartbeat_ack: crossbeam_channel::Sender, + // 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, } @@ -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, + // 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, + // 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 diff --git a/shpool/tests/regression.rs b/shpool/tests/regression.rs index e8a8e14d..286f9a50 100644 --- a/shpool/tests/regression.rs +++ b/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.