[3.12] gh-129185: Fix PyTraceMalloc_Untrack() at Python exit (#129191) (#129217) (#129221)

[3.13] gh-129185: Fix PyTraceMalloc_Untrack() at Python exit (#129191) (#129217)

gh-129185: Fix PyTraceMalloc_Untrack() at Python exit (#129191)

Support calling PyTraceMalloc_Track() and PyTraceMalloc_Untrack()
during late Python finalization.

* Call _PyTraceMalloc_Fini() later in Python finalization.
* Test also PyTraceMalloc_Untrack() without the GIL
* PyTraceMalloc_Untrack() now gets the GIL.
* Test also PyTraceMalloc_Untrack() in test_tracemalloc_track_race().

(cherry picked from commit 46c7e13c05)
(cherry picked from commit e3b3e01d6a)

Lib/test/test_tracemalloc.py

@@ -1,6 +1,7 @@
 import contextlib
 import os
 import sys
+import textwrap
 import tracemalloc
 import unittest
 from unittest.mock import patch
@@ -16,6 +17,7 @@ except ImportError:
     _testcapi = None
 
 
+DEFAULT_DOMAIN = 0
 EMPTY_STRING_SIZE = sys.getsizeof(b'')
 INVALID_NFRAME = (-1, 2**30)
 
@@ -1020,8 +1022,8 @@ class TestCAPI(unittest.TestCase):
                                     release_gil)
         return frames
 
-    def untrack(self):
-        _testcapi.tracemalloc_untrack(self.domain, self.ptr)
+    def untrack(self, release_gil=False):
+        _testcapi.tracemalloc_untrack(self.domain, self.ptr, release_gil)
 
     def get_traced_memory(self):
         # Get the traced size in the domain
@@ -1063,7 +1065,7 @@ class TestCAPI(unittest.TestCase):
         self.assertEqual(self.get_traceback(),
                          tracemalloc.Traceback(frames))
 
-    def test_untrack(self):
+    def check_untrack(self, release_gil):
         tracemalloc.start()
 
         self.track()
@@ -1071,13 +1073,19 @@ class TestCAPI(unittest.TestCase):
         self.assertEqual(self.get_traced_memory(), self.size)
 
         # untrack must remove the trace
-        self.untrack()
+        self.untrack(release_gil)
         self.assertIsNone(self.get_traceback())
         self.assertEqual(self.get_traced_memory(), 0)
 
         # calling _PyTraceMalloc_Untrack() multiple times must not crash
-        self.untrack()
-        self.untrack()
+        self.untrack(release_gil)
+        self.untrack(release_gil)
+
+    def test_untrack(self):
+        self.check_untrack(False)
+
+    def test_untrack_without_gil(self):
+        self.check_untrack(True)
 
     def test_stop_track(self):
         tracemalloc.start()
@@ -1103,6 +1111,29 @@ class TestCAPI(unittest.TestCase):
         # gh-128679: Test fix for tracemalloc.stop() race condition
         _testcapi.tracemalloc_track_race()
 
+    def test_late_untrack(self):
+        code = textwrap.dedent(f"""
+            from test import support
+            import tracemalloc
+            import _testcapi
+
+            class Tracked:
+                def __init__(self, domain, size):
+                    self.domain = domain
+                    self.ptr = id(self)
+                    self.size = size
+                    _testcapi.tracemalloc_track(self.domain, self.ptr, self.size)
+
+                def __del__(self, untrack=_testcapi.tracemalloc_untrack):
+                    untrack(self.domain, self.ptr, 1)
+
+            domain = {DEFAULT_DOMAIN}
+            tracemalloc.start()
+            obj = Tracked(domain, 1024 * 1024)
+            support.late_deletion(obj)
+        """)
+        assert_python_ok("-c", code)
+
 
 if __name__ == "__main__":
     unittest.main()

Modules/_testcapi/mem.c

@@ -637,8 +637,9 @@ tracemalloc_untrack(PyObject *self, PyObject *args)
 {
     unsigned int domain;
     PyObject *ptr_obj;
+    int release_gil = 0;
 
-    if (!PyArg_ParseTuple(args, "IO", &domain, &ptr_obj)) {
+    if (!PyArg_ParseTuple(args, "IO|i", &domain, &ptr_obj, &release_gil)) {
         return NULL;
     }
     void *ptr = PyLong_AsVoidPtr(ptr_obj);
@@ -646,7 +647,15 @@ tracemalloc_untrack(PyObject *self, PyObject *args)
         return NULL;
     }
 
-    int res = PyTraceMalloc_Untrack(domain, (uintptr_t)ptr);
+    int res;
+    if (release_gil) {
+        Py_BEGIN_ALLOW_THREADS
+        res = PyTraceMalloc_Untrack(domain, (uintptr_t)ptr);
+        Py_END_ALLOW_THREADS
+    }
+    else {
+        res = PyTraceMalloc_Untrack(domain, (uintptr_t)ptr);
+    }
     if (res < 0) {
         PyErr_SetString(PyExc_RuntimeError, "PyTraceMalloc_Untrack error");
         return NULL;

Modules/_testcapimodule.c

@@ -3231,6 +3231,7 @@ static void
 tracemalloc_track_race_thread(void *data)
 {
     PyTraceMalloc_Track(123, 10, 1);
+    PyTraceMalloc_Untrack(123, 10);
 
     PyThread_type_lock lock = (PyThread_type_lock)data;
     PyThread_release_lock(lock);

Python/pylifecycle.c

@@ -1933,7 +1933,7 @@ Py_FinalizeEx(void)
 
     /* Disable tracemalloc after all Python objects have been destroyed,
        so it is possible to use tracemalloc in objects destructor. */
-    _PyTraceMalloc_Fini();
+    _PyTraceMalloc_Stop();
 
     /* Finalize any remaining import state */
     // XXX Move these up to where finalize_modules() is currently.
@@ -1986,6 +1986,8 @@ Py_FinalizeEx(void)
 
     finalize_interp_clear(tstate);
 
+    _PyTraceMalloc_Fini();
+
 #ifdef WITH_PYMALLOC
     if (malloc_stats) {
         _PyObject_DebugMallocStats(stderr);

Python/tracemalloc.c

@@ -1313,29 +1313,48 @@ PyTraceMalloc_Track(unsigned int domain, uintptr_t ptr,
                     size_t size)
 {
     PyGILState_STATE gil_state = PyGILState_Ensure();
+    int result;
+
+    // gh-129185: Check before TABLES_LOCK() to support calls after
+    // _PyTraceMalloc_Fini().
+    if (!tracemalloc_config.tracing) {
+        result = -2;
+        goto done;
+    }
+
     TABLES_LOCK();
 
-    int res;
     if (tracemalloc_config.tracing) {
-        res = tracemalloc_add_trace(domain, ptr, size);
+        result = tracemalloc_add_trace(domain, ptr, size);
     }
     else {
         // gh-128679: tracemalloc.stop() was called by another thread
-        res = -2;
+        result = -2;
     }
 
     TABLES_UNLOCK();
+done:
     PyGILState_Release(gil_state);
-    return res;
+    return result;
 }
 
 
 int
 PyTraceMalloc_Untrack(unsigned int domain, uintptr_t ptr)
 {
+    // Need the GIL to prevent races on the first 'tracing' test
+    PyGILState_STATE gil_state = PyGILState_Ensure();
+    int result;
+
+    // gh-129185: Check before TABLES_LOCK() to support calls after
+    // _PyTraceMalloc_Fini()
+    if (!tracemalloc_config.tracing) {
+        result = -2;
+        goto done;
+    }
+
     TABLES_LOCK();
 
-    int result;
     if (tracemalloc_config.tracing) {
         tracemalloc_remove_trace(domain, ptr);
         result = 0;
@@ -1346,6 +1365,8 @@ PyTraceMalloc_Untrack(unsigned int domain, uintptr_t ptr)
     }
 
     TABLES_UNLOCK();
+done:
+    PyGILState_Release(gil_state);
     return result;
 }