Allow attaching to existing macOS process

samply/resources/entitlement.xml

@@ -0,0 +1,10 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0">
+<dict>
+	<key>com.apple.security.cs.debugger</key>
+	<true/>
+	<key>com.apple.security.cs.disable-library-validation</key>
+	<true/>
+</dict>
+</plist>

samply/src/mac/process_launcher.rs

@@ -2,27 +2,97 @@ use std::collections::BTreeMap;
 use std::ffi::{OsStr, OsString};
 use std::fs::File;
 use std::io::Write;
-use std::mem;
 use std::os::unix::prelude::OsStrExt;
 use std::path::PathBuf;
-use std::process::{Child, Command};
+use std::process::{Child, Command, ExitStatus};
 use std::sync::Arc;
 use std::time::Duration;
 
 use flate2::write::GzDecoder;
+use mach::task::{task_resume, task_suspend};
+use mach::traps::task_for_pid;
 use tempfile::tempdir;
 
+use crate::shared::ctrl_c::CtrlC;
+
 pub use super::mach_ipc::{mach_port_t, MachError, OsIpcSender};
-use super::mach_ipc::{BlockingMode, OsIpcMultiShotServer, MACH_PORT_NULL};
+use super::mach_ipc::{mach_task_self, BlockingMode, OsIpcMultiShotServer, MACH_PORT_NULL};
+
+pub trait RootTaskRunner {
+    fn run_root_task(&self) -> Result<ExitStatus, MachError>;
+}
 
 pub struct TaskLauncher {
     program: OsString,
     args: Vec<OsString>,
     child_env: Vec<(OsString, OsString)>,
-    _temp_dir: Arc<tempfile::TempDir>,
+    iteration_count: u32,
+}
+
+impl RootTaskRunner for TaskLauncher {
+    fn run_root_task(&self) -> Result<ExitStatus, MachError> {
+        // Ignore Ctrl+C while the subcommand is running. The signal still reaches the process
+        // under observation while we continue to record it. (ctrl+c will send the SIGINT signal
+        // to all processes in the foreground process group).
+        let mut ctrl_c_receiver = CtrlC::observe_oneshot();
+
+        let mut root_child = self.launch_child();
+        let mut exit_status = root_child.wait().expect("couldn't wait for child");
+
+        for i in 2..=self.iteration_count {
+            if !exit_status.success() {
+                eprintln!(
+                    "Skipping remaining iterations due to non-success exit status: \"{}\"",
+                    exit_status
+                );
+                break;
+            }
+            eprintln!("Running iteration {i} of {}...", self.iteration_count);
+            let mut root_child = self.launch_child();
+            exit_status = root_child.wait().expect("couldn't wait for child");
+        }
+
+        // From now on, we want to terminate if the user presses Ctrl+C.
+        ctrl_c_receiver.close();
+
+        Ok(exit_status)
+    }
 }
 
 impl TaskLauncher {
+    pub fn new<I, S>(
+        program: S,
+        args: I,
+        iteration_count: u32,
+        env_vars: &[(OsString, OsString)],
+        extra_env_vars: &[(OsString, OsString)],
+    ) -> Result<TaskLauncher, MachError>
+    where
+        I: IntoIterator<Item = S>,
+        S: Into<OsString>,
+    {
+        // Take this process's environment variables and add DYLD_INSERT_LIBRARIES
+        // and SAMPLY_BOOTSTRAP_SERVER_NAME.
+        let mut child_env: BTreeMap<OsString, OsString> = std::env::vars_os().collect();
+        for (name, val) in env_vars {
+            child_env.insert(name.to_owned(), val.to_owned());
+        }
+        for (name, val) in extra_env_vars {
+            child_env.insert(name.to_owned(), val.to_owned());
+        }
+        let child_env: Vec<(OsString, OsString)> = child_env.into_iter().collect();
+
+        let args: Vec<OsString> = args.into_iter().map(|a| a.into()).collect();
+        let program: OsString = program.into();
+
+        Ok(TaskLauncher {
+            program,
+            args,
+            child_env,
+            iteration_count,
+        })
+    }
+
     pub fn launch_child(&self) -> Child {
         match Command::new(&self.program)
             .args(&self.args)
@@ -49,22 +119,16 @@ impl TaskLauncher {
 
 pub struct TaskAccepter {
     server: OsIpcMultiShotServer,
+    added_env: Vec<(OsString, OsString)>,
+    queue: Vec<ReceivedStuff>,
     _temp_dir: Arc<tempfile::TempDir>,
 }
 
 static PRELOAD_LIB_CONTENTS: &[u8] =
     include_bytes!("../../resources/libsamply_mac_preload.dylib.gz");
 
 impl TaskAccepter {
-    pub fn new<I, S>(
-        program: S,
-        args: I,
-        env_vars: &[(OsString, OsString)],
-    ) -> Result<(Self, TaskLauncher), MachError>
-    where
-        I: IntoIterator<Item = S>,
-        S: Into<OsString>,
-    {
+    pub fn new() -> Result<Self, MachError> {
         let (server, server_name) = OsIpcMultiShotServer::new()?;
 
         // Launch the child with DYLD_INSERT_LIBRARIES set to libsamply_mac_preload.dylib.
@@ -84,42 +148,38 @@ impl TaskAccepter {
             .finish()
             .expect("Couldn't write libsamply_mac_preload.dylib (error during finish)");
 
-        // Take this process's environment variables and add DYLD_INSERT_LIBRARIES
-        // and SAMPLY_BOOTSTRAP_SERVER_NAME.
-        let mut child_env: BTreeMap<OsString, OsString> = std::env::vars_os().collect();
-        for (name, val) in env_vars {
-            child_env.insert(name.to_owned(), val.to_owned());
-        }
+        let mut added_env: Vec<(OsString, OsString)> = vec![];
         let mut add_env = |name: &str, val: &OsStr| {
-            child_env.insert(name.into(), val.to_owned());
+            added_env.push((name.into(), val.to_owned()));
             // Also set the same variable with an `__XPC_` prefix, so that it gets applied
             // to services launched via XPC. XPC strips the prefix when setting these environment
             // variables on the launched process.
-            child_env.insert(format!("__XPC_{name}").into(), val.to_owned());
+            added_env.push((format!("__XPC_{name}").into(), val.to_owned()));
         };
         add_env("DYLD_INSERT_LIBRARIES", preload_lib_path.as_os_str());
         add_env("SAMPLY_BOOTSTRAP_SERVER_NAME", OsStr::new(&server_name));
-        let child_env: Vec<(OsString, OsString)> = child_env.into_iter().collect();
 
-        let args: Vec<OsString> = args.into_iter().map(|a| a.into()).collect();
-        let program: OsString = program.into();
-        let dir = Arc::new(dir);
-
-        Ok((
-            TaskAccepter {
-                server,
-                _temp_dir: dir.clone(),
-            },
-            TaskLauncher {
-                program,
-                args,
-                child_env,
-                _temp_dir: dir,
-            },
-        ))
+        Ok(TaskAccepter {
+            server,
+            added_env,
+            queue: vec![],
+            _temp_dir: Arc::new(dir),
+        })
+    }
+
+    pub fn extra_env_vars(&self) -> &[(OsString, OsString)] {
+        &self.added_env
+    }
+
+    pub fn queue_received_stuff(&mut self, rs: ReceivedStuff) {
+        self.queue.push(rs);
     }
 
     pub fn next_message(&mut self, timeout: Duration) -> Result<ReceivedStuff, MachError> {
+        if let Some(rs) = self.queue.pop() {
+            return Ok(rs);
+        }
+
         // Wait until the child is ready
         let (res, mut channels, _) = self
             .server
@@ -138,7 +198,7 @@ impl TaskAccepter {
                 ReceivedStuff::AcceptedTask(AcceptedTask {
                     task,
                     pid,
-                    sender_channel,
+                    sender_channel: Some(sender_channel),
                 })
             }
             (b"Jitdump", jitdump_info) => {
@@ -174,19 +234,64 @@ pub enum ReceivedStuff {
 pub struct AcceptedTask {
     task: mach_port_t,
     pid: u32,
-    sender_channel: OsIpcSender,
+    sender_channel: Option<OsIpcSender>,
 }
 
 impl AcceptedTask {
-    pub fn take_task(&mut self) -> mach_port_t {
-        mem::replace(&mut self.task, MACH_PORT_NULL)
+    pub fn task(&self) -> mach_port_t {
+        self.task
     }
 
     pub fn get_id(&self) -> u32 {
         self.pid
     }
 
     pub fn start_execution(&self) {
-        self.sender_channel.send(b"Proceed", vec![]).unwrap();
+        if let Some(sender_channel) = &self.sender_channel {
+            sender_channel.send(b"Proceed", vec![]).unwrap();
+        } else {
+            unsafe { task_resume(self.task) };
+        }
+    }
+}
+
+pub struct ExistingProcessRunner {
+    pid: u32,
+}
+
+impl RootTaskRunner for ExistingProcessRunner {
+    fn run_root_task(&self) -> Result<ExitStatus, MachError> {
+        let ctrl_c_receiver = CtrlC::observe_oneshot();
+
+        eprintln!("Profiling {}, press Ctrl-C to stop...", self.pid);
+
+        ctrl_c_receiver
+            .blocking_recv()
+            .expect("Ctrl+C receiver failed");
+
+        eprintln!("Done.");
+
+        Ok(ExitStatus::default())
+    }
+}
+
+impl ExistingProcessRunner {
+    pub fn new(pid: u32, task_accepter: &mut TaskAccepter) -> ExistingProcessRunner {
+        let task = unsafe {
+            let mut task = MACH_PORT_NULL;
+            let kr = task_for_pid(mach_task_self(), pid as i32, &mut task);
+            if kr != 0 {
+                eprintln!("Error: task_for_pid failed with error code {kr}. Does the profiler have entitlements?");
+                std::process::exit(1);
+            }
+            task_suspend(task);
+            task
+        };
+        task_accepter.queue_received_stuff(ReceivedStuff::AcceptedTask(AcceptedTask {
+            task,
+            pid,
+            sender_channel: None,
+        }));
+        ExistingProcessRunner { pid }
     }
 }