diff --git a/Lib/test/test_tracemalloc.py b/Lib/test/test_tracemalloc.py index 5755f7697de91ae..da2db28775578a2 100644 --- a/Lib/test/test_tracemalloc.py +++ b/Lib/test/test_tracemalloc.py @@ -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 @@ -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): @@ -1101,6 +1101,12 @@ def test_stop_untrack(self): with self.assertRaises(RuntimeError): self.untrack() + @unittest.skipIf(_testcapi is None, 'need _testcapi') + @threading_helper.requires_working_threading() + def test_tracemalloc_track_race(self): + # gh-128679: Test fix for tracemalloc.stop() race condition + _testcapi.tracemalloc_track_race() + if __name__ == "__main__": unittest.main() diff --git a/Modules/_testcapimodule.c b/Modules/_testcapimodule.c index a0a1f8af6710a3d..7d304add5999d1f 100644 --- a/Modules/_testcapimodule.c +++ b/Modules/_testcapimodule.c @@ -3435,6 +3435,104 @@ code_offset_to_line(PyObject* self, PyObject* const* args, Py_ssize_t nargsf) return PyLong_FromInt32(PyCode_Addr2Line(code, offset)); } + +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}, @@ -3578,6 +3676,7 @@ static PyMethodDef TestMethods[] = { {"type_freeze", type_freeze, METH_VARARGS}, {"test_atexit", test_atexit, METH_NOARGS}, {"code_offset_to_line", _PyCFunction_CAST(code_offset_to_line), METH_FASTCALL}, + {"tracemalloc_track_race", tracemalloc_track_race, METH_NOARGS}, {NULL, NULL} /* sentinel */ }; diff --git a/Python/tracemalloc.c b/Python/tracemalloc.c index 68b641c51868b95..919c564ee729670 100644 --- a/Python/tracemalloc.c +++ b/Python/tracemalloc.c @@ -567,11 +567,14 @@ tracemalloc_alloc(int need_gil, int use_calloc, } TABLES_LOCK(); - if (ADD_TRACE(ptr, nelem * elsize) < 0) { - // Failed to allocate a trace for the new memory block - alloc->free(alloc->ctx, ptr); - ptr = NULL; + if (tracemalloc_config.tracing) { + if (ADD_TRACE(ptr, nelem * elsize) < 0) { + // Failed to allocate a trace for the new memory block + alloc->free(alloc->ctx, ptr); + ptr = NULL; + } } + // else: gh-128679: tracemalloc.stop() was called by another thread TABLES_UNLOCK(); if (need_gil) { @@ -614,6 +617,11 @@ tracemalloc_realloc(int need_gil, void *ctx, void *ptr, size_t new_size) } TABLES_LOCK(); + if (!tracemalloc_config.tracing) { + // gh-128679: tracemalloc.stop() was called by another thread + goto unlock; + } + if (ptr != NULL) { // An existing memory block has been resized @@ -646,6 +654,7 @@ tracemalloc_realloc(int need_gil, void *ctx, void *ptr, size_t new_size) } } +unlock: TABLES_UNLOCK(); if (need_gil) { PyGILState_Release(gil_state); @@ -674,7 +683,12 @@ tracemalloc_free(void *ctx, void *ptr) } TABLES_LOCK(); - REMOVE_TRACE(ptr); + + if (tracemalloc_config.tracing) { + REMOVE_TRACE(ptr); + } + // else: gh-128679: tracemalloc.stop() was called by another thread + TABLES_UNLOCK(); } @@ -1312,8 +1326,9 @@ _PyTraceMalloc_TraceRef(PyObject *op, PyRefTracerEvent event, assert(PyGILState_Check()); TABLES_LOCK(); - int result = -1; - assert(tracemalloc_config.tracing); + if (!tracemalloc_config.tracing) { + goto done; + } PyTypeObject *type = Py_TYPE(op); const size_t presize = _PyType_PreHeaderSize(type); @@ -1325,13 +1340,13 @@ _PyTraceMalloc_TraceRef(PyObject *op, PyRefTracerEvent event, traceback_t *traceback = traceback_new(); if (traceback != NULL) { trace->traceback = traceback; - result = 0; } } /* else: cannot track the object, its memory block size is unknown */ +done: TABLES_UNLOCK(); - return result; + return 0; } @@ -1472,13 +1487,19 @@ int _PyTraceMalloc_GetTracebackLimit(void) size_t _PyTraceMalloc_GetMemory(void) { - size_t size = _Py_hashtable_size(tracemalloc_tracebacks); - size += _Py_hashtable_size(tracemalloc_filenames); - TABLES_LOCK(); - size += _Py_hashtable_size(tracemalloc_traces); - _Py_hashtable_foreach(tracemalloc_domains, - tracemalloc_get_tracemalloc_memory_cb, &size); + size_t size; + if (tracemalloc_config.tracing) { + size = _Py_hashtable_size(tracemalloc_tracebacks); + size += _Py_hashtable_size(tracemalloc_filenames); + + size += _Py_hashtable_size(tracemalloc_traces); + _Py_hashtable_foreach(tracemalloc_domains, + tracemalloc_get_tracemalloc_memory_cb, &size); + } + else { + size = 0; + } TABLES_UNLOCK(); return size; }