refactor(profiling): remove redundant locks from memalloc

We added locking to memalloc, the memory profiler, in #11460 in order to
address crashes. These locks made the crashes go away, but significantly
increased the baseline overhead of the profiler and introduced subtle
bugs. The locks we added turned out to be fundamentally incompatible
with the global interpreter lock (GIL), at least with the implementation
from #11460. This PR refactors the profiler to use the GIL exclusively
for locking.

First, we should acknowledge no-GIL and subinterpreters. As of right
now, our module does not support either. A module has to explicitly
opt-in to support either, so there is no risk of those modes being
enabled under our feet. Supporting either mode is likely a repo-wide
project. For now, we can assume the GIL exists.

This work was motivated by overhead. We currently acquire and release
locks in every memory allocation and free. Even when the locks aren't
contended, allocations and frees are very frequent, and the extra works
adds up. We add about ~8x overhead to the baselien cost of allocation
just with our locking, not including the cost of actually sampling an
allocation. We can't get rid of this overhead just by reducing sampling
frequency.

There are a few rules to follow in order to use the GIL correctly for
locking:

1) The GIL is held when a C extension function is called, _except_
   possibly in the raw allocator, which we do not profile
2) The GIL may be released during C Python API calls. Even if it is
   released, though, it will be held again after the call
3) Thus, the GIL creates critical sections only between C Python API
   calls, and the beginning and end of C extension functions. Modifications
   to shared state across those points are not atomic.
4) If we take a lock of our own in a C extension code (i.e. a
   pthread_mutex), and the extension code releases the GIL, then the
   program will deadlock due to lock order inversion. We can only safely
   take locks in C extension when the GIL is released.

The crashes that #11460 addresed were due to breaking the first three
rules. In particular, we could race on accessing the shared scratch
buffer used when collecting tracebacks, which lead to double-frees.
See #13185 for more details.

Our mitigation involved using C locks around any access to the shared
profiler state. We nearly broke rule 4 in the process. However, we used
try-locks specifically out of a fear of introducing deadlocks. Try-locks
mean that we attempt to acquire the lock, but return a failure if the
lock is already held. This stopped deadlocks, but introduced bugs: For
example:

- If we failed to take the lock when trying to report allocation
  profile events, we'd raise an exception when it was in fact not
  reasonable for doing that to fail. See #12075.
- memalloc_heap_untrack, which removes tracked allocations, was guarded
  with a try-lock. If we couldn't acquire the lock, we would fail to
  remove a record for an allocation and effectively leak memory.
  See #13317
- We attempted to make our locking fork-safe. The first attempt was
  inefficient; we made it less inefficient but the fix only "worked"
  because of try-locks. See #11848

Try-locks hide concurrency problems and we shouldn't use them. Using our
own locks requires releasing the GIL before acquisition, and then
re-acquiring the GIL. That adds unnecessary overhead. We don't
inherently need to do any off-GIL work. So, we should try to just use
the GIL as long as it is available.

The basic refactor is actually pretty simple. In a nutshell, we
rearrange the memalloc_add_event and memalloc_heap_track functions so
that they make the sampling decision, then take a traceback, then insert
the traceback into the appropriate data structure. Collecting a
traceback can release the GIL, so we make sure that modifying the data
structure happens completely after the traceback is collected. We also
safeguard against the possibility that the profiler was stopped during
sampling, if the GIL was released. This requires a small rearrangement
of memalloc_stop to make sure that the sampling functions don't see
partially-freed profiler data structures.

For testing, I have mainly used the code from test_memealloc_data_race_regression.
I also added a debug mode, enabled by compiling with
MEMALLOC_TESTING_GIL_RELEASE, which releases the GIL at places where it
would be expected. For performance I examined the overhead of profiling
on a basic flask application.

Author: nsrip-dd

ddtrace/profiling/collector/_memalloc.c

@@ -31,9 +31,9 @@ static memalloc_context_t global_memalloc_ctx;
 /* Allocation tracker */
 typedef struct
 {
-    /* List of traceback */
+    /* List of tracebacks for sampled allocations */
     traceback_array_t allocs;
-    /* Total number of allocations */
+    /* Total number of observed allocations, sampled or not */
     uint64_t alloc_count;
 } alloc_tracker_t;
 
@@ -42,113 +42,81 @@ static PyObject* object_string = NULL;
 
 #define ALLOC_TRACKER_MAX_COUNT UINT64_MAX
 
-/* This lock protects access to global_alloc_tracker. The GIL is NOT sufficient
-   to protect our data structures from concurrent access. For one, the GIL is an
-   implementation detail and may go away in the future. Additionally, even if the
-   GIL is held on _entry_ to our C extension functions, making it safe to call
-   Python C API functions, the GIL can be released during Python C API calls if
-   we call back into interpreter code. This can happen if we allocate a Python
-   object (such as frame info), trigger garbage collection, and run arbitrary
-   destructors. When this happens, other threads can run python code, such as the
-   thread that aggregates and uploads the profile data and mutates the global
-   data structures. The GIL does not create critical sections for C extension
-   functions!
- */
-static memlock_t g_memalloc_lock;
-
 static alloc_tracker_t* global_alloc_tracker;
 
-// This is a multiplatform way to define an operation to happen at static initialization time
-static void
-memalloc_init(void);
-
-static void
-memalloc_prefork(void)
-{
-    // Lock the mutex prior to forking. This ensures that the memory profiler
-    // data structures will be in a consistent state in the child process.
-    // The rest of the memalloc calls do trylock so we don't run the risk
-    // of deadlocking if some other fork handler allocates
-    memlock_lock(&g_memalloc_lock);
-}
-
-static void
-memalloc_postfork_parent(void)
-{
-    memlock_unlock(&g_memalloc_lock);
-}
-
-static void
-memalloc_postfork_child(void)
-{
-    memlock_unlock(&g_memalloc_lock);
-}
-
-#ifdef _MSC_VER
-#pragma section(".CRT$XCU", read)
-__declspec(allocate(".CRT$XCU")) void (*memalloc_init_func)(void) = memalloc_init;
-
-#elif defined(__GNUC__) || defined(__clang__)
-__attribute__((constructor))
-#else
-#error Unsupported compiler
-#endif
-static void
-memalloc_init()
-{
-    memlock_init(&g_memalloc_lock);
-#ifndef _WIN32
-    pthread_atfork(memalloc_prefork, memalloc_postfork_parent, memalloc_postfork_child);
-#endif
-}
-
 static void
 memalloc_add_event(memalloc_context_t* ctx, void* ptr, size_t size)
 {
-    uint64_t alloc_count = atomic_add_clamped(&global_alloc_tracker->alloc_count, 1, ALLOC_TRACKER_MAX_COUNT);
+    /* Safety check: is profiling still enabled? */
+    if (!global_alloc_tracker)
+        return;
 
     /* Return if we've reached the maximum number of allocations */
-    if (alloc_count == 0)
+    if (global_alloc_tracker->alloc_count == ALLOC_TRACKER_MAX_COUNT)
         return;
 
-    // Return if we can't take the guard
-    if (!memalloc_take_guard()) {
-        return;
-    }
+    /* Incrementing the allocation count before taking the guard gives
+     * an accurate count, but means we'll overweight the allocations we
+     * do sample relative to their actual size since our count necessarily
+     * includes allocations we won't sample. We could consider moving this
+     * below the guard and instead adding a placeholder "self allocation"
+     * entry to represent the allocations we intentionally don't sample.
+     */
+    uint64_t alloc_count = global_alloc_tracker->alloc_count++;
 
-    // In this implementation, the `global_alloc_tracker` isn't intrinsically protected.  Before we read or modify,
-    // take the lock.  The count of allocations is already forward-attributed elsewhere, so if we can't take the lock
-    // there's nothing to do.
-    if (!memlock_trylock(&g_memalloc_lock)) {
+    /* Return if we can't take the guard, since we're not going to be able to sample */
+    if (!memalloc_take_guard()) {
         return;
     }
 
     /* Determine if we can capture or if we need to sample */
-    if (global_alloc_tracker->allocs.count < ctx->max_events) {
+    bool should_sample = false;
+    if (alloc_count < ctx->max_events) {
         /* Buffer is not full, fill it */
-        traceback_t* tb = memalloc_get_traceback(ctx->max_nframe, ptr, size, ctx->domain);
-        if (tb) {
-            traceback_array_append(&global_alloc_tracker->allocs, tb);
-        }
+        should_sample = true;
     } else {
         /* Sampling mode using a reservoir sampling algorithm: replace a random
-         * traceback with this one */
+         * traceback with this one
+         * NB: this just decides whether we sample. See comment below;
+         * we will probably have to recompute the index to replace */
         uint64_t r = random_range(alloc_count);
+        should_sample = r < ctx->max_events;
+    }
 
-        // In addition to event size, need to check that the tab is in a good state
-        if (r < ctx->max_events && global_alloc_tracker->allocs.tab != NULL) {
-            /* Replace a random traceback with this one */
-            traceback_t* tb = memalloc_get_traceback(ctx->max_nframe, ptr, size, ctx->domain);
+    if (!should_sample) {
+        goto done;
+    }
 
-            // Need to check not only that the tb returned
-            if (tb) {
-                traceback_free(global_alloc_tracker->allocs.tab[r]);
-                global_alloc_tracker->allocs.tab[r] = tb;
-            }
-        }
+    traceback_t* tb = memalloc_get_traceback(ctx->max_nframe, ptr, size, ctx->domain);
+    if (!tb) {
+        goto done;
     }
 
-    memlock_unlock(&g_memalloc_lock);
+    if (!global_alloc_tracker) {
+        /* If getting a traceback lead to a GIL release, it is possible
+         * that the profiler was stopped and the traceback array was freed.
+         * Don't add the sample.
+         */
+        traceback_free(tb);
+        goto done;
+    }
+
+    /* The GIL may have been released while collecting a traceback above.
+     * If so, the aggregation/uploading thread may have rotated the tracker.
+     * We need to recompute the index where the sample will go.
+     */
+    if (global_alloc_tracker->allocs.count < ctx->max_events) {
+        traceback_array_append(&global_alloc_tracker->allocs, tb);
+    } else {
+        uint64_t r = random_range(ctx->max_events);
+        traceback_t* old = global_alloc_tracker->allocs.tab[r];
+        global_alloc_tracker->allocs.tab[r] = tb;
+        /* Free the old traceback only after modifying the table, in case freeing
+         * releases the GIL */
+        traceback_free(old);
+    }
+
+done:
     memalloc_yield_guard();
 }
 
@@ -301,11 +269,14 @@ memalloc_start(PyObject* Py_UNUSED(module), PyObject* args)
 
     global_memalloc_ctx.domain = PYMEM_DOMAIN_OBJ;
 
-    if (memlock_trylock(&g_memalloc_lock)) {
-        global_alloc_tracker = alloc_tracker_new();
-        memlock_unlock(&g_memalloc_lock);
+    alloc_tracker_t* tracker = alloc_tracker_new();
+    if (!tracker) {
+        PyErr_SetString(PyExc_RuntimeError, "failed to allocate profiler state");
+        return NULL;
     }
 
+    global_alloc_tracker = tracker;
+
     PyMem_GetAllocator(PYMEM_DOMAIN_OBJ, &global_memalloc_ctx.pymem_allocator_obj);
     PyMem_SetAllocator(PYMEM_DOMAIN_OBJ, &alloc);
 
@@ -327,16 +298,24 @@ memalloc_stop(PyObject* Py_UNUSED(module), PyObject* Py_UNUSED(args))
         return NULL;
     }
 
+    /* First, uninstall our wrappers. There may still be calls to our wrapper in progress,
+     * if they happened to release the GIL.
+     * NB: We're assuming here that this is not called concurrently with iter_events
+     * or memalloc_heap. The higher-level collector deals with this. */
     PyMem_SetAllocator(PYMEM_DOMAIN_OBJ, &global_memalloc_ctx.pymem_allocator_obj);
-    memalloc_tb_deinit();
-    if (memlock_trylock(&g_memalloc_lock)) {
-        alloc_tracker_free(global_alloc_tracker);
-        global_alloc_tracker = NULL;
-        memlock_unlock(&g_memalloc_lock);
-    }
+
+    alloc_tracker_t* tracker = global_alloc_tracker;
+    /* Setting this to NULL indicates that in-progress sampling shouldn't add a sample */
+    global_alloc_tracker = NULL;
+    /* Now free the tracker, in case GIL release while deallocating tracebacks lets
+     * in-process sampling continue and observe a partially de-initialized tracker */
+    alloc_tracker_free(tracker);
 
     memalloc_heap_tracker_deinit();
 
+    /* Finally, we know in-progress sampling won't use the buffer pool, so clear it out */
+    memalloc_tb_deinit();
+
     Py_RETURN_NONE;
 }
 
@@ -396,10 +375,8 @@ iterevents_new(PyTypeObject* type, PyObject* Py_UNUSED(args), PyObject* Py_UNUSE
         return NULL;
     }
 
-    memlock_lock(&g_memalloc_lock);
     iestate->alloc_tracker = global_alloc_tracker;
     global_alloc_tracker = tracker;
-    memlock_unlock(&g_memalloc_lock);
 
     iestate->seq_index = 0;
 
@@ -414,11 +391,8 @@ iterevents_new(PyTypeObject* type, PyObject* Py_UNUSED(args), PyObject* Py_UNUSE
 static void
 iterevents_dealloc(IterEventsState* iestate)
 {
-    if (memlock_trylock(&g_memalloc_lock)) {
-        alloc_tracker_free(iestate->alloc_tracker);
-        Py_TYPE(iestate)->tp_free(iestate);
-        memlock_unlock(&g_memalloc_lock);
-    }
+    alloc_tracker_free(iestate->alloc_tracker);
+    Py_TYPE(iestate)->tp_free(iestate);
 }
 
 static PyObject*

ddtrace/profiling/collector/_memalloc_debug.h

@@ -0,0 +1,15 @@
+#ifndef _DDTRACE_MEMALLOC_DEBUG_H
+#define _DDTRACE_MEMALLOC_DEBUG_H
+
+/* Release the GIL. For debugging when GIL release allows memory profiling functions
+ * to interleave from different threads. Call near C Python API calls. */
+static inline void
+memalloc_debug_gil_release(void)
+{
+#ifdef MEMALLOC_TESTING_GIL_RELEASE
+    Py_BEGIN_ALLOW_THREADS;
+    Py_END_ALLOW_THREADS;
+#endif
+}
+
+#endif