python · vstinner · Jan 18, 2025 · Jan 15, 2025 · Jan 16, 2025 · Jan 16, 2025
diff --git a/Include/internal/pycore_tracemalloc.h b/Include/internal/pycore_tracemalloc.h
@@ -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);

diff --git 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,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()
diff --git a/Misc/NEWS.d/next/Library/2025-01-10-15-43-52.gh-issue-128679.KcfVVR.rst b/Misc/NEWS.d/next/Library/2025-01-10-15-43-52.gh-issue-128679.KcfVVR.rst
@@ -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.
diff --git a/Modules/_testcapimodule.c b/Modules/_testcapimodule.c
@@ -3386,6 +3386,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},
@@ -3532,6 +3630,7 @@ static PyMethodDef TestMethods[] = {
    {"test_critical_sections", test_critical_sections, METH_NOARGS},
    {"pyeval_getlocals", pyeval_getlocals, METH_NOARGS},
    {"test_atexit", test_atexit, METH_NOARGS},
+    {"tracemalloc_track_race", tracemalloc_track_race, METH_NOARGS},
    {NULL, NULL} /* sentinel */
 };


diff --git a/Modules/_tracemalloc.c b/Modules/_tracemalloc.c
@@ -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;
 }
@@ -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) {