Skip to content

fix: keep stalled clients from wedging session attach #399

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
crystal150 wants to merge 1 commit into
base: master
Choose a base branch
from
Open

fix: keep stalled clients from wedging session attach #399

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
166 changes: 138 additions & 28 deletions libshpool/src/daemon/shell.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,11 @@ const SHELL_TO_CLIENT_POLL_MS: u16 = 50;
// shell->client thread.
const SHELL_TO_CLIENT_CTL_TIMEOUT: time::Duration = time::Duration::from_millis(300);

// Keep socket backpressure out of the shell->client control loop. At the
// largest protocol chunk size this allows a few MiB of temporary client lag.
const CLIENT_OUTPUT_QUEUE_DEPTH: usize = 256;
const CLIENT_OUTPUT_DRAIN_TIMEOUT: time::Duration = time::Duration::from_millis(100);

/// Lifecycle state tracking when sessions were last connected/disconnected and
/// by whom. The fields update in lockstep, so they live behind a single private
/// lock that only this type's methods take, exactly once per update or read.
Expand Down Expand Up @@ -168,7 +173,7 @@ pub struct SessionInner {
pub struct ClientConnection {
/// All output data should be written to this sink rather than
/// directly to the unix stream.
sink: io::BufWriter<UnixStream>,
sink: io::BufWriter<ClientWriter>,
/// The size of the client tty.
size: TtySize,
/// The raw unix socket stream. The shell->client thread should
Expand All @@ -177,6 +182,72 @@ pub struct ClientConnection {
stream: UnixStream,
}

struct ClientWriter {
queue: crossbeam_channel::Sender<ClientOutputMsg>,
}

enum ClientOutputMsg {
Data(Vec<u8>),
Flush(crossbeam_channel::Sender<()>),
}

impl ClientWriter {
fn new(mut stream: UnixStream) -> io::Result<Self> {
let (queue, pending) =
crossbeam_channel::bounded::<ClientOutputMsg>(CLIENT_OUTPUT_QUEUE_DEPTH);
thread::Builder::new().name(String::from("client-output-writer")).spawn(move || {
test_hooks::emit("daemon-client-output-writer-before-read");
while let Ok(msg) = pending.recv() {
match msg {
ClientOutputMsg::Data(buf) => {
if let Err(err) = stream.write_all(&buf).and_then(|_| stream.flush()) {
trace!("client output writer stopped: {}", err);
return;
}
}
ClientOutputMsg::Flush(ack) => {
let _ = ack.send(());
}
}
}
})?;
Ok(Self { queue })
}

fn wait_flushed(&self) -> bool {
let (ack, flushed) = crossbeam_channel::bounded(0);
if self
.queue
.send_timeout(ClientOutputMsg::Flush(ack), CLIENT_OUTPUT_DRAIN_TIMEOUT)
.is_err()
{
return false;
}
flushed.recv_timeout(CLIENT_OUTPUT_DRAIN_TIMEOUT).is_ok()
}
}

impl io::Write for ClientWriter {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
if buf.is_empty() {
return Ok(0);
}
match self.queue.try_send(ClientOutputMsg::Data(buf.to_vec())) {
Ok(()) => Ok(buf.len()),
Err(crossbeam_channel::TrySendError::Full(_)) => {
Err(io::Error::new(io::ErrorKind::WouldBlock, "client output queue is full"))
}
Err(crossbeam_channel::TrySendError::Disconnected(_)) => {
Err(io::Error::new(io::ErrorKind::BrokenPipe, "client output writer stopped"))
}
}
}

fn flush(&mut self) -> io::Result<()> {

@ethanpailes ethanpailes Jul 10, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is lying. flush must not return before the data is actually written to the file descriptor and flushed.

Ok(())
}
}

#[derive(Debug)]
pub enum ClientConnectionStatus {
/// The new session replaced an existing session client.
Expand Down Expand Up @@ -411,25 +482,29 @@ impl SessionInner {
}
}
recv(args.heartbeat) -> _ => {
let mut write_error = None;
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()) {
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")?;
write_error = Some(e);
false
}
}
} else {
false
};
if let Some(err) = write_error.as_ref() {
Self::disconnect_client_after_write_error(
&mut client_conn,
"heartbeat",
err,
);
}

args.heartbeat_ack.send(client_present)
.context("sending heartbeat ack")?;
Expand Down Expand Up @@ -457,17 +532,22 @@ impl SessionInner {
}

let chunk = Chunk { kind: ChunkKind::MaybeSwitch, buf: &encoded[..] };
match chunk.write_to(&mut conn.sink).and_then(|_| conn.sink.flush()) {
let write_error = match chunk
.write_to(&mut conn.sink)
.and_then(|_| conn.sink.flush())
{
Ok(_) => {
trace!("wrote MaybeSwitch");
None
}
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")?;
}
Err(e) => Some(e),
};
if let Some(err) = write_error.as_ref() {
Self::disconnect_client_after_write_error(
&mut client_conn,
"MaybeSwitch",
err,
);
}
}

Expand Down Expand Up @@ -498,6 +578,7 @@ impl SessionInner {
if do_reattach {
info!("executing reattach protocol");
let restore_buf = output_spool.restore_buffer();
let mut write_error = None;
if let (true, ClientConnectionMsg::New(conn)) =
(!restore_buf.is_empty(), &mut client_conn)
{
Expand All @@ -508,13 +589,23 @@ impl SessionInner {
let chunk = Chunk { kind: ChunkKind::Data, buf: block };

if let Err(err) = chunk.write_to(&mut conn.sink) {
warn!("err writing session-restore buf: {:?}", err);
write_error = Some(err);
break;
}
}
if let Err(err) = conn.sink.flush() {
warn!("err flushing session-restore: {:?}", err);
if write_error.is_none() {
if let Err(err) = conn.sink.flush() {
write_error = Some(err);
}
}
}
if let Some(err) = write_error.as_ref() {
Self::disconnect_client_after_write_error(
&mut client_conn,
"session restore",
err,
);
}
}

// TODO(ethan): what if poll times out on a tick when we have just
Expand Down Expand Up @@ -588,7 +679,7 @@ impl SessionInner {
output_spool.process(buf);
}

let mut reset_client_conn = false;
let mut write_error = None;
if let (ClientConnectionMsg::New(conn), true) =
(&mut client_conn, has_seen_prompt_sentinel)
{
Expand All @@ -607,14 +698,13 @@ impl SessionInner {
let write_result =
chunk.write_to(&mut conn.sink).and_then(|_| conn.sink.flush());
if let Err(err) = write_result {
info!("client_stream write err, assuming hangup: {:?}", err);
reset_client_conn = true;
write_error = Some(err);
} else {
test_hooks::emit("daemon-wrote-s2c-chunk");
}
}
if reset_client_conn {
client_conn = ClientConnectionMsg::Disconnect;
if let Some(err) = write_error.as_ref() {
Self::disconnect_client_after_write_error(&mut client_conn, "PTY output", err);
}
}
};
Expand All @@ -624,12 +714,16 @@ impl SessionInner {
.spawn(move || log_if_error("error in shell->client", closure()))?)
}

fn write_exit_chunk<W: io::Write>(mut sink: W, status: i32) {
fn write_exit_chunk(sink: &mut io::BufWriter<ClientWriter>, status: i32) {
let status_buf: [u8; 4] = status.to_le_bytes();
let chunk = Chunk { kind: ChunkKind::ExitStatus, buf: status_buf.as_slice() };
match chunk.write_to(&mut sink).and_then(|_| sink.flush()) {
match chunk.write_to(&mut *sink).and_then(|_| sink.flush()) {
Ok(_) => {
trace!("wrote exit status chunk");
if sink.get_ref().wait_flushed() {
trace!("wrote exit status chunk");
} else {
trace!("timed out draining exit status chunk");
}
}
Err(e) if e.kind() == io::ErrorKind::BrokenPipe => {
trace!("client hangup: {:?}", e);
Expand All @@ -640,6 +734,20 @@ impl SessionInner {
};
}

fn disconnect_client_after_write_error(
client_conn: &mut ClientConnectionMsg,
operation: &str,
err: &dyn std::fmt::Display,
) {
info!("{} write failed; detaching stalled client: {}", operation, err);
if let ClientConnectionMsg::New(conn) = client_conn {
if let Err(shutdown_err) = shutdown_socket(&conn.stream, net::Shutdown::Both) {
warn!("failed to shut down stalled client socket: {}", shutdown_err);
}
}
*client_conn = ClientConnectionMsg::Disconnect;
}

/// bidi_stream shuffles bytes between the subprocess and
/// the client connection. It returns true if the subprocess
/// has exited, and false if it is still running.
Expand All @@ -665,8 +773,10 @@ impl SessionInner {
client_stream.try_clone().context("creating client->shell client stream")?;
let shell_to_client_client_stream =
client_stream.try_clone().context("creating shell->client client stream handle")?;
let output_sink =
io::BufWriter::new(client_stream.try_clone().context("wrapping stream in bufwriter")?);
let output_stream = client_stream.try_clone().context("cloning client output stream")?;
let output_sink = io::BufWriter::new(
ClientWriter::new(output_stream).context("starting client output writer")?,
);

{
let _s = span!(Level::INFO, "initial_attach_lock(shell_to_client_ctl)").entered();
Expand Down
44 changes: 44 additions & 0 deletions shpool/tests/regression.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -215,6 +215,50 @@ fn client_eof_does_not_spin() -> anyhow::Result<()> {
Ok(())
}

/// A client can stay connected while its terminal stops consuming output. The
/// daemon must drop only that client before its socket write blocks the
/// shell-to-client control loop, then allow a fresh client to reattach.
#[test]
#[timeout(30000)]
fn stalled_output_client_does_not_wedge_session() -> anyhow::Result<()> {
let mut daemon_proc = support::daemon::Proc::new(
"norc.toml",
DaemonArgs { verbosity: 1, ..DaemonArgs::default() },
)
.context("starting daemon proc")?;
let marker = daemon_proc.tmp_dir.path().join("burst-finished");

daemon_proc.send_event_command("pause-at daemon-client-output-writer-before-read")?;

let mut stalled =
daemon_proc.attach("sh1", Default::default()).context("starting stalled attach proc")?;
daemon_proc.await_event("daemon-bidi-stream-enter")?;
daemon_proc.await_event("paused-at daemon-client-output-writer-before-read")?;

let burst = format!(
"i=0; while [ \"$i\" -lt 131072 ]; do printf \
'0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef\\n'; \
i=$((i + 1)); done; touch {}",
marker.display()
);
stalled.run_cmd(&burst)?;

daemon_proc.wait_until_list_matches(|out| out.contains("disconnected"))?;
support::wait_until(|| Ok(marker.exists()))?;

let mut reattached = daemon_proc
.attach("sh1", Default::default())
.context("reattaching after stalled client")?;
daemon_proc.await_event("daemon-bidi-stream-enter")?;
daemon_proc.await_event("paused-at daemon-client-output-writer-before-read")?;
daemon_proc.send_event_command("release daemon-client-output-writer-before-read")?;
let mut line_matcher = reattached.line_matcher()?;
reattached.run_cmd("echo transport-recovered")?;
line_matcher.scan_until_re("transport-recovered$")?;

Ok(())
}

/// Regression test for an EOF spin loop in the pager display thread.
/// If the client abruptly disconnects while viewing the MOTD pager, the daemon
/// should quickly detect EOF and abort the pager display.
Expand Down