functiontrace 0.5.3__tar.gz → 0.6.0__tar.gz

{functiontrace-0.5.3 → functiontrace-0.6.0}/Cargo.toml

@@ -10,7 +10,7 @@ functiontrace-server = { version = "0.8.4", default-features = false, features =
 jiff = { version = "0.2", default-features = false, features = ["std", "serde"] }
 libc = "0.2"
 owo-colors = "4.2"
-pyo3 = { version = "0.29", features = ["eyre"] }
+pyo3 = { version = "0.30", features = ["eyre"] }
 reqwest = { version = "0.13", default-features = false, features = ["blocking", "json", "rustls"] }
 rmp-serde = "1.3"
 serde = "1"
@@ -20,7 +20,7 @@ tracing-subscriber = { version = "0.3", default-features = false, features = ["a
 
 [build-dependencies]
 toml = { version = "1.1", default-features = false, features = ["parse", "serde"] }
-pyo3-build-config = { version = "0.29" }
+pyo3-build-config = { version = "0.30" }
 
 [lib]
 # See https://github.com/PyO3/pyo3 for details

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

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

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

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

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

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

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

@@ -6,7 +6,7 @@ mod _functiontrace_rs {
     use pyo3::prelude::*;
     use serde::{Deserialize, Serialize};
     use std::borrow::Cow;
-    use std::cell::Cell;
+    use std::cell::RefCell;
     use std::ffi::{CStr, OsString};
     use std::io::Write;
     use std::os::unix::net::UnixStream;
@@ -50,7 +50,7 @@ mod _functiontrace_rs {
         /// 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
+                // We haven't started yet, so we shouldn't write anything out. This avoids logging
                 // various internal setup things before we're ready.
                 return Ok(src.len());
             }
@@ -58,8 +58,8 @@ mod _functiontrace_rs {
             let bytes = src.len();
             let tail = self.head + bytes;
 
-            if tail >= self.buffer.len() {
-                if bytes >= self.buffer.len() {
+            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,
@@ -67,7 +67,7 @@ mod _functiontrace_rs {
                     ));
                 }
 
-                // We can't fit any more data, so need to flush.  The caller should retry after
+                // We can't fit any more data, so need to flush. The caller should retry after
                 // this.
                 self.flush()?;
 
@@ -85,13 +85,25 @@ mod _functiontrace_rs {
 
         /// 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
+            // 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 = self.head;
-            self.head = 0;
+            let head = std::mem::take(&mut self.head);
 
-            self.socket.write_all(&self.buffer[..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();
         }
     }
 
@@ -105,16 +117,15 @@ mod _functiontrace_rs {
     /// `sys.argv` for our process, which is used when registering new threads.
     static ARGV: OnceLock<String> = OnceLock::new();
 
-    /// True iff we've started tracing and haven't been marked as terminated.  When this is
+    /// True iff we've started tracing and haven't been marked as terminated. When this is
     /// set, we're allowed to send messages to the profile generation server.
     static STARTED: AtomicBool = AtomicBool::new(false);
 
     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()) };
+        // NOTE: The `const` block allows the macro to avoid lazy initialization.
+        static THREAD_STATE: RefCell<Option<Box<ThreadState>>> = const { RefCell::new(None) };
     }
 
     ////////////////////////////////////////////////////////////////////////////////////////////////
@@ -140,22 +151,17 @@ mod _functiontrace_rs {
         Duration::new(time.tv_sec as u64, time.tv_nsec as u32)
     }
 
-    /// Return the [`ThreadState`] for the current thread.  This should either be None (during
-    /// startup), or a valid [`ThreadState`] that has been allocated on the heap and leaked via
-    /// `Box::into_raw`.
+    /// Run the given callback under the [`ThreadState`] for the current thread.
+    /// Returns `None` if the thread is not registered or the state is currently borrowed.
+    ///
+    /// NOTE: The latter should only happen if we're running Python code inside a tracing event,
+    /// which is something we don't actually want to trace (it would be tracing FunctionTrace
+    /// implementation rather than the user's code).
     #[inline(always)]
-    pub fn thread_state() -> Option<&'static mut ThreadState> {
-        let state = THREAD_STATE.get();
-
-        if state.is_null() {
-            // We shouldn't be asked to reset a thread that hasn't yet been registered.
-            return None;
-        }
-
-        Some(unsafe {
-            // SAFETY: state is a valid pointer to a ThreadState, since we always remove from TLS
-            // before freeing.
-            state.as_mut().expect("state is non-null")
+    pub fn thread_state<T>(access: impl FnOnce(&mut ThreadState) -> T) -> Option<T> {
+        THREAD_STATE.with(|state| {
+            let mut state = state.try_borrow_mut().ok()?;
+            Some(access(state.as_deref_mut()?))
         })
     }
 
@@ -288,7 +294,7 @@ mod _functiontrace_rs {
                 // We'll display this in the UI, so ensure it's under (the rather
                 // arbitrary) 100 chars.
                 if full_args.chars().count() > 95 {
-                    full_args.truncate(95);
+                    full_args = full_args.chars().take(95).collect();
                     full_args.push_str(" ...");
                 }
 
@@ -446,6 +452,7 @@ mod _functiontrace_rs {
             };
 
             // Create a multiprocessing wrapper for the current executable if it doesn't exist yet.
+            // TODO: This is slow. Better to just hash the executable.
             std::fs::create_dir_all(
                 wrapper
                     .parent()
@@ -465,6 +472,8 @@ mod _functiontrace_rs {
                         &format!("'{}'", executable.replace('\'', "'\"'\"'")),
                     );
 
+                    // TODO: this needs to be atomic, since multiple children could hit this case at
+                    // once.
                     file.write_all(script.as_bytes())
                         .wrap_err("Failed to write multiprocessing wrapper")?;
                     file.set_permissions(std::fs::Permissions::from_mode(0o700))
@@ -486,18 +495,17 @@ mod _functiontrace_rs {
         // Ensure we're properly handling teardown scenarios, including flushing data in both
         // single and multi-threaded scenarios, and reseting anything that's needed on forks.
 
-        // Mark that we'll need to forget some information on forks.  In
-        // particular, we shouldn't think that we have a thread that's sent any
-        // information.
+        // Mark that we'll need to forget some information on forks. In particular, we shouldn't
+        // think that we have a thread that's sent any information.
         if unsafe { libc::pthread_atfork(None, None, Some(c_thread_reset)) } != 0 {
             return Err(eyre!("Failed to register pthread_atfork() handler"));
         }
 
         // Things can get into a weird state during shutdown due to GC
-        // (bugs.python.org/issue21512).  Halt our tracing instead to avoid odd behaviour and
-        // ensure we've flushed all our data.
+        // (bugs.python.org/issue21512). Halt our tracing instead to avoid odd behavior and ensure
+        // we've flushed all our data.
         //
-        // NOTE: We specifically do this rather than `Py_AtExit`  since we want to run before all
+        // NOTE: We specifically do this rather than `Py_AtExit` since we want to run before all
         // the Python teardown occurs.
         Python::attach(|py| -> Result<()> {
             py.import("atexit")
@@ -562,7 +570,8 @@ mod _functiontrace_rs {
         _: &Bound<'_, pyo3::types::PyTuple>,
         _kwargs: Option<&Bound<'_, pyo3::types::PyDict>>,
     ) -> Result<()> {
-        // Stop profiling and don't allow further logging, then flush any remaining messages.
+        // Stop profiling and don't allow further logging.
+        // The outstanding messages will be flushed on exit when the ThreadState is dropped.
         STARTED.store(false, Ordering::Relaxed);
 
         unsafe {
@@ -570,23 +579,19 @@ mod _functiontrace_rs {
             pyo3::ffi::PyEval_SetProfile(None, std::ptr::null_mut());
         }
 
-        // TODO: We could easily keep track of all the outstanding Writers and flush them all too.
-        if let Some(writer) = thread_state() {
-            thread_teardown(writer).wrap_err("Failed to teardown main thread")?;
-        }
+        // TLS destruction isn't guaranteed to run on the main thread's shutdown.
+        thread_teardown();
 
         Ok(())
     }
 
     /// Register a new thread, including creating the UnixStream to log messages for this thread
-    /// on.
-    ///
-    /// A new [`ThreadState`] will be returned, which must freed once the thread has terminated.
+    /// on. The new [`ThreadState`] is stored in TLS until until [`thread_teardown`] is run.
     fn thread_register() -> Result<()> {
-        if thread_state().is_some() {
-            // We shouldn't have a ThreadState yet, as we're the ones that create it.
-            return Err(eyre!("Thread has already been registered"));
-        }
+        assert!(
+            THREAD_STATE.with(|state| { state.borrow().is_none() }),
+            "We can't have a ThreadState yet since we're responsible for registering it"
+        );
 
         let register =
             function_trace::FunctionTrace::RegisterThread(function_trace::ThreadRegistration {
@@ -610,37 +615,28 @@ mod _functiontrace_rs {
             .wrap_err("Failed to emit RegisterThread message")?;
 
         // Initialize the `ThreadState` for this thread, then associate the socket with it.
-        //
-        // 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 {
+        let state = 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.
-        THREAD_STATE.set(state);
+        THREAD_STATE.set(Some(state));
 
         Ok(())
     }
 
-    /// Tear down the given thread by flushing any outstanding messages.
+    /// 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, including in multithread/process situations.
-    pub fn thread_teardown(writer: &mut ThreadState) -> Result<()> {
-        writer
-            .flush()
-            .wrap_err("Failed to flush remaining messages")?;
-
-        // TODO: We should tear down the socket to avoid leaking O(# threads) resources.
-        // Historically doing this has crashed for some reason...
-        //
-        // TODO: And we should probably free state since it's pretty large...
-        Ok(())
+    /// This is called when the thread has shut down but we can't rely on `ThreadState::drop`
+    /// running, including in multiprocess situations.
+    pub fn thread_teardown() {
+        let _ = THREAD_STATE.take();
     }
 
     #[unsafe(no_mangle)]
@@ -654,26 +650,21 @@ 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<()> {
         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.
+            // We haven't actually started yet, but are for some reason being asked to fork. This
+            // appears to be OS dependent.
             return Ok(());
         }
 
-        // Load the old state so we can free it.
-        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.
-        THREAD_STATE.set(std::ptr::null_mut());
+        // Remove the old state without flushing the parent's socket.
+        let mut state = 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);
 
-        // Free the thread's state, including the socket
-        unsafe {
-            // SAFETY: `state` came from `Box::into_raw` previously and can only be freed once.
-            let _ = Box::from_raw(state);
-        }
-
-        // 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.
+        // 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.
         thread_register().wrap_err("Failed to register new thread")?;
 
         Ok(())
@@ -766,7 +757,7 @@ mod _functiontrace_rs {
         Ok(())
     }
 
-    /// The core logging function for FunctionTrace.  Configured by `PyEval_SetProfile`, and used
+    /// 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,
@@ -774,142 +765,141 @@ mod _functiontrace_rs {
         what: i32,
         arg: *mut pyo3::ffi::PyObject,
     ) -> i32 {
-        let Some(mut writer) = thread_state() else {
-            // No reason to look at the event since we can't log it yet.
-            return 0;
-        };
-
         if !STARTED.load(Ordering::Relaxed) {
             // We aren't running (either starting up or shutting down), so shouldn't write
             // anything.
             return 0;
         }
 
-        let time = trace_time();
+        // We hold a mutable reference to the current ThreadState throughout this entire function.
+        // This ensures that we don't log anything from Python code that we call as part of the
+        // function (such as a `str()` implementation), since that's not part of the user's code.
+        thread_state(|mut writer| {
+            let time = trace_time();
 
-        // Return the function name when we're in a C_* context.
-        let c_func_name = || {
-            let func = arg as *mut pyo3::ffi::PyCFunctionObject;
+            // Return the function name when we're in a C_* context.
+            let c_func_name = || {
+                let func = arg as *mut pyo3::ffi::PyCFunctionObject;
 
-            unsafe {
-                // SAFETY: In the `C_*` cases, `arg` is a non-NULL `*PyCFunctionObject`, which
-                // always contains a valid method.
-                let method = (*func).m_ml;
-                (*method).ml_name
-            }
-        };
-        // Given a C string sourced from Python, convert it to a Cow that can be serialized.
-        let cstr_to_cow = |ptr: *const i8| {
-            if !ptr.is_null() {
                 unsafe {
-                    // SAFETY: `ptr` is non-NULL and contains a valid C string.
-                    // TODO: Can this change underneath us?
-                    CStr::from_ptr(ptr)
+                    // SAFETY: In the `C_*` cases, `arg` is a non-NULL `*PyCFunctionObject`, which
+                    // always contains a valid method.
+                    let method = (*func).m_ml;
+                    (*method).ml_name
                 }
-                .to_str()
-                .unwrap_or("<invalid string>")
-                .into()
-            } else {
-                Cow::Borrowed("<null>")
-            }
-        };
-        match what {
-            pyo3::ffi::PyTrace_CALL | pyo3::ffi::PyTrace_RETURN => {
-                // 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::attach(|py| {
+            };
+            // Given a C string sourced from Python, convert it to a Cow that can be serialized.
+            let cstr_to_cow = |ptr: *const i8| {
+                if !ptr.is_null() {
+                    unsafe {
+                        // SAFETY: `ptr` is non-NULL and contains a valid C string.
+                        // TODO: Can this change underneath us?
+                        CStr::from_ptr(ptr)
+                    }
+                    .to_str()
+                    .unwrap_or("<invalid string>")
+                    .into()
+                } else {
+                    Cow::Borrowed("<null>")
+                }
+            };
+            match what {
+                pyo3::ffi::PyTrace_CALL | pyo3::ffi::PyTrace_RETURN => {
+                    // 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: code is guaranteed to be a valid PyCodeObject since it
-                            // came from `PyFrame_GetCode`.
-                            Bound::from_borrowed_ptr(py, code)
-                        };
+                            // SAFETY: frame is guaranteed to be a valid PyFrameObject
 
-                        // 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.cast::<pyo3::types::PyString>() {
-                                        Cow::from(filename.to_string_lossy().into_owned())
-                                    } else {
-                                        Cow::Borrowed("<UNKNOWN>")
-                                    }
-                                })
-                                .unwrap_or_else(|_| Cow::Borrowed("<NONE>"))
+                            // 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
                         };
 
-                        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")
-                        }));
+                        Python::attach(|py| {
+                            let code = unsafe {
+                                // SAFETY: code is guaranteed to be a valid PyCodeObject since it
+                                // came from `PyFrame_GetCode`.
+                                Bound::from_owned_ptr(py, code)
+                            };
 
-                        (filename, func_name)
-                    })
-                });
+                            // 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.cast::<pyo3::types::PyString>() {
+                                            Cow::from(filename.to_string_lossy().into_owned())
+                                        } else {
+                                            Cow::Borrowed("<UNKNOWN>")
+                                        }
+                                    })
+                                    .unwrap_or_else(|_| Cow::Borrowed("<NONE>"))
+                            };
 
-                if what == pyo3::ffi::PyTrace_CALL {
-                    let linenumber = unsafe { pyo3::ffi::PyFrame_GetLineNumber(frame) }
-                        .try_into()
-                        .expect("Invalid line number for executing frame");
+                            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")
+                            }));
 
-                    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,
+                            (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")
                     }
-                    .serialize(&mut rmp_serde::encode::Serializer::new(&mut writer))
-                    .expect("Failed to emit Return message")
                 }
-            }
-            pyo3::ffi::PyTrace_C_CALL => {
-                let func = unsafe {
-                    // SAFETY: In the `C_CALL` case, `arg` is guaranteed to be a non-NULL pointer
-                    // to a `PyCFunctionObject`.
-                    &*(arg as *mut pyo3::ffi::PyCFunctionObject)
-                };
+                pyo3::ffi::PyTrace_C_CALL => {
+                    let func = unsafe {
+                        // SAFETY: In the `C_CALL` case, `arg` is guaranteed to be a non-NULL pointer
+                        // to a `PyCFunctionObject`.
+                        &*(arg as *mut pyo3::ffi::PyCFunctionObject)
+                    };
 
-                let func_name = cstr_to_cow(c_func_name());
+                    let func_name = cstr_to_cow(c_func_name());
 
-                // We either belong to a module or are a method on a class.
-                if !func.m_module.is_null() {
+                    // We either belong to a module or are a method on a class.
+                    if !func.m_module.is_null() {
                     // NOTE: _Typically_ `m_module` is a PyModule, but there are no requirements on
-                    // what type it needs to be!  In practice, we only see PyModules and PyStrings.
+                    // what type it needs to be! In practice, we only see PyModules and PyStrings.
                     Python::attach(|py| {
                         let m_module = unsafe {
                             // SAFETY: m_module is a non-NULL pointer to a valid Python object.
@@ -973,35 +963,35 @@ mod _functiontrace_rs {
                     .serialize(&mut rmp_serde::encode::Serializer::new(&mut writer))
                 }
                 .expect("Failed to emit NativeCall message");
-            }
-            pyo3::ffi::PyTrace_C_RETURN | pyo3::ffi::PyTrace_C_EXCEPTION => {
-                if what == pyo3::ffi::PyTrace_C_EXCEPTION {
-                    // C exceptions aren't exposed in a particularly interesting fashion, since they're
-                    // done using the `PyErr_*` functions plus a NULL return value.
-                    // However, if a `C_EXCEPTION` event is sent, we won't receive the corresponding
-                    // `C_RETURN` event, so in practice we treat both of these events the same.
-                    //
-                    // TODO: Should we try exposing information about them?  We'd probably need to
-                    // change the `Exception` format in `functiontrace-server`.
                 }
+                pyo3::ffi::PyTrace_C_RETURN | pyo3::ffi::PyTrace_C_EXCEPTION => {
+                    if what == pyo3::ffi::PyTrace_C_EXCEPTION {
+                        // C exceptions aren't exposed in a particularly interesting fashion, since they're
+                        // done using the `PyErr_*` functions plus a NULL return value.
+                        // However, if a `C_EXCEPTION` event is sent, we won't receive the corresponding
+                        // `C_RETURN` event, so in practice we treat both of these events the same.
+                        //
+                        // TODO: Should we try exposing information about them?  We'd probably need to
+                        // change the `Exception` format in `functiontrace-server`.
+                    }
 
-                function_trace::FunctionTrace::NativeReturn {
-                    func_name: cstr_to_cow(c_func_name()),
-                    time,
+                    function_trace::FunctionTrace::NativeReturn {
+                        func_name: cstr_to_cow(c_func_name()),
+                        time,
+                    }
+                    .serialize(&mut rmp_serde::encode::Serializer::new(&mut writer))
+                    .expect("Failed to emit NativeReturn message");
                 }
-                .serialize(&mut rmp_serde::encode::Serializer::new(&mut writer))
-                .expect("Failed to emit NativeReturn message");
-            }
-            pyo3::ffi::PyTrace_EXCEPTION => {
-                // TODO: Actually handle exceptions and do something interesting.
-                // We unfortunately can't get these from `PyEval_SetProfile`, so we'd need to use a
-                // different method to capture them.
-            }
-            v => {
-                tracing::error!(event = v, "Unexpected PyTrace message!");
-            }
-        };
-
+                pyo3::ffi::PyTrace_EXCEPTION => {
+                    // TODO: Actually handle exceptions and do something interesting.
+                    // We unfortunately can't get these from `PyEval_SetProfile`, so we'd need to use a
+                    // different method to capture them.
+                }
+                v => {
+                    tracing::error!(event = v, "Unexpected PyTrace message!");
+                }
+            };
+        });
         0
     }
 
@@ -1051,15 +1041,17 @@ mod _functiontrace_rs {
             }
         };
 
-        let config = if let Ok(config) = std::fs::read(&config_path)
-            .wrap_err("Failed to read config file")
-            .and_then(|x| toml::from_slice::<Config>(&x).wrap_err("Failed to parse config file"))
-        {
+        let config = if let Ok(contents) = std::fs::read(&config_path) {
             // We successfully loaded an existing config file.
-            config
+            toml::from_slice::<Config>(&contents).wrap_err_with(|| {
+                format!(
+                    "Failed to load existing Functiontrace config file at {}",
+                    config_path.display()
+                )
+            })?
         } else {
-            // There's no config file, so assume this is the first time FunctionTrace has been
-            // used and generate a new one.
+            // There's no config file, so assume this is the first time FunctionTrace has been used
+            // and generate a new one.
             let logo = "\
 ███████╗██╗   ██╗███╗   ██╗ ██████╗████████╗██╗ ██████╗ ███╗   ██╗████████╗██████╗  █████╗  ██████╗███████╗\n\
 ██╔════╝██║   ██║████╗  ██║██╔════╝╚══██╔══╝██║██╔═══██╗████╗  ██║╚══██╔══╝██╔══██╗██╔══██╗██╔════╝██╔════╝\n\
@@ -1097,7 +1089,7 @@ mod _functiontrace_rs {
             println!();
             println!(
                 "{}",
-                "Welcome to FunctionTrace, a low overheard tracing profiler for Python".bold()
+                "Welcome to FunctionTrace, a low overhead tracing profiler for Python".bold()
             );
             println!();
             println!(
@@ -1140,8 +1132,8 @@ mod _functiontrace_rs {
             config
         };
 
-        // We need to write out the config if it's changed.  Track the new and old configs so we
-        // can tell.
+        // We need to write out the config if it's changed. Track the new and old configs so we can
+        // tell.
         let mut new_config = config.clone();
 
         ////////////////////////////////////////////////////////////////////////
@@ -1292,7 +1284,7 @@ mod _functiontrace_rs {
             if let Ok(valid_expiration) = Timestamp::strptime(DEMO_KEY_FORMAT, demo_key) {
                 if now
                     .until(valid_expiration)
-                    .map_or(false, |remaining| remaining.is_positive())
+                    .is_ok_and(|remaining| remaining.is_positive())
                 {
                     // The license is still valid
                     eprintln!(
@@ -1453,14 +1445,14 @@ mod allocation_wrappers {
 
     /// Log the given allocation details.
     fn allocation_log<F: FnOnce() -> function_trace::AllocationDetails>(msg: F) {
-        if let Some(mut writer) = crate::_functiontrace_rs::thread_state() {
+        crate::_functiontrace_rs::thread_state(|mut writer| {
             function_trace::FunctionTrace::Allocation {
                 time: crate::_functiontrace_rs::trace_time(),
                 details: msg(),
             }
             .serialize(&mut rmp_serde::encode::Serializer::new(&mut writer))
             .expect("Failed to emit Allocation message");
-        }
+        });
     }
 
     /// Given an allocator->ctx for a logging allocator, convert it to point to the original
@@ -1477,33 +1469,36 @@ mod allocation_wrappers {
     #[unsafe(no_mangle)]
     pub extern "C" fn log_malloc(ctx: *mut c_void, bytes: usize) -> *mut c_void {
         let wrapped_allocator = allocator_ctx(ctx);
-        let addr = wrapped_allocator
+        if let Some(addr) = wrapped_allocator
             .malloc
             .map(|malloc| malloc(wrapped_allocator.ctx, bytes))
-            .unwrap_or(std::ptr::null_mut());
-
-        allocation_log(|| function_trace::AllocationDetails::Alloc {
-            bytes,
-            addr: addr as usize,
-        });
-
-        addr
+        {
+            allocation_log(|| function_trace::AllocationDetails::Alloc {
+                bytes,
+                addr: addr as usize,
+            });
+            addr
+        } else {
+            std::ptr::null_mut()
+        }
     }
 
     #[unsafe(no_mangle)]
     pub extern "C" fn log_calloc(ctx: *mut c_void, nelems: usize, elsize: usize) -> *mut c_void {
         let wrapped_allocator = allocator_ctx(ctx);
-        let addr = wrapped_allocator
+        if let Some(addr) = wrapped_allocator
             .calloc
             .map(|calloc| calloc(wrapped_allocator.ctx, nelems, elsize))
-            .unwrap_or(std::ptr::null_mut());
-
-        allocation_log(|| function_trace::AllocationDetails::Alloc {
-            bytes: nelems * elsize,
-            addr: addr as usize,
-        });
+        {
+            allocation_log(|| function_trace::AllocationDetails::Alloc {
+                bytes: nelems * elsize,
+                addr: addr as usize,
+            });
 
-        addr
+            addr
+        } else {
+            std::ptr::null_mut()
+        }
     }
 
     #[unsafe(no_mangle)]
@@ -1513,18 +1508,20 @@ mod allocation_wrappers {
         new_size: usize,
     ) -> *mut c_void {
         let wrapped_allocator = allocator_ctx(ctx);
-        let addr = wrapped_allocator
+        if let Some(addr) = wrapped_allocator
             .realloc
             .map(|realloc| realloc(wrapped_allocator.ctx, old_addr, new_size))
-            .unwrap_or(std::ptr::null_mut());
-
-        allocation_log(|| function_trace::AllocationDetails::Realloc {
-            bytes: new_size,
-            old_addr: old_addr as usize,
-            new_addr: addr as usize,
-        });
+        {
+            allocation_log(|| function_trace::AllocationDetails::Realloc {
+                bytes: new_size,
+                old_addr: old_addr as usize,
+                new_addr: addr as usize,
+            });
 
-        addr
+            addr
+        } else {
+            std::ptr::null_mut()
+        }
     }
 
     #[unsafe(no_mangle)]
@@ -1608,52 +1605,28 @@ mod hooks {
     // Hooking Framework
     ////////////////////////////////////////////////////////////////////////////
 
-    /// Implementation details behind [`hooks`].
-    ///
-    /// Keep a:
-    /// - counter that is incremented for each hook tuple we're passed, allowing us to know the
-    ///   index of that hook
-    /// - accumulator that gathers the counted hooks, allowing us to emit a single expression at
-    ///   the end (an array consisting of all hooks with their index associated)
-    ///
-    /// NOTE: I suspect this is grosser than it needs to be, but it works and really isn't fun to
-    /// modify.
-    macro_rules! hook_counter {
-        // The accumulator has fully collected all of our hooks, so we can now emit a vec of them.
-        ( [$( ($target:literal, $wrapper:ident, $counter:expr) ),*  $(,)?], $_counter:expr) => {
-            [$( Hook {
+    /// Generate a static slice of [`Hook`]s with info about the target functions.
+    macro_rules! hooks {
+        // Entry: setup a counter an accumulator array
+        ($(($target:literal, $wrapper:ident)),* $(,)?) => {
+            hooks!(@count 0; []; $(($target, $wrapper),)*)
+        };
+        // Recursive case: consume one function and put info into the accumulator array.
+        (@count $index:expr; [$($hooks:expr,)*]; ($target:literal, $wrapper:ident), $($rest:tt)*) => {
+            hooks!(@count $index + 1; [$($hooks,)* Hook {
                 target: $target,
-                wrapper: $wrapper::<{ $counter }>,
+                wrapper: $wrapper::<{ $index }>,
                 original: OnceLock::new(),
-            }),*]
+            },]; $($rest)*)
         };
-        // Handle the first iteration specifically
-        ( [ ], $counter:expr, $target:literal, $wrapper:ident, $($rest:tt),*) => {
-            hook_counter!( [
-                ($target, $wrapper, $counter)
-            ], $counter +1, $($rest),*)
-        };
-        // Handle the last iteration specifically, since we'll switch out of the normal recursive
-        // case and into our accumulator case.
-        ( [$($acc:tt),*  $(,)?], $counter:expr, $target:literal, $wrapper:ident) => {
-            hook_counter!( [
-                $($acc),*, ($target, $wrapper, $counter)
-            ], $counter +1)
-        };
-        // The main recursive case
-        ( [ $( $acc:tt ),* $(,)?], $counter:expr, $target:literal, $wrapper:ident, $($rest:tt),*) => {
-            hook_counter!( [
-                $($acc),*, ($target, $wrapper, $counter)
-            ], $counter +1, $($rest),*)
-        };
-    }
-
-    /// Generate an array of [`Hook`]s.  In particular, we need to generate unique references into
-    /// [`HOOKS`] that each wrapper is aware of, meaning we need to use a custom push-down
-    /// accumulator and counting macro.
-    macro_rules! hooks {
-        ($(($target:literal, $wrapper:ident)),* $(,)?) => {
-            hook_counter!([], 0, $($target, $wrapper),*)
+        // Base case: we have no more functions to hook, so emit the hooks.
+        (@count $index:expr; [$($hooks:expr,)*];) => {
+            {
+                // We need a static array with a fixed size to get a reference to. Generate this in
+                // the macro so the caller doesn't need to manually update the size.
+                static ENTRIES: [Hook; $index] = [$($hooks,)*];
+                &ENTRIES
+            }
         };
     }
 
@@ -1662,18 +1635,19 @@ mod hooks {
     /// NOTE: This must be created by [`hooks`], since we actually generate new variants of the
     /// wrapper function for each hook, allowing the function to determine which hook triggered it
     /// at runtime.
-    static HOOKS: [Hook; 11] = hooks!(
+    static HOOKS: &[Hook] = hooks!(
+        ("builtins.__import__", logging_import),
         ("builtins.print", logging_print),
+        ("logging.critical", logging_print),
         ("logging.debug", logging_print),
-        ("logging.log", logging_print),
-        ("logging.info", logging_print),
-        ("logging.warning", logging_print),
         ("logging.error", logging_print),
-        ("logging.critical", logging_print),
-        ("logging.fatal", logging_print),
         ("logging.exception", logging_print),
-        ("multiprocessing.util._exit_function", multiprocessing_exit),
-        ("builtins.__import__", logging_import),
+        ("logging.fatal", logging_print),
+        ("logging.info", logging_print),
+        ("logging.log", logging_print),
+        ("logging.warning", logging_print),
+        ("multiprocessing.util._exit_function", process_exit),
+        ("os._exit", process_exit),
     );
 
     /// Information, used during both runtime and at [`install`] time, about a specific function
@@ -1704,7 +1678,7 @@ mod hooks {
             .expect("Wrapper functions can't be called until hooked");
 
         let proxied = Python::attach(|py| -> Result<()> {
-            if let Some(mut writer) = crate::_functiontrace_rs::thread_state() {
+            crate::_functiontrace_rs::thread_state(|mut writer| -> Result<()> {
                 let args = unsafe {
                     // SAFETY: `args` represents a valid argument tuple
                     Bound::from_borrowed_ptr(py, args)
@@ -1718,7 +1692,9 @@ mod hooks {
                     msg.serialize(&mut rmp_serde::encode::Serializer::new(&mut writer))
                         .wrap_err("Failed to emit Allocation message")?;
                 }
-            }
+                Ok(())
+            })
+            .transpose()?;
 
             Ok(())
         })
@@ -1809,19 +1785,17 @@ mod hooks {
     /// Though this is mixed in with the rest of the hooks, this is actually load-bearing for
     /// general FunctionTrace functionality.
     ///
-    /// `atexit()` will only run in the main process, so we need to register a separate
-    /// multiprocessesing exit hook to do similar teardown on any multiprocessing workers.
-    extern "C" fn multiprocessing_exit<const HOOK: usize>(
+    /// 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.
+    extern "C" fn process_exit<const HOOK: usize>(
         _self: *mut PyObject,
         args: *mut PyObject,
         kwargs: *mut PyObject,
     ) -> *mut PyObject {
         let hook = &HOOKS[HOOK];
 
-        if let Some(writer) = crate::_functiontrace_rs::thread_state() {
-            crate::_functiontrace_rs::thread_teardown(writer)
-                .expect("Multiprocessing teardown should be successful");
-        }
+        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)