GIL monitoring for Python 3.10.12

Author: cyme

Include/cpython/ceval.h

@@ -4,6 +4,7 @@
 
 PyAPI_FUNC(void) PyEval_SetProfile(Py_tracefunc, PyObject *);
 PyAPI_DATA(int) _PyEval_SetProfile(PyThreadState *tstate, Py_tracefunc func, PyObject *arg);
+PyAPI_FUNC(void) PyEval_SetProfileAllThreads(Py_tracefunc func, PyObject *arg);
 PyAPI_FUNC(void) PyEval_SetTrace(Py_tracefunc, PyObject *);
 PyAPI_FUNC(int) _PyEval_SetTrace(PyThreadState *tstate, Py_tracefunc func, PyObject *arg);
 PyAPI_FUNC(int) _PyEval_GetCoroutineOriginTrackingDepth(void);

Include/cpython/pystate.h

@@ -25,7 +25,14 @@ typedef int (*Py_tracefunc)(PyObject *, PyFrameObject *, int, PyObject *);
 #define PyTrace_C_EXCEPTION 5
 #define PyTrace_C_RETURN 6
 #define PyTrace_OPCODE 7
-
+#define PyTrace_LOCK_ACQUIRE 8
+#define PyTrace_LOCK_RELEASE 9
+#define PyTrace_THREAD_PREEMPT 10
+#define PyTrace_THREAD_RESUME 11
+#define PyTrace_THREAD_ACQUIRE 12
+#define PyTrace_THREAD_RELEASE 13
+#define PyTrace_THREAD_SAVE 14
+#define PyTrace_THREAD_RESTORE 15
 
 typedef struct _cframe {
     /* This struct will be threaded through the C stack

Include/internal/pycore_ceval.h

@@ -31,6 +31,8 @@ PyAPI_FUNC(void) _PyEval_SetCoroutineOriginTrackingDepth(
     PyThreadState *tstate,
     int new_depth);
 
+extern int _PyEval_NotifyThreadStateChange(PyThreadState *tstate, int event);
+
 void _PyEval_Fini(void);
 
 

Python/ceval.c

@@ -81,6 +81,7 @@ static int maybe_call_line_trace(Py_tracefunc, PyObject *,
 static void maybe_dtrace_line(PyFrameObject *, PyTraceInfo *, int);
 static void dtrace_function_entry(PyFrameObject *);
 static void dtrace_function_return(PyFrameObject *);
+static int notify_thread_state_change(PyThreadState *tstate, int event, long arg);
 
 static PyObject * import_name(PyThreadState *, PyFrameObject *,
                               PyObject *, PyObject *, PyObject *);
@@ -415,18 +416,31 @@ _PyEval_Fini(void)
 void
 PyEval_AcquireLock(void)
 {
+    struct timespec start, end;
+    clock_gettime(CLOCK_MONOTONIC, &start);
+
     _PyRuntimeState *runtime = &_PyRuntime;
     PyThreadState *tstate = _PyRuntimeState_GetThreadState(runtime);
     _Py_EnsureTstateNotNULL(tstate);
 
     take_gil(tstate);
+
+    clock_gettime(CLOCK_MONOTONIC, &end);
+    long timediff = ((long)end.tv_sec - (long)start.tv_sec) * (long)1000000
+         + ((long)end.tv_nsec - (long)start.tv_nsec) / 1000;
+    notify_thread_state_change(tstate, PyTrace_LOCK_ACQUIRE, timediff);
 }
 
 void
 PyEval_ReleaseLock(void)
 {
     _PyRuntimeState *runtime = &_PyRuntime;
     PyThreadState *tstate = _PyRuntimeState_GetThreadState(runtime);
+
+    if (tstate != NULL) {
+        notify_thread_state_change(tstate, PyTrace_LOCK_RELEASE, -1);
+    }
+
     /* This function must succeed when the current thread state is NULL.
        We therefore avoid PyThreadState_Get() which dumps a fatal error
        in debug mode. */
@@ -438,6 +452,8 @@ PyEval_ReleaseLock(void)
 void
 _PyEval_ReleaseLock(PyThreadState *tstate)
 {
+    notify_thread_state_change(tstate, PyTrace_LOCK_RELEASE, -1);
+
     struct _ceval_runtime_state *ceval = &tstate->interp->runtime->ceval;
     struct _ceval_state *ceval2 = &tstate->interp->ceval;
     drop_gil(ceval, ceval2, tstate);
@@ -448,6 +464,9 @@ PyEval_AcquireThread(PyThreadState *tstate)
 {
     _Py_EnsureTstateNotNULL(tstate);
 
+    struct timespec start, end;
+    clock_gettime(CLOCK_MONOTONIC, &start);
+
     take_gil(tstate);
 
     struct _gilstate_runtime_state *gilstate = &tstate->interp->runtime->gilstate;
@@ -458,13 +477,20 @@ PyEval_AcquireThread(PyThreadState *tstate)
         Py_FatalError("non-NULL old thread state");
     }
 #endif
+
+    clock_gettime(CLOCK_MONOTONIC, &end);
+    long timediff = ((long)end.tv_sec - (long)start.tv_sec) * (long)1000000
+         + ((long)end.tv_nsec - (long)start.tv_nsec) / 1000;
+    notify_thread_state_change(tstate, PyTrace_THREAD_ACQUIRE, timediff);
 }
 
 void
 PyEval_ReleaseThread(PyThreadState *tstate)
 {
     assert(is_tstate_valid(tstate));
 
+    notify_thread_state_change(tstate, PyTrace_THREAD_RELEASE, -1);
+
     _PyRuntimeState *runtime = tstate->interp->runtime;
     PyThreadState *new_tstate = _PyThreadState_Swap(&runtime->gilstate, NULL);
     if (new_tstate != tstate) {
@@ -519,12 +545,16 @@ _PyEval_SignalAsyncExc(PyInterpreterState *interp)
 PyThreadState *
 PyEval_SaveThread(void)
 {
+    PyThreadState *tstate = PyGILState_GetThisThreadState();
+    if (tstate != NULL) {
+        notify_thread_state_change(tstate, PyTrace_THREAD_SAVE, -1);
+    }
     _PyRuntimeState *runtime = &_PyRuntime;
 #ifdef EXPERIMENTAL_ISOLATED_SUBINTERPRETERS
     PyThreadState *old_tstate = _PyThreadState_GET();
-    PyThreadState *tstate = _PyThreadState_Swap(&runtime->gilstate, old_tstate);
+    tstate = _PyThreadState_Swap(&runtime->gilstate, old_tstate);
 #else
-    PyThreadState *tstate = _PyThreadState_Swap(&runtime->gilstate, NULL);
+    tstate = _PyThreadState_Swap(&runtime->gilstate, NULL);
 #endif
     _Py_EnsureTstateNotNULL(tstate);
 
@@ -544,10 +574,18 @@ PyEval_RestoreThread(PyThreadState *tstate)
 {
     _Py_EnsureTstateNotNULL(tstate);
 
+    struct timespec start, end;
+    clock_gettime(CLOCK_MONOTONIC, &start);
+
     take_gil(tstate);
 
     struct _gilstate_runtime_state *gilstate = &tstate->interp->runtime->gilstate;
     _PyThreadState_Swap(gilstate, tstate);
+
+    clock_gettime(CLOCK_MONOTONIC, &end);
+    long timediff = ((long)end.tv_sec - (long)start.tv_sec) * (long)1000000
+         + ((long)end.tv_nsec - (long)start.tv_nsec) / 1000;
+    notify_thread_state_change(tstate, PyTrace_THREAD_RESTORE, timediff);
 }
 
 
@@ -1177,7 +1215,9 @@ eval_frame_handle_pending(PyThreadState *tstate)
     }
 
     /* GIL drop request */
-    if (_Py_atomic_load_relaxed(&ceval2->gil_drop_request)) {
+    if (_Py_atomic_load_relaxed(&ceval2->gil_drop_request) && !tstate->tracing) {
+
+        notify_thread_state_change(tstate, PyTrace_THREAD_PREEMPT, -1);
         /* Give another thread a chance */
         if (_PyThreadState_Swap(&runtime->gilstate, NULL) != tstate) {
             Py_FatalError("tstate mix-up");
@@ -1186,6 +1226,9 @@ eval_frame_handle_pending(PyThreadState *tstate)
 
         /* Other threads may run now */
 
+        struct timespec start, end;
+        clock_gettime(CLOCK_MONOTONIC, &start);
+
         take_gil(tstate);
 
 #ifdef EXPERIMENTAL_ISOLATED_SUBINTERPRETERS
@@ -1195,6 +1238,11 @@ eval_frame_handle_pending(PyThreadState *tstate)
             Py_FatalError("orphan tstate");
         }
 #endif
+
+        clock_gettime(CLOCK_MONOTONIC, &end);
+        long timediff = ((long)end.tv_sec - (long)start.tv_sec) * (long)1000000
+            + ((long)end.tv_nsec - (long)start.tv_nsec) / 1000;
+        notify_thread_state_change(tstate, PyTrace_THREAD_RESUME, timediff);
     }
 
     /* Check for asynchronous exception. */
@@ -5572,6 +5620,27 @@ PyEval_SetProfile(Py_tracefunc func, PyObject *arg)
     }
 }
 
+void
+PyEval_SetProfileAllThreads(Py_tracefunc func, PyObject *arg)
+{
+    PyThreadState *this_tstate = _PyThreadState_GET();
+    PyInterpreterState* interp = this_tstate->interp;
+
+    /* unused var _PyRuntimeState *runtime = &_PyRuntime; */
+    /* missing call to HEAD_LOCK(runtime); */
+    PyThreadState* ts = PyInterpreterState_ThreadHead(interp);
+    /* missing call to HEAD_UNLOCK(runtime); */
+
+    while (ts) {
+        if (_PyEval_SetProfile(ts, func, arg) < 0) {
+            _PyErr_WriteUnraisableMsg("Exception ignored in PyEval_SetProfileAllThreads", NULL);
+        }
+        /* missing call to HEAD_LOCK(runtime); */
+        ts = PyThreadState_Next(ts);
+        /* missing call to HEAD_UNLOCK(runtime); */
+    }
+}
+
 int
 _PyEval_SetTrace(PyThreadState *tstate, Py_tracefunc func, PyObject *arg)
 {
@@ -6512,6 +6581,31 @@ maybe_dtrace_line(PyFrameObject *frame,
     }
 }
 
+static int
+notify_thread_state_change(PyThreadState *tstate, int event, long arg)
+{
+    if (tstate->c_profilefunc == NULL)
+        return 0;
+    PyObject *arg_obj = Py_None;
+    if (arg >= 0) {
+        arg_obj = PyLong_FromLong(arg);
+    }
+    if (arg_obj == NULL) {
+        return -1;
+    }
+    PyTraceInfo trace_info;
+    /* Mark trace_info as uninitialized */
+    trace_info.code = NULL;
+    int res = call_trace_protected(tstate->c_profilefunc,
+                                    tstate->c_profileobj,
+                                    tstate, tstate->frame,
+                                    &trace_info,
+                                    event, arg_obj);
+    if (arg_obj != Py_None) {
+        Py_DECREF(arg_obj);
+    }
+    return res;
+}
 
 /* Implement Py_EnterRecursiveCall() and Py_LeaveRecursiveCall() as functions
    for the limited API. */

Python/sysmodule.c

@@ -924,19 +924,24 @@ sys_intern_impl(PyObject *module, PyObject *s)
  * Cached interned string objects used for calling the profile and
  * trace functions.  Initialized by trace_init().
  */
-static PyObject *whatstrings[8] = {NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL};
+static PyObject *whatstrings[16] = {
+    NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL,
+    NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL
+};
 
 static int
 trace_init(void)
 {
-    static const char * const whatnames[8] = {
+    static const char * const whatnames[16] = {
         "call", "exception", "line", "return",
         "c_call", "c_exception", "c_return",
-        "opcode"
+        "opcode", "lock_acquire", "lock_release",
+        "thread_preempt", "thread_resume", "thread_acquire",
+        "thread_release", "thread_save", "thread_restore"
     };
     PyObject *name;
     int i;
-    for (i = 0; i < 8; ++i) {
+    for (i = 0; i < 16; ++i) {
         if (whatstrings[i] == NULL) {
             name = PyUnicode_InternFromString(whatnames[i]);
             if (name == NULL)
@@ -952,21 +957,23 @@ static PyObject *
 call_trampoline(PyThreadState *tstate, PyObject* callback,
                 PyFrameObject *frame, int what, PyObject *arg)
 {
-    if (PyFrame_FastToLocalsWithError(frame) < 0) {
+    if (frame != NULL && PyFrame_FastToLocalsWithError(frame) < 0) {
         return NULL;
     }
 
     PyObject *stack[3];
-    stack[0] = (PyObject *)frame;
+    stack[0] = (frame != NULL) ? (PyObject *)frame : Py_None;
     stack[1] = whatstrings[what];
     stack[2] = (arg != NULL) ? arg : Py_None;
 
     /* call the Python-level function */
     PyObject *result = _PyObject_FastCallTstate(tstate, callback, stack, 3);
 
-    PyFrame_LocalsToFast(frame, 1);
-    if (result == NULL) {
-        PyTraceBack_Here(frame);
+    if (frame != NULL) {
+        PyFrame_LocalsToFast(frame, 1);
+        if (result == NULL) {
+            PyTraceBack_Here(frame);
+        }
     }
 
     return result;
@@ -1100,6 +1107,34 @@ Set the profiling function.  It will be called on each function call\n\
 and return.  See the profiler chapter in the library manual."
 );
 
+static PyObject *
+sys__setprofileallthreads(PyObject *module, PyObject *arg)
+{
+    PyObject* argument = NULL;
+    Py_tracefunc func = NULL;
+
+    if (trace_init() == -1) {
+        return NULL;
+    }
+
+    if (arg != Py_None) {
+        func = profile_trampoline;
+        argument = arg;
+    }
+
+    PyEval_SetProfileAllThreads(func, argument);
+
+    Py_RETURN_NONE;
+}
+
+PyDoc_STRVAR(_setprofileallthreads_doc,
+"_setprofileallthreads(function)\n\
+\n\
+Set the profiling function for current and all other running threads.\n\
+It will be called on each function calland return.  See the profiler\n\
+chapter in the library manual."
+);
+
 /*[clinic input]
 sys.getprofile
 
@@ -2055,6 +2090,7 @@ static PyMethodDef sys_methods[] = {
     SYS_GETSWITCHINTERVAL_METHODDEF
     SYS_SETDLOPENFLAGS_METHODDEF
     {"setprofile",      sys_setprofile, METH_O, setprofile_doc},
+    {"_setprofileallthreads", sys__setprofileallthreads, METH_O, _setprofileallthreads_doc},
     SYS_GETPROFILE_METHODDEF
     SYS_SETRECURSIONLIMIT_METHODDEF
     {"settrace",        sys_settrace, METH_O, settrace_doc},