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().

1 files changed, 37 insertions, 6 deletions

diff --git a/Lib/test/test_tracemalloc.py b/Lib/test/test_tracemalloc.py
index 238ae14b388..0220a83d24b 100644
--- a/Lib/test/test_tracemalloc.py
+++ b/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
@@ -19,6 +20,7 @@ except ImportError:
     _testinternalcapi = None
 
 
+DEFAULT_DOMAIN = 0
 EMPTY_STRING_SIZE = sys.getsizeof(b'')
 INVALID_NFRAME = (-1, 2**30)
 
@@ -1027,8 +1029,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
@@ -1070,7 +1072,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()
@@ -1078,13 +1080,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()
@@ -1110,6 +1118,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()