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, 23 insertions, 2 deletions

diff --git a/Python/tracemalloc.c b/Python/tracemalloc.c
index b398ea3..62065e8 100644
--- a/Python/tracemalloc.c
+++ b/Python/tracemalloc.c
@@ -1256,9 +1256,17 @@ 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 result;
     if (tracemalloc_config.tracing) {
         result = tracemalloc_add_trace_unlocked(domain, ptr, size);
     }
@@ -1268,6 +1276,7 @@ PyTraceMalloc_Track(unsigned int domain, uintptr_t ptr,
     }
 
     TABLES_UNLOCK();
+done:
     PyGILState_Release(gil_state);
 
     return result;
@@ -1277,9 +1286,19 @@ PyTraceMalloc_Track(unsigned int domain, uintptr_t ptr,
 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_unlocked(domain, ptr);
         result = 0;
@@ -1290,6 +1309,8 @@ PyTraceMalloc_Untrack(unsigned int domain, uintptr_t ptr)
     }
 
     TABLES_UNLOCK();
+done:
+    PyGILState_Release(gil_state);
     return result;
 }