Skip to content

Commit

Permalink
[3.13] gh-128679: Fix tracemalloc.stop() race conditions (#128897)
Browse files Browse the repository at this point in the history
tracemalloc_alloc(), tracemalloc_realloc(), PyTraceMalloc_Track(),
PyTraceMalloc_Untrack() and _PyTraceMalloc_TraceRef() now check
tracemalloc_config.tracing after calling TABLES_LOCK().

_PyTraceMalloc_Stop() now protects more code with TABLES_LOCK(),
especially setting tracemalloc_config.tracing to 1.

Add a test using PyTraceMalloc_Track() to test tracemalloc.stop()
race condition.

Call _PyTraceMalloc_Init() at Python startup.
  • Loading branch information
vstinner authored Jan 18, 2025
1 parent ef99618 commit 6b47499
Show file tree
Hide file tree
Showing 7 changed files with 250 additions and 117 deletions.
2 changes: 1 addition & 1 deletion Include/internal/pycore_tracemalloc.h
Original file line number Diff line number Diff line change
Expand Up @@ -144,7 +144,7 @@ extern PyObject* _PyTraceMalloc_GetTraces(void);
extern PyObject* _PyTraceMalloc_GetObjectTraceback(PyObject *obj);

/* Initialize tracemalloc */
extern int _PyTraceMalloc_Init(void);
extern PyStatus _PyTraceMalloc_Init(void);

/* Start tracemalloc */
extern int _PyTraceMalloc_Start(int max_nframe);
Expand Down
12 changes: 10 additions & 2 deletions Lib/test/test_tracemalloc.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,9 @@
from test.support.script_helper import (assert_python_ok, assert_python_failure,
interpreter_requires_environment)
from test import support
from test.support import os_helper
from test.support import force_not_colorized
from test.support import os_helper
from test.support import threading_helper

try:
import _testcapi
Expand Down Expand Up @@ -952,7 +953,6 @@ def check_env_var_invalid(self, nframe):
return
self.fail(f"unexpected output: {stderr!a}")


def test_env_var_invalid(self):
for nframe in INVALID_NFRAME:
with self.subTest(nframe=nframe):
Expand Down Expand Up @@ -1101,6 +1101,14 @@ def test_stop_untrack(self):
with self.assertRaises(RuntimeError):
self.untrack()

@unittest.skipIf(_testcapi is None, 'need _testcapi')
@threading_helper.requires_working_threading()
# gh-128679: Test crash on a debug build (especially on FreeBSD).
@unittest.skipIf(support.Py_DEBUG, 'need release build')
def test_tracemalloc_track_race(self):
# gh-128679: Test fix for tracemalloc.stop() race condition
_testcapi.tracemalloc_track_race()


if __name__ == "__main__":
unittest.main()
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
Fix :func:`tracemalloc.stop` race condition. Fix :mod:`tracemalloc` to
support calling :func:`tracemalloc.stop` in one thread, while another thread
is tracing memory allocations. Patch by Victor Stinner.
99 changes: 99 additions & 0 deletions Modules/_testcapimodule.c
Original file line number Diff line number Diff line change
Expand Up @@ -3322,6 +3322,104 @@ test_atexit(PyObject *self, PyObject *Py_UNUSED(args))
Py_RETURN_NONE;
}


static void
tracemalloc_track_race_thread(void *data)
{
PyTraceMalloc_Track(123, 10, 1);

PyThread_type_lock lock = (PyThread_type_lock)data;
PyThread_release_lock(lock);
}

// gh-128679: Test fix for tracemalloc.stop() race condition
static PyObject *
tracemalloc_track_race(PyObject *self, PyObject *args)
{
#define NTHREAD 50
PyObject *tracemalloc = NULL;
PyObject *stop = NULL;
PyThread_type_lock locks[NTHREAD];
memset(locks, 0, sizeof(locks));

// Call tracemalloc.start()
tracemalloc = PyImport_ImportModule("tracemalloc");
if (tracemalloc == NULL) {
goto error;
}
PyObject *start = PyObject_GetAttrString(tracemalloc, "start");
if (start == NULL) {
goto error;
}
PyObject *res = PyObject_CallNoArgs(start);
Py_DECREF(start);
if (res == NULL) {
goto error;
}
Py_DECREF(res);

stop = PyObject_GetAttrString(tracemalloc, "stop");
Py_CLEAR(tracemalloc);
if (stop == NULL) {
goto error;
}

// Start threads
for (size_t i = 0; i < NTHREAD; i++) {
PyThread_type_lock lock = PyThread_allocate_lock();
if (!lock) {
PyErr_NoMemory();
goto error;
}
locks[i] = lock;
PyThread_acquire_lock(lock, 1);

unsigned long thread;
thread = PyThread_start_new_thread(tracemalloc_track_race_thread,
(void*)lock);
if (thread == (unsigned long)-1) {
PyErr_SetString(PyExc_RuntimeError, "can't start new thread");
goto error;
}
}

// Call tracemalloc.stop() while threads are running
res = PyObject_CallNoArgs(stop);
Py_CLEAR(stop);
if (res == NULL) {
goto error;
}
Py_DECREF(res);

// Wait until threads complete with the GIL released
Py_BEGIN_ALLOW_THREADS
for (size_t i = 0; i < NTHREAD; i++) {
PyThread_type_lock lock = locks[i];
PyThread_acquire_lock(lock, 1);
PyThread_release_lock(lock);
}
Py_END_ALLOW_THREADS

// Free threads locks
for (size_t i=0; i < NTHREAD; i++) {
PyThread_type_lock lock = locks[i];
PyThread_free_lock(lock);
}
Py_RETURN_NONE;

error:
Py_CLEAR(tracemalloc);
Py_CLEAR(stop);
for (size_t i=0; i < NTHREAD; i++) {
PyThread_type_lock lock = locks[i];
if (lock) {
PyThread_free_lock(lock);
}
}
return NULL;
#undef NTHREAD
}

static PyMethodDef TestMethods[] = {
{"set_errno", set_errno, METH_VARARGS},
{"test_config", test_config, METH_NOARGS},
Expand Down Expand Up @@ -3464,6 +3562,7 @@ static PyMethodDef TestMethods[] = {
{"function_set_warning", function_set_warning, METH_NOARGS},
{"test_critical_sections", test_critical_sections, METH_NOARGS},
{"test_atexit", test_atexit, METH_NOARGS},
{"tracemalloc_track_race", tracemalloc_track_race, METH_NOARGS},
{NULL, NULL} /* sentinel */
};

Expand Down
5 changes: 0 additions & 5 deletions Modules/_tracemalloc.c
Original file line number Diff line number Diff line change
Expand Up @@ -223,10 +223,5 @@ PyInit__tracemalloc(void)
PyUnstable_Module_SetGIL(m, Py_MOD_GIL_NOT_USED);
#endif

if (_PyTraceMalloc_Init() < 0) {
Py_DECREF(m);
return NULL;
}

return m;
}
5 changes: 5 additions & 0 deletions Python/pylifecycle.c
Original file line number Diff line number Diff line change
Expand Up @@ -706,6 +706,11 @@ pycore_create_interpreter(_PyRuntimeState *runtime,
return _PyStatus_NO_MEMORY();
}

status = _PyTraceMalloc_Init();
if (_PyStatus_EXCEPTION(status)) {
return status;
}

PyThreadState *tstate = _PyThreadState_New(interp,
_PyThreadState_WHENCE_INIT);
if (tstate == NULL) {
Expand Down
Loading

0 comments on commit 6b47499

Please sign in to comment.