Hook open+fopen and forward jitdump paths.

On macOS, we now interpose the open and fopen functions in the profiled
process. If we detect the opening of a file which matches the jitdump
path syntax (`**/jit-*.dump`), we send this path to the sampler and
forward it to the right TaskProfiler.

samply-mac-preload/Cargo.toml

@@ -5,6 +5,9 @@ authors = ["Markus Stange <[email protected]>"]
 edition = "2021"
 license = "MIT OR Apache-2.0"
 
+[workspace]
+# This crate is not part of the samply workspace.
+
 [lib]
 crate_type = ["cdylib"]
 
@@ -17,4 +20,5 @@ panic = 'abort'
 
 [dependencies]
 libc = { version = "0.2.70", default-features = false }
+spin = "0.9.6"
 

samply-mac-preload/README.md

@@ -5,12 +5,4 @@ rustup target add x86_64-apple-darwin
 rustup target add aarch64-apple-darwin
 ```
 
-```
-export SDKROOT=/Applications/Xcode.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX.sdk
-MACOSX_DEPLOYMENT_TARGET=10.7 cargo build --release --target=x86_64-apple-darwin
-mv target/x86_64-apple-darwin/release/libsamply_mac_preload.dylib binaries/libsamply_mac_preload_x86_64.dylib
-MACOSX_DEPLOYMENT_TARGET=11.0 cargo build --release --target=aarch64-apple-darwin
-mv target/aarch64-apple-darwin/release/libsamply_mac_preload.dylib binaries/libsamply_mac_preload_arm64.dylib
-lipo binaries/libsamply_mac_preload_* -create -output binaries/libsamply_mac_preload.dylib
-gzip -cvf binaries/libsamply_mac_preload.dylib > ../samply/resources/libsamply_mac_preload.dylib.gz
-```
+Run `build.sh` from inside this directory to update the files inside `binaries/` and to copy the updated dylib into `../samply/resources/`.

samply-mac-preload/build.sh

@@ -0,0 +1,7 @@
+export SDKROOT=/Applications/Xcode.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX.sdk
+MACOSX_DEPLOYMENT_TARGET=10.7 cargo build --release --target=x86_64-apple-darwin
+mv target/x86_64-apple-darwin/release/libsamply_mac_preload.dylib binaries/libsamply_mac_preload_x86_64.dylib
+MACOSX_DEPLOYMENT_TARGET=11.0 cargo build --release --target=aarch64-apple-darwin
+mv target/aarch64-apple-darwin/release/libsamply_mac_preload.dylib binaries/libsamply_mac_preload_arm64.dylib
+lipo binaries/libsamply_mac_preload_* -create -output binaries/libsamply_mac_preload.dylib
+gzip -cvf binaries/libsamply_mac_preload.dylib > ../samply/resources/libsamply_mac_preload.dylib.gz

samply-mac-preload/src/lib.rs

@@ -1,17 +1,28 @@
 #![no_main]
 #![no_std]
 
-#[cfg(not(test))]
-#[panic_handler]
-fn panic(_panic: &core::panic::PanicInfo<'_>) -> ! {
-    unsafe { libc::abort() }
-}
+use libc::{c_char, c_int, mode_t, FILE};
+
+use core::ffi::CStr;
 
 mod mach_ipc;
 mod mach_sys;
 
 use mach_ipc::{channel, mach_task_self, OsIpcChannel, OsIpcSender};
 
+extern "C" {
+    fn open(path: *const c_char, flags: c_int, mode: mode_t) -> c_int;
+    fn fopen(filename: *const c_char, mode: *const c_char) -> *mut FILE;
+}
+
+static CHANNEL_SENDER: spin::Mutex<Option<OsIpcSender>> = spin::Mutex::new(None);
+
+#[cfg(not(test))]
+#[panic_handler]
+fn panic(_panic: &core::panic::PanicInfo<'_>) -> ! {
+    unsafe { libc::abort() }
+}
+
 // Run our code as early as possible, by pretending to be a global constructor.
 // This code was taken from https://github.com/neon-bindings/neon/blob/2277e943a619579c144c1da543874f4a7ec39879/src/lib.rs#L40-L44
 #[used]
@@ -48,9 +59,90 @@ fn set_up_samply_connection() -> Option<()> {
     message_bytes[7..11].copy_from_slice(&pid.to_le_bytes());
     tx1.send(&message_bytes, [OsIpcChannel::Sender(tx0), c])
         .ok()?;
+    *CHANNEL_SENDER.lock() = Some(tx1);
     // Wait for the parent to tell us to proceed, in case it wants to do any more setup with our task.
-    let mut recv_buf = [0; 32];
+    let mut recv_buf = [0; 256];
     let result = rx0.recv(&mut recv_buf).ok()?;
     assert_eq!(b"Proceed", &result);
     Some(())
 }
+
+// Override the `open` function, in order to be able to observe the file
+// paths of opened files.
+//
+// We use this to detect jitdump files.
+#[no_mangle]
+extern "C" fn samply_hooked_open(path: *const c_char, flags: c_int, mode: mode_t) -> c_int {
+    // unsafe {
+    //     libc::printf(b"open(%s, %d, %u)\n\0".as_ptr() as *const i8, path, flags, mode as c_uint);
+    // }
+
+    if let Ok(path) = unsafe { CStr::from_ptr(path) }.to_str() {
+        detect_and_send_jitdump_path(path);
+    }
+
+    // Call the original. Do this at the end, so that this is compiled as a tail call.
+    //
+    // WARNING: What we are doing here is even sketchier than it seems. The `open` function
+    // is variadic: It can be called with or without the mode parameter. I have not found
+    // the right way to forward those variadic args properly. So by using a tail call, we
+    // can hope that the compiled code leaves the arguments completely untouched and just
+    // jumps to the called function, and everything should work out fine in terms of the
+    // call ABI.
+    unsafe { open(path, flags, mode) }
+}
+
+// Override fopen for the same reason.
+#[no_mangle]
+extern "C" fn samply_hooked_fopen(path: *const c_char, mode: *const c_char) -> *mut FILE {
+    // unsafe {
+    //     libc::printf(b"fopen(%s, %s\n\0".as_ptr() as *const i8, path, mode);
+    // }
+
+    if let Ok(path) = unsafe { CStr::from_ptr(path) }.to_str() {
+        detect_and_send_jitdump_path(path);
+    }
+
+    // Call the original.
+    unsafe { fopen(path, mode) }
+}
+
+fn detect_and_send_jitdump_path(path: &str) {
+    if path.len() > 256 - 12 || !path.ends_with(".dump") || !path.contains("/jit-") {
+        return;
+    }
+
+    let channel_sender = CHANNEL_SENDER.lock();
+    let Some(sender) = channel_sender.as_ref() else { return };
+    let pid = unsafe { libc::getpid() };
+    let mut message_bytes = [0; 256];
+    message_bytes[0..7].copy_from_slice(b"Jitdump");
+    message_bytes[7..11].copy_from_slice(&pid.to_le_bytes());
+    message_bytes[11] = path.len() as u8;
+    message_bytes[12..][..path.len()].copy_from_slice(path.as_bytes());
+    let _ = sender.send(&message_bytes, []);
+}
+
+#[allow(non_camel_case_types)]
+pub struct InterposeEntry {
+    _new: *const (),
+    _old: *const (),
+}
+
+#[used]
+#[allow(dead_code)]
+#[allow(non_upper_case_globals)]
+#[link_section = "__DATA,__interpose"]
+pub static mut _interpose_open: InterposeEntry = InterposeEntry {
+    _new: samply_hooked_open as *const (),
+    _old: open as *const (),
+};
+
+#[used]
+#[allow(dead_code)]
+#[allow(non_upper_case_globals)]
+#[link_section = "__DATA,__interpose"]
+pub static mut _interpose_fopen: InterposeEntry = InterposeEntry {
+    _new: samply_hooked_fopen as *const (),
+    _old: fopen as *const (),
+};

samply-mac-preload/src/mach_ipc.rs

@@ -233,7 +233,7 @@ impl OsIpcReceiver {
                 message as *mut _,
                 flags,
                 0,
-                (*message).header.msgh_size,
+                buf.len() as u32,
                 port,
                 timeout,
                 MACH_PORT_NULL,

samply/src/mac/process_launcher.rs

@@ -2,6 +2,8 @@ use std::ffi::OsStr;
 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::time::Duration;
 
@@ -86,30 +88,49 @@ impl TaskAccepter {
         ))
     }
 
-    pub fn try_accept(&mut self, timeout: Duration) -> Result<AcceptedTask, MachError> {
+    pub fn next_message(&mut self, timeout: Duration) -> Result<ReceivedStuff, MachError> {
         // Wait until the child is ready
         let (res, mut channels, _) = self
             .server
             .accept(BlockingMode::BlockingWithTimeout(timeout))?;
-        assert_eq!(res.len(), 11);
-        assert!(&res[0..7] == b"My task");
-        let mut pid_bytes: [u8; 4] = Default::default();
-        pid_bytes.copy_from_slice(&res[7..11]);
-        let pid = u32::from_le_bytes(pid_bytes);
-        let task_channel = channels.pop().unwrap();
-        let sender_channel = channels.pop().unwrap();
-        let sender_channel = sender_channel.into_sender();
-
-        let task = task_channel.into_port();
-
-        Ok(AcceptedTask {
-            task,
-            pid,
-            sender_channel,
-        })
+        let received_stuff = match res.split_at(7) {
+            (b"My task", pid_bytes) => {
+                assert!(pid_bytes.len() == 4);
+                let pid =
+                    u32::from_le_bytes([pid_bytes[0], pid_bytes[1], pid_bytes[2], pid_bytes[3]]);
+                let task_channel = channels.pop().unwrap();
+                let sender_channel = channels.pop().unwrap();
+                let sender_channel = sender_channel.into_sender();
+
+                let task = task_channel.into_port();
+
+                ReceivedStuff::AcceptedTask(AcceptedTask {
+                    task,
+                    pid,
+                    sender_channel,
+                })
+            }
+            (b"Jitdump", jitdump_info) => {
+                let pid_bytes = &jitdump_info[0..4];
+                let pid =
+                    u32::from_le_bytes([pid_bytes[0], pid_bytes[1], pid_bytes[2], pid_bytes[3]]);
+                let len = jitdump_info[4] as usize;
+                let path = &jitdump_info[5..][..len];
+                ReceivedStuff::JitdumpPath(pid, OsStr::from_bytes(path).into())
+            }
+            (other, _) => {
+                panic!("Unexpected message: {:?}", other);
+            }
+        };
+        Ok(received_stuff)
     }
 }
 
+pub enum ReceivedStuff {
+    AcceptedTask(AcceptedTask),
+    JitdumpPath(u32, PathBuf),
+}
+
 pub struct AcceptedTask {
     task: mach_port_t,
     pid: u32,

samply/src/mac/profiler.rs

@@ -1,6 +1,8 @@
 use crossbeam_channel::unbounded;
 use serde_json::to_writer;
 
+use std::collections::hash_map::Entry;
+use std::collections::HashMap;
 use std::ffi::OsString;
 use std::fs::File;
 use std::io::BufWriter;
@@ -12,7 +14,7 @@ use std::thread;
 use std::time::{Duration, Instant};
 
 use super::error::SamplingError;
-use super::process_launcher::{MachError, TaskAccepter};
+use super::process_launcher::{MachError, ReceivedStuff, TaskAccepter};
 use super::sampler::{Sampler, TaskInit};
 use crate::server::{start_server_main, ServerProps};
 
@@ -58,28 +60,57 @@ pub fn start_recording(
         TaskAccepter::create_and_launch_root_task(&command_name, command_args)?;
 
     let (accepter_sender, accepter_receiver) = unbounded();
-    let accepter_thread = thread::spawn(move || loop {
-        if let Ok(()) = accepter_receiver.try_recv() {
-            break;
-        }
-        let timeout = Duration::from_secs_f64(1.0);
-        match task_accepter.try_accept(timeout) {
-            Ok(mut accepted_task) => {
-                let send_result = task_sender.send(TaskInit {
-                    start_time: Instant::now(),
-                    task: accepted_task.take_task(),
-                    pid: accepted_task.get_id(),
-                });
-                if send_result.is_err() {
-                    // The sampler has already shut down. This task arrived too late.
-                }
-                accepted_task.start_execution();
-            }
-            Err(MachError::RcvTimedOut) => {
-                // TODO: give status back via task_sender
+    let accepter_thread = thread::spawn(move || {
+        // Loop while accepting messages from the spawned process tree.
+
+        // A map of pids to channel senders, to notify existing tasks of Jitdump
+        // paths. Having the mapping here lets us deliver the path to the right
+        // task even in cases where a process execs into a new task with the same pid.
+        let mut jitdump_path_senders_per_pid = HashMap::new();
+
+        loop {
+            if let Ok(()) = accepter_receiver.try_recv() {
+                break;
             }
-            Err(err) => {
-                eprintln!("Encountered error while waiting for task port: {err:?}");
+            let timeout = Duration::from_secs_f64(1.0);
+            match task_accepter.next_message(timeout) {
+                Ok(ReceivedStuff::AcceptedTask(mut accepted_task)) => {
+                    let pid = accepted_task.get_id();
+                    let (jitdump_path_sender, jitdump_path_receiver) = unbounded();
+                    let send_result = task_sender.send(TaskInit {
+                        start_time: Instant::now(),
+                        task: accepted_task.take_task(),
+                        pid,
+                        jitdump_path_receiver,
+                    });
+                    jitdump_path_senders_per_pid.insert(pid, jitdump_path_sender);
+                    if send_result.is_err() {
+                        // The sampler has already shut down. This task arrived too late.
+                    }
+                    accepted_task.start_execution();
+                }
+                Ok(ReceivedStuff::JitdumpPath(pid, path)) => {
+                    match jitdump_path_senders_per_pid.entry(pid) {
+                        Entry::Occupied(mut entry) => {
+                            let send_result = entry.get_mut().send(path);
+                            if send_result.is_err() {
+                                // The task is probably already dead. The path arrived too late.
+                                entry.remove();
+                            }
+                        }
+                        Entry::Vacant(_entry) => {
+                            eprintln!(
+                                "Received a Jitdump path for pid {pid} which I don't have a task for."
+                            );
+                        }
+                    }
+                }
+                Err(MachError::RcvTimedOut) => {
+                    // TODO: give status back via task_sender
+                }
+                Err(err) => {
+                    eprintln!("Encountered error while waiting for task port: {err:?}");
+                }
             }
         }
     });