gh-128679: Redesign tracemalloc locking (#128888)

* Use TABLES_LOCK() to protect 'tracemalloc_config.tracing'.
* Hold TABLES_LOCK() longer while accessing tables.
* tracemalloc_realloc() and tracemalloc_free() no longer
  remove the trace on reentrant call.
* _PyTraceMalloc_Stop() unregisters _PyTraceMalloc_TraceRef().
* _PyTraceMalloc_GetTraces() sets the reentrant flag.
* tracemalloc_clear_traces_unlocked() sets the reentrant flag.

2 files changed, 283 insertions, 335 deletions

diff --git a/Python/pylifecycle.c b/Python/pylifecycle.c
index 9a8a92aaceb..f1ecee6a92e 100644
--- a/Python/pylifecycle.c
+++ b/Python/pylifecycle.c
@@ -707,7 +707,12 @@ pycore_create_interpreter(_PyRuntimeState *runtime,
     // the settings are loaded (so that feature_flags are set) but before
     // any calls are made to obmalloc functions.
     if (_PyMem_init_obmalloc(interp) < 0) {
-        return  _PyStatus_NO_MEMORY();
+        return _PyStatus_NO_MEMORY();
+    }
+
+    status = _PyTraceMalloc_Init();
+    if (_PyStatus_EXCEPTION(status)) {
+        return status;
     }
 
     PyThreadState *tstate = _PyThreadState_New(interp,
diff --git a/Python/tracemalloc.c b/Python/tracemalloc.c
index f661d69c031..68b641c5186 100644
--- a/Python/tracemalloc.c
+++ b/Python/tracemalloc.c
@@ -2,6 +2,7 @@
 #include "pycore_fileutils.h"     // _Py_write_noraise()
 #include "pycore_gc.h"            // PyGC_Head
 #include "pycore_hashtable.h"     // _Py_hashtable_t
+#include "pycore_initconfig.h"    // _PyStatus_NO_MEMORY()
 #include "pycore_object.h"        // _PyType_PreHeaderSize()
 #include "pycore_pymem.h"         // _Py_tracemalloc_config
 #include "pycore_runtime.h"       // _Py_ID()
@@ -19,6 +20,8 @@ _Py_DECLARE_STR(anon_unknown, "<unknown>");
 /* Forward declaration */
 static void* raw_malloc(size_t size);
 static void raw_free(void *ptr);
+static int _PyTraceMalloc_TraceRef(PyObject *op, PyRefTracerEvent event,
+                                   void* Py_UNUSED(ignore));
 
 #ifdef Py_DEBUG
 #  define TRACE_DEBUG
@@ -106,16 +109,16 @@ tracemalloc_error(const char *format, ...)
 static int
 get_reentrant(void)
 {
-    void *ptr;
-
     assert(PyThread_tss_is_created(&tracemalloc_reentrant_key));
-    ptr = PyThread_tss_get(&tracemalloc_reentrant_key);
+
+    void *ptr = PyThread_tss_get(&tracemalloc_reentrant_key);
     if (ptr != NULL) {
         assert(ptr == REENTRANT);
         return 1;
     }
-    else
+    else {
         return 0;
+    }
 }
 
 static void
@@ -252,6 +255,7 @@ tracemalloc_get_frame(_PyInterpreterFrame *pyframe, frame_t *frame)
 {
     assert(PyStackRef_CodeCheck(pyframe->f_executable));
     frame->filename = &_Py_STR(anon_unknown);
+
     int lineno = PyUnstable_InterpreterFrame_GetLine(pyframe);
     if (lineno < 0) {
         lineno = 0;
@@ -259,7 +263,6 @@ tracemalloc_get_frame(_PyInterpreterFrame *pyframe, frame_t *frame)
     frame->lineno = (unsigned int)lineno;
 
     PyObject *filename = filename = _PyFrame_GetCode(pyframe)->co_filename;
-
     if (filename == NULL) {
 #ifdef TRACE_DEBUG
         tracemalloc_error("failed to get the filename of the code object");
@@ -275,7 +278,7 @@ tracemalloc_get_frame(_PyInterpreterFrame *pyframe, frame_t *frame)
     }
     if (!PyUnicode_IS_READY(filename)) {
         /* Don't make a Unicode string ready to avoid reentrant calls
-           to tracemalloc_malloc() or tracemalloc_realloc() */
+           to tracemalloc_alloc() or tracemalloc_realloc() */
 #ifdef TRACE_DEBUG
         tracemalloc_error("filename is not a ready unicode string");
 #endif
@@ -309,7 +312,7 @@ tracemalloc_get_frame(_PyInterpreterFrame *pyframe, frame_t *frame)
 static Py_uhash_t
 traceback_hash(traceback_t *traceback)
 {
-    /* code based on tuplehash() of Objects/tupleobject.c */
+    /* code based on tuple_hash() of Objects/tupleobject.c */
     Py_uhash_t x, y;  /* Unsigned for defined overflow behavior. */
     int len = traceback->nframe;
     Py_uhash_t mult = PyHASH_MULTIPLIER;
@@ -440,7 +443,7 @@ tracemalloc_get_traces_table(unsigned int domain)
 
 
 static void
-tracemalloc_remove_trace(unsigned int domain, uintptr_t ptr)
+tracemalloc_remove_trace_unlocked(unsigned int domain, uintptr_t ptr)
 {
     assert(tracemalloc_config.tracing);
 
@@ -459,12 +462,12 @@ tracemalloc_remove_trace(unsigned int domain, uintptr_t ptr)
 }
 
 #define REMOVE_TRACE(ptr) \
-            tracemalloc_remove_trace(DEFAULT_DOMAIN, (uintptr_t)(ptr))
+    tracemalloc_remove_trace_unlocked(DEFAULT_DOMAIN, (uintptr_t)(ptr))
 
 
 static int
-tracemalloc_add_trace(unsigned int domain, uintptr_t ptr,
-                      size_t size)
+tracemalloc_add_trace_unlocked(unsigned int domain, uintptr_t ptr,
+                               size_t size)
 {
     assert(tracemalloc_config.tracing);
 
@@ -519,82 +522,138 @@ tracemalloc_add_trace(unsigned int domain, uintptr_t ptr,
 }
 
 #define ADD_TRACE(ptr, size) \
-            tracemalloc_add_trace(DEFAULT_DOMAIN, (uintptr_t)(ptr), size)
+    tracemalloc_add_trace_unlocked(DEFAULT_DOMAIN, (uintptr_t)(ptr), size)
 
 
 static void*
-tracemalloc_alloc(int use_calloc, void *ctx, size_t nelem, size_t elsize)
+tracemalloc_alloc(int need_gil, int use_calloc,
+                  void *ctx, size_t nelem, size_t elsize)
 {
-    PyMemAllocatorEx *alloc = (PyMemAllocatorEx *)ctx;
-    void *ptr;
-
     assert(elsize == 0 || nelem <= SIZE_MAX / elsize);
 
-    if (use_calloc)
+    int reentrant = get_reentrant();
+
+    // Ignore reentrant call.
+    //
+    // For example, PyObjet_Malloc() calls
+    // PyMem_Malloc() for allocations larger than 512 bytes: don't trace the
+    // same memory allocation twice.
+    //
+    // If reentrant calls are not ignored, PyGILState_Ensure() can call
+    // PyMem_RawMalloc() which would call PyGILState_Ensure() again in a loop.
+    if (!reentrant) {
+        set_reentrant(1);
+    }
+
+    PyMemAllocatorEx *alloc = (PyMemAllocatorEx *)ctx;
+    void *ptr;
+    if (use_calloc) {
         ptr = alloc->calloc(alloc->ctx, nelem, elsize);
-    else
+    }
+    else {
         ptr = alloc->malloc(alloc->ctx, nelem * elsize);
-    if (ptr == NULL)
-        return NULL;
+    }
 
+    if (ptr == NULL) {
+        goto done;
+    }
+    if (reentrant) {
+        goto done;
+    }
+
+    PyGILState_STATE gil_state;
+    if (need_gil) {
+        gil_state = PyGILState_Ensure();
+    }
     TABLES_LOCK();
+
     if (ADD_TRACE(ptr, nelem * elsize) < 0) {
-        /* Failed to allocate a trace for the new memory block */
-        TABLES_UNLOCK();
+        // Failed to allocate a trace for the new memory block
         alloc->free(alloc->ctx, ptr);
-        return NULL;
+        ptr = NULL;
     }
+
     TABLES_UNLOCK();
+    if (need_gil) {
+        PyGILState_Release(gil_state);
+    }
+
+done:
+    if (!reentrant) {
+        set_reentrant(0);
+    }
     return ptr;
 }
 
 
 static void*
-tracemalloc_realloc(void *ctx, void *ptr, size_t new_size)
+tracemalloc_realloc(int need_gil, void *ctx, void *ptr, size_t new_size)
 {
+    int reentrant = get_reentrant();
+
+    // Ignore reentrant call. PyObjet_Realloc() calls PyMem_Realloc() for
+    // allocations larger than 512 bytes: don't trace the same memory block
+    // twice.
+    if (!reentrant) {
+        set_reentrant(1);
+    }
+
     PyMemAllocatorEx *alloc = (PyMemAllocatorEx *)ctx;
-    void *ptr2;
+    void *ptr2 = alloc->realloc(alloc->ctx, ptr, new_size);
 
-    ptr2 = alloc->realloc(alloc->ctx, ptr, new_size);
-    if (ptr2 == NULL)
-        return NULL;
+    if (ptr2 == NULL) {
+        goto done;
+    }
+    if (reentrant) {
+        goto done;
+    }
 
-    if (ptr != NULL) {
-        /* an existing memory block has been resized */
+    PyGILState_STATE gil_state;
+    if (need_gil) {
+        gil_state = PyGILState_Ensure();
+    }
+    TABLES_LOCK();
 
-        TABLES_LOCK();
+    if (ptr != NULL) {
+        // An existing memory block has been resized
 
-        /* tracemalloc_add_trace() updates the trace if there is already
-           a trace at address ptr2 */
+        // tracemalloc_add_trace_unlocked() updates the trace if there is
+        // already a trace at address ptr2.
         if (ptr2 != ptr) {
             REMOVE_TRACE(ptr);
         }
 
         if (ADD_TRACE(ptr2, new_size) < 0) {
-            /* Memory allocation failed. The error cannot be reported to
-               the caller, because realloc() may already have shrunk the
-               memory block and so removed bytes.
-
-               This case is very unlikely: a hash entry has just been
-               released, so the hash table should have at least one free entry.
-
-               The GIL and the table lock ensures that only one thread is
-               allocating memory. */
+            // Memory allocation failed. The error cannot be reported to the
+            // caller, because realloc() already have shrunk the memory block
+            // and so removed bytes.
+            //
+            // This case is very unlikely: a hash entry has just been released,
+            // so the hash table should have at least one free entry.
+            //
+            // The GIL and the table lock ensures that only one thread is
+            // allocating memory.
             Py_FatalError("tracemalloc_realloc() failed to allocate a trace");
         }
-        TABLES_UNLOCK();
     }
     else {
-        /* new allocation */
+        // New allocation
 
-        TABLES_LOCK();
         if (ADD_TRACE(ptr2, new_size) < 0) {
-            /* Failed to allocate a trace for the new memory block */
-            TABLES_UNLOCK();
+            // Failed to allocate a trace for the new memory block
             alloc->free(alloc->ctx, ptr2);
-            return NULL;
+            ptr2 = NULL;
         }
-        TABLES_UNLOCK();
+    }
+
+    TABLES_UNLOCK();
+    if (need_gil) {
+        PyGILState_Release(gil_state);
+    }
+
+done:
+    if (!reentrant) {
+        set_reentrant(0);
     }
     return ptr2;
 }
@@ -603,168 +662,63 @@ tracemalloc_realloc(void *ctx, void *ptr, size_t new_size)
 static void
 tracemalloc_free(void *ctx, void *ptr)
 {
-    PyMemAllocatorEx *alloc = (PyMemAllocatorEx *)ctx;
-
-    if (ptr == NULL)
+    if (ptr == NULL) {
         return;
+    }
 
-     /* GIL cannot be locked in PyMem_RawFree() because it would introduce
-        a deadlock in _PyThreadState_DeleteCurrent(). */
-
+    PyMemAllocatorEx *alloc = (PyMemAllocatorEx *)ctx;
     alloc->free(alloc->ctx, ptr);
 
-    TABLES_LOCK();
-    REMOVE_TRACE(ptr);
-    TABLES_UNLOCK();
-}
-
-
-static void*
-tracemalloc_alloc_gil(int use_calloc, void *ctx, size_t nelem, size_t elsize)
-{
-    void *ptr;
-
     if (get_reentrant()) {
-        PyMemAllocatorEx *alloc = (PyMemAllocatorEx *)ctx;
-        if (use_calloc)
-            return alloc->calloc(alloc->ctx, nelem, elsize);
-        else
-            return alloc->malloc(alloc->ctx, nelem * elsize);
+        return;
     }
 
-    /* Ignore reentrant call. PyObjet_Malloc() calls PyMem_Malloc() for
-       allocations larger than 512 bytes, don't trace the same memory
-       allocation twice. */
-    set_reentrant(1);
-
-    ptr = tracemalloc_alloc(use_calloc, ctx, nelem, elsize);
-
-    set_reentrant(0);
-    return ptr;
+    TABLES_LOCK();
+    REMOVE_TRACE(ptr);
+    TABLES_UNLOCK();
 }
 
 
 static void*
 tracemalloc_malloc_gil(void *ctx, size_t size)
 {
-    return tracemalloc_alloc_gil(0, ctx, 1, size);
+    return tracemalloc_alloc(0, 0, ctx, 1, size);
 }
 
 
 static void*
 tracemalloc_calloc_gil(void *ctx, size_t nelem, size_t elsize)
 {
-    return tracemalloc_alloc_gil(1, ctx, nelem, elsize);
+    return tracemalloc_alloc(0, 1, ctx, nelem, elsize);
 }
 
 
 static void*
 tracemalloc_realloc_gil(void *ctx, void *ptr, size_t new_size)
 {
-    void *ptr2;
-
-    if (get_reentrant()) {
-        /* Reentrant call to PyMem_Realloc() and PyMem_RawRealloc().
-           Example: PyMem_RawRealloc() is called internally by pymalloc
-           (_PyObject_Malloc() and  _PyObject_Realloc()) to allocate a new
-           arena (new_arena()). */
-        PyMemAllocatorEx *alloc = (PyMemAllocatorEx *)ctx;
-
-        ptr2 = alloc->realloc(alloc->ctx, ptr, new_size);
-        if (ptr2 != NULL && ptr != NULL) {
-            TABLES_LOCK();
-            REMOVE_TRACE(ptr);
-            TABLES_UNLOCK();
-        }
-        return ptr2;
-    }
-
-    /* Ignore reentrant call. PyObjet_Realloc() calls PyMem_Realloc() for
-       allocations larger than 512 bytes. Don't trace the same memory
-       allocation twice. */
-    set_reentrant(1);
-
-    ptr2 = tracemalloc_realloc(ctx, ptr, new_size);
-
-    set_reentrant(0);
-    return ptr2;
+    return tracemalloc_realloc(0, ctx, ptr, new_size);
 }
 
 
 #ifdef TRACE_RAW_MALLOC
 static void*
-tracemalloc_raw_alloc(int use_calloc, void *ctx, size_t nelem, size_t elsize)
-{
-    PyGILState_STATE gil_state;
-    void *ptr;
-
-    if (get_reentrant()) {
-        PyMemAllocatorEx *alloc = (PyMemAllocatorEx *)ctx;
-        if (use_calloc)
-            return alloc->calloc(alloc->ctx, nelem, elsize);
-        else
-            return alloc->malloc(alloc->ctx, nelem * elsize);
-    }
-
-    /* Ignore reentrant call. PyGILState_Ensure() may call PyMem_RawMalloc()
-       indirectly which would call PyGILState_Ensure() if reentrant are not
-       disabled. */
-    set_reentrant(1);
-
-    gil_state = PyGILState_Ensure();
-    ptr = tracemalloc_alloc(use_calloc, ctx, nelem, elsize);
-    PyGILState_Release(gil_state);
-
-    set_reentrant(0);
-    return ptr;
-}
-
-
-static void*
 tracemalloc_raw_malloc(void *ctx, size_t size)
 {
-    return tracemalloc_raw_alloc(0, ctx, 1, size);
+    return tracemalloc_alloc(1, 0, ctx, 1, size);
 }
 
 
 static void*
 tracemalloc_raw_calloc(void *ctx, size_t nelem, size_t elsize)
 {
-    return tracemalloc_raw_alloc(1, ctx, nelem, elsize);
+    return tracemalloc_alloc(1, 1, ctx, nelem, elsize);
 }
 
 
 static void*
 tracemalloc_raw_realloc(void *ctx, void *ptr, size_t new_size)
 {
-    PyGILState_STATE gil_state;
-    void *ptr2;
-
-    if (get_reentrant()) {
-        /* Reentrant call to PyMem_RawRealloc(). */
-        PyMemAllocatorEx *alloc = (PyMemAllocatorEx *)ctx;
-
-        ptr2 = alloc->realloc(alloc->ctx, ptr, new_size);
-
-        if (ptr2 != NULL && ptr != NULL) {
-            TABLES_LOCK();
-            REMOVE_TRACE(ptr);
-            TABLES_UNLOCK();
-        }
-        return ptr2;
-    }
-
-    /* Ignore reentrant call. PyGILState_Ensure() may call PyMem_RawMalloc()
-       indirectly which would call PyGILState_Ensure() if reentrant calls are
-       not disabled. */
-    set_reentrant(1);
-
-    gil_state = PyGILState_Ensure();
-    ptr2 = tracemalloc_realloc(ctx, ptr, new_size);
-    PyGILState_Release(gil_state);
-
-    set_reentrant(0);
-    return ptr2;
+    return tracemalloc_realloc(1, ctx, ptr, new_size);
 }
 #endif   /* TRACE_RAW_MALLOC */
 
@@ -777,48 +731,36 @@ tracemalloc_clear_filename(void *value)
 }
 
 
-/* reentrant flag must be set to call this function and GIL must be held */
 static void
-tracemalloc_clear_traces(void)
+tracemalloc_clear_traces_unlocked(void)
 {
-    /* The GIL protects variables against concurrent access */
+    // Clearing tracemalloc_filenames requires the GIL to call Py_DECREF()
     assert(PyGILState_Check());
 
-    TABLES_LOCK();
+    set_reentrant(1);
+
     _Py_hashtable_clear(tracemalloc_traces);
     _Py_hashtable_clear(tracemalloc_domains);
+    _Py_hashtable_clear(tracemalloc_tracebacks);
+    _Py_hashtable_clear(tracemalloc_filenames);
+
     tracemalloc_traced_memory = 0;
     tracemalloc_peak_traced_memory = 0;
-    TABLES_UNLOCK();
-
-    _Py_hashtable_clear(tracemalloc_tracebacks);
 
-    _Py_hashtable_clear(tracemalloc_filenames);
+    set_reentrant(0);
 }
 
 
-int
+PyStatus
 _PyTraceMalloc_Init(void)
 {
-    if (tracemalloc_config.initialized == TRACEMALLOC_FINALIZED) {
-        PyErr_SetString(PyExc_RuntimeError,
-                        "the tracemalloc module has been unloaded");
-        return -1;
-    }
-
-    if (tracemalloc_config.initialized == TRACEMALLOC_INITIALIZED)
-        return 0;
+    assert(tracemalloc_config.initialized == TRACEMALLOC_NOT_INITIALIZED);
 
     PyMem_GetAllocator(PYMEM_DOMAIN_RAW, &allocators.raw);
 
 #ifdef REENTRANT_THREADLOCAL
     if (PyThread_tss_create(&tracemalloc_reentrant_key) != 0) {
-#ifdef MS_WINDOWS
-        PyErr_SetFromWindowsErr(0);
-#else
-        PyErr_SetFromErrno(PyExc_OSError);
-#endif
-        return -1;
+        return _PyStatus_NO_MEMORY();
     }
 #endif
 
@@ -826,8 +768,7 @@ _PyTraceMalloc_Init(void)
     if (tables_lock == NULL) {
         tables_lock = PyThread_allocate_lock();
         if (tables_lock == NULL) {
-            PyErr_SetString(PyExc_RuntimeError, "cannot allocate lock");
-            return -1;
+            return _PyStatus_NO_MEMORY();
         }
     }
 #endif
@@ -844,9 +785,9 @@ _PyTraceMalloc_Init(void)
     tracemalloc_domains = tracemalloc_create_domains_table();
 
     if (tracemalloc_filenames == NULL || tracemalloc_tracebacks == NULL
-       || tracemalloc_traces == NULL || tracemalloc_domains == NULL) {
-        PyErr_NoMemory();
-        return -1;
+       || tracemalloc_traces == NULL || tracemalloc_domains == NULL)
+    {
+        return _PyStatus_NO_MEMORY();
     }
 
     tracemalloc_empty_traceback.nframe = 1;
@@ -857,7 +798,7 @@ _PyTraceMalloc_Init(void)
     tracemalloc_empty_traceback.hash = traceback_hash(&tracemalloc_empty_traceback);
 
     tracemalloc_config.initialized = TRACEMALLOC_INITIALIZED;
-    return 0;
+    return _PyStatus_OK();
 }
 
 
@@ -892,9 +833,6 @@ tracemalloc_deinit(void)
 int
 _PyTraceMalloc_Start(int max_nframe)
 {
-    PyMemAllocatorEx alloc;
-    size_t size;
-
     if (max_nframe < 1 || (unsigned long) max_nframe > MAX_NFRAME) {
         PyErr_Format(PyExc_ValueError,
                      "the number of frames must be in range [1; %lu]",
@@ -902,23 +840,15 @@ _PyTraceMalloc_Start(int max_nframe)
         return -1;
     }
 
-    if (_PyTraceMalloc_Init() < 0) {
-        return -1;
-    }
-
-    if (PyRefTracer_SetTracer(_PyTraceMalloc_TraceRef, NULL) < 0) {
-        return -1;
-    }
-
-    if (tracemalloc_config.tracing) {
-        /* hook already installed: do nothing */
+    if (_PyTraceMalloc_IsTracing()) {
+        /* hooks already installed: do nothing */
         return 0;
     }
 
     tracemalloc_config.max_nframe = max_nframe;
 
     /* allocate a buffer to store a new traceback */
-    size = TRACEBACK_SIZE(max_nframe);
+    size_t size = TRACEBACK_SIZE(max_nframe);
     assert(tracemalloc_traceback == NULL);
     tracemalloc_traceback = raw_malloc(size);
     if (tracemalloc_traceback == NULL) {
@@ -926,6 +856,7 @@ _PyTraceMalloc_Start(int max_nframe)
         return -1;
     }
 
+    PyMemAllocatorEx alloc;
 #ifdef TRACE_RAW_MALLOC
     alloc.malloc = tracemalloc_raw_malloc;
     alloc.calloc = tracemalloc_raw_calloc;
@@ -950,8 +881,14 @@ _PyTraceMalloc_Start(int max_nframe)
     PyMem_GetAllocator(PYMEM_DOMAIN_OBJ, &allocators.obj);
     PyMem_SetAllocator(PYMEM_DOMAIN_OBJ, &alloc);
 
+    if (PyRefTracer_SetTracer(_PyTraceMalloc_TraceRef, NULL) < 0) {
+        return -1;
+    }
+
     /* everything is ready: start tracing Python memory allocations */
+    TABLES_LOCK();
     tracemalloc_config.tracing = 1;
+    TABLES_UNLOCK();
 
     return 0;
 }
@@ -960,8 +897,11 @@ _PyTraceMalloc_Start(int max_nframe)
 void
 _PyTraceMalloc_Stop(void)
 {
-    if (!tracemalloc_config.tracing)
-        return;
+    TABLES_LOCK();
+
+    if (!tracemalloc_config.tracing) {
+        goto done;
+    }
 
     /* stop tracing Python memory allocations */
     tracemalloc_config.tracing = 0;
@@ -973,11 +913,16 @@ _PyTraceMalloc_Stop(void)
     PyMem_SetAllocator(PYMEM_DOMAIN_MEM, &allocators.mem);
     PyMem_SetAllocator(PYMEM_DOMAIN_OBJ, &allocators.obj);
 
-    tracemalloc_clear_traces();
+    tracemalloc_clear_traces_unlocked();
 
     /* release memory */
     raw_free(tracemalloc_traceback);
     tracemalloc_traceback = NULL;
+
+    (void)PyRefTracer_SetTracer(NULL, NULL);
+
+done:
+    TABLES_UNLOCK();
 }
 
 
@@ -985,15 +930,16 @@ _PyTraceMalloc_Stop(void)
 static PyObject*
 frame_to_pyobject(frame_t *frame)
 {
-    PyObject *frame_obj, *lineno_obj;
+    assert(get_reentrant());
 
-    frame_obj = PyTuple_New(2);
-    if (frame_obj == NULL)
+    PyObject *frame_obj = PyTuple_New(2);
+    if (frame_obj == NULL) {
         return NULL;
+    }
 
     PyTuple_SET_ITEM(frame_obj, 0, Py_NewRef(frame->filename));
 
-    lineno_obj = PyLong_FromUnsignedLong(frame->lineno);
+    PyObject *lineno_obj = PyLong_FromUnsignedLong(frame->lineno);
     if (lineno_obj == NULL) {
         Py_DECREF(frame_obj);
         return NULL;
@@ -1008,7 +954,6 @@ static PyObject*
 traceback_to_pyobject(traceback_t *traceback, _Py_hashtable_t *intern_table)
 {
     PyObject *frames;
-
     if (intern_table != NULL) {
         frames = _Py_hashtable_get(intern_table, (const void *)traceback);
         if (frames) {
@@ -1017,8 +962,9 @@ traceback_to_pyobject(traceback_t *traceback, _Py_hashtable_t *intern_table)
     }
 
     frames = PyTuple_New(traceback->nframe);
-    if (frames == NULL)
+    if (frames == NULL) {
         return NULL;
+    }
 
     for (int i=0; i < traceback->nframe; i++) {
         PyObject *frame = frame_to_pyobject(&traceback->frames[i]);
@@ -1046,14 +992,14 @@ static PyObject*
 trace_to_pyobject(unsigned int domain, const trace_t *trace,
                   _Py_hashtable_t *intern_tracebacks)
 {
-    PyObject *trace_obj = NULL;
-    PyObject *obj;
+    assert(get_reentrant());
 
-    trace_obj = PyTuple_New(4);
-    if (trace_obj == NULL)
+    PyObject *trace_obj = PyTuple_New(4);
+    if (trace_obj == NULL) {
         return NULL;
+    }
 
-    obj = PyLong_FromSize_t(domain);
+    PyObject *obj = PyLong_FromSize_t(domain);
     if (obj == NULL) {
         Py_DECREF(trace_obj);
         return NULL;
@@ -1100,7 +1046,6 @@ tracemalloc_copy_trace(_Py_hashtable_t *traces,
                        void *user_data)
 {
     _Py_hashtable_t *traces2 = (_Py_hashtable_t *)user_data;
-
     trace_t *trace = (trace_t *)value;
 
     trace_t *trace2 = raw_malloc(sizeof(trace_t));
@@ -1141,7 +1086,6 @@ tracemalloc_copy_domain(_Py_hashtable_t *domains,
                         void *user_data)
 {
     _Py_hashtable_t *domains2 = (_Py_hashtable_t *)user_data;
-
     unsigned int domain = (unsigned int)FROM_PTR(key);
     _Py_hashtable_t *traces = (_Py_hashtable_t *)value;
 
@@ -1182,7 +1126,6 @@ tracemalloc_get_traces_fill(_Py_hashtable_t *traces,
                             void *user_data)
 {
     get_traces_t *get_traces = user_data;
-
     const trace_t *trace = (const trace_t *)value;
 
     PyObject *tuple = trace_to_pyobject(get_traces->domain, trace,
@@ -1196,7 +1139,6 @@ tracemalloc_get_traces_fill(_Py_hashtable_t *traces,
     if (res < 0) {
         return 1;
     }
-
     return 0;
 }
 
@@ -1207,7 +1149,6 @@ tracemalloc_get_traces_domain(_Py_hashtable_t *domains,
                               void *user_data)
 {
     get_traces_t *get_traces = user_data;
-
     unsigned int domain = (unsigned int)FROM_PTR(key);
     _Py_hashtable_t *traces = (_Py_hashtable_t *)value;
 
@@ -1227,27 +1168,21 @@ tracemalloc_pyobject_decref(void *value)
 
 
 static traceback_t*
-tracemalloc_get_traceback(unsigned int domain, uintptr_t ptr)
+tracemalloc_get_traceback_unlocked(unsigned int domain, uintptr_t ptr)
 {
-
-    if (!tracemalloc_config.tracing)
+    if (!tracemalloc_config.tracing) {
         return NULL;
+    }
 
-    trace_t *trace;
-    TABLES_LOCK();
     _Py_hashtable_t *traces = tracemalloc_get_traces_table(domain);
-    if (traces) {
-        trace = _Py_hashtable_get(traces, TO_PTR(ptr));
-    }
-    else {
-        trace = NULL;
+    if (!traces) {
+        return NULL;
     }
-    TABLES_UNLOCK();
 
+    trace_t *trace = _Py_hashtable_get(traces, TO_PTR(ptr));
     if (!trace) {
         return NULL;
     }
-
     return trace->traceback;
 }
 
@@ -1269,24 +1204,28 @@ _PyMem_DumpFrame(int fd, frame_t * frame)
 void
 _PyMem_DumpTraceback(int fd, const void *ptr)
 {
-    traceback_t *traceback;
-    int i;
-
+    TABLES_LOCK();
     if (!tracemalloc_config.tracing) {
         PUTS(fd, "Enable tracemalloc to get the memory block "
                  "allocation traceback\n\n");
-        return;
+        goto done;
     }
 
-    traceback = tracemalloc_get_traceback(DEFAULT_DOMAIN, (uintptr_t)ptr);
-    if (traceback == NULL)
-        return;
+    traceback_t *traceback;
+    traceback = tracemalloc_get_traceback_unlocked(DEFAULT_DOMAIN,
+                                                   (uintptr_t)ptr);
+    if (traceback == NULL) {
+        goto done;
+    }
 
     PUTS(fd, "Memory block allocated at (most recent call first):\n");
-    for (i=0; i < traceback->nframe; i++) {
+    for (int i=0; i < traceback->nframe; i++) {
         _PyMem_DumpFrame(fd, &traceback->frames[i]);
     }
     PUTS(fd, "\n");
+
+done:
+    TABLES_UNLOCK();
 }
 
 #undef PUTS
@@ -1307,38 +1246,42 @@ int
 PyTraceMalloc_Track(unsigned int domain, uintptr_t ptr,
                     size_t size)
 {
-    int res;
-    PyGILState_STATE gil_state;
+    PyGILState_STATE gil_state = PyGILState_Ensure();
+    TABLES_LOCK();
 
-    if (!tracemalloc_config.tracing) {
+    int result;
+    if (tracemalloc_config.tracing) {
+        result = tracemalloc_add_trace_unlocked(domain, ptr, size);
+    }
+    else {
         /* tracemalloc is not tracing: do nothing */
-        return -2;
+        result = -2;
     }
 
-    gil_state = PyGILState_Ensure();
-
-    TABLES_LOCK();
-    res = tracemalloc_add_trace(domain, ptr, size);
     TABLES_UNLOCK();
-
     PyGILState_Release(gil_state);
-    return res;
+
+    return result;
 }
 
 
 int
 PyTraceMalloc_Untrack(unsigned int domain, uintptr_t ptr)
 {
-    if (!tracemalloc_config.tracing) {
+    TABLES_LOCK();
+
+    int result;
+    if (tracemalloc_config.tracing) {
+        tracemalloc_remove_trace_unlocked(domain, ptr);
+        result = 0;
+    }
+    else {
         /* tracemalloc is not tracing: do nothing */
-        return -2;
+        result = -2;
     }
 
-    TABLES_LOCK();
-    tracemalloc_remove_trace(domain, ptr);
     TABLES_UNLOCK();
-
-    return 0;
+    return result;
 }
 
 
@@ -1355,87 +1298,101 @@ _PyTraceMalloc_Fini(void)
 
    Do nothing if tracemalloc is not tracing memory allocations
    or if the object memory block is not already traced. */
-int
-_PyTraceMalloc_TraceRef(PyObject *op, PyRefTracerEvent event, void* Py_UNUSED(ignore))
+static int
+_PyTraceMalloc_TraceRef(PyObject *op, PyRefTracerEvent event,
+                        void* Py_UNUSED(ignore))
 {
     if (event != PyRefTracer_CREATE) {
         return 0;
     }
+    if (get_reentrant()) {
+        return 0;
+    }
 
     assert(PyGILState_Check());
+    TABLES_LOCK();
 
-    if (!tracemalloc_config.tracing) {
-        /* tracemalloc is not tracing: do nothing */
-        return -1;
-    }
+    int result = -1;
+    assert(tracemalloc_config.tracing);
 
     PyTypeObject *type = Py_TYPE(op);
     const size_t presize = _PyType_PreHeaderSize(type);
     uintptr_t ptr = (uintptr_t)((char *)op - presize);
 
-    int res = -1;
-
-    TABLES_LOCK();
     trace_t *trace = _Py_hashtable_get(tracemalloc_traces, TO_PTR(ptr));
     if (trace != NULL) {
         /* update the traceback of the memory block */
         traceback_t *traceback = traceback_new();
         if (traceback != NULL) {
             trace->traceback = traceback;
-            res = 0;
+            result = 0;
         }
     }
     /* else: cannot track the object, its memory block size is unknown */
-    TABLES_UNLOCK();
 
-    return res;
+    TABLES_UNLOCK();
+    return result;
 }
 
 
 PyObject*
 _PyTraceMalloc_GetTraceback(unsigned int domain, uintptr_t ptr)
 {
-    traceback_t *traceback;
+    TABLES_LOCK();
 
-    traceback = tracemalloc_get_traceback(domain, ptr);
-    if (traceback == NULL)
-        Py_RETURN_NONE;
+    traceback_t *traceback = tracemalloc_get_traceback_unlocked(domain, ptr);
+    PyObject *result;
+    if (traceback) {
+        set_reentrant(1);
+        result = traceback_to_pyobject(traceback, NULL);
+        set_reentrant(0);
+    }
+    else {
+        result = Py_NewRef(Py_None);
+    }
 
-    return traceback_to_pyobject(traceback, NULL);
+    TABLES_UNLOCK();
+    return result;
 }
 
 int
 _PyTraceMalloc_IsTracing(void)
 {
-    return tracemalloc_config.tracing;
+    TABLES_LOCK();
+    int tracing = tracemalloc_config.tracing;
+    TABLES_UNLOCK();
+    return tracing;
 }
 
 void
 _PyTraceMalloc_ClearTraces(void)
 {
-
-    if (!tracemalloc_config.tracing) {
-        return;
+    TABLES_LOCK();
+    if (tracemalloc_config.tracing) {
+        tracemalloc_clear_traces_unlocked();
     }
-    set_reentrant(1);
-    tracemalloc_clear_traces();
-    set_reentrant(0);
+    TABLES_UNLOCK();
 }
 
 PyObject *
 _PyTraceMalloc_GetTraces(void)
 {
+    TABLES_LOCK();
+    set_reentrant(1);
+
     get_traces_t get_traces;
     get_traces.domain = DEFAULT_DOMAIN;
     get_traces.traces = NULL;
     get_traces.domains = NULL;
     get_traces.tracebacks = NULL;
     get_traces.list = PyList_New(0);
-    if (get_traces.list == NULL)
-        goto error;
+    if (get_traces.list == NULL) {
+        goto finally;
+    }
 
-    if (!tracemalloc_config.tracing)
-        return get_traces.list;
+    if (!tracemalloc_config.tracing) {
+        goto finally;
+    }
 
     /* the traceback hash table is used temporarily to intern traceback tuple
        of (filename, lineno) tuples */
@@ -1449,24 +1406,17 @@ _PyTraceMalloc_GetTraces(void)
     // Copy all traces so tracemalloc_get_traces_fill() doesn't have to disable
     // temporarily tracemalloc which would impact other threads and so would
     // miss allocations while get_traces() is called.
-    TABLES_LOCK();
     get_traces.traces = tracemalloc_copy_traces(tracemalloc_traces);
-    TABLES_UNLOCK();
-
     if (get_traces.traces == NULL) {
         goto no_memory;
     }
 
-    TABLES_LOCK();
     get_traces.domains = tracemalloc_copy_domains(tracemalloc_domains);
-    TABLES_UNLOCK();
-
     if (get_traces.domains == NULL) {
         goto no_memory;
     }
 
     // Convert traces to a list of tuples
-    set_reentrant(1);
     int err = _Py_hashtable_foreach(get_traces.traces,
                                     tracemalloc_get_traces_fill,
                                     &get_traces);
@@ -1475,20 +1425,22 @@ _PyTraceMalloc_GetTraces(void)
                                     tracemalloc_get_traces_domain,
                                     &get_traces);
     }
-    set_reentrant(0);
+
     if (err) {
-        goto error;
+        Py_CLEAR(get_traces.list);
+        goto finally;
     }
-
     goto finally;
 
 no_memory:
     PyErr_NoMemory();
-
-error:
     Py_CLEAR(get_traces.list);
+    goto finally;
 
 finally:
+    set_reentrant(0);
+    TABLES_UNLOCK();
+
     if (get_traces.tracebacks != NULL) {
         _Py_hashtable_destroy(get_traces.tracebacks);
     }
@@ -1506,31 +1458,21 @@ PyObject *
 _PyTraceMalloc_GetObjectTraceback(PyObject *obj)
 /*[clinic end generated code: output=41ee0553a658b0aa input=29495f1b21c53212]*/
 {
-    PyTypeObject *type;
-    traceback_t *traceback;
-
-    type = Py_TYPE(obj);
+    PyTypeObject *type = Py_TYPE(obj);
     const size_t presize = _PyType_PreHeaderSize(type);
     uintptr_t ptr = (uintptr_t)((char *)obj - presize);
-
-    traceback = tracemalloc_get_traceback(DEFAULT_DOMAIN, ptr);
-    if (traceback == NULL) {
-        Py_RETURN_NONE;
-    }
-
-    return traceback_to_pyobject(traceback, NULL);
+    return _PyTraceMalloc_GetTraceback(DEFAULT_DOMAIN, ptr);
 }
 
-int _PyTraceMalloc_GetTracebackLimit(void) {
+int _PyTraceMalloc_GetTracebackLimit(void)
+{
     return tracemalloc_config.max_nframe;
 }
 
 size_t
-_PyTraceMalloc_GetMemory(void) {
-
-    size_t size;
-
-    size = _Py_hashtable_size(tracemalloc_tracebacks);
+_PyTraceMalloc_GetMemory(void)
+{
+    size_t size = _Py_hashtable_size(tracemalloc_tracebacks);
     size += _Py_hashtable_size(tracemalloc_filenames);
 
     TABLES_LOCK();
@@ -1545,26 +1487,27 @@ _PyTraceMalloc_GetMemory(void) {
 PyObject *
 _PyTraceMalloc_GetTracedMemory(void)
 {
-    Py_ssize_t size, peak_size;
-
-    if (!tracemalloc_config.tracing)
-        return Py_BuildValue("ii", 0, 0);
-
     TABLES_LOCK();
-    size = tracemalloc_traced_memory;
-    peak_size = tracemalloc_peak_traced_memory;
+    Py_ssize_t traced, peak;
+    if (tracemalloc_config.tracing) {
+        traced = tracemalloc_traced_memory;
+        peak = tracemalloc_peak_traced_memory;
+    }
+    else {
+        traced = 0;
+        peak = 0;
+    }
     TABLES_UNLOCK();
 
-    return Py_BuildValue("nn", size, peak_size);
+    return Py_BuildValue("nn", traced, peak);
 }
 
 void
 _PyTraceMalloc_ResetPeak(void)
 {
-    if (!tracemalloc_config.tracing) {
-        return;
-    }
     TABLES_LOCK();
-    tracemalloc_peak_traced_memory = tracemalloc_traced_memory;
+    if (tracemalloc_config.tracing) {
+        tracemalloc_peak_traced_memory = tracemalloc_traced_memory;
+    }
     TABLES_UNLOCK();
 }