functiontrace 0.4.3__tar.gz → 0.4.4__tar.gz

{functiontrace-0.4.3 → functiontrace-0.4.4}/Cargo.toml

@@ -5,6 +5,7 @@ edition = "2021"
 
 [dependencies]
 color-eyre = { version = "0.6", default-features = false }
+dashmap = "6"
 functiontrace-server = { version = "0.8.4", default-features = false, features = ["serialize"] }
 libc = "0.2"
 pyo3 = { version = "0.25", features = ["eyre"] }

{functiontrace-0.4.3 → functiontrace-0.4.4}/MANIFEST.in

@@ -1,4 +1,3 @@
-graft mpack
 graft rust
 include Cargo.toml
 include build.rs

{functiontrace-0.4.3/functiontrace.egg-info → functiontrace-0.4.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: functiontrace
-Version: 0.4.3
+Version: 0.4.4
 Summary: A profiling tool that traces your application's execution
 Author: Alistair Dobke
 Author-email: Matt Bryant <functiontrace@programsareproofs.com>

{functiontrace-0.4.3 → functiontrace-0.4.4/functiontrace.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: functiontrace
-Version: 0.4.3
+Version: 0.4.4
 Summary: A profiling tool that traces your application's execution
 Author: Alistair Dobke
 Author-email: Matt Bryant <functiontrace@programsareproofs.com>

{functiontrace-0.4.3 → functiontrace-0.4.4}/functiontrace.egg-info/SOURCES.txt

@@ -2,7 +2,6 @@ Cargo.toml
 LICENSE.md
 MANIFEST.in
 README.rst
-_functiontrace.c
 build.rs
 functiontrace.py
 pyproject.toml
@@ -12,7 +11,4 @@ functiontrace.egg-info/SOURCES.txt
 functiontrace.egg-info/dependency_links.txt
 functiontrace.egg-info/entry_points.txt
 functiontrace.egg-info/top_level.txt
-mpack/LICENSE
-mpack/mpack.c
-mpack/mpack.h
 rust/lib.rs

functiontrace-0.4.4/functiontrace.egg-info/top_level.txt

@@ -0,0 +1 @@
+functiontrace

{functiontrace-0.4.3 → functiontrace-0.4.4}/functiontrace.py

@@ -5,7 +5,6 @@ import sys
 import tempfile
 from importlib import machinery, util
 
-import _functiontrace  # noqa: F401 - this import is load bearing since we need to run code on import
 import _functiontrace_rs as rust
 
 PYTHON_TEMPLATE = """#!/bin/sh

{functiontrace-0.4.3 → functiontrace-0.4.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "functiontrace"
-version = "0.4.3"
+version = "0.4.4"
 description = "A profiling tool that traces your application's execution"
 authors = [
     { name = "Matt Bryant", email = "functiontrace@programsareproofs.com"},

{functiontrace-0.4.3 → functiontrace-0.4.4}/rust/lib.rs

@@ -5,11 +5,12 @@ mod _functiontrace_rs {
     use pyo3::prelude::*; // TODO: What's in here?
     use serde::Serialize;
     use std::borrow::Cow;
+    use std::cell::Cell;
     use std::ffi::{c_void, CStr, OsString};
     use std::io::Write;
     use std::os::unix::net::UnixStream;
-    use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
-    use std::sync::OnceLock;
+    use std::sync::atomic::{AtomicBool, Ordering};
+    use std::sync::{LazyLock, OnceLock};
     use std::time::Duration;
 
     ////////////////////////////////////////////////////////////////////////////////////////////////
@@ -18,34 +19,22 @@ mod _functiontrace_rs {
     /// The size of messagepack buffers.  This is picked somewhat arbitrarily, but should be
     /// reasonably small since we'll need one per thread, and should be large enough that we
     /// can fit many messages in before needing to do expensive UnixStream syscalls.
-    const MPACK_BUFFER: usize = 1 << 17;
+    const MSGPACK_BUFFER: usize = 1 << 17;
 
     /// Storage specific to each thread used for buffering/transmitting data.
-    type ThreadState = MpackWriter;
-
-    /// A `mpack_writer_t` struct, which `mpack.h` will use to pass us our per-thread context in.
-    ///
-    /// NOTE: This must directly match the `mpack_writer_t` definition in C, since we're going to
-    /// be creating and modifying it ourselves!
-    #[repr(C)]
-    pub struct MpackWriter {
-        flush: extern "C" fn(writer: *const MpackWriter, buffer: *const u8, bytes: usize),
-        error_fn: *const (),
-        teardown: *const (),
-        context: *mut UnixStream,
-
-        buffer: *mut u8,
-        current: *mut u8,
-        end: *mut u8,
-
-        // XXX: This is an enum that we never use.  Since it doesn't have a clearly-defined size,
-        // we allocate 8 bytes for it.  We MUST NOT access this field!
-        _error: u64,
+    pub struct ThreadState {
+        socket: UnixStream,
+
+        // TODO: Conceptually this should just be a `BufWriter`, but in practice using that
+        // triggers 0 byte writes that cause `functiontrace-server` to close the connection.
+        // Implementing this ourselves is pretty trivial, so...
+        buffer: [u8; MSGPACK_BUFFER],
+        head: usize,
     }
 
-    impl Write for MpackWriter {
-        /// A small [`Write`] implementation that allows us to write to an
-        /// `mpack_writer_t`'s interior buffer while keeping its state consistent.
+    impl Write for ThreadState {
+        /// A small [`Write`] implementation that writes to a buffer and only flushes to the socket
+        /// once full.
         fn write(&mut self, src: &[u8]) -> std::io::Result<usize> {
             if !STARTED.load(Ordering::Relaxed) {
                 // We haven't started yet, so we shouldn't write anything out.  This avoids logging
@@ -53,21 +42,15 @@ mod _functiontrace_rs {
                 return Ok(src.len());
             }
 
-            let len = unsafe {
-                // SAFETY: writer.(current|end) are both within writer.buffer.
-                self.end.offset_from(self.current)
-            }
-            .try_into()
-            .expect("end > current");
-
             let bytes = src.len();
+            let tail = self.head + bytes;
 
-            if bytes > len {
-                if bytes >= MPACK_BUFFER {
+            if tail >= self.buffer.len() {
+                if bytes >= self.buffer.len() {
                     // We'll never be able to fit this much.
                     return Err(std::io::Error::new(
                         std::io::ErrorKind::StorageFull,
-                        "mpack buffer too small",
+                        "msgpack buffer too small",
                     ));
                 }
 
@@ -77,44 +60,25 @@ mod _functiontrace_rs {
 
                 return Err(std::io::Error::new(
                     std::io::ErrorKind::Interrupted,
-                    "mpack buffer is full, flushing",
+                    "msgpack buffer is full, flushing",
                 ));
             }
 
-            let buf = unsafe { std::slice::from_raw_parts_mut(self.current, len) };
-
-            buf[..bytes].copy_from_slice(src);
-
-            unsafe {
-                // SAFETY: end = current + len, and bytes < len, so current + bytes < end
-                self.current = self.current.add(bytes)
-            };
+            self.buffer[self.head..tail].copy_from_slice(src);
+            self.head = tail;
 
             Ok(bytes)
         }
 
-        /// This roughly implements `Mpack_Flush`.
+        /// Flush our local buffer to the socket.
         fn flush(&mut self) -> std::io::Result<()> {
-            let len = unsafe {
-                // SAFETY: writer.current is within writer.buffer.
-                self.current.offset_from(self.buffer)
-            }
-            .try_into()
-            .expect("current > end");
-
-            let data = unsafe { std::slice::from_raw_parts_mut(self.buffer, len) };
-
             // We're about to send all of the data, so reset our buffer.  Do this now to ensure we
             // don't miss it, since the buffer will be left in an invalid state if this doesn't
             // happen.
-            self.current = self.buffer;
+            let head = self.head;
+            self.head = 0;
 
-            let mut socket = unsafe {
-                // SAFETY: `writer` is valid and `writer->context` contains a `UnixStream` emitted from
-                // `thread_register`.
-                self.context.as_ref().expect("context is non-null")
-            };
-            socket.write_all(data)
+            self.socket.write_all(&self.buffer[..head])
         }
     }
 
@@ -125,13 +89,6 @@ mod _functiontrace_rs {
     /// The UnixStream address we're using to communicate with the profile server.
     static SOCKADDR: OnceLock<OsString> = OnceLock::new();
 
-    /// The pthread_key we use for accessing [`ThreadState`], acquired from the C extension.
-    ///
-    /// NOTE: We store this as an atomic rather than use a [`OnceLock`] since we need to access it
-    /// in the critical path.  There's some overhead to `OnceLock` acquisitions, while simple loads
-    /// and stores of [`AtomicU32`]s have no overhead (at least on x86).
-    static TSS_KEY: AtomicU32 = AtomicU32::new(0);
-
     /// `sys.argv` for our process, which is used when registering new threads.
     static ARGV: OnceLock<String> = OnceLock::new();
 
@@ -139,53 +96,12 @@ mod _functiontrace_rs {
     /// set, we're allowed to send messages to the profile generation server.
     static STARTED: AtomicBool = AtomicBool::new(false);
 
-    /// The `Fprofile_FunctionTrace` function defined in the C extension.  This is needed to finish
-    /// moving setup to the Rust code.
-    static TRACEFUNC: OnceLock<pyo3::ffi::Py_tracefunc> = OnceLock::new();
-
-    ////////////////////////////////////////////////////////////////////////////////////////////////
-    // Rust <-> C Communication
-    ////////////////////////////////////////////////////////////////////////////////////////////////
-
-    /// We expose all of our C functions in a single struct that's loaded once at the start of the
-    /// program, avoiding most of the overhead from the previous `dlopen-esque` interface.
-    #[derive(Debug)]
-    #[repr(C)]
-    struct RustFunctions {
-        set_config: usize,
-    }
-
-    /// It's easiest to return pointer-sized types over the Python API, so allocate our
-    /// [`RustFunctions`] and leak it.  This will be used by the C extension and never freed.
-    #[pyfunction]
-    fn c_api() -> usize {
-        let api = Box::new(RustFunctions {
-            set_config: set_config as usize,
-        });
-
-        Box::into_raw(api) as usize
-    }
-
-    /// Called by the C module to pass the `c_functiontrace` handler to us.
-    #[unsafe(no_mangle)]
-    extern "C" fn set_config(c_functiontrace: pyo3::ffi::Py_tracefunc) -> libc::pthread_key_t {
-        // Allocate a Tss_Key for storing thread state
-        let mut tss_key = 0;
-        if unsafe {
-            // SAFETY: tss_key is live
-            libc::pthread_key_create(&raw mut tss_key, Some(c_thread_teardown))
-        } != 0
-        {
-            panic!("Failed to create Tss_Key");
-        }
-
-        TRACEFUNC
-            .set(c_functiontrace)
-            .expect("set_config can only be called once");
-
-        TSS_KEY.store(tss_key, Ordering::Relaxed);
-
-        tss_key
+    thread_local! {
+        /// A per-thread [`ThreadState`] for writing events to.
+        //
+        // NOTE: The `const` block appears to be necessary in order to generate reasonable
+        // assembly.  Unclear why...
+        static THREAD_STATE: Cell<*mut ThreadState> = const { Cell::new(std::ptr::null_mut()) };
     }
 
     ////////////////////////////////////////////////////////////////////////////////////////////////
@@ -216,13 +132,7 @@ mod _functiontrace_rs {
     /// `Box::into_raw`.
     #[inline(always)]
     pub fn thread_state() -> Option<&'static mut ThreadState> {
-        // Check that we're able to properly initialize
-        let tss_key = TSS_KEY.load(Ordering::Relaxed);
-
-        let state = unsafe {
-            // SAFETY: tss_key is a valid `pthread_key_t`
-            libc::pthread_getspecific(tss_key) as *mut ThreadState
-        };
+        let state = THREAD_STATE.get();
 
         if state.is_null() {
             // We shouldn't be asked to reset at thread that hasn't yet been registered.
@@ -241,7 +151,7 @@ mod _functiontrace_rs {
     ////////////////////////////////////////////////////////////////////////////////////////////////
 
     /// Initializes and returns the [`UnixStream`] we'll communicate to `functiontrace-server` on.
-    fn message_initialize() -> Box<UnixStream> {
+    fn message_initialize() -> UnixStream {
         let sockaddr = SOCKADDR
             .get()
             .expect("Must be initialized before setting up messaging");
@@ -252,7 +162,7 @@ mod _functiontrace_rs {
         loop {
             match UnixStream::connect(sockaddr) {
                 Ok(s) => {
-                    break Box::new(s);
+                    break s;
                 }
                 _ => {
                     if start.elapsed() > Duration::from_millis(1000) {
@@ -265,30 +175,6 @@ mod _functiontrace_rs {
         }
     }
 
-    /// Flush events to this socket once the given buffer is full according to `mpack`.
-    #[unsafe(no_mangle)]
-    extern "C" fn Mpack_Flush(writer: *const MpackWriter, buffer: *const u8, bytes: usize) {
-        let data = unsafe {
-            // SAFETY: `bytes` represents the initialized number of bytes in `buffer` when called
-            // by `mpack_writer_flush_message`.
-            std::slice::from_raw_parts(buffer, bytes)
-        };
-
-        let mut socket = unsafe {
-            // SAFETY: `writer` is valid and `writer->context` contains a `UnixStream` emitted from
-            // `register_thread`.
-            writer
-                .as_ref()
-                .expect("writer is non-null")
-                .context
-                .as_ref()
-                .expect("context is non-null")
-        };
-        if let Err(e) = socket.write_all(data) {
-            panic!("Socket send failed: {e}");
-        }
-    }
-
     ////////////////////////////////////////////////////////////////////////////////////////////////
     // Tracing Implementation
     ////////////////////////////////////////////////////////////////////////////////////////////////
@@ -568,11 +454,10 @@ mod _functiontrace_rs {
         Ok(())
     }
 
-    /// Register a new thread, including creating the UnixStream for mpack to log messages for this
-    /// thread on.
+    /// Register a new thread, including creating the UnixStream to log messages for this thread
+    /// on.
     ///
-    /// A new [`ThreadState`] will be returned, which contains a fully initialized `MpackWriter`.
-    /// This state must be returned to Rust code later to be freed.
+    /// A new [`ThreadState`] will be returned, which must freed once the thread has terminated.
     fn thread_register() -> Result<()> {
         if thread_state().is_some() {
             // We shouldn't have a ThreadState yet, as we're the ones that create it.
@@ -600,46 +485,21 @@ mod _functiontrace_rs {
             .serialize(&mut rmp_serde::encode::Serializer::new(&mut socket))
             .wrap_err("Failed to emit RegisterThread message")?;
 
-        // Allocate a new buffer for mpack to write messages into.
-        // This will be stored in ThreadState, and will need to be freed by Rust code.
-        let buf = Box::into_raw(Box::new([0u8; MPACK_BUFFER])) as *mut u8;
-
-        // Initialize the `MpackWriter` for this thread, then associate the socket with it.
+        // Initialize the `ThreadState` for this thread, then associate the socket with it.
         //
-        // We'll return this state to the C code, which must pass it back to us to avoid leaking.
-        let state = Box::into_raw(Box::new(
-            // XXX: This must match `mpack_writer_init`, since mpack.c will be using this structure
-            // later!
-            MpackWriter {
-                context: Box::into_raw(socket),
-                flush: Mpack_Flush,
-
-                buffer: buf as _,
-                current: buf as _,
-                end: unsafe {
-                    // SAFETY: buf is an allocated array of MPACK_BUFFER bytes
-                    buf.add(MPACK_BUFFER)
-                } as _,
-
-                teardown: std::ptr::null(),
-                error_fn: std::ptr::null(),
-                _error: 0,
-            },
-        ));
+        // We'll move this state into thread local storage, and will need to recover it on
+        // termination in order to free it.
+        let state = Box::into_raw(Box::new(ThreadState {
+            socket,
+
+            buffer: [0u8; MSGPACK_BUFFER],
+            head: 0,
+        }));
 
         // Store our state in thread-specific storage so we can find it in the future.
         // Tracing begins for the thread at this point, since we've now published `state` for other
         // functions to start using.
-        let tss_key = TSS_KEY.load(Ordering::Relaxed);
-
-        if unsafe {
-            // SAFETY: tss_key is a valid `pthread_key_t`, and `state` is a valid pointer to a
-            // heap-allocated `ThreadState` object
-            libc::pthread_setspecific(tss_key, state as *const _)
-        } != 0
-        {
-            return Err(eyre!("Failed to set tss_key on thread startup"));
-        }
+        THREAD_STATE.set(state);
 
         Ok(())
     }
@@ -652,7 +512,7 @@ mod _functiontrace_rs {
         }
 
         let writer = unsafe {
-            // SAFETY: writer is a valid pointer to an initialized `MpackWriter`.
+            // SAFETY: writer is a valid pointer to an initialized `ThreadState`.
             &mut (state as *mut ThreadState)
                 .as_mut()
                 .expect("state is non-null")
@@ -664,7 +524,7 @@ mod _functiontrace_rs {
     /// Tear down the given thread by flushing any outstanding messages.
     ///
     /// This is called when the thread has shut down, including in multithread/process situations.
-    pub fn thread_teardown(writer: &mut MpackWriter) -> Result<()> {
+    pub fn thread_teardown(writer: &mut ThreadState) -> Result<()> {
         writer
             .flush()
             .wrap_err("Failed to flush remaining messages")?;
@@ -686,8 +546,6 @@ mod _functiontrace_rs {
     ///
     /// This is useful when we've just forked and want to ensure we don't reuse an existing socket.
     fn thread_reset() -> Result<()> {
-        let tss_key = TSS_KEY.load(Ordering::Relaxed);
-
         if !STARTED.load(Ordering::Relaxed) {
             // We haven't actually started yet, but are for some reason being asked
             // to fork.  This appears to be OS dependent.
@@ -698,30 +556,14 @@ mod _functiontrace_rs {
         let state = thread_state().ok_or_eyre(eyre!("Thread wasn't yet registered"))?;
 
         // Remove the TLS reference so no one else can load state after we've freed it.
-        if unsafe {
-            // SAFETY: tss_key is a valid `pthread_key_t`, and `state` is a valid pointer to a
-            // heap-allocated `ThreadState` object
-            libc::pthread_setspecific(tss_key, std::ptr::null())
-        } != 0
-        {
-            return Err(eyre!("Failed to clear TLS"));
-        }
+        THREAD_STATE.set(std::ptr::null_mut());
 
-        let buf = unsafe {
-            // SAFETY: writer.buffer contains a valid pointer to a boxed buffer
-            (state.buffer as *mut [u8; MPACK_BUFFER])
-                .as_mut()
-                .expect("buffer is initialized when state is and is never modified")
-        };
-
-        // Free the thread's mpack state
+        // Free the thread's state, including the socket
         unsafe {
-            let _ = Box::from_raw(buf);
+            // SAFETY: `state` came from `Box::into_raw` previously and can only be freed once.
             let _ = Box::from_raw(state);
         }
 
-        // TODO: Should we free the socket too?
-
         // Create a new thread for this process.  We don't need to start an entire new trace like
         // with subprocess calls since we were forked and therefore already share our
         // configuration.
@@ -816,7 +658,7 @@ mod _functiontrace_rs {
     /// The core logging function for FunctionTrace.  Configured by `PyEval_SetProfile`, and used
     /// to record all function calls/returns.
     extern "C" fn functiontrace(
-        obj: *mut pyo3::ffi::PyObject,
+        _obj: *mut pyo3::ffi::PyObject,
         frame: *mut pyo3::ffi::PyFrameObject,
         what: i32,
         arg: *mut pyo3::ffi::PyObject,
@@ -833,7 +675,6 @@ mod _functiontrace_rs {
         }
 
         let time = trace_time();
-        let c_functiontrace = *TRACEFUNC.get().expect("tracefunc is set in set_config");
 
         // Return the function name when we're in a C_* context.
         let c_func_name = || {
@@ -861,14 +702,88 @@ mod _functiontrace_rs {
                 Cow::Borrowed("<null>")
             }
         };
-
         match what {
             pyo3::ffi::PyTrace_CALL | pyo3::ffi::PyTrace_RETURN => {
-                // For now we delegate these to the C implementation, since we can't access the
-                // `PyCodeObject` fields that we require via Rust.
-                unsafe {
-                    // SAFETY: tracefunc is a valid function pointer to `Fprofile_FunctionTrace`.
-                    return c_functiontrace(obj, frame, what, arg);
+                // Cache information about each frame we see so that we don't need to look it up
+                // every time, which involves expensive Python operations.
+                static FRAMES: LazyLock<
+                    dashmap::DashMap<usize, (Cow<'static, str>, Cow<'static, str>)>,
+                > = LazyLock::new(|| {
+                    dashmap::DashMap::with_capacity(
+                        // Even small programs will have a decent amount of function calls in the
+                        // interpreter setup.
+                        512,
+                    )
+                });
+
+                let entry = FRAMES.entry(frame.addr()).or_insert_with(|| {
+                    // We've never seen this frame before, so we'll need to load its
+                    // information.  This is fairly slow, since it'll involve going via Python
+                    // rather than fast struct accesses.
+                    let code = unsafe {
+                        // SAFETY: frame is guaranteed to be a valid PyFrameObject
+
+                        // We'll be storing `frame` in our hashmap, so we don't want anyone
+                        // else freeing it.
+                        pyo3::ffi::Py_IncRef(frame as *mut pyo3::ffi::PyObject);
+
+                        pyo3::ffi::PyFrame_GetCode(frame) as *mut pyo3::ffi::PyObject
+                    };
+
+                    Python::with_gil(|py| {
+                        let code = unsafe {
+                            // SAFETY: code is guaranteed to be a valid PyCodeObject since it
+                            // came from `PyFrame_GetCode`.
+                            Bound::from_borrowed_ptr(py, code)
+                        };
+
+                        // Given a Python object that _should_ be a string, convert it into an
+                        // a Rust string.
+                        let unicode_converter = |pyobj: PyResult<Bound<'_, PyAny>>| {
+                            pyobj
+                                .map(|obj| {
+                                    if let Ok(filename) = obj.downcast::<pyo3::types::PyString>() {
+                                        Cow::from(filename.to_string_lossy().into_owned())
+                                    } else {
+                                        Cow::Borrowed("<UNKNOWN>")
+                                    }
+                                })
+                                .unwrap_or_else(|_| Cow::Borrowed("<NONE>"))
+                        };
+
+                        let filename =
+                            unicode_converter(code.getattr(pyo3::intern!(py, "co_filename")));
+                        let func_name = unicode_converter(code.getattr(if cfg!(Py_3_11) {
+                            // co_qualname is better but was only added in 3.11.
+                            pyo3::intern!(py, "co_qualname")
+                        } else {
+                            pyo3::intern!(py, "co_name")
+                        }));
+
+                        (filename, func_name)
+                    })
+                });
+
+                if what == pyo3::ffi::PyTrace_CALL {
+                    let linenumber = unsafe { pyo3::ffi::PyFrame_GetLineNumber(frame) }
+                        .try_into()
+                        .expect("Invalid line number for executing frame");
+
+                    function_trace::FunctionTrace::Call {
+                        filename: entry.0.clone(),
+                        func_name: entry.1.clone(),
+                        linenumber,
+                        time,
+                    }
+                    .serialize(&mut rmp_serde::encode::Serializer::new(&mut writer))
+                    .expect("Failed to emit Call message");
+                } else {
+                    function_trace::FunctionTrace::Return {
+                        func_name: entry.1.clone(),
+                        time,
+                    }
+                    .serialize(&mut rmp_serde::encode::Serializer::new(&mut writer))
+                    .expect("Failed to emit Return message")
                 }
             }
             pyo3::ffi::PyTrace_C_CALL => {

functiontrace-0.4.4/setup.py

@@ -0,0 +1,5 @@
+from setuptools import setup
+
+
+if __name__ == "__main__":
+    setup(py_modules=["functiontrace"])