Crash fix in case of forks: remove usage of pthread_once

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 pthread_once_t _key_once; // ensures we call key creation a single time
-  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

@@ -3,7 +3,7 @@
 // developed at Datadog (https://www.datadoghq.com/). Copyright 2021-Present
 // Datadog, Inc.
 
-#include "allocation_tracker.hpp"
+#include "lib/allocation_tracker.hpp"
 
 #include "allocation_event.hpp"
 #include "ddprof_perf_event.hpp"
@@ -23,16 +23,14 @@
 #include <chrono>
 #include <cstdint>
 #include <cstdlib>
+#include <thread>
 #include <unistd.h>
 
 namespace ddprof {
 
-// Static declarations
-pthread_once_t AllocationTracker::_key_once = PTHREAD_ONCE_INIT;
-
-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,
@@ -50,15 +48,28 @@ DDPROF_NOINLINE auto sleep_and_retry_reserve(MPSCRingBufferWriter &writer,
   }
   return Buffer{};
 }
+
+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();
+}
 } // namespace
 
 TrackerThreadLocalState *AllocationTracker::get_tl_state() {
   // In shared libraries, TLS access requires a call to tls_get_addr,
   // 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
-  pthread_once(&_key_once, make_key);
-  auto *tl_state = static_cast<TrackerThreadLocalState *>(
-      pthread_getspecific(_tl_state_key));
+
+  // The pthread key is initialized during allocation_tracking_init() and
+  // maintained by the fork handler, so we can directly use it here (hot path)
+  const 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(key));
   return tl_state;
 }
 
@@ -91,15 +102,43 @@ 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)
+  const 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);
+    }
+  }
+}
+
+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/CMakeLists.txt

