gh-129185: Simplify PyTraceMalloc_Track() (#129256)

Since tracemalloc uses PyMutex, it becomes safe to use TABLES_LOCK()
even after _PyTraceMalloc_Fini(): remove the "pre-check" in
PyTraceMalloc_Track() and PyTraceMalloc_Untrack().

PyTraceMalloc_Untrack() no longer needs to acquire the GIL.

_PyTraceMalloc_Fini() can be called earlier during Python
finalization.

2 files changed, 3 insertions, 26 deletions

diff --git a/Python/pylifecycle.c b/Python/pylifecycle.c
index 52890cf..f357ddf 100644
--- a/Python/pylifecycle.c
+++ b/Python/pylifecycle.c
@@ -2113,7 +2113,7 @@ _Py_Finalize(_PyRuntimeState *runtime)
 
     /* Disable tracemalloc after all Python objects have been destroyed,
        so it is possible to use tracemalloc in objects destructor. */
-    _PyTraceMalloc_Stop();
+    _PyTraceMalloc_Fini();
 
     /* Finalize any remaining import state */
     // XXX Move these up to where finalize_modules() is currently.
@@ -2166,7 +2166,6 @@ _Py_Finalize(_PyRuntimeState *runtime)
 
     finalize_interp_clear(tstate);
 
-    _PyTraceMalloc_Fini();
 
 #ifdef Py_TRACE_REFS
     /* Display addresses (& refcnts) of all objects still alive.
diff --git a/Python/tracemalloc.c b/Python/tracemalloc.c
index d27c2f9..d69b0eb 100644
--- a/Python/tracemalloc.c
+++ b/Python/tracemalloc.c
@@ -1203,17 +1203,9 @@ 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);
     }
@@ -1223,9 +1215,7 @@ PyTraceMalloc_Track(unsigned int domain, uintptr_t ptr,
     }
 
     TABLES_UNLOCK();
-done:
     PyGILState_Release(gil_state);
-
     return result;
 }
 
@@ -1233,19 +1223,9 @@ done:
 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;
@@ -1256,8 +1236,6 @@ PyTraceMalloc_Untrack(unsigned int domain, uintptr_t ptr)
     }
 
     TABLES_UNLOCK();
-done:
-    PyGILState_Release(gil_state);
     return result;
 }