functiontrace 0.3.10__tar.gz → 0.4.0__tar.gz

functiontrace-0.4.0/Cargo.toml

@@ -0,0 +1,19 @@
+[package]
+name = "_functiontrace-rs"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+pyo3 = "0.24"
+
+[profile.release-lto]
+inherits = "release"
+lto = true
+
+[lib]
+# See https://github.com/PyO3/pyo3 for details
+name = "_lib"  # private module to be nested into Python package
+crate-type = ["cdylib"]
+path = "rust/lib.rs"

functiontrace-0.4.0/MANIFEST.in

@@ -0,0 +1,3 @@
+graft mpack
+graft rust
+include Cargo.toml

{functiontrace-0.3.10/functiontrace.egg-info → functiontrace-0.4.0}/PKG-INFO

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

{functiontrace-0.3.10 → functiontrace-0.4.0}/_functiontrace.c

@@ -10,6 +10,9 @@
 
 #include "mpack/mpack.h"
 
+// We use the Rust implementation for all data transport.
+#define RUST_SOCKETS
+
 //#include <assert.h>
 // By default, python c extensions compile with asserts disable and it is not
 // super easy to disable this behavior (AFAIK). Use this custom assert for now.
@@ -56,7 +59,9 @@ typedef struct {
 // Prototypes
 ////////////////////////////////////////////////////////////////////////////////
 static Writer* Fprofile_Writer(void);
+#if !defined(RUST_SOCKETS)
 static void    Mpack_Flush(Writer* writer, const char* buffer, size_t bytes);
+#endif
 
 static Timestamp Fprofile_Time(void);
 
@@ -117,6 +122,11 @@ static struct {
     // The Unix socket address we're using to communicate with the profile
     // server.
     struct sockaddr_un socket;
+
+#if defined(RUST_SOCKETS)
+    // The Rust extension.
+    PyObject* rust;
+#endif
 } moduleConfiguration;
 
 // The set of methods exposed by this module to Python.
@@ -146,7 +156,7 @@ static PyMethodDef methods[] = {
         .ml_flags = METH_VARARGS,
         .ml_doc   = "INTERNAL USE ONLY"
     },
-    [METHOD_TERMINATE ] = {
+    [METHOD_TERMINATE] = {
         .ml_name  = "terminate",
         .ml_meth  = (PyCFunction) Fprofile_Terminate,
         .ml_flags = METH_NOARGS,
@@ -181,6 +191,7 @@ static Writer* Fprofile_Writer(void) {
     return &state->writer;
 }
 
+#if !defined(RUST_SOCKETS)
 static void Mpack_Flush(Writer* writer, const char* buffer, size_t bytes) {
     // The context passed to Mpack is the socket to send on.
     intptr_t socket = (intptr_t) writer->context;
@@ -190,6 +201,7 @@ static void Mpack_Flush(Writer* writer, const char* buffer, size_t bytes) {
         exit(-1);
     }
 }
+#endif
 
 ////////////////////////////////////////////////////////////////////////////////
 // Misc
@@ -732,6 +744,9 @@ static void Fprofile_ThreadTeardown(void* key) {
     }
 
     mpack_writer_flush_message(&state->writer);
+
+    // TODO: We should tear down the socket to avoid leaking O(# threads)
+    // resources.
 }
 
 static PyObject* Fprofile_MpexitHook(HookedFunction* hook, PyObject* args, PyObject* kwargs) {
@@ -836,10 +851,12 @@ static PyObject* Fprofile_StartTrace(PyObject* obj, PyObject* args, PyObject* kw
         Timestamp   time                      = Fprofile_Time();
         char        traceInitialization[1024] = { 0 };
         Writer      initWriter                = { 0 };
-        intptr_t    sock                      = -1;
         const char* version                   = Py_GetVersion();
         const char* platform                  = Py_GetPlatform();
         char*       outputDirectory           = NULL;
+#if !defined(RUST_SOCKETS)
+        intptr_t    sock                      = -1;
+#endif
 
         // Parse the args to figure out what our output directory is.
         if (!PyArg_ParseTuple(args, "s", &outputDirectory)) {
@@ -848,7 +865,22 @@ static PyObject* Fprofile_StartTrace(PyObject* obj, PyObject* args, PyObject* kw
             exit(-1);
         }
 
-        // Spawn the server and wait for it to connect.
+#if defined(RUST_SOCKETS)
+        // NOTE: We only import this module once.  Importing while tracing is
+        // live will be recorded, which exposes internals users don't care
+        // about and breaks the invariant that we send a `RegisterThread`
+        // message first when threads are traced.
+        moduleConfiguration.rust = PyImport_ImportModule("_functiontrace_rs._lib");
+
+        // Verify that the rust extension loads if we're using it.
+        if (moduleConfiguration.rust == NULL) {
+            perror("Failed to load internal Rust extension");
+            exit(-1);
+        }
+#endif
+
+        // Spawn the server.  We'll wait a bit for it to start accepting
+        // sockets.
         pid_t server = fork();
         if (server == 0) {
             // Detach ourselves as a daemon
@@ -876,6 +908,17 @@ static PyObject* Fprofile_StartTrace(PyObject* obj, PyObject* args, PyObject* kw
                 sizeof(moduleConfiguration.socket.sun_path),
                 "/tmp/functiontrace-server.sock.%d", server);
 
+        // Register an env var so subprocesses know how to connect to our
+        // profile.
+        if (setenv(breadcrumb, (char*) &moduleConfiguration.socket.sun_path, 0)) {
+            perror("Failed to set recursive breadcrumb");
+            exit(-1);
+        }
+
+        // Setup a local Mpack context for this first message.
+        mpack_writer_init(&initWriter, (char*) &traceInitialization, sizeof(traceInitialization));
+
+#if !defined(RUST_SOCKETS)
         // Generate the TraceInitialization message to the profiler server telling
         // them about the profile we'll be creating.
         if ((sock = socket(AF_UNIX, SOCK_STREAM, 0)) == -1) {
@@ -902,17 +945,26 @@ static PyObject* Fprofile_StartTrace(PyObject* obj, PyObject* args, PyObject* kw
             usleep(1000);
         }
 
-        // Register an env var so subprocesses know how to connect to our
-        // profile.
-        if (setenv(breadcrumb, (char*) &moduleConfiguration.socket.sun_path, 0)) {
-            perror("Failed to set recursive breadcrumb");
-            exit(-1);
-        }
-
-        // Setup a local Mpack context for this first message.
-        mpack_writer_init(&initWriter, (char*) &traceInitialization, sizeof(traceInitialization));
         mpack_writer_set_context(&initWriter, (void*) sock);
         mpack_writer_set_flush(&initWriter, &Mpack_Flush);
+#else
+        {
+            PyObject* msg_init = PyObject_GetAttrString(moduleConfiguration.rust, "message_initialization");
+            assert(msg_init != NULL);
+
+            PyObject* arg = PyLong_FromVoidPtr(&moduleConfiguration.socket.sun_path);
+            PyObject* context = PyObject_CallFunctionObjArgs(msg_init, arg, NULL);
+            assert(context != NULL);
+
+            PyObject* msg_flush = PyObject_GetAttrString(moduleConfiguration.rust, "message_flush");
+            assert(msg_flush != NULL);
+            PyObject* mpack_flush = PyObject_CallFunctionObjArgs(msg_flush, NULL);
+            assert(mpack_flush != NULL);
+
+            mpack_writer_set_context(&initWriter, PyLong_AsVoidPtr(context));
+            mpack_writer_set_flush(&initWriter, PyLong_AsVoidPtr(mpack_flush));
+        }
+#endif
 
         mpack_start_map(&initWriter, 5);
         {
@@ -947,10 +999,22 @@ static PyObject* Fprofile_StartTrace(PyObject* obj, PyObject* args, PyObject* kw
         // out first.
         mpack_writer_flush_message(&initWriter);
 
+        // Close the initialization socket to trigger `functiontrace-server` to
+        // read our message, since otherwise it'll wait forever for more data.
+#if !defined(RUST_SOCKETS)
         if (shutdown(sock, SHUT_WR) != 0) {
             perror("Failed to close TraceInitialization socket");
             exit(-1);
         }
+#else
+        {
+            PyObject* msg_shutdown = PyObject_GetAttrString(moduleConfiguration.rust, "message_shutdown");
+            assert(msg_shutdown != NULL);
+
+            PyObject* socket = PyLong_FromVoidPtr(mpack_writer_context(&initWriter));
+            assert(PyObject_CallFunctionObjArgs(msg_shutdown, socket, NULL) != NULL);
+        }
+#endif
     } else {
         // Read the socket addr out of `existingProfile` so our threads know
         // how to talk to the profile server.
@@ -1033,6 +1097,14 @@ static ThreadState* Fprofile_CreateThreadState(void) {
         exit(-1);
     }
 
+    // Setup msgpack and a socket for this thread to report tracing info on.
+    //
+    // We need to do this before any allocations or Python function calls
+    // occur, since consumers can't tell whether mpack is fully initialized or
+    // not.
+    mpack_writer_init(&state->writer, (char*) &state->buffer, sizeof(state->buffer));
+
+#if !defined(RUST_SOCKETS)
     // Allocate a socket to report this thread's information back to the
     // profile server.
     if ((state->socket = socket(AF_UNIX, SOCK_STREAM, 0)) == -1) {
@@ -1046,11 +1118,33 @@ static ThreadState* Fprofile_CreateThreadState(void) {
         exit(-1);
     }
 
-    // Register a messagepack writer for this thread, including telling it to
-    // flush on this thread's socket once the buffer fills up.
-    mpack_writer_init(&state->writer, (char*) &state->buffer, sizeof(state->buffer));
+    // Configure Mpack to flush events to this thread's socket once the buffer
+    // fills up.
     mpack_writer_set_context(&state->writer, (void*) state->socket);
     mpack_writer_set_flush(&state->writer, &Mpack_Flush);
+#else
+    {
+        PyObject* msg_init = PyObject_GetAttrString(moduleConfiguration.rust, "message_initialization");
+        assert(msg_init != NULL);
+
+        // TODO: Convert this to a shared Rust configuration object so we don't
+        // need to pass it back and forth.
+        PyObject* arg = PyLong_FromVoidPtr(&moduleConfiguration.socket.sun_path);
+        PyObject* context = PyObject_CallFunctionObjArgs(msg_init, arg, NULL);
+        assert(context != NULL);
+
+        PyObject* msg_flush = PyObject_GetAttrString(moduleConfiguration.rust, "message_flush");
+        assert(msg_flush != NULL);
+        PyObject* mpack_flush = PyObject_CallFunctionObjArgs(msg_flush, NULL);
+        assert(mpack_flush != NULL);
+
+        // Re-initialize msgpack, since we've done some allocations and Python
+        // calls that we don't want to actually record.
+        mpack_writer_init(&state->writer, (char*) &state->buffer, sizeof(state->buffer));
+        mpack_writer_set_context(&state->writer, PyLong_AsVoidPtr(context));
+        mpack_writer_set_flush(&state->writer, PyLong_AsVoidPtr(mpack_flush));
+    }
+#endif
 
     // Send a message about our thread's registration.
     Timestamp time = Fprofile_Time();

{functiontrace-0.3.10 → functiontrace-0.4.0/functiontrace.egg-info}/PKG-INFO

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

{functiontrace-0.3.10 → functiontrace-0.4.0}/functiontrace.egg-info/SOURCES.txt

@@ -1,3 +1,4 @@
+Cargo.toml
 MANIFEST.in
 README.rst
 _functiontrace.c
@@ -11,4 +12,5 @@ functiontrace.egg-info/entry_points.txt
 functiontrace.egg-info/top_level.txt
 mpack/LICENSE
 mpack/mpack.c
-mpack/mpack.h
+mpack/mpack.h
+rust/lib.rs

{functiontrace-0.3.10 → functiontrace-0.4.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "functiontrace"
-version = "0.3.10"
+version = "0.4.0"
 description = "A profiling tool that traces your application's execution"
 authors = [
     { name = "Matt Bryant", email = "functiontrace@programsareproofs.com"},
@@ -41,5 +41,9 @@ functiontrace = "functiontrace:main"
 functiontrace = "functiontrace"
 
 [build-system]
-requires = ["setuptools", "toml"]
+requires = ["setuptools", "toml", "setuptools-rust"]
 build-backend = "setuptools.build_meta"
+
+[[tool.setuptools-rust.ext-modules]]
+target = "_functiontrace_rs._lib" # Must match `lib.name` in Cargo.toml
+binding = "PyO3"

functiontrace-0.4.0/rust/lib.rs

@@ -0,0 +1,100 @@
+#[pyo3::pymodule]
+mod _lib {
+    use pyo3::prelude::*;
+    use std::ffi::{CStr, OsStr};
+    use std::os::raw::c_char;
+    use std::os::unix::ffi::OsStrExt;
+    use std::os::unix::net::UnixStream;
+
+    ////////////////////////////////////////////////////////////////////////////////////////////////
+    // Socket Communication
+    ////////////////////////////////////////////////////////////////////////////////////////////////
+
+    /// Initializes and returns a pointer to the [`UnixStream`] we'll communicate to
+    /// `functiontrace-server` on.
+    ///
+    /// The returned value will be stored as [`MpackWriter.context`] and used by [`Mpack_Flush`].
+    #[pyfunction]
+    fn message_initialization(socket: usize) -> usize {
+        let c_socket = unsafe {
+            // SAFETY: socket is a valid `char*` converted via `PyLong_FromVoidPtr`.
+            CStr::from_ptr(socket as *const c_char)
+        };
+        let sockaddr = OsStr::from_bytes(c_socket.to_bytes());
+
+        // The functiontrace-server might not be ready to receive connections yet, so we retry for
+        // a bit.
+        let start = std::time::Instant::now();
+        let socket = loop {
+            match UnixStream::connect(sockaddr) {
+                Ok(s) => {
+                    break Box::new(s);
+                }
+                _ => {
+                    if start.elapsed() > std::time::Duration::from_millis(1000) {
+                        panic!("Timed out trying to connect to functiontrace-server");
+                    }
+
+                    std::thread::sleep(std::time::Duration::from_millis(10));
+                }
+            }
+        };
+
+        // This must be freed by `message_shutdown`
+        Box::into_raw(socket) as usize
+    }
+
+    /// Shutdown a socket that was opened by [`message_initialize`], triggering
+    /// `functiontrace-server` to read the data sent over it.
+    ///
+    /// If we forget this, we'll leak sockets and a bit of memory.
+    #[pyfunction]
+    fn message_shutdown(socket: usize) {
+        let socket = unsafe {
+            // SAFETY: socket is a `UnixStream` emitted from `message_initialization` encoded via
+            // `PyLong_FromVoidPtr`.
+            Box::from_raw(socket as *mut UnixStream)
+        };
+
+        if let Err(e) = socket.shutdown(std::net::Shutdown::Both) {
+            panic!("Failed to close socket: {e}");
+        }
+    }
+
+    /// The initial subset of the `mpack_writer_t` struct, which `mpack.h` will use to pass us our
+    /// per-thread context in.
+    #[repr(C)]
+    struct MpackWriter {
+        flush: *const (),
+        error_fn: *const (),
+        teardown: *const (),
+        context: *mut UnixStream,
+    }
+
+    /// 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) {
+        use std::io::Write;
+
+        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
+            // `message_initialization`.
+            &*(*writer).context
+        };
+        if let Err(e) = socket.write_all(data) {
+            panic!("Socket send failed: {e}");
+        }
+    }
+
+    /// C accessor for [`Mpack_Flush`]
+    #[pyfunction]
+    fn message_flush() -> PyResult<usize> {
+        Ok(Mpack_Flush as usize)
+    }
+}

functiontrace-0.3.10/MANIFEST.in

@@ -1 +0,0 @@
-graft mpack