@@ -393,6 +393,38 @@ add_unit_test(address_bitset-ut address_bitset-ut.cc ../src/lib/address_bitset.c
 
 add_unit_test(lib_logger-ut ./lib_logger-ut.cc)
 
+# Manual test for allocation_tracker fork behavior (no gtest due to TLS issues)
+add_exe(
+  allocation_tracker_fork_test
+  allocation_tracker_fork_test.c
+  allocation_tracker_fork_test_wrapper.cc
+  ../src/lib/allocation_tracker.cc
+  ../src/lib/address_bitset.cc
+  ../src/logger.cc
+  ../src/lib/pthread_fixes.cc
+  ../src/lib/savecontext.cc
+  ../src/lib/saveregisters.cc
+  ../src/procutils.cc
+  ../src/ratelimiter.cc
+  ../src/ringbuffer_utils.cc
+  ../src/sys_utils.cc
+  ../src/tsc_clock.cc
+  ../src/perf_clock.cc
+  ../src/perf.cc
+  ../src/ddres_list.cc
+  ../src/perf_ringbuffer.cc
+  ../src/pevent_lib.cc
+  ../src/user_override.cc
+  LIBRARIES llvm-demangle)
+target_include_directories(
+  allocation_tracker_fork_test PRIVATE ${CMAKE_SOURCE_DIR}/include ${CMAKE_SOURCE_DIR}/include/lib
+                                       ${CMAKE_SOURCE_DIR}/src)
+
+add_test(
+  NAME allocation_tracker_fork_test
+  COMMAND allocation_tracker_fork_test
+  WORKING_DIRECTORY ${CMAKE_BINARY_DIR})
+
 add_unit_test(
   create_elf-ut
   create_elf-ut.cc

test/allocation_tracker_fork_test.c

@@ -0,0 +1,144 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0. This product includes software
+// developed at Datadog (https://www.datadoghq.com/). Copyright 2021-Present
+// Datadog, Inc.
+
+// Simple test to verify pthread key initialization works across fork()
+// This tests the atomic-based pthread_once replacement
+
+#include <pthread.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <sys/types.h>
+#include <sys/wait.h>
+#include <unistd.h>
+
+// Forward declare the C++ functions we need
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+void *AllocationTracker_get_tl_state(void);
+void *AllocationTracker_init_tl_state(void);
+void AllocationTracker_ensure_key_initialized(void);
+
+#ifdef __cplusplus
+}
+#endif
+
+// Thread function to test TLS in new thread
+static void *thread_test_func(void *arg) {
+  int *result = (int *)arg;
+
+  void *thread_state = AllocationTracker_get_tl_state();
+
+  if (thread_state != NULL) {
+    fprintf(stderr, "FAIL: New thread expected NULL TLS, got %p\n",
+            thread_state);
+    *result = 1;
+    return NULL;
+  }
+
+  *result = 0;
+  return NULL;
+}
+
+int main(void) {
+  // 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");
+    return 1;
+  }
+
+  void *retrieved = AllocationTracker_get_tl_state();
+  if (retrieved != parent_state) {
+    fprintf(stderr, "FAIL: Parent get_tl_state() mismatch\n");
+    return 1;
+  }
+
+  // Test fork - this is where the old pthread_once bug would crash
+  fflush(stdout);
+  fflush(stderr);
+
+  pid_t pid = fork();
+  if (pid == -1) {
+    perror("FAIL: fork");
+    return 1;
+  }
+
+  if (pid == 0) {
+    // Child process - test pthread key works after fork
+    void *child_state = AllocationTracker_init_tl_state();
+    if (child_state == NULL) {
+      fprintf(stderr, "FAIL: Child init_tl_state() returned NULL\n");
+      _exit(1);
+    }
+
+    void *child_retrieved = AllocationTracker_get_tl_state();
+    if (child_retrieved != child_state) {
+      fprintf(stderr, "FAIL: Child pthread key broken after fork\n");
+      _exit(2);
+    }
+
+    // Test that a new thread starts with NULL TLS
+    pthread_t thread;
+    int thread_result = -1;
+    if (pthread_create(&thread, NULL, thread_test_func, &thread_result) != 0) {
+      fprintf(stderr, "FAIL: pthread_create failed\n");
+      _exit(3);
+    }
+
+    pthread_join(thread, NULL);
+    if (thread_result != 0) {
+      fprintf(stderr, "FAIL: Thread test failed - new thread TLS not NULL\n");
+      _exit(4);
+    }
+
+    _exit(0);
+  }
+
+  // Parent waits for child
+  int status;
+  waitpid(pid, &status, 0);
+
+  if (!WIFEXITED(status)) {
+    if (WIFSIGNALED(status)) {
+      fprintf(stderr,
+              "FAIL: Child crashed with signal %d (pthread_once ABI bug!)\n",
+              WTERMSIG(status));
+    } else {
+      fprintf(stderr, "FAIL: Child did not exit normally\n");
+    }
+    return 1;
+  }
+
+  int exit_code = WEXITSTATUS(status);
+  if (exit_code != 0) {
+    const char *reason = "unknown";
+    if (exit_code == 1)
+      reason = "init_tl_state failed";
+    else if (exit_code == 2)
+      reason = "pthread key broken after fork";
+    else if (exit_code == 3)
+      reason = "pthread_create failed";
+    else if (exit_code == 4)
+      reason = "new thread TLS not NULL";
+    fprintf(stderr, "FAIL: Child exited with code %d (%s)\n", exit_code,
+            reason);
+    return 1;
+  }
+
+  // Verify parent still works after fork
+  void *parent_after = AllocationTracker_get_tl_state();
+  if (parent_after != parent_state) {
+    fprintf(stderr, "FAIL: Parent TLS corrupted after fork\n");
+    return 1;
+  }
+
+  // All tests passed - silent success for CI
+  return 0;
+}

test/allocation_tracker_fork_test_wrapper.cc

@@ -0,0 +1,23 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0. This product includes software
+// developed at Datadog (https://www.datadoghq.com/). Copyright 2021-Present
+// Datadog, Inc.
+
+// C++ wrapper to expose AllocationTracker functions to C test
+
+#include "lib/allocation_tracker.hpp"
+
+extern "C" {
+
+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();
+}
+}