Use atomics for the _PyRuntime.tracemalloc.config.tracing so we can bail

early in `PyTraceMalloc_Track` and `PyTraceMalloc_Untrack` if tracemalloc
isn't enabled. This avoids a significant scaling issue when tracemalloc
isn't enabled, in extension modules that properly use the
PyTraceMalloc_Track API (like numpy).

Author: Yhg1s

Include/internal/pycore_tracemalloc.h

@@ -20,8 +20,7 @@ struct _PyTraceMalloc_Config {
         TRACEMALLOC_FINALIZED
     } initialized;
 
-    /* Is tracemalloc tracing memory allocations?
-       Variable protected by the TABLES_LOCK(). */
+    /* Is tracemalloc tracing memory allocations? (Atomic.) */
     int tracing;
 
     /* limit of the number of frames in a traceback, 1 by default.

Python/tracemalloc.c

@@ -850,7 +850,7 @@ _PyTraceMalloc_Start(int max_nframe)
 
     /* everything is ready: start tracing Python memory allocations */
     TABLES_LOCK();
-    tracemalloc_config.tracing = 1;
+    _Py_atomic_store_int_release(&tracemalloc_config.tracing, 1);
     TABLES_UNLOCK();
 
     return 0;
@@ -867,7 +867,7 @@ _PyTraceMalloc_Stop(void)
     }
 
     /* stop tracing Python memory allocations */
-    tracemalloc_config.tracing = 0;
+    _Py_atomic_store_int_release(&tracemalloc_config.tracing, 0);
 
     /* unregister the hook on memory allocators */
     PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &allocators.raw);
@@ -1207,6 +1207,12 @@ int
 PyTraceMalloc_Track(unsigned int domain, uintptr_t ptr,
                     size_t size)
 {
+    if (!_Py_atomic_load_int_relaxed(&tracemalloc_config.tracing)) {
+        // Exit early without attempting to lock if it doesn't look like
+        // we're tracing.
+        return -2;
+    }
+
     PyGILState_STATE gil_state = PyGILState_Ensure();
     TABLES_LOCK();
 
@@ -1215,7 +1221,7 @@ PyTraceMalloc_Track(unsigned int domain, uintptr_t ptr,
         result = tracemalloc_add_trace_unlocked(domain, ptr, size);
     }
     else {
-        /* tracemalloc is not tracing: do nothing */
+        /* tracemalloc was disabled before we locked tables; do nothing. */
         result = -2;
     }
 
@@ -1228,6 +1234,11 @@ PyTraceMalloc_Track(unsigned int domain, uintptr_t ptr,
 int
 PyTraceMalloc_Untrack(unsigned int domain, uintptr_t ptr)
 {
+    if (!_Py_atomic_load_int_relaxed(&tracemalloc_config.tracing)) {
+        // Exit early without attempting to lock if it doesn't look like
+        // we're tracing.
+        return -2;
+    }
     TABLES_LOCK();
 
     int result;
@@ -1236,7 +1247,7 @@ PyTraceMalloc_Untrack(unsigned int domain, uintptr_t ptr)
         result = 0;
     }
     else {
-        /* tracemalloc is not tracing: do nothing */
+        /* tracemalloc was disabled before we locked tables; do nothing. */
         result = -2;
     }