Merge branch 'master' into set_gha_workflow_permissions

Author: benfred

src/dump.rs

@@ -1,3 +1,5 @@
+use std::io::Write;
+
 use anyhow::{Context, Error};
 use console::{style, Term};
 
@@ -8,64 +10,91 @@ use crate::stack_trace::StackTrace;
 use remoteprocess::Pid;
 
 pub fn print_traces(pid: Pid, config: &Config, parent: Option<Pid>) -> Result<(), Error> {
+    let stdout = std::io::stdout();
+    let mut out = stdout.lock();
+    write_traces(&mut out, pid, config, parent)
+}
+
+pub fn write_traces<W: Write>(
+    out: &mut W,
+    pid: Pid,
+    config: &Config,
+    parent: Option<Pid>,
+) -> Result<(), Error> {
     let mut process = PythonSpy::new(pid, config)?;
     if config.dump_json {
         let traces = process
             .get_stack_traces()
             .context("Failed to get stack traces")?;
-        println!("{}", serde_json::to_string_pretty(&traces)?);
+        writeln!(out, "{}", serde_json::to_string_pretty(&traces)?)?;
         return Ok(());
     }
 
-    println!(
+    writeln!(
+        out,
         "Process {}: {}",
         style(process.pid).bold().yellow(),
         process.process.cmdline()?.join(" ")
-    );
+    )?;
 
-    println!(
+    writeln!(
+        out,
         "Python v{} ({})",
         style(&process.version).bold(),
         style(process.process.exe()?).dim()
-    );
+    )?;
 
     if let Some(parentpid) = parent {
         let parentprocess = remoteprocess::Process::new(parentpid)?;
-        println!(
+        writeln!(
+            out,
             "Parent Process {}: {}",
             style(parentpid).bold().yellow(),
             parentprocess.cmdline()?.join(" ")
-        );
+        )?;
     }
-    println!();
+    writeln!(out)?;
     let traces = process
         .get_stack_traces()
         .context("Failed to get stack traces")?;
     for trace in traces.iter().rev() {
-        print_trace(trace, true);
-        if config.subprocesses {
-            for (childpid, parentpid) in process
-                .process
-                .child_processes()
-                .expect("failed to get subprocesses")
-            {
-                let term = Term::stdout();
-                let (_, width) = term.size();
-
-                println!("\n{}", &style("-".repeat(width as usize)).dim());
-                // child_processes() returns the whole process tree, since we're recursing here
-                // though we could end up printing grandchild processes multiple times. Limit down
-                // to just once
-                if parentpid == pid {
-                    print_traces(childpid, config, Some(parentpid))?;
-                }
+        write_trace(out, trace, true)?;
+    }
+
+    if config.subprocesses {
+        for (childpid, parentpid) in process
+            .process
+            .child_processes()
+            .expect("failed to get subprocesses")
+        {
+            let term = Term::stdout();
+            let (_, width) = term.size();
+
+            writeln!(out, "\n{}", &style("-".repeat(width as usize)).dim())?;
+            // child_processes() returns the whole process tree, since we're recursing here
+            // though we could end up printing grandchild processes multiple times. Limit down
+            // to just once
+            if parentpid == pid {
+                write_traces(out, childpid, config, Some(parentpid))?;
             }
         }
     }
     Ok(())
 }
 
+#[cfg(target_os = "linux")]
 pub fn print_trace(trace: &StackTrace, include_activity: bool) {
+    let stdout = std::io::stdout();
+    let mut out = stdout.lock();
+    // Swallow write errors here to preserve the old println! behaviour.
+    let _ = write_trace(&mut out, trace, include_activity);
+}
+
+pub fn write_trace<W: Write>(
+    out: &mut W,
+    trace: &StackTrace,
+    include_activity: bool,
+) -> Result<(), Error> {
     let thread_id = trace.format_threadid();
 
     let status = if include_activity {
@@ -78,15 +107,16 @@ pub fn print_trace(trace: &StackTrace, include_activity: bool) {
 
     match trace.thread_name.as_ref() {
         Some(name) => {
-            println!(
+            writeln!(
+                out,
                 "Thread {}{}: \"{}\"",
                 style(thread_id).bold().yellow(),
                 status,
                 name
-            );
+            )?;
         }
         None => {
-            println!("Thread {}{}", style(thread_id).bold().yellow(), status);
+            writeln!(out, "Thread {}{}", style(thread_id).bold().yellow(), status)?;
         }
     };
 
@@ -96,35 +126,38 @@ pub fn print_trace(trace: &StackTrace, include_activity: bool) {
             None => &frame.filename,
         };
         if frame.line != 0 {
-            println!(
+            writeln!(
+                out,
                 "    {} ({}:{})",
                 style(&frame.name).green(),
                 style(&filename).cyan(),
                 style(frame.line).dim()
-            );
+            )?;
         } else {
-            println!(
+            writeln!(
+                out,
                 "    {} ({})",
                 style(&frame.name).green(),
                 style(&filename).cyan()
-            );
+            )?;
         }
 
         if let Some(locals) = &frame.locals {
             let mut shown_args = false;
             let mut shown_locals = false;
             for local in locals {
                 if local.arg && !shown_args {
-                    println!("        {}", style("Arguments:").dim());
+                    writeln!(out, "        {}", style("Arguments:").dim())?;
                     shown_args = true;
                 } else if !local.arg && !shown_locals {
-                    println!("        {}", style("Locals:").dim());
+                    writeln!(out, "        {}", style("Locals:").dim())?;
                     shown_locals = true;
                 }
 
                 let repr = local.repr.as_deref().unwrap_or("?");
-                println!("            {}: {}", local.name, repr);
+                writeln!(out, "            {}: {}", local.name, repr)?;
             }
         }
     }
+    Ok(())
 }

tests/integration_test.rs

@@ -554,3 +554,84 @@ fn test_delayed_subprocess() {
         break;
     }
 }
+
+#[cfg(not(target_os = "freebsd"))]
+#[test]
+fn test_dump_subprocesses_no_duplication() {
+    #[cfg(target_os = "macos")]
+    {
+        if unsafe { libc::geteuid() } != 0 {
+            return;
+        }
+    }
+
+    use std::io::{BufRead, BufReader};
+
+    // Regression test for the dup-per-thread bug in dump.rs: spawning the
+    // target as a child of this test process makes the test the ptrace
+    // ancestor, so we don't need elevated privileges on Linux.
+    console::set_colors_enabled(false);
+
+    let mut child = std::process::Command::new("python")
+        .arg("./tests/scripts/dump_subprocesses_threads.py")
+        .stdout(std::process::Stdio::piped())
+        .stderr(std::process::Stdio::null())
+        .spawn()
+        .expect("failed to spawn parent script");
+
+    struct Killer(std::process::Child);
+    impl Drop for Killer {
+        fn drop(&mut self) {
+            let _ = self.0.kill();
+        }
+    }
+    let stdout = child.stdout.take().unwrap();
+    let mut child = Killer(child);
+
+    let mut parent_pid: Option<Pid> = None;
+    let mut child_pid: Option<Pid> = None;
+    let reader = BufReader::new(stdout);
+    for line in reader.lines() {
+        let line = line.expect("failed to read from parent stdout");
+        if let Some(rest) = line.strip_prefix("PID_PARENT=") {
+            parent_pid = Some(rest.trim().parse().unwrap());
+        } else if let Some(rest) = line.strip_prefix("PID_CHILD=") {
+            child_pid = Some(rest.trim().parse().unwrap());
+        } else if line.trim() == "READY" {
+            break;
+        }
+    }
+    let parent_pid = parent_pid.expect("parent script never reported PID_PARENT");
+    let child_pid = child_pid.expect("parent script never reported PID_CHILD");
+
+    // Give the child subprocess a moment to finish Python initialization.
+    std::thread::sleep(std::time::Duration::from_millis(500));
+
+    let config = Config {
+        subprocesses: true,
+        ..Default::default()
+    };
+    let mut buf: Vec<u8> = Vec::new();
+    py_spy::dump::write_traces(&mut buf, parent_pid, &config, None).expect("write_traces failed");
+    let out = String::from_utf8(buf).expect("dump output was not utf-8");
+
+    // Match at line start only, so "Parent Process <pid>:" headers don't count.
+    let parent_header = format!("Process {}:", parent_pid);
+    let child_header = format!("Process {}:", child_pid);
+    let count_line_start =
+        |needle: &str| -> usize { out.lines().filter(|line| line.starts_with(needle)).count() };
+    let parent_occurrences = count_line_start(&parent_header);
+    let child_occurrences = count_line_start(&child_header);
+    assert_eq!(
+        parent_occurrences, 1,
+        "expected parent process header once, got {}:\n{}",
+        parent_occurrences, out
+    );
+    assert_eq!(
+        child_occurrences, 1,
+        "expected child process header once, got {}:\n{}",
+        child_occurrences, out
+    );
+
+    let _ = child.0.kill();
+}

tests/scripts/dump_subprocesses_threads.py

@@ -0,0 +1,35 @@
+import os
+import subprocess
+import sys
+import threading
+import time
+
+
+def _worker():
+    while True:
+        time.sleep(0.05)
+
+
+def main():
+    for _ in range(4):
+        threading.Thread(target=_worker, daemon=True).start()
+
+    child = subprocess.Popen(
+        [sys.executable, "-c", "import time\nwhile True: time.sleep(1)"],
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+    )
+
+    print("PID_PARENT=%d" % os.getpid(), flush=True)
+    print("PID_CHILD=%d" % child.pid, flush=True)
+    print("READY", flush=True)
+
+    try:
+        child.wait()
+    finally:
+        if child.poll() is None:
+            child.kill()
+
+
+if __name__ == "__main__":
+    main()