Allocation profiling crash - Fork safety

- Use the key as synchronization
- Avoid init during a get of the thread local state
- Add a fork handler to ensure we have the key initialized

Author: r1viollet

include/lib/allocation_tracker.hpp

@@ -80,6 +80,11 @@ class AllocationTracker {
   // can return null (does not init)
   static TrackerThreadLocalState *get_tl_state();
 
+  // Initialize pthread key for TLS (idempotent, fork-safe)
+  static void ensure_key_initialized();
+  // Register pthread_atfork handler (called during init)
+  static void register_fork_handler();
+
 private:
   static constexpr unsigned k_ratio_max_elt_to_bitset_size = 16;
 
@@ -120,8 +125,6 @@ class AllocationTracker {
 
   static void delete_tl_state(void *tl_state);
 
-  static void make_key();
-
   void track_allocation(uintptr_t addr, size_t size,
                         TrackerThreadLocalState &tl_state, bool is_large_alloc);
   void track_deallocation(uintptr_t addr, TrackerThreadLocalState &tl_state,
@@ -160,8 +163,8 @@ class AllocationTracker {
 
   // These can not be tied to the internal state of the instance.
   // The creation of the instance depends on this
-  static std::atomic<int> _key_init_state;
-  static pthread_key_t _tl_state_key;
+  static constexpr pthread_key_t kInvalidKey = static_cast<pthread_key_t>(-1);
+  static std::atomic<pthread_key_t> _tl_state_key;
 
   static AllocationTracker *_instance;
 };

src/lib/allocation_tracker.cc

@@ -28,19 +28,9 @@
 
 namespace ddprof {
 
-namespace {
-// TLS key initialization states
-constexpr int kKeyNotInitialized = 0;
-constexpr int kKeyInitializing = 1;
-constexpr int kKeyInitialized = 2;
-} // namespace
-
-// Static declarations
-std::atomic<int> AllocationTracker::_key_init_state{kKeyNotInitialized};
-
-pthread_key_t AllocationTracker::_tl_state_key;
-
-AllocationTracker *AllocationTracker::_instance;
+AllocationTracker *AllocationTracker::_instance = nullptr;
+std::atomic<pthread_key_t> AllocationTracker::_tl_state_key{
+    AllocationTracker::kInvalidKey};
 
 namespace {
 DDPROF_NOINLINE auto sleep_and_retry_reserve(MPSCRingBufferWriter &writer,
@@ -65,36 +55,15 @@ TrackerThreadLocalState *AllocationTracker::get_tl_state() {
   // tls_get_addr can call into malloc, which can create a recursive loop
   // instead we call pthread APIs to control the creation of TLS objects
 
-  // Thread-safe initialization using atomics (avoids pthread_once ABI issues)
-
-  // Fast path: relaxed load is sufficient since we only care if it's
-  // initialized Once initialized, this value never changes, so no
-  // synchronization needed
-  if (_key_init_state.load(std::memory_order_relaxed) != kKeyInitialized) {
-    // Slow path: need proper synchronization for initialization
-    int expected = kKeyNotInitialized;
-    if (_key_init_state.compare_exchange_strong(expected, kKeyInitializing,
-                                                std::memory_order_acq_rel)) {
-      // We won the race, do the initialization
-      make_key();
-      _key_init_state.store(kKeyInitialized, std::memory_order_release);
-    } else {
-      // Another thread is initializing or already done, wait until complete
-      constexpr int k_max_init_wait_attempts = 100;
-      int attempts = 0;
-      while (_key_init_state.load(std::memory_order_acquire) !=
-             kKeyInitialized) {
-        if (++attempts >= k_max_init_wait_attempts) {
-          return nullptr;
-        }
-        std::this_thread::yield();
-        std::this_thread::sleep_for(std::chrono::microseconds(1));
-      }
-    }
-  }
+  // The pthread key is initialized during allocation_tracking_init() and
+  // maintained by the fork handler, so we can directly use it here (hot path)
+  pthread_key_t key = _tl_state_key.load(std::memory_order_relaxed);
+
+  // In debug builds, verify our assumption
+  assert(key != kInvalidKey && "pthread key should be initialized before use");
 
-  auto *tl_state = static_cast<TrackerThreadLocalState *>(
-      pthread_getspecific(_tl_state_key));
+  auto *tl_state =
+      static_cast<TrackerThreadLocalState *>(pthread_getspecific(key));
   return tl_state;
 }
 
@@ -127,15 +96,49 @@ void AllocationTracker::delete_tl_state(void *tl_state) {
   delete static_cast<TrackerThreadLocalState *>(tl_state);
 }
 
-void AllocationTracker::make_key() {
-  // delete is called on all key objects
-  pthread_key_create(&_tl_state_key, delete_tl_state);
+void AllocationTracker::ensure_key_initialized() {
+  // Ensure pthread key is initialized (idempotent)
+  pthread_key_t key = _tl_state_key.load(std::memory_order_acquire);
+
+  if (key == kInvalidKey) {
+    pthread_key_t new_key;
+    if (pthread_key_create(&new_key, delete_tl_state) != 0) {
+      return; // Failed, will be retried later
+    }
+
+    pthread_key_t expected = kInvalidKey;
+    if (!_tl_state_key.compare_exchange_strong(expected, new_key,
+                                               std::memory_order_release)) {
+      // Another thread beat us, clean up our key
+      pthread_key_delete(new_key);
+    }
+  }
+}
+
+static void child_after_fork_handler() {
+  // After fork in child, verify the pthread key is still valid.
+  // The key itself survives fork, but we want to ensure it's properly set.
+  AllocationTracker::ensure_key_initialized();
+}
+
+void AllocationTracker::register_fork_handler() {
+  static bool registered = false;
+  if (!registered) {
+    pthread_atfork(nullptr, nullptr, child_after_fork_handler);
+    registered = true;
+  }
 }
 
 DDRes AllocationTracker::allocation_tracking_init(
     uint64_t allocation_profiling_rate, uint32_t flags,
     uint32_t stack_sample_size, const RingBufferInfo &ring_buffer,
     const IntervalTimerCheck &timer_check) {
+  // Register fork handler to ensure key is valid after fork
+  register_fork_handler();
+
+  // Ensure pthread key is initialized before we try to use it
+  ensure_key_initialized();
+
   TrackerThreadLocalState *tl_state = get_tl_state();
   if (!tl_state) {
     // This is the time at which the init_tl_state should not fail

test/allocation_tracker_fork_test.c

@@ -20,6 +20,7 @@ extern "C" {
 
 void *AllocationTracker_get_tl_state(void);
 void *AllocationTracker_init_tl_state(void);
+void AllocationTracker_ensure_key_initialized(void);
 
 #ifdef __cplusplus
 }
@@ -43,9 +44,10 @@ static void *thread_test_func(void *arg) {
 }
 
 int main(void) {
-  // Initialize pthread key and create TLS data in parent
-  void *state_before = AllocationTracker_get_tl_state();
+  // Initialize pthread key before using it
+  AllocationTracker_ensure_key_initialized();
 
+  // Initialize pthread key and create TLS data in parent
   void *parent_state = AllocationTracker_init_tl_state();
   if (parent_state == NULL) {
     fprintf(stderr, "FAIL: Parent init_tl_state() returned NULL\n");
@@ -70,8 +72,6 @@ int main(void) {
 
   if (pid == 0) {
     // Child process - test pthread key works after fork
-    void *child_inherited = AllocationTracker_get_tl_state();
-
     void *child_state = AllocationTracker_init_tl_state();
     if (child_state == NULL) {
       fprintf(stderr, "FAIL: Child init_tl_state() returned NULL\n");
@@ -94,6 +94,7 @@ int main(void) {
 
     pthread_join(thread, NULL);
     if (thread_result != 0) {
+      fprintf(stderr, "FAIL: Thread test failed - new thread TLS not NULL\n");
       _exit(4);
     }
 
@@ -121,11 +122,11 @@ int main(void) {
     if (exit_code == 1)
       reason = "init_tl_state failed";
     else if (exit_code == 2)
-      reason = "pthread key broken";
+      reason = "pthread key broken after fork";
     else if (exit_code == 3)
       reason = "pthread_create failed";
     else if (exit_code == 4)
-      reason = "thread test failed";
+      reason = "new thread TLS not NULL";
     fprintf(stderr, "FAIL: Child exited with code %d (%s)\n", exit_code,
             reason);
     return 1;

test/allocation_tracker_fork_test_wrapper.cc

@@ -13,6 +13,10 @@ void *AllocationTracker_get_tl_state(void) {
   return ddprof::AllocationTracker::get_tl_state();
 }
 
+void AllocationTracker_ensure_key_initialized(void) {
+  ddprof::AllocationTracker::ensure_key_initialized();
+}
+
 void *AllocationTracker_init_tl_state(void) {
   return ddprof::AllocationTracker::init_tl_state();
 }