vmprof · mattip · Dec 20, 2023 · Jan 21, 2023 · Jan 25, 2023 · Jan 25, 2023
diff --git a/setup.py b/setup.py
@@ -78,6 +78,8 @@ def run(self):
         raise NotImplementedError("platform '%s' is not supported!" % sys.platform)
     extra_compile_args.append('-I src/')
     extra_compile_args.append('-I src/libbacktrace')
+    if sys.version_info[:2] == (3,11):
+        extra_source_files += ['src/populate_frames.c']
     ext_modules = [Extension('_vmprof',
                            sources=[
                                'src/_vmprof.c',
@@ -121,7 +123,7 @@ def run(self):
         'pytz',
         'colorama',
     ] + extra_install_requires,
-    python_requires='>=3.6, <3.11',
+    python_requires='>=3.6, <3.12',
     tests_require=['pytest','cffi','hypothesis'],
     entry_points = {
         'console_scripts': [

diff --git a/src/_vmprof.c b/src/_vmprof.c
@@ -11,6 +11,12 @@
 #include "_vmprof.h"
 #include "vmprof_common.h"
 
+
+#if PY_VERSION_HEX >= 0x030b00f0 /* >= 3.11 */
+#include "internal/pycore_frame.h"
+#include "populate_frames.h"
+#endif
+
 static destructor Original_code_dealloc = 0;
 static PyObject* (*_default_eval_loop)(PyFrameObject *, int) = 0;
 
@@ -112,7 +118,7 @@ static int _look_for_code_object(PyObject *o, void * param)
         /* as a special case, recursively look for and add code
            objects found in the co_consts.  The problem is that code
            objects are not created as GC-aware in CPython, so we need
-           to hack like this to hope to find most of them. 
+           to hack like this to hope to find most of them.
         */
         i = PyTuple_Size(co->co_consts);
         while (i > 0) {
@@ -267,6 +273,15 @@ write_all_code_objects(PyObject *module, PyObject * seen_code_ids)
 }
 
 
+#if PY_VERSION_HEX < 0x030900B1 /* < 3.9 */
+static inline PyFrameObject* PyThreadState_GetFrame(PyThreadState *tstate)
+{
+    Py_XINCREF(tstate->frame);
+    return tstate->frame;
+}
+#endif
+
+
 
 static PyObject *
 sample_stack_now(PyObject *module, PyObject * args)
@@ -298,7 +313,17 @@ sample_stack_now(PyObject *module, PyObject * args)
         vmprof_ignore_signals(0);
         return NULL;
     }
-    entry_count = vmp_walk_and_record_stack(tstate->frame, m, SINGLE_BUF_SIZE/sizeof(void*)-1, (int)skip, 0);
+
+#if PY_VERSION_HEX >= 0x030B0000 /* < 3.11, no pypy 3.11 at the moment*/ 
+    _PyInterpreterFrame *  frame = unsafe_PyThreadState_GetInterpreterFrame(tstate);
+#else
+    PyFrameObject* frame = PyThreadState_GetFrame(tstate);
+#endif
+
+    entry_count = vmp_walk_and_record_stack(frame, m, SINGLE_BUF_SIZE/sizeof(void*)-1, (int)skip, 0);
+
+    Py_XDECREF(frame);
+
 
     for (i = 0; i < entry_count; i++) {
         routine_ip = m[i];

diff --git a/src/populate_frames.c b/src/populate_frames.c
@@ -0,0 +1,76 @@
+/* This code was taken from https://github.com/GoogleCloudPlatform/cloud-profiler-python/blob/main/googlecloudprofiler/src/populate_frames.cc */ 
+
+#include "populate_frames.h"
+
+#include <Python.h>
+
+
+// 0x030B0000 is 3.11.
+#define PY_311 0x030B0000
+#if PY_VERSION_HEX >= PY_311
+
+/**
+ * The PyFrameObject structure members have been removed from the public C API
+ * in 3.11:
+https://docs.python.org/3/whatsnew/3.11.html#pyframeobject-3-11-hiding.
+ *
+ * Instead, getters are provided which participate in reference counting; since
+ * this code runs as part of the SIGPROF handler, it cannot modify Python
+ * objects (including their refcounts) and the getters can't be used. Instead,
+ * we expose the internal _PyInterpreterFrame and use that directly.
+ *
+ */
+
+#define Py_BUILD_CORE
+#include "internal/pycore_frame.h"
+#undef Py_BUILD_CORE
+
+// Modified from
+// https://github.com/python/cpython/blob/v3.11.4/Python/pystate.c#L1278-L1285
+_PyInterpreterFrame *unsafe_PyThreadState_GetInterpreterFrame(
+    PyThreadState *tstate) {
+  assert(tstate != NULL);
+  _PyInterpreterFrame *f = tstate->cframe->current_frame;
+  while (f && _PyFrame_IsIncomplete(f)) {
+    f = f->previous;
+  }
+  if (f == NULL) {
+    return NULL;
+  }
+  return f;
+}
+
+// Modified from
+// https://github.com/python/cpython/blob/v3.11.4/Objects/frameobject.c#L1310-L1315
+// with refcounting removed
+PyCodeObject *unsafe_PyInterpreterFrame_GetCode(
+    _PyInterpreterFrame *frame) {
+  assert(frame != NULL);
+  assert(!_PyFrame_IsIncomplete(frame));
+  PyCodeObject *code = frame->f_code;
+  assert(code != NULL);
+  return code;
+}
+
+// Modified from
+// https://github.com/python/cpython/blob/v3.11.4/Objects/frameobject.c#L1326-L1329
+// with refcounting removed
+_PyInterpreterFrame *unsafe_PyInterpreterFrame_GetBack(
+    _PyInterpreterFrame *frame) {
+  assert(frame != NULL);
+  assert(!_PyFrame_IsIncomplete(frame));
+  _PyInterpreterFrame *prev = frame->previous;
+  while (prev && _PyFrame_IsIncomplete(prev)) {
+    prev = prev->previous;
+  }
+  return prev;
+}
+
+// Copied from
+// https://github.com/python/cpython/blob/v3.11.4/Python/frame.c#L165-L170 as
+// this function is not available in libpython
+int _PyInterpreterFrame_GetLine(_PyInterpreterFrame *frame) {
+  int addr = _PyInterpreterFrame_LASTI(frame) * sizeof(_Py_CODEUNIT);
+  return PyCode_Addr2Line(frame->f_code, addr);
+}
+#endif  // PY_VERSION_HEX >= PY_311
diff --git a/src/populate_frames.h b/src/populate_frames.h
@@ -0,0 +1,22 @@
+/* This code was taken from https://github.com/GoogleCloudPlatform/cloud-profiler-python/blob/main/googlecloudprofiler/src/populate_frames.h */ 
+
+#ifndef pp_frames
+#define pp_frames
+
+#include <Python.h>
+
+#include <frameobject.h>
+
+#define Py_BUILD_CORE
+#include "internal/pycore_frame.h"
+#undef Py_BUILD_CORE
+
+_PyInterpreterFrame *unsafe_PyThreadState_GetInterpreterFrame(PyThreadState *tstate);
+
+PyCodeObject *unsafe_PyInterpreterFrame_GetCode(_PyInterpreterFrame *frame);
+
+_PyInterpreterFrame *unsafe_PyInterpreterFrame_GetBack(_PyInterpreterFrame *frame);
+
+int _PyInterpreterFrame_GetLine(_PyInterpreterFrame *frame);
+
+#endif