Get GIL status on Python 3.7+

Cargo.toml

@@ -20,7 +20,6 @@ goblin = "0.0.21"
 lazy_static = "1.1.0"
 libc = "0.2.34"
 log = "0.4"
-memoffset = "0.3"
 regex = "1"
 tempdir = "0.3"
 tempfile = "3.0.3"

generate_bindings.py

@@ -9,6 +9,7 @@
 import argparse
 import os
 import sys
+import tempfile
 
 
 def build_python(cpython_path, version):
@@ -35,6 +36,59 @@ def build_python(cpython_path, version):
     return os.system(f"{pip} install setuptools_rust wheel")
 
 
+def calculate_pyruntime_offsets(cpython_path, version, configure=False):
+    ret = os.system(f"""
+        cd {cpython_path}
+        git checkout {version}
+
+        # need to run configure on the current branch to generate pyconfig.h sometimes
+        {("./configure prefix=" + os.path.join(cpython_path, version)) if configure else ""}
+    """)
+    if ret:
+        return ret
+
+    # simple little c program to get the offsets we need from the pyruntime struct
+    # (using rust bindgen here is more complicated than necessary)
+    program = r"""
+        #include <stddef.h>
+        #include <stdio.h>
+        #define Py_BUILD_CORE 1
+        #include "Include/Python.h"
+        #include "Include/internal/pystate.h"
+
+        int main(int argc, const char * argv[]) {
+            size_t interp_head = offsetof(_PyRuntimeState, interpreters.head);
+            printf("pub static INTERP_HEAD_OFFSET: usize = %i;\n", interp_head);
+
+            size_t tstate_current = offsetof(_PyRuntimeState, gilstate.tstate_current);
+            printf("pub static TSTATE_CURRENT: usize = %i;\n", tstate_current);
+        }
+    """
+
+    if not os.path.isfile(os.path.join(cpython_path, "Include", "internal", "pystate.h")):
+        if os.path.isfile(os.path.join(cpython_path, "Include", "internal", "pycore_pystate.h")):
+            program = program.replace("pystate.h", "pycore_pystate.h")
+        else:
+            print("failed to find Include/internal/pystate.h in cpython directory =(")
+            return
+
+    with tempfile.TemporaryDirectory() as path:
+        source_filename = os.path.join(path, "pyruntime_offsets.c")
+        exe = os.path.join(path, "pyruntime_offsets")
+
+        with open(source_filename, "w") as o:
+            o.write(program)
+        ret = os.system(f"""gcc {source_filename} -I {cpython_path} -I {cpython_path}/include -o {exe}""")
+        if ret:
+            print("Failed to compile""")
+            return ret
+
+        ret = os.system(exe)
+        if ret:
+            print("Failed to run pyruntime file")
+            return ret
+
+
 def extract_bindings(cpython_path, version, configure=False):
     print("Generating bindings for python %s from repo at %s" % (version, cpython_path))
     ret = os.system(f"""
@@ -46,12 +100,10 @@ def extract_bindings(cpython_path, version, configure=False):
 
         cat Include/Python.h > bindgen_input.h
         cat Include/frameobject.h >> bindgen_input.h
-        cat Include/internal/pystate.h >> bindgen_input.h
 
         bindgen  bindgen_input.h -o bindgen_output.rs \
             --with-derive-default \
             --no-layout-tests --no-doc-comments \
-            --whitelist-type _PyRuntimeState \
             --whitelist-type PyInterpreterState \
             --whitelist-type PyFrameObject \
             --whitelist-type PyThreadState \
@@ -62,7 +114,7 @@ def extract_bindings(cpython_path, version, configure=False):
             --whitelist-type PyUnicodeObject \
             --whitelist-type PyCompactUnicodeObject \
             --whitelist-type PyStringObject \
-             -- -I . -I ./Include -DPy_BUILD_CORE
+             -- -I . -I ./Include
     """)
     if ret:
         return ret
@@ -77,7 +129,7 @@ def extract_bindings(cpython_path, version, configure=False):
         o.write("#![allow(clippy::useless_transmute)]\n")
         o.write("#![allow(clippy::default_trait_access)]\n")
         o.write("#![allow(clippy::cast_lossless)]\n")
-        o.write("#![allow(clippy::trivially_copy_pass_by_ref)]\n")
+        o.write("#![allow(clippy::trivially_copy_pass_by_ref)]\n\n")
         o.write(open(os.path.join(cpython_path, "bindgen_output.rs")).read())
 
 
@@ -91,6 +143,9 @@ def extract_bindings(cpython_path, version, configure=False):
     parser.add_argument("--configure",
                         help="Run configure script prior to generating bindings",
                         action="store_true")
+    parser.add_argument("--pyruntime",
+                        help="generate offsets for pyruntime",
+                        action="store_true")
     parser.add_argument("--build",
                         help="Build python for this version",
                         action="store_true")
@@ -120,6 +175,9 @@ def extract_bindings(cpython_path, version, configure=False):
             # todo: this probably shoudl be a separate script
             if build_python(args.cpython, version):
                 print("Failed to build python")
+        elif args.pyruntime:
+            calculate_pyruntime_offsets(args.cpython, version, configure=args.configure)
+
         else:
             if extract_bindings(args.cpython, version, configure=args.configure):
                 print("Failed to generate bindings")

src/main.rs

@@ -17,8 +17,6 @@ extern crate libc;
 #[macro_use]
 extern crate log;
 extern crate memmap;
-#[macro_use]
-extern crate memoffset;
 extern crate proc_maps;
 extern crate benfred_read_process_memory as read_process_memory;
 extern crate regex;
@@ -42,6 +40,7 @@ mod stack_trace;
 mod console_viewer;
 mod flamegraph;
 mod utils;
+mod version;
 
 use std::io::Read;
 use std::sync::atomic::{AtomicBool, Ordering};

src/python_bindings/mod.rs

@@ -3,3 +3,76 @@ pub mod v3_3_7;
 pub mod v3_5_5;
 pub mod v3_6_6;
 pub mod v3_7_0;
+
+// currently the PyRuntime struct used from Python 3.7 on really can't be
+// exposed in a cross platform way using bindgen. PyRuntime has several mutex's
+// as member variables, and these have different sizes depending on the operating
+// system and system architecture.
+// Instead we will define some constants here that define valid offsets for the
+// member variables we care about here
+// (note 'generate_bindings.py' has code to figure out these offsets)
+pub mod pyruntime {
+    use version::Version;
+
+    // There aren't any OS specific members of PyRuntime before pyinterpreters.head,
+    // so these offsets should be valid for all OS'es
+    #[cfg(target_pointer_width = "32")]
+    pub static INTERP_HEAD_OFFSET: usize = 16;
+
+    #[cfg(target_pointer_width = "64")]
+    pub static INTERP_HEAD_OFFSET: usize = 24;
+
+    // getting gilstate.tstate_current is different for all OS
+    // and is also different for each python version, and even
+    // between v3.8.0a1 and v3.8.0a2 =(
+    #[cfg(target_os="macos")]
+    pub fn get_tstate_current_offset(version: &Version) -> Option<usize> {
+        match version {
+             Version{major: 3, minor: 7, ..} => Some(1440),
+             Version{major: 3, minor: 8, patch: 0, ..} => {
+                 match version.release_flags.as_ref() {
+                    "a1" => Some(1432),
+                    "a2" => Some(888),
+                    _ => None
+                }
+             },
+             _ => None
+        }
+    }
+
+    #[cfg(all(target_os="linux", target_pointer_width = "32"))]
+    pub fn get_tstate_current_offset(version: &Version) -> Option<usize> {
+        match version {
+             Version{major: 3, minor: 7, ..} => Some(796),
+             Version{major: 3, minor: 8, patch: 0, ..} => {
+                 match version.release_flags.as_ref() {
+                    "a1" => Some(792),
+                    "a2" => Some(512),
+                    _ => None
+                }
+             },
+             _ => None
+        }
+    }
+
+    #[cfg(all(target_os="linux", target_pointer_width = "64"))]
+    pub fn get_tstate_current_offset(version: &Version) -> Option<usize> {
+        match version {
+             Version{major: 3, minor: 7, ..} => Some(1392),
+             Version{major: 3, minor: 8, patch: 0, ..} => {
+                 match version.release_flags.as_ref() {
+                    "a1" => Some(1384),
+                    "a2" => Some(840),
+                    _ => None
+                }
+             },
+             _ => None
+        }
+    }
+
+    #[cfg(windows)]
+    pub fn get_tstate_current_offset(version: &Version) -> Option<usize> {
+        // TODO: compute offsets for windows
+        None
+    }
+}