8000 gh-128679: Fix tracemalloc.stop() race conditions (#128893) · python/cpython@3193cb5 · GitHub
[go: up one dir, main page]

Skip to content

Commit 3193cb5

Browse files
authored
gh-128679: Fix tracemalloc.stop() race conditions (#128893)
tracemalloc_alloc(), tracemalloc_realloc(), tracemalloc_free(), _PyTraceMalloc_TraceRef() and _PyTraceMalloc_GetMemory() now check 'tracemalloc_config.tracing' after calling TABLES_LOCK(). _PyTraceMalloc_TraceRef() now always returns 0.
1 parent 313b96e commit 3193cb5

File tree

4 files changed

+146
-17
lines changed

4 files changed

+146
-17
lines changed

Lib/test/test_tracemalloc.py

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,9 @@
77
from test.support.script_helper import (assert_python_ok, assert_python_failure,
88
interpreter_requires_environment)
99
from test import support
10-
from test.support import os_helper
1110
from test.support import force_not_colorized
11+
from test.support import os_helper
12+
from test.support import threading_helper
1213

1314
try:
1415
import _testcapi
@@ -952,7 +953,6 @@ def check_env_var_invalid(self, nframe):
952953
return
953954
self.fail(f"unexpected output: {stderr!a}")
954955

955-
956956
def test_env_var_invalid(self):
957957
for nframe in INVALID_NFRAME:
958958
with self.subTest(nframe=nframe):
@@ -1101,6 +1101,12 @@ def test_stop_untrack(self):
11011101
with self.assertRaises(RuntimeError):
11021102
self.untrack()
11031103

1104+
@unittest.skipIf(_testcapi is None, 'need _testcapi')
1105+
@threading_helper.requires_working_threading()
1106+
def test_tracemalloc_track_race(self):
1107+
# gh-128679: Test fix for tracemalloc.stop() race condition
1108+
_testcapi.tracemalloc_track_race()
1109+
11041110

11051111
if __name__ == "__main__":
11061112
unittest.main()
Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
:mod:`tracemalloc`: Fix race conditions when :func:`tracemalloc.stop` is
2+
called by a thread, while other threads are tracing memory allocations.
3+
Patch by Victor Stinner.

Modules/_testcapimodule.c

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3435,6 +3435,104 @@ code_offset_to_line(PyObject* self, PyObject* const* args, Py_ssize_t nargsf)
34353435
return PyLong_FromInt32(PyCode_Addr2Line(code, offset));
34363436
}
34373437

3438+
3439+
static void
3440+
tracemalloc_track_race_thread(void *data)
3441+
{
3442+
PyTraceMalloc_Track(123, 10, 1);
3443+
3444+
PyThread_type_lock lock = (PyThread_type_lock)data;
3445+
PyThread_release_lock(lock);
3446+
}
3447+
3448+
// gh-128679: Test fix for tracemalloc.stop() race condition
3449+
static PyObject *
3450+
tracemalloc_track_race(PyObject *self, PyObject *args)
3451+
{
3452+
#define NTHREAD 50
3453+
PyObject *tracemalloc = NULL;
3454+
PyObject *stop = NULL;
3455+
PyThread_type_lock locks[NTHREAD];
3456+
memset(locks, 0, sizeof(locks));
3457+
3458+
// Call tracemalloc.start()
3459+
tracemalloc = PyImport_ImportModule("tracemalloc");
3460+
if (tracemalloc == NULL) {
3461+
goto error;
3462+
}
3463+
PyObject *start = PyObject_GetAttrString(tracemalloc, "start");
3464+
if (start == NULL) {
3465+
goto error;
3466+
}
3467+
PyObject *res = PyObject_CallNoArgs(start);
3468+
Py_DECREF(start);
3469+
if (res == NULL) {
3470+
goto error;
3471+
}
3472+
Py_DECREF(res);
3473+
3474+
stop = PyObject_GetAttrString(tracemalloc, "stop");
3475+
Py_CLEAR(tracemalloc);
3476+
if (stop == NULL) {
3477+
goto error;
3478+
}
3479+
3480+
// Start threads
3481+
for (size_t i = 0; i < NTHREAD; i++) {
3482+
PyThread_type_lock lock = PyThread_allocate_lock();
3483+
if (!lock) {
3484+
PyErr_NoMemory();
3485+
goto error;
3486+
}
3487+
locks[i] = lock;
3488+
PyThread_acquire_lock(lock, 1);
3489+
3490+
unsigned long thread;
3491+
thread = PyThread_start_new_thread(tracemalloc F438 _track_race_thread,
3492+
(void*)lock);
3493+
if (thread == (unsigned long)-1) {
3494+
PyErr_SetString(PyExc_RuntimeError, "can't start new thread");
3495+
goto error;
3496+
}
3497+
}
3498+
3499+
// Call tracemalloc.stop() while threads are running
3500+
res = PyObject_CallNoArgs(stop);
3501+
Py_CLEAR(stop);
3502+
if (res == NULL) {
3503+
goto error;
3504+
}
3505+
Py_DECREF(res);
3506+
3507+
// Wait until threads complete with the GIL released
3508+
Py_BEGIN_ALLOW_THREADS
3509+
for (size_t i = 0; i < NTHREAD; i++) {
3510+
PyThread_type_lock lock = locks[i];
3511+
PyThread_acquire_lock(lock, 1);
3512+
PyThread_release_lock(lock);
3513+
}
3514+
Py_END_ALLOW_THREADS
3515+
3516+
// Free threads locks
3517+
for (size_t i=0; i < NTHREAD; i++) {
3518+
PyThread_type_lock lock = locks[i];
3519+
PyThread_free_lock(lock);
3520+
}
3521+
Py_RETURN_NONE;
3522+
3523+
error:
3524+
Py_CLEAR(tracemalloc);
3525+
Py_CLEAR(stop);
3526+
for (size_t i=0; i < NTHREAD; i++) {
3527+
PyThread_type_lock lock = locks[i];
3528+
if (lock) {
3529+
PyThread_free_lock(lock);
3530+
}
3531+
}
3532+
return NULL;
3533+
#undef NTHREAD
3534+
}
3535+
34383536
static PyMethodDef TestMethods[] = {
34393537
{"set_errno", set_errno, METH_VARARGS},
34403538
{"test_config", test_config, METH_NOARGS},
@@ -3578,6 +3676,7 @@ static PyMethodDef TestMethods[] = {
35783676
{"type_freeze", type_freeze, METH_VARARGS},
35793677
{"test_atexit", test_atexit, METH_NOARGS},
35803678
{"code_offset_to_line", _PyCFunction_CAST(code_offset_to_line), METH_FASTCALL},
3679+
{"tracemalloc_track_race", tracemalloc_track_race, METH_NOARGS},
35813680
{NULL, NULL} /* sentinel */
35823681
};
35833682

Python/tracemalloc.c

Lines changed: 36 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -567,11 +567,14 @@ tracemalloc_alloc(int need_gil, int use_calloc,
567567
}
568568
TABLES_LOCK();
569569

570-
if (ADD_TRACE(ptr, nelem * elsize) < 0) {
571-
// Failed to allocate a trace for the new memory block
572-
alloc->free(alloc->ctx, ptr);
573-
ptr = NULL;
570+
if (tracemalloc_config.tracing) {
571+
if (ADD_TRACE(ptr, nelem * elsize) < 0) {
572+
// Failed to allocate a trace for the new memory block
573+
alloc->free(alloc->ctx, ptr);
574+
ptr = NULL;
575+
}
574576
}
577+
// else: gh-128679: tracemalloc.stop() was called by another thread
575578

576579
TABLES_UNLOCK();
577580
if (need_gil) {
@@ -614,6 +617,11 @@ tracemalloc_realloc(int need_gil, void *ctx, void *ptr, size_t new_size)
614617
}
615618
TABLES_LOCK();
616619

620+
if (!tracemalloc_config.tracing) {
621+
// gh-128679: tracemalloc.stop() was called by another thread
622+
goto unlock;
623+
}
624+
617625
if (ptr != NULL) {
618626
// An existing memory block has been resized
619627

@@ -646,6 +654,7 @@ tracemalloc_realloc(int need_gil, void *ctx, void *ptr, size_t new_size)
646654
}
647655
}
648656

657+
unlock:
649658
TABLES_UNLOCK();
650659
if (need_gil) {
651660
PyGILState_Release(gil_state);
@@ -674,7 +683,12 @@ tracemalloc_free(void *ctx, void *ptr)
674683
}
675684

676685
TABLES_LOCK();
677-
REMOVE_TRACE(ptr);
686+
687+
if (tracemalloc_config.tracing) {
688+
REMOVE_TRACE(ptr);
689+
}
690+
// else: gh-128679: tracemalloc.stop() was called by another thread
691+
678692
TABLES_UNLOCK();
679693
}
680694

@@ -1312,8 +1326,9 @@ _PyTraceMalloc_TraceRef(PyObject *op, PyRefTracerEvent event,
13121326
assert(PyGILState_Check());
13131327
TABLES_LOCK();
13141328

1315-
int result = -1;
1316-
assert(tracemalloc_config.tracing);
1329+
if (!tracemalloc_config.tracing) {
1330+
goto done;
1331+
}
13171332

13181333
PyTypeObject *type = Py_TYPE(op);
13191334
const size_t presize = _PyType_PreHeaderSize(type);
@@ -1325,13 +1340,13 @@ _PyTraceMalloc_TraceRef(PyObject *op, PyRefTracerEvent event,
13251340
traceback_t *traceback = traceback_new();
13261341
if (traceback != NULL) {
13271342
trace->traceback = traceback;
1328-
result = 0;
13291343
}
13301344
}
13311345
/* else: cannot track the object, its memory block size is unknown */
13321346

1347+
done:
13331348
TABLES_UNLOCK();
1334-
return result;
1349+
return 0;
13351350
}
13361351

13371352

@@ -1472,13 +1487,19 @@ int _PyTraceMalloc_GetTracebackLimit(void)
14721487
size_t
14731488
_PyTraceMalloc_GetMemory(void)
14741489
{
1475-
size_t size = _Py_hashtable_size(tracemalloc_tracebacks);
1476-
size += _Py_hashtable_size(tracemalloc_filenames);
1477-
14781490
TABLES_LOCK();
1479-
size += _Py_hashtable_size(tracemalloc_traces);
1480-
_Py_hashtable_foreach(tracemalloc_domains,
1481-
tracemalloc_get_tracemalloc_memory_cb, &size);
1491+
size_t size;
1492+
if (tracemalloc_config.tracing) {
1493+
size = _Py_hashtable_size(tracemalloc_tracebacks);
1494+
size += _Py_hashtable_size(tracemalloc_filenames);
1495+
1496+
size += _Py_hashtable_size(tracemalloc_traces);
1497+
_Py_hashtable_foreach(tracemalloc_domains,
1498+
tracemalloc_get_tracemalloc_memory_cb, &size);
1499+
}
1500+
else {
1501+
size = 0;
1502+
}
14821503
TABLES_UNLOCK();
14831504
return size;
14841505
}

0 commit comments

Comments
 (0)
0