(arigo, plan_rich) a signal interrupting a thread will now always record the stack trace of the same thread (not others)

Author: planrich

src/vmprof_common.h

@@ -143,13 +143,3 @@ void *volatile _PyThreadState_Current;
 #  define _Py_atomic_load_relaxed(pp)  (*(pp))
 #endif
 
-PyThreadState * get_current_thread_state(void)
-{
-#if PY_MAJOR_VERSION < 3
-    return _PyThreadState_Current;
-#elif PY_VERSION_HEX < 0x03050200
-    return (PyThreadState*) _Py_atomic_load_relaxed(&_PyThreadState_Current);
-#else
-    return _PyThreadState_UncheckedGet();
-#endif
-}

src/vmprof_main.h

@@ -98,36 +98,14 @@ static char atfork_hook_installed = 0;
  * *************************************************************
  */
 
-static int get_stack_trace(void** result, int max_depth, ucontext_t *ucontext)
+static int get_stack_trace(PyThreadState * current, void** result, int max_depth, ucontext_t *ucontext)
 {
-    PyThreadState* current = get_current_thread_state();
-
     if (!current)
         return 0;
     PyFrameObject *frame = current->frame;
     return read_trace_from_cpy_frame(current->frame, result, max_depth);
 }
 
-static void *get_current_thread_id(void)
-{
-    /* xxx This function is a hack on two fronts:
-
-       - It assumes that pthread_self() is async-signal-safe.  This
-         should be true on Linux and OS X.  I hope it is also true elsewhere.
-
-       - It abuses pthread_self() by assuming it just returns an
-         integer.  According to comments in CPython's source code, the
-         platforms where it is not the case are rare nowadays.
-
-       An alternative would be to try to look if the information is
-       available in the ucontext_t in the caller.
-    */
-#ifdef __APPLE__
-    return (void *)get_current_thread_state();
-#else
-    return (void *)pthread_self();
-#endif
-}
 
 
 /* *************************************************************
@@ -148,31 +126,32 @@ static void segfault_handler(int arg)
 
 static void sigprof_handler(int sig_nr, siginfo_t* info, void *ucontext)
 {
-#ifdef __APPLE__
+    PyThreadState * tstate = NULL;
+    void (*prevhandler)(int);
     // TERRIBLE HACK AHEAD
     // on OS X, the thread local storage is sometimes uninitialized
     // when the signal handler runs - it means it's impossible to read errno
     // or call any syscall or read PyThread_Current or pthread_self. Additionally,
     // it seems impossible to read the register gs.
     // here we register segfault handler (all guarded by a spinlock) and call
     // longjmp in case segfault happens while reading a thread local
+    //
+    // We do the same error detection for linux to ensure that
+    // get_current_thread_state returns a sane result
     while (__sync_lock_test_and_set(&spinlock, 1)) {
     }
-    signal(SIGSEGV, &segfault_handler);
+    prevhandler = signal(SIGSEGV, &segfault_handler);
     int fault_code = setjmp(restore_point);
     if (fault_code == 0) {
         pthread_self();
-        get_current_thread_state();
+        tstate = PyGILState_GetThisThreadState();
     } else {
-        signal(SIGSEGV, SIG_DFL);
-        __sync_synchronize();
-        spinlock = 0;
+        signal(SIGSEGV, prevhandler);
+        __sync_lock_release(&spinlock);
         return;    
     }
-    signal(SIGSEGV, SIG_DFL);
-    __sync_synchronize();
-    spinlock = 0;
-#endif
+    signal(SIGSEGV, prevhandler);
+    __sync_lock_release(&spinlock);
     long val = __sync_fetch_and_add(&signal_handler_value, 2L);
 
     if ((val & 1) == 0) {
@@ -189,11 +168,11 @@ static void sigprof_handler(int sig_nr, siginfo_t* info, void *ucontext)
             struct prof_stacktrace_s *st = (struct prof_stacktrace_s *)p->data;
             st->marker = MARKER_STACKTRACE;
             st->count = 1;
-            depth = get_stack_trace(st->stack, MAX_STACK_DEPTH-1, ucontext);
+            depth = get_stack_trace(tstate, st->stack, MAX_STACK_DEPTH-1, ucontext);
             //st->stack[0] = GetPC((ucontext_t*)ucontext);
             // we gonna need that for pypy
             st->depth = depth;
-            st->stack[depth++] = get_current_thread_id();
+            st->stack[depth++] = tstate;
             long rss = get_current_proc_rss();
             if (rss >= 0)
                 st->stack[depth++] = (void*)rss;

src/vmprof_main_win32.h

@@ -76,6 +76,18 @@ int vmprof_snapshot_thread(DWORD thread_id, PyThreadState *tstate, prof_stacktra
     return depth;
 }
 
+static
+PyThreadState * get_current_thread_state(void)
+{
+#if PY_MAJOR_VERSION < 3
+    return _PyThreadState_Current;
+#elif PY_VERSION_HEX < 0x03050200
+    return (PyThreadState*) _Py_atomic_load_relaxed(&_PyThreadState_Current);
+#else
+    return _PyThreadState_UncheckedGet();
+#endif
+}
+
 long __stdcall vmprof_mainloop(void *arg)
 {   
     prof_stacktrace_s *stack = (prof_stacktrace_s*)malloc(SINGLE_BUF_SIZE);

vmprof/test/test_run.py

@@ -112,8 +112,8 @@ def test_nested_call():
         assert len(t[''].children) == 0
 
 def test_multithreaded():
-    if '__pypy__' in sys.builtin_module_names or PY3K:
-        py.test.skip("not supported on pypy and python3 just yet")
+    if '__pypy__' in sys.builtin_module_names:
+        py.test.skip("not supported on pypy just yet")
     import threading
     finished = []