diff --git a/include/lib/allocation_tracker.hpp b/include/lib/allocation_tracker.hpp index ef7c9705..17fc04a5 100644 --- a/include/lib/allocation_tracker.hpp +++ b/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(-1); + static std::atomic _tl_state_key; static AllocationTracker *_instance; }; diff --git a/src/lib/allocation_tracker.cc b/src/lib/allocation_tracker.cc index e7a2981d..625496ac 100644 --- a/src/lib/allocation_tracker.cc +++ b/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 #include #include +#include #include 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 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( - 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(pthread_getspecific(key)); return tl_state; } @@ -91,15 +102,43 @@ void AllocationTracker::delete_tl_state(void *tl_state) { delete static_cast(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 diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt index 01d4e625..fb26e7b7 100644 --- a/test/CMakeLists.txt +++ b/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 diff --git a/test/allocation_tracker_fork_test.c b/test/allocation_tracker_fork_test.c new file mode 100644 index 00000000..dfaf5902 --- /dev/null +++ b/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 +#include +#include +#include +#include +#include + +// 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; +} diff --git a/test/allocation_tracker_fork_test_wrapper.cc b/test/allocation_tracker_fork_test_wrapper.cc new file mode 100644 index 00000000..97d63019 --- /dev/null +++ b/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(); +} +}