functiontrace 0.6.0__tar.gz → 0.6.2__tar.gz

{functiontrace-0.6.0/functiontrace.egg-info → functiontrace-0.6.2}/PKG-INFO

@@ -1,9 +1,9 @@
 Metadata-Version: 2.4
 Name: functiontrace
-Version: 0.6.0
+Version: 0.6.2
 Summary: A profiling tool that traces your application's execution
-Author: Alistair Dobke
-Author-email: Matt Bryant <functiontrace@programsareproofs.com>
+Author: Matt Bryant, Alistair Dobke
+Maintainer-email: Matt Bryant <functiontrace@programsareproofs.com>
 License-Expression: LicenseRef-Prosperity-Public-License-3.0.0
 Project-URL: homepage, https://functiontrace.com
 Project-URL: repository, https://gitlab.com/mbryant/functiontrace

{functiontrace-0.6.0 → functiontrace-0.6.2/functiontrace.egg-info}/PKG-INFO

@@ -1,9 +1,9 @@
 Metadata-Version: 2.4
 Name: functiontrace
-Version: 0.6.0
+Version: 0.6.2
 Summary: A profiling tool that traces your application's execution
-Author: Alistair Dobke
-Author-email: Matt Bryant <functiontrace@programsareproofs.com>
+Author: Matt Bryant, Alistair Dobke
+Maintainer-email: Matt Bryant <functiontrace@programsareproofs.com>
 License-Expression: LicenseRef-Prosperity-Public-License-3.0.0
 Project-URL: homepage, https://functiontrace.com
 Project-URL: repository, https://gitlab.com/mbryant/functiontrace

{functiontrace-0.6.0 → functiontrace-0.6.2}/pyproject.toml

@@ -1,11 +1,14 @@
 [project]
 name = "functiontrace"
-version = "0.6.0"
+version = "0.6.2"
 description = "A profiling tool that traces your application's execution"
 authors = [
-    { name = "Matt Bryant", email = "functiontrace@programsareproofs.com"},
+    { name = "Matt Bryant" },
     { name = "Alistair Dobke" }
 ]
+maintainers = [
+    { name = "Matt Bryant", email = "functiontrace@programsareproofs.com" }
+]
 license = "LicenseRef-Prosperity-Public-License-3.0.0"
 license-files = ["LICENSE.md"]
 readme = "README.rst"

{functiontrace-0.6.0 → functiontrace-0.6.2}/rust/lib.rs

@@ -8,7 +8,7 @@ mod _functiontrace_rs {
     use std::borrow::Cow;
     use std::cell::RefCell;
     use std::ffi::{CStr, OsString};
-    use std::io::Write;
+    use std::io::{BufWriter, Read, Write};
     use std::os::unix::net::UnixStream;
     use std::sync::atomic::{AtomicBool, Ordering};
     use std::sync::{LazyLock, OnceLock};
@@ -29,25 +29,18 @@ mod _functiontrace_rs {
         license_key: Option<String>,
     }
 
-    /// The size of messagepack buffers.  This is picked somewhat arbitrarily, but should be
+    /// 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 MSGPACK_BUFFER: usize = 1 << 17;
 
     /// Storage specific to each thread used for buffering/transmitting data.
     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,
+        writer: BufWriter<UnixStream>,
     }
 
     impl Write for ThreadState {
-        /// A small [`Write`] implementation that writes to a buffer and only flushes to the socket
-        /// once full.
+        /// Buffer trace data for the socket while tracing is enabled.
         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
@@ -55,55 +48,12 @@ mod _functiontrace_rs {
                 return Ok(src.len());
             }
 
-            let bytes = src.len();
-            let tail = self.head + bytes;
-
-            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,
-                        "msgpack buffer too small",
-                    ));
-                }
-
-                // We can't fit any more data, so need to flush. The caller should retry after
-                // this.
-                self.flush()?;
-
-                return Err(std::io::Error::new(
-                    std::io::ErrorKind::Interrupted,
-                    "msgpack buffer is full, flushing",
-                ));
-            }
-
-            self.buffer[self.head..tail].copy_from_slice(src);
-            self.head = tail;
-
-            Ok(bytes)
+            self.writer.write(src)
         }
 
         /// Flush our local buffer to the socket.
         fn flush(&mut self) -> std::io::Result<()> {
-            // We're about to send all 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.
-            let head = std::mem::take(&mut self.head);
-
-            if head != 0 {
-                self.socket.write_all(&self.buffer[..head])
-            } else {
-                Ok(())
-            }
-        }
-    }
-
-    impl Drop for ThreadState {
-        // Attempt to flush the buffer on drop. This should be run whenever TLS is torn down for
-        // threads, which should allow us to collect data from short-lived threads that don't call
-        // `thread_teardown` themselves.
-        fn drop(&mut self) {
-            let _ = self.flush();
+            self.writer.flush()
         }
     }
 
@@ -427,8 +377,28 @@ mod _functiontrace_rs {
             // Close the initialization socket to trigger `functiontrace-server` to
             // read our message, since otherwise it'll wait forever for more data.
             socket
-                .shutdown(std::net::Shutdown::Both)
-                .wrap_err("Failed to close initialization socket")?;
+                .shutdown(std::net::Shutdown::Write)
+                .wrap_err("Failed to finish initialization message")?;
+
+            // Modern `functiontrace-server`s send their version, allowing us to ensure
+            // protocols match.
+            let version = {
+                let mut reply = Vec::with_capacity(16);
+                socket
+                    .read_to_end(&mut reply)
+                    .wrap_err("Failed to read server version")?;
+
+                if reply.is_empty() {
+                    // Old servers will close the socket without sending anything. We know they're
+                    // older than 0.8.6, which is sufficient for us to handle.
+                    "0.8.6".to_string()
+                } else {
+                    rmp_serde::from_slice(&reply)
+                        .wrap_err("Failed to parse server version as string")?
+                }
+            };
+
+            tracing::info!("Connected to functiontrace-server {version}");
         }
 
         // `multiprocessing` uses an absolute path to an executable to launch forkserver workers
@@ -616,10 +586,7 @@ mod _functiontrace_rs {
 
         // Initialize the `ThreadState` for this thread, then associate the socket with it.
         let state = Box::new(ThreadState {
-            socket,
-
-            buffer: [0u8; MSGPACK_BUFFER],
-            head: 0,
+            writer: BufWriter::with_capacity(MSGPACK_BUFFER, socket),
         });
 
         // Store our state in thread-specific storage so we can find it in the future.
@@ -633,8 +600,8 @@ mod _functiontrace_rs {
     /// Tear down the given thread by dropping its [`ThreadState`], which will also flush any
     /// outstanding messages.
     ///
-    /// This is called when the thread has shut down but we can't rely on `ThreadState::drop`
-    /// running, including in multiprocess situations.
+    /// This is called when the thread has shut down but we can't rely on TLS destructors running,
+    /// including in multiprocess situations.
     pub fn thread_teardown() {
         let _ = THREAD_STATE.take();
     }
@@ -656,12 +623,11 @@ mod _functiontrace_rs {
         }
 
         // Remove the old state without flushing the parent's socket.
-        let mut state = THREAD_STATE
+        let ThreadState { writer } = *THREAD_STATE
             .take()
             .ok_or_eyre(eyre!("Thread wasn't yet registered"))?;
-        // Empty the buffer so the `Drop` impl doesn't flush it.
-        state.head = 0;
-        drop(state);
+        let (socket, _buffer) = writer.into_parts();
+        drop(socket);
 
         // 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.
@@ -1664,11 +1630,13 @@ mod hooks {
 
     /// A generic wrapper that parses out the arguments, calls the given handler function to handle
     /// custom hooking logic, then proxies the arguments out to the original wrapped function.
+    /// The handler emits events synchronously so they can borrow from its local Python objects.
     fn proxy(
         hook: &Hook,
         handler: impl FnOnce(
             &Bound<'_, pyo3::types::PyTuple>,
-        ) -> Option<function_trace::FunctionTrace<'static>>,
+            &mut dyn FnMut(function_trace::FunctionTrace<'_>),
+        ),
         args: *mut PyObject,
         kwargs: *mut PyObject,
     ) -> *mut PyObject {
@@ -1688,10 +1656,14 @@ mod hooks {
                     .map_err(Into::<PyErr>::into)
                     .wrap_err("Python args must be a tuple")?;
 
-                if let Some(msg) = handler(args) {
+                handler(args, &mut |msg| {
                     msg.serialize(&mut rmp_serde::encode::Serializer::new(&mut writer))
-                        .wrap_err("Failed to emit Allocation message")?;
-                }
+                        .expect(
+                            // TODO: This isn't entirely true, but you'd have to be actively
+                            // malicious to hit this case.
+                            "Hooks emit capped size events so can't fail to serialize",
+                        );
+                });
                 Ok(())
             })
             .transpose()?;
@@ -1730,18 +1702,27 @@ mod hooks {
 
         proxy(
             hook,
-            |args| {
+            |args, emit| {
                 // Logging functions should roughly call `str()` on their arguments
-                // TODO: We shouldn't need to explicitly allocate for this.
-                let args = args
-                    .str()
-                    .and_then(|x| x.to_str().map(|x| Cow::Owned(x.to_string())))
-                    .unwrap_or(Cow::Borrowed("<invalid string>"));
+                let args = args.str().ok();
+                let value = args
+                    .as_ref()
+                    .and_then(|args| args.to_str().ok())
+                    .unwrap_or("<invalid string>");
 
-                Some(function_trace::FunctionTrace::Log {
+                let value = if value.chars().count() > 1024 {
+                    let mut value: String = value.chars().take(1021).collect();
+                    value.push_str("...");
+
+                    Cow::Owned(value)
+                } else {
+                    Cow::Borrowed(value)
+                };
+
+                emit(function_trace::FunctionTrace::Log {
                     time: trace_time(),
                     log_type: Cow::Borrowed(hook.target),
-                    log_value: args,
+                    log_value: value,
                 })
             },
             args,
@@ -1759,22 +1740,22 @@ mod hooks {
 
         proxy(
             hook,
-            |args| {
+            |args, emit| {
                 // `import` takes rather complicated arguments, but in practice we only care about
                 // the first argument (name)
-                let module = if let Ok(arg) = args.get_item(0) {
+                let arg = args.get_item(0);
+                let module = if let Ok(arg) = &arg {
                     arg.cast()
                         .map_err(Into::into)
                         .and_then(|module| module.to_str())
-                        .map(|module| Cow::Owned(module.to_string()))
-                        .unwrap_or(Cow::Borrowed("<unknown module>"))
+                        .unwrap_or("<unknown module>")
                 } else {
-                    Cow::Borrowed("<no module specified>")
+                    "<no module specified>"
                 };
 
-                Some(function_trace::FunctionTrace::Import {
+                emit(function_trace::FunctionTrace::Import {
                     time: trace_time(),
-                    module_name: module,
+                    module_name: Cow::Borrowed(module),
                 })
             },
             args,
@@ -1788,6 +1769,9 @@ mod hooks {
     /// Anything that exits with `os._exit` (including `multiprocessing` by default) will
     /// immediately terminate the process without running cleanup/`atexit()`. Hook this to ensure we
     /// can flush as needed.
+    ///
+    /// TODO: We should also track any other outstanding threads and flush their buffers, since they
+    /// won't get the chance to cleanly exit.
     extern "C" fn process_exit<const HOOK: usize>(
         _self: *mut PyObject,
         args: *mut PyObject,
@@ -1798,6 +1782,6 @@ mod hooks {
         crate::_functiontrace_rs::thread_teardown();
 
         // Directly pass the command on, since we don't need to log/etc for this.
-        proxy(hook, |_args| None, args, kwargs)
+        proxy(hook, |_args, _emit| {}, args, kwargs)
     }
 }