functiontrace 0.5.2__tar.gz → 0.5.3__tar.gz

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

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

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

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

{functiontrace-0.5.2 → functiontrace-0.5.3}/functiontrace.egg-info/SOURCES.txt

@@ -11,4 +11,5 @@ functiontrace.egg-info/SOURCES.txt
 functiontrace.egg-info/dependency_links.txt
 functiontrace.egg-info/entry_points.txt
 functiontrace.egg-info/top_level.txt
-rust/lib.rs
+rust/lib.rs
+rust/multiprocessing.sh.tpl

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

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

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

@@ -241,6 +241,8 @@ mod _functiontrace_rs {
         module: &Bound<'_, PyModule>,
         output_directory: OsString,
     ) -> Result<()> {
+        use std::os::unix::fs::{DirBuilderExt, PermissionsExt};
+
         // The name of the env var we use to notify subprocesses that there's a trace in progress.
         const BREADCRUMB: &str = "FUNCTIONTRACE_LIVE";
 
@@ -337,7 +339,16 @@ mod _functiontrace_rs {
                     e => e.wrap_err("Failed to launch functiontrace-server")?,
                 };
 
-                format!("/tmp/functiontrace-server.sock.{}", server.id()).into()
+                let sockaddr = format!("/tmp/functiontrace-server.sock.{}", server.id());
+
+                // Also allocate a directory for storing wrappers (multiprocessing.set_executable,
+                // etc) we may need.
+                std::fs::DirBuilder::new()
+                    .recursive(true)
+                    .mode(0o700)
+                    .create(format!("{sockaddr}.wrappers"))
+                    .wrap_err("Failed to create wrapper directory")?;
+                sockaddr.into()
             };
 
             Python::attach(|_| {
@@ -414,6 +425,61 @@ mod _functiontrace_rs {
                 .wrap_err("Failed to close initialization socket")?;
         }
 
+        // `multiprocessing` uses an absolute path to an executable to launch forkserver workers
+        // using `-c` for the actual code,
+        // Override this with our own wrapper that also sets up functiontrace.
+        Python::attach(|py| -> Result<()> {
+            let executable: String = py.import("sys")?.getattr("executable")?.extract()?;
+            let wrapper = {
+                // We'll store the wrappers next to the socket.
+                // TODO: Ideally we'd have these all in the same tmpdir, but that requires changing
+                // client and server code so would be a breaking change.
+                let mut directory = SOCKADDR
+                    .get()
+                    .expect("sockaddr is set during server initialization")
+                    .clone();
+                directory.push(".wrappers");
+
+                std::path::PathBuf::from(directory)
+                    .join("multiprocessing")
+                    .join(executable.strip_prefix("/").unwrap_or(&executable))
+            };
+
+            // Create a multiprocessing wrapper for the current executable if it doesn't exist yet.
+            std::fs::create_dir_all(
+                wrapper
+                    .parent()
+                    .expect("Wrappers are always in a subdirectory"),
+            )
+            .wrap_err("Failed to create multiprocessing wrapper directories")?;
+
+            match std::fs::OpenOptions::new()
+                .write(true)
+                .create_new(true)
+                .open(&wrapper)
+            {
+                Ok(mut file) => {
+                    let script = include_str!("multiprocessing.sh.tpl").replace(
+                        "@PYTHON_INTERPRETER@",
+                        // NOTE: This is equivalent to shlex.quote
+                        &format!("'{}'", executable.replace('\'', "'\"'\"'")),
+                    );
+
+                    file.write_all(script.as_bytes())
+                        .wrap_err("Failed to write multiprocessing wrapper")?;
+                    file.set_permissions(std::fs::Permissions::from_mode(0o700))
+                        .wrap_err("Failed to `chmod +x` multiprocessing wrapper")?;
+                }
+                Err(err) if err.kind() == std::io::ErrorKind::AlreadyExists => {}
+                Err(err) => return Err(err).wrap_err("Failed to create multiprocessing wrapper"),
+            }
+            py.import("multiprocessing")?
+                .getattr("set_executable")?
+                .call1((wrapper,))?;
+            Ok(())
+        })
+        .wrap_err("Failed to configure multiprocessing tracing")?;
+
         ////////////////////////////////////////////////////////////////////////////
         // Teardown Configuration
         ////////////////////////////////////////////////////////////////////////////

functiontrace-0.5.3/rust/multiprocessing.sh.tpl

@@ -0,0 +1,53 @@
+#!/bin/sh
+
+# Multiprocessing passes the code to execute in workers as `-c <code>`.
+# Strip that out and run tracing setup first, maintaining the order of arguments just in case.
+
+remaining=$#
+while [ "$remaining" -gt 0 ]; do
+    arg=$1
+    shift
+    remaining=$((remaining - 1))
+    case "$arg" in
+        -W|-X)
+            # These arguments can consume other arguments. Skip them so we don't parse the wrong `-c`.
+            value=$1
+            shift
+            remaining=$((remaining - 1))
+            set -- "$@" "$arg" "$value"
+            ;;
+        -c)
+            code=$1
+            shift
+            remaining=$((remaining - 1))
+            set -- "$@" -c '
+exec(
+    (
+        compile(
+            __import__("sys").argv.pop(1),
+            "<string>",
+            "exec",
+            dont_inherit=True,
+        ),
+        # Start tracing before exec runs the compiled code.
+        __import__("_functiontrace_rs").begin_tracing("."),
+    )[0]
+)
+' "$code"
+            break
+            ;;
+        *)
+            set -- "$@" "$arg"
+            ;;
+    esac
+done
+
+# Consume the rest of the args without looking.
+while [ "$remaining" -gt 0 ]; do
+    arg=$1
+    shift
+    remaining=$((remaining - 1))
+    set -- "$@" "$arg"
+done
+
+exec @PYTHON_INTERPRETER@ "$@"