Skip to content

Commit 45eb346

Browse files
arigoarigo
authored
Merge pull request #205 from timpalpant/insert-thread-by-id
Insert real time thread by ID
2 parents e4e99e5 + bf1e2db commit 45eb346

4 files changed

Lines changed: 81 additions & 17 deletions

File tree

‎src/_vmprof.c‎

Lines changed: 34 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -383,8 +383,22 @@ static PyObject * vmp_get_profile_path(PyObject *module, PyObject *noargs) {
383383

384384
#ifdef VMPROF_UNIX
385385
static PyObject *
386-
insert_real_time_thread(PyObject *module, PyObject * noargs) {
386+
insert_real_time_thread(PyObject *module, PyObject * args) {
387387
ssize_t thread_count;
388+
unsigned long thread_id = 0;
389+
pthread_t th = pthread_self();
390+
391+
if (!PyArg_ParseTuple(args, "|k", &thread_id)) {
392+
return NULL;
393+
}
394+
395+
if (thread_id) {
396+
#if SIZEOF_LONG <= SIZEOF_PTHREAD_T
397+
th = (pthread_t) thread_id;
398+
#else
399+
th = (pthread_t) *(unsigned long *) &thread_id;
400+
#endif
401+
}
388402

389403
if (!vmprof_is_enabled()) {
390404
PyErr_SetString(PyExc_ValueError, "vmprof is not enabled");
@@ -397,15 +411,29 @@ insert_real_time_thread(PyObject *module, PyObject * noargs) {
397411
}
398412

399413
vmprof_aquire_lock();
400-
thread_count = insert_thread(pthread_self(), -1);
414+
thread_count = insert_thread(th, -1);
401415
vmprof_release_lock();
402416

403417
return PyLong_FromSsize_t(thread_count);
404418
}
405419

406420
static PyObject *
407-
remove_real_time_thread(PyObject *module, PyObject * noargs) {
421+
remove_real_time_thread(PyObject *module, PyObject * args) {
408422
ssize_t thread_count;
423+
unsigned long thread_id = 0;
424+
pthread_t th = pthread_self();
425+
426+
if (!PyArg_ParseTuple(args, "|k", &thread_id)) {
427+
return NULL;
428+
}
429+
430+
if (thread_id) {
431+
#if SIZEOF_LONG <= SIZEOF_PTHREAD_T
432+
th = (pthread_t) thread_id;
433+
#else
434+
th = (pthread_t) *(unsigned long *) &thread_id;
435+
#endif
436+
}
409437

410438
if (!vmprof_is_enabled()) {
411439
PyErr_SetString(PyExc_ValueError, "vmprof is not enabled");
@@ -418,7 +446,7 @@ remove_real_time_thread(PyObject *module, PyObject * noargs) {
418446
}
419447

420448
vmprof_aquire_lock();
421-
thread_count = remove_thread(pthread_self(), -1);
449+
thread_count = remove_thread(th, -1);
422450
vmprof_release_lock();
423451

424452
return PyLong_FromSsize_t(thread_count);
@@ -445,9 +473,9 @@ static PyMethodDef VMProfMethods[] = {
445473
#ifdef VMPROF_UNIX
446474
{"get_profile_path", vmp_get_profile_path, METH_NOARGS,
447475
"Profile path the profiler logs to."},
448-
{"insert_real_time_thread", insert_real_time_thread, METH_NOARGS,
476+
{"insert_real_time_thread", insert_real_time_thread, METH_VARARGS,
449477
"Insert a thread into the real time profiling list."},
450-
{"remove_real_time_thread", remove_real_time_thread, METH_NOARGS,
478+
{"remove_real_time_thread", remove_real_time_thread, METH_VARARGS,
451479
"Remove a thread from the real time profiling list."},
452480
#endif
453481
{NULL, NULL, 0, NULL} /* Sentinel */

‎src/vmprof_unix.c‎

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -244,11 +244,7 @@ void sigprof_handler(int sig_nr, siginfo_t* info, void *ucontext)
244244
if (commit) {
245245
commit_buffer(fd, p);
246246
} else {
247-
#ifndef RPYTHON_VMPROF
248-
fprintf(stderr, "WARNING: canceled buffer, no stack trace was written\n");
249-
#else
250247
fprintf(stderr, "WARNING: canceled buffer, no stack trace was written\n");
251-
#endif
252248
cancel_buffer(p);
253249
}
254250
}

‎vmprof/__init__.py‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -102,19 +102,21 @@ def resolve_addr(addr):
102102
"""
103103
return _vmprof.resolve_addr(addr)
104104

105-
def insert_real_time_thread():
105+
def insert_real_time_thread(thread_id=0):
106106
""" Inserts a thread into the list of threads to be sampled in real time mode.
107107
When enabling real time mode, the caller thread is inserted automatically.
108108
Returns the number of registered threads, or -1 if we can't insert thread.
109+
Inserts the current thread if thread_id is not provided.
109110
"""
110-
return _vmprof.insert_real_time_thread()
111+
return _vmprof.insert_real_time_thread(thread_id)
111112

112-
def remove_real_time_thread():
113+
def remove_real_time_thread(thread_id=0):
113114
""" Removes a thread from the list of threads to be sampled in real time mode.
114115
When disabling in real time mode, *all* threads are removed automatically.
115116
Returns the number of registered threads, or -1 if we can't remove thread.
117+
Removes the current thread if thread_id is not provided.
116118
"""
117-
return _vmprof.remove_real_time_thread()
119+
return _vmprof.remove_real_time_thread(thread_id)
118120

119121

120122
def is_enabled():

‎vmprof/test/test_run.py‎

Lines changed: 41 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -81,13 +81,22 @@ def function_bar():
8181
def functime_foo(t=0.05, insert=False):
8282
if (insert):
8383
vmprof.insert_real_time_thread()
84-
return time.sleep(t)
84+
sleep_retry_eintr(t)
8585

8686

8787
def functime_bar(t=0.05, remove=False):
8888
if (remove):
8989
vmprof.remove_real_time_thread()
90-
return time.sleep(t)
90+
sleep_retry_eintr(t)
91+
92+
93+
def sleep_retry_eintr(t):
94+
start = time.time()
95+
remaining = t
96+
while remaining > 0:
97+
time.sleep(remaining)
98+
elapsed = time.time() - start
99+
remaining = t - elapsed
91100

92101

93102
foo_full_name = "py:function_foo:%d:%s" % (function_foo.__code__.co_firstlineno,
@@ -255,7 +264,6 @@ def test_vmprof_real_time():
255264
assert d[foo_time_name] > 0
256265

257266

258-
@py.test.mark.xfail()
259267
@py.test.mark.skipif("'__pypy__' in sys.builtin_module_names")
260268
@py.test.mark.skipif("sys.platform == 'win32'")
261269
@py.test.mark.parametrize("insert_foo,remove_bar", [
@@ -280,6 +288,36 @@ def test_vmprof_real_time_threaded(insert_foo, remove_bar):
280288
assert remove_bar != (bar_time_name in d)
281289

282290

291+
@py.test.mark.skipif("'__pypy__' in sys.builtin_module_names")
292+
@py.test.mark.skipif("sys.platform == 'win32'")
293+
@py.test.mark.parametrize("insert_foo,remove_bar", [
294+
(False, False),
295+
(False, True),
296+
( True, False),
297+
( True, True),
298+
])
299+
def test_insert_other_real_time_thread(insert_foo, remove_bar):
300+
import threading
301+
prof = vmprof.Profiler()
302+
wait = 0.5
303+
# This test is the same as above, except that we manually add/remove
304+
# all threads explicitly by id from the main thread.
305+
thread = threading.Thread(target=functime_foo, args=[wait, False])
306+
with prof.measure(period=0.25, real_time=True):
307+
thread.start()
308+
if insert_foo:
309+
vmprof.insert_real_time_thread(thread.ident)
310+
if remove_bar:
311+
vmprof.remove_real_time_thread(threading.current_thread().ident)
312+
functime_bar(wait, False)
313+
thread.join()
314+
stats = prof.get_stats()
315+
tprof = stats.top_profile()
316+
d = dict(tprof)
317+
assert insert_foo == (foo_time_name in d)
318+
assert remove_bar != (bar_time_name in d)
319+
320+
283321

284322
if GZIP:
285323
def test_gzip_problem():

0 commit comments

Comments
 (0)