diff --git a/src/_vmprof.c b/src/_vmprof.c index bdefdc28..34ff0e47 100644 --- a/src/_vmprof.c +++ b/src/_vmprof.c @@ -383,8 +383,22 @@ static PyObject * vmp_get_profile_path(PyObject *module, PyObject *noargs) { #ifdef VMPROF_UNIX static PyObject * -insert_real_time_thread(PyObject *module, PyObject * noargs) { +insert_real_time_thread(PyObject *module, PyObject * args) { ssize_t thread_count; + unsigned long thread_id = 0; + pthread_t th = pthread_self(); + + if (!PyArg_ParseTuple(args, "|k", &thread_id)) { + return NULL; + } + + if (thread_id) { +#if SIZEOF_LONG <= SIZEOF_PTHREAD_T + th = (pthread_t) thread_id; +#else + th = (pthread_t) *(unsigned long *) &thread_id; +#endif + } if (!vmprof_is_enabled()) { PyErr_SetString(PyExc_ValueError, "vmprof is not enabled"); @@ -397,15 +411,29 @@ insert_real_time_thread(PyObject *module, PyObject * noargs) { } vmprof_aquire_lock(); - thread_count = insert_thread(pthread_self(), -1); + thread_count = insert_thread(th, -1); vmprof_release_lock(); return PyLong_FromSsize_t(thread_count); } static PyObject * -remove_real_time_thread(PyObject *module, PyObject * noargs) { +remove_real_time_thread(PyObject *module, PyObject * args) { ssize_t thread_count; + unsigned long thread_id = 0; + pthread_t th = pthread_self(); + + if (!PyArg_ParseTuple(args, "|k", &thread_id)) { + return NULL; + } + + if (thread_id) { +#if SIZEOF_LONG <= SIZEOF_PTHREAD_T + th = (pthread_t) thread_id; +#else + th = (pthread_t) *(unsigned long *) &thread_id; +#endif + } if (!vmprof_is_enabled()) { PyErr_SetString(PyExc_ValueError, "vmprof is not enabled"); @@ -418,7 +446,7 @@ remove_real_time_thread(PyObject *module, PyObject * noargs) { } vmprof_aquire_lock(); - thread_count = remove_thread(pthread_self(), -1); + thread_count = remove_thread(th, -1); vmprof_release_lock(); return PyLong_FromSsize_t(thread_count); @@ -445,9 +473,9 @@ static PyMethodDef VMProfMethods[] = { #ifdef VMPROF_UNIX {"get_profile_path", vmp_get_profile_path, METH_NOARGS, "Profile path the profiler logs to."}, - {"insert_real_time_thread", insert_real_time_thread, METH_NOARGS, + {"insert_real_time_thread", insert_real_time_thread, METH_VARARGS, "Insert a thread into the real time profiling list."}, - {"remove_real_time_thread", remove_real_time_thread, METH_NOARGS, + {"remove_real_time_thread", remove_real_time_thread, METH_VARARGS, "Remove a thread from the real time profiling list."}, #endif {NULL, NULL, 0, NULL} /* Sentinel */ diff --git a/src/vmprof_unix.c b/src/vmprof_unix.c index 1a85a10c..595fff55 100644 --- a/src/vmprof_unix.c +++ b/src/vmprof_unix.c @@ -244,11 +244,7 @@ void sigprof_handler(int sig_nr, siginfo_t* info, void *ucontext) if (commit) { commit_buffer(fd, p); } else { -#ifndef RPYTHON_VMPROF - fprintf(stderr, "WARNING: canceled buffer, no stack trace was written\n"); -#else fprintf(stderr, "WARNING: canceled buffer, no stack trace was written\n"); -#endif cancel_buffer(p); } } diff --git a/vmprof/__init__.py b/vmprof/__init__.py index 1a9c7a58..8b0c5bb2 100644 --- a/vmprof/__init__.py +++ b/vmprof/__init__.py @@ -102,19 +102,21 @@ def resolve_addr(addr): """ return _vmprof.resolve_addr(addr) -def insert_real_time_thread(): +def insert_real_time_thread(thread_id=0): """ Inserts a thread into the list of threads to be sampled in real time mode. When enabling real time mode, the caller thread is inserted automatically. Returns the number of registered threads, or -1 if we can't insert thread. + Inserts the current thread if thread_id is not provided. """ - return _vmprof.insert_real_time_thread() + return _vmprof.insert_real_time_thread(thread_id) -def remove_real_time_thread(): +def remove_real_time_thread(thread_id=0): """ Removes a thread from the list of threads to be sampled in real time mode. When disabling in real time mode, *all* threads are removed automatically. Returns the number of registered threads, or -1 if we can't remove thread. + Removes the current thread if thread_id is not provided. """ - return _vmprof.remove_real_time_thread() + return _vmprof.remove_real_time_thread(thread_id) def is_enabled(): diff --git a/vmprof/test/test_run.py b/vmprof/test/test_run.py index f1790cea..f97b3e07 100644 --- a/vmprof/test/test_run.py +++ b/vmprof/test/test_run.py @@ -81,13 +81,22 @@ def function_bar(): def functime_foo(t=0.05, insert=False): if (insert): vmprof.insert_real_time_thread() - return time.sleep(t) + sleep_retry_eintr(t) def functime_bar(t=0.05, remove=False): if (remove): vmprof.remove_real_time_thread() - return time.sleep(t) + sleep_retry_eintr(t) + + +def sleep_retry_eintr(t): + start = time.time() + remaining = t + while remaining > 0: + time.sleep(remaining) + elapsed = time.time() - start + remaining = t - elapsed foo_full_name = "py:function_foo:%d:%s" % (function_foo.__code__.co_firstlineno, @@ -255,7 +264,6 @@ def test_vmprof_real_time(): assert d[foo_time_name] > 0 -@py.test.mark.xfail() @py.test.mark.skipif("'__pypy__' in sys.builtin_module_names") @py.test.mark.skipif("sys.platform == 'win32'") @py.test.mark.parametrize("insert_foo,remove_bar", [ @@ -280,6 +288,36 @@ def test_vmprof_real_time_threaded(insert_foo, remove_bar): assert remove_bar != (bar_time_name in d) +@py.test.mark.skipif("'__pypy__' in sys.builtin_module_names") +@py.test.mark.skipif("sys.platform == 'win32'") +@py.test.mark.parametrize("insert_foo,remove_bar", [ + (False, False), + (False, True), + ( True, False), + ( True, True), +]) +def test_insert_other_real_time_thread(insert_foo, remove_bar): + import threading + prof = vmprof.Profiler() + wait = 0.5 + # This test is the same as above, except that we manually add/remove + # all threads explicitly by id from the main thread. + thread = threading.Thread(target=functime_foo, args=[wait, False]) + with prof.measure(period=0.25, real_time=True): + thread.start() + if insert_foo: + vmprof.insert_real_time_thread(thread.ident) + if remove_bar: + vmprof.remove_real_time_thread(threading.current_thread().ident) + functime_bar(wait, False) + thread.join() + stats = prof.get_stats() + tprof = stats.top_profile() + d = dict(tprof) + assert insert_foo == (foo_time_name in d) + assert remove_bar != (bar_time_name in d) + + if GZIP: def test_gzip_problem():