fix: keep script shutdown independent of stdout writes

e2e-tests/tests/script_output_lifecycle.rs

@@ -0,0 +1,205 @@
+mod common;
+
+use anyhow::{ensure, Context, Result};
+use common::sandbox::SandboxHandle;
+use std::ffi::OsString;
+use std::os::fd::AsRawFd;
+use std::process::Stdio;
+use std::time::Duration;
+use tokio::io::{AsyncBufReadExt, BufReader};
+
+enum PipeShutdown {
+    Signal { signal: i32, merge_stderr: bool },
+    CloseReader { fill_pipe: bool },
+}
+
+async fn exercise_pipe_shutdown(shutdown: PipeShutdown) -> Result<()> {
+    common::init();
+    let sandbox = SandboxHandle::default_ghostscope()?;
+    if !sandbox.is_host_backend() {
+        eprintln!("Host pipe lifecycle test is covered by Standard E2E and host-to-private E2E");
+        return Ok(());
+    }
+    let binary = common::FIXTURES.get_test_binary("scalar_types_program")?;
+    let target = common::targets::TargetLauncher::binary(&binary)
+        .spawn()
+        .await?;
+    let stderr = tempfile::NamedTempFile::new()?;
+    let result = async {
+        let script = format!("trace scalar_anchor {{ print \"{}\"; }}", "x".repeat(1024));
+        let args: Vec<OsString> = [
+            "--no-log",
+            "--no-status",
+            "--debuginfod",
+            "off",
+            "--no-save-llvm-ir",
+            "--no-save-ebpf",
+            "--no-save-ast",
+            "--script-output",
+            "plain",
+            "--script-output-events-per-sec",
+            "0",
+            "--emit-ready-marker",
+            "PIPE_READY",
+            "-p",
+            &target.visible_pid_from(&sandbox)?.to_string(),
+            "-s",
+            &script,
+        ]
+        .into_iter()
+        .map(OsString::from)
+        .collect();
+        let launch = sandbox.ghostscope_runner_command(&args)?;
+        let mut command = tokio::process::Command::new(&launch.program);
+        command
+            .args(&launch.args)
+            .stdout(Stdio::piped())
+            .stderr(stderr.reopen()?)
+            .kill_on_drop(true);
+        if matches!(
+            shutdown,
+            PipeShutdown::Signal {
+                merge_stderr: true,
+                ..
+            }
+        ) {
+            // SAFETY: the child only calls async-signal-safe dup2 before exec;
+            // descriptors 1 and 2 are configured by Command above.
+            unsafe {
+                command.pre_exec(|| {
+                    if libc::dup2(libc::STDOUT_FILENO, libc::STDERR_FILENO) == -1 {
+                        return Err(std::io::Error::last_os_error());
+                    }
+                    Ok(())
+                });
+            }
+        }
+        let mut child = command.spawn()?;
+        let mut reader = BufReader::new(child.stdout.take().unwrap());
+        tokio::time::timeout(Duration::from_secs(30), async {
+            loop {
+                let mut line = String::new();
+                ensure!(
+                    reader.read_line(&mut line).await? != 0,
+                    "exited before ready marker"
+                );
+                if line.contains("PIPE_READY") {
+                    return Ok::<_, anyhow::Error>(());
+                }
+            }
+        })
+        .await
+        .context("waiting for tracing to become ready")??;
+
+        if !matches!(shutdown, PipeShutdown::CloseReader { fill_pipe: false }) {
+            // Poll the pipe's write readiness: fragmented pipe buffers can block
+            // writes before FIONREAD reaches the nominal byte capacity.
+            let pid = child
+                .id()
+                .context("GhostScope exited before pipe inspection")?;
+            let pipe = std::fs::OpenOptions::new()
+                .write(true)
+                .open(format!("/proc/{pid}/fd/1"))?;
+            tokio::time::timeout(Duration::from_secs(10), async {
+                loop {
+                    ensure!(
+                        child.try_wait()?.is_none(),
+                        "exited while filling stdout pipe"
+                    );
+                    let mut descriptor = libc::pollfd {
+                        fd: pipe.as_raw_fd(),
+                        events: libc::POLLOUT,
+                        revents: 0,
+                    };
+                    // SAFETY: descriptor is one initialized pollfd for the live pipe.
+                    let status = unsafe { libc::poll(&mut descriptor, 1, 0) };
+                    ensure!(status >= 0, "failed to inspect stdout pipe");
+                    if descriptor.revents & libc::POLLOUT == 0 {
+                        return Ok::<_, anyhow::Error>(());
+                    }
+                    tokio::time::sleep(Duration::from_millis(10)).await;
+                }
+            })
+            .await
+            .context("stdout pipe never filled")??;
+            // Leave enough events for the bounded userspace queue to fill too.
+            tokio::time::sleep(Duration::from_millis(300)).await;
+        }
+        let status = match shutdown {
+            PipeShutdown::Signal {
+                signal,
+                merge_stderr,
+            } => {
+                let pid = child
+                    .id()
+                    .context("GhostScope exited before shutdown signal")?;
+                // SAFETY: pid belongs to the live child retained above; kill takes scalar arguments.
+                ensure!(
+                    unsafe { libc::kill(pid as libc::pid_t, signal) } == 0,
+                    "failed to send signal {signal}"
+                );
+                // Keep the read end open without draining it until the child has exited.
+                let status = tokio::time::timeout(Duration::from_secs(2), child.wait())
+                    .await
+                    .context("shutdown signal was blocked by full stdout")??;
+                if !merge_stderr {
+                    ensure!(
+                        std::fs::read_to_string(stderr.path())?
+                            .contains("discarding pending script output"),
+                        "shutdown must disclose output abandoned after the drain deadline"
+                    );
+                }
+                drop(reader);
+                status
+            }
+            PipeShutdown::CloseReader { .. } => {
+                drop(reader);
+                tokio::time::timeout(Duration::from_secs(2), child.wait())
+                    .await
+                    .context("tracing continued after the stdout reader closed")??
+            }
+        };
+        ensure!(status.success(), "unexpected shutdown status: {status}");
+        Ok(())
+    }
+    .await;
+    target.terminate().await?;
+    result.with_context(|| std::fs::read_to_string(stderr.path()).unwrap_or_default())
+}
+
+#[tokio::test]
+async fn test_sigterm_stops_tracing_with_a_full_stdout_pipe() -> Result<()> {
+    exercise_pipe_shutdown(PipeShutdown::Signal {
+        signal: libc::SIGTERM,
+        merge_stderr: false,
+    })
+    .await
+}
+
+#[tokio::test]
+async fn test_sigint_stops_tracing_with_a_full_stdout_pipe() -> Result<()> {
+    exercise_pipe_shutdown(PipeShutdown::Signal {
+        signal: libc::SIGINT,
+        merge_stderr: false,
+    })
+    .await
+}
+
+#[tokio::test]
+async fn test_sigterm_stops_tracing_with_a_full_combined_stdout_stderr_pipe() -> Result<()> {
+    exercise_pipe_shutdown(PipeShutdown::Signal {
+        signal: libc::SIGTERM,
+        merge_stderr: true,
+    })
+    .await
+}
+
+#[tokio::test]
+async fn test_closed_stdout_reader_stops_tracing() -> Result<()> {
+    exercise_pipe_shutdown(PipeShutdown::CloseReader { fill_pipe: false }).await
+}
+
+#[tokio::test]
+async fn test_closed_stdout_reader_stops_tracing_with_a_full_queue() -> Result<()> {
+    exercise_pipe_shutdown(PipeShutdown::CloseReader { fill_pipe: true }).await
+}

ghostscope/src/cli/mod.rs

@@ -5,6 +5,7 @@ mod docs;
 mod dry_run;
 mod loading_reporter;
 pub mod script_output;
+mod script_output_writer;
 pub mod script_runtime;
 
 use crate::config::BpffsPruneArgs;

ghostscope/src/cli/script_output_writer.rs

@@ -0,0 +1,166 @@
+//! Bounded output delivery independent of the tracing and signal tasks.
+
+use std::fs::File;
+use std::io::{self, Write};
+use std::os::fd::AsFd;
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::Arc;
+use tokio::sync::{mpsc, oneshot};
+
+const CHUNK_BYTES: usize = 64 * 1024;
+const QUEUED_CHUNKS: usize = 8;
+
+pub(super) struct ScriptOutputWriter {
+    sender: Option<mpsc::Sender<Vec<u8>>>,
+    completed: oneshot::Receiver<io::Result<()>>,
+    cancelled: Arc<AtomicBool>,
+}
+
+impl ScriptOutputWriter {
+    pub(super) fn stdout() -> io::Result<Self> {
+        // Own a duplicate descriptor, never StdoutLock: Rust's exit-time stdout
+        // cleanup must remain free to run while a pipe write is blocked.
+        let fd = io::stdout().as_fd().try_clone_to_owned()?;
+        Self::new(File::from(fd))
+    }
+
+    pub(super) fn stderr() -> io::Result<Self> {
+        let fd = io::stderr().as_fd().try_clone_to_owned()?;
+        Self::new(File::from(fd))
+    }
+
+    fn new(mut writer: impl Write + Send + 'static) -> io::Result<Self> {
+        let (sender, mut receiver) = mpsc::channel::<Vec<u8>>(QUEUED_CHUNKS);
+        let (done, completed) = oneshot::channel();
+        let cancelled = Arc::new(AtomicBool::new(false));
+        let worker_cancelled = Arc::clone(&cancelled);
+        // A dedicated thread owns only this descriptor and bounded byte chunks.
+        // Blocking OS writes cannot be cancelled. Do not put them in Tokio's
+        // blocking pool, whose shutdown waits for outstanding writes to finish.
+        std::thread::Builder::new()
+            .name("ghostscope-output".into())
+            .spawn(move || {
+                let result = (|| {
+                    while let Some(bytes) = receiver.blocking_recv() {
+                        if worker_cancelled.load(Ordering::Acquire) {
+                            return Ok(());
+                        }
+                        writer.write_all(&bytes)?;
+                    }
+                    if !worker_cancelled.load(Ordering::Acquire) {
+                        writer.flush()?;
+                    }
+                    Ok(())
+                })();
+                let _ = done.send(result);
+            })?;
+        Ok(Self {
+            sender: Some(sender),
+            completed,
+            cancelled,
+        })
+    }
+
+    /// Backpressure is asynchronous and cancellable by the caller's signal select.
+    /// The queue holds at most 512 KiB, plus one chunk in the blocking writer.
+    pub(super) async fn write(&mut self, bytes: &[u8]) -> io::Result<()> {
+        for chunk in bytes.chunks(CHUNK_BYTES) {
+            if self
+                .sender
+                .as_ref()
+                .expect("output sender is present until drop")
+                .send(chunk.to_vec())
+                .await
+                .is_err()
+            {
+                return self.completed().await;
+            }
+        }
+        Ok(())
+    }
+
+    pub(super) async fn completed(&mut self) -> io::Result<()> {
+        (&mut self.completed)
+            .await
+            .map_err(|_| io::Error::other("script output worker stopped unexpectedly"))?
+    }
+
+    /// Deliver accepted chunks and close the descriptor. The caller must impose
+    /// a deadline: cancelling this future falls back to the non-blocking Drop.
+    pub(super) async fn finish(mut self) -> io::Result<()> {
+        self.sender.take();
+        self.completed().await
+    }
+}
+
+impl Drop for ScriptOutputWriter {
+    fn drop(&mut self) {
+        self.cancelled.store(true, Ordering::Release);
+        self.sender.take();
+        // Do not join a thread that may be blocked in the consumer's pipe.
+        // Session teardown releases all probes independently of this writer.
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use tokio::io::AsyncReadExt;
+
+    #[tokio::test]
+    async fn graceful_finish_delivers_all_queued_output() {
+        let (writer, reader) = std::os::unix::net::UnixStream::pair().unwrap();
+        reader.set_nonblocking(true).unwrap();
+        let mut reader = tokio::net::UnixStream::from_std(reader).unwrap();
+        let mut output = ScriptOutputWriter::new(writer).unwrap();
+        let expected: Vec<u8> = (0..CHUNK_BYTES * QUEUED_CHUNKS)
+            .map(|index| (index % 251) as u8)
+            .collect();
+        // Accept a full queue before allowing the consumer to read it.
+        output.write(&expected).await.unwrap();
+        let received = tokio::spawn(async move {
+            let mut bytes = Vec::new();
+            reader.read_to_end(&mut bytes).await.unwrap();
+            bytes
+        });
+        tokio::time::timeout(std::time::Duration::from_secs(2), output.finish())
+            .await
+            .expect("a healthy consumer must drain during shutdown")
+            .unwrap();
+        let received = tokio::time::timeout(std::time::Duration::from_secs(2), received)
+            .await
+            .expect("the worker must close its descriptor after draining")
+            .unwrap();
+        assert_eq!(
+            received.len(),
+            expected.len(),
+            "shutdown lost accepted output"
+        );
+        assert_eq!(received, expected);
+    }
+
+    #[tokio::test]
+    async fn finish_propagates_non_pipe_write_failures() {
+        let writer = File::options().write(true).open("/dev/full").unwrap();
+        let mut output = ScriptOutputWriter::new(writer).unwrap();
+        output.write(b"event\n").await.unwrap();
+        let error = tokio::time::timeout(std::time::Duration::from_secs(1), output.finish())
+            .await
+            .expect("a failed writer must finish shutdown")
+            .unwrap_err();
+        assert_eq!(error.raw_os_error(), Some(libc::ENOSPC));
+    }
+
+    #[tokio::test]
+    async fn closed_consumer_notifies_the_control_task() {
+        let (writer, reader) = std::os::unix::net::UnixStream::pair().unwrap();
+        drop(reader);
+        let mut output = ScriptOutputWriter::new(writer).unwrap();
+        output.write(b"event\n").await.unwrap();
+        let error = tokio::time::timeout(std::time::Duration::from_secs(1), output.completed())
+            .await
+            .expect("a failed writer must wake the control task")
+            .unwrap_err();
+        assert_eq!(error.kind(), io::ErrorKind::BrokenPipe);
+    }
+}