An alternative TLS implementation to fix a crash on musl/aarch64

ddprof-lib/src/main/cpp/livenessTracker.cpp

@@ -19,6 +19,7 @@
 #include "os.h"
 #include "profiler.h"
 #include "thread.h"
+#include "threadLocal.h"
 #include "tsc.h"
 #include <jni.h>
 #include <string.h>
@@ -276,6 +277,29 @@ Error LivenessTracker::initialize(Arguments &args) {
  return _stored_error = Error::OK;
 }
 
+static void* create_mt19937() {
+ // std::mt19937 itself is noexcept, but std::random_device and `new` may throw.
+ // If that happens we let the failure terminate the process (same outcome as
+ // failing thread_local initialization previously).
+ return static_cast<void*>(new std::mt19937(std::random_device{}()));
+}
+
+static void* create_uniform_real_distribution() {
+ // std::uniform_real_distribution<> construction is noexcept, but `new` may throw.
+ // If allocation fails the process is likely to abort anyway.
+ return static_cast<void*>(new std::uniform_real_distribution<>(0, 1.0));
+}
+
+static void free_mt19937(void* p) {
+ std::mt19937* mt = static_cast<std::mt19937*>(p);
+ delete mt;
+}
+
+static void free_uniform_real_distribution(void* p) {
+ std::uniform_real_distribution<>* urd = static_cast<std::uniform_real_distribution<>*>(p);
+ delete urd;
+}
+
 void LivenessTracker::track(JNIEnv *env, AllocEvent &event, jint tid,
  jobject object, u64 call_trace_id) {
  if (!_enabled) {
@@ -287,13 +311,17 @@ void LivenessTracker::track(JNIEnv *env, AllocEvent &event, jint tid,
  return;
  }
 
- static thread_local std::mt19937gen(std::random_device{}());
- static thread_local std::uniform_real_distribution<> dis(0, 1.0);
- static thread_local double skipped = 0;
+ static ThreadLocal<std::mt19937*, create_mt19937, free_mt19937> gen;
+ static ThreadLocal<std::uniform_real_distribution<>*, create_uniform_real_distribution, free_uniform_real_distribution> dis;
+ static ThreadLocal<double> skipped;
 
- if (_subsample_ratio < 1.0 && dis(gen) > _subsample_ratio) {
- skipped += static_cast<double>(event._weight) * event._size;
- return;
+ if (_subsample_ratio < 1.0) {
+ std::mt19937* genp = gen.get();
+ std::uniform_real_distribution<>* disp = dis.get();
+ if (disp->operator()(*genp) > _subsample_ratio) {
+ skipped.set(skipped.get() + static_cast<double>(event._weight) * event._size);
+ return;
+ }
  }
 
  jweak ref = env->NewWeakGlobalRef(object);
@@ -322,7 +350,8 @@ void LivenessTracker::track(JNIEnv *env, AllocEvent &event, jint tid,
  _table[idx].time = TSC::ticks();
  _table[idx].ref = ref;
  _table[idx].alloc = event;
- _table[idx].skipped = skipped;
+ _table[idx].skipped = skipped.get();
+ skipped.set(0);
  _table[idx].age = 0;
  _table[idx].call_trace_id = call_trace_id;
  _table[idx].ctx = ContextApi::snapshot();
@@ -375,8 +404,8 @@ void LivenessTracker::track(JNIEnv *env, AllocEvent &event, jint tid,
  } else {
  env->DeleteWeakGlobalRef(ref);
  }
+ skipped.set(0); // reset the subsampling skipped bytes
  }
- skipped = 0; // reset the subsampling skipped bytes
 }
 
 void JNICALL LivenessTracker::GarbageCollectionFinish(jvmtiEnv *jvmti_env) {

ddprof-lib/src/main/cpp/threadLocal.h

@@ -0,0 +1,168 @@
+/*
+ * Copyright 2026, Datadog, Inc.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#ifndef _THREADLOCAL_H
+#define _THREADLOCAL_H
+
+#include <cassert>
+#include <cstring>
+#include <pthread.h>
+#include "arch.h"
+
+/**
+ * This file implements an alternative to C/C++ thread local.
+ * Due to some restrictions of the language implementations, especially, on musl/aarch64,
+ * they cannot be safely used in profiler.
+ * 
+ * How to use?
+ * A ThreadLocal should be declared as a static variable, e.g.
+ * 
+ * static void* create_my_object() {
+ * return new MyObject();
+ * }
+ * 
+ * static void delete_my_object(void* p) {
+ * MyObject* obj = (MyObject*)p;
+ * delete obj; 
+ * }
+ * 
+ * static ThreadLocal<MyObject*, create_my_object, delete_my_object> var;
+ * MyObject* value = var.get();
+ * ...
+ * var.clear();
+ * ... 
+ * MyObject* new_value = new MyObject();
+ * var.set(new_value);
+ * ...
+ * var.clear();
+ * 
+ */
+
+// The function to create value if it does not exist
+typedef void* (*CREATE_FUNC)(void);
+// Cleanup the value when deleting the key
+typedef void (*CLEAN_FUNC)(void*);
+template <typename T, CREATE_FUNC C = nullptr, CLEAN_FUNC F = nullptr>
+class ThreadLocal {
+protected:
+ pthread_key_t _key;
+ bool _key_valid;
+
+public:
+ ThreadLocal(const ThreadLocal&) = delete;
+ ThreadLocal& operator=(const ThreadLocal&) = delete;
+
+ ThreadLocal() {
+ static_assert(sizeof(T) == sizeof(void*),
+ "ThreadLocal<T> requires sizeof(T)==sizeof(void*); use a pointer type or add a specialization");
+ _key_valid = pthread_key_create(&_key, F) == 0;
+ // What to do if we can not create a key?
+ // We probably want to shutdown profiler gracefully, instead of
+ // aborting user application - We will need this mechanism globally,
+ // defer to a separate task.
+ assert(_key_valid);
+ }
+
+ ~ThreadLocal() {
+ if(_key_valid) {
+ pthread_key_delete(_key);
+ } else {
+ assert(false && "Invalid pthread key");
+ }
+ }
+
+ /**
+ * set(nullptr) will result in the value being recreated when get() is called
+ * when CREATE_FUNC is not nullptr.
+ * Note: caller is responsible to free old value, which mirrors thread_local
+ */
+ void set(T value) {
+ assert(_key_valid && "Invalid pthread key");
+ int err = pthread_setspecific(_key, reinterpret_cast<const void*>(value));
+ assert(err == 0);
+ }
+
+ T get() {
+ assert(_key_valid && "Invalid pthread key");
+ void* p = pthread_getspecific(_key);
+ if (p == nullptr && C != nullptr) {
+ p = C();
+ set((T)p);
+ }
+ return reinterpret_cast<T>(p);
+ }
+
+ // Clear the value
+ void clear() {
+ assert(_key_valid && "Invalid pthread key");
+ void* p = pthread_getspecific(_key);
+ if (p == nullptr) return;
+ int err = pthread_setspecific(_key, nullptr);
+ // Safety: if reset the value failed, get() can see staled value if
+ // it is freed.
+ if (err == 0 && F != nullptr) {
+ F(p);
+ }
+ }
+};
+
+template <>
+class ThreadLocal<double> {
+protected:
+ pthread_key_t _key;
+ bool _key_valid;
+
+public:
+ ThreadLocal(const ThreadLocal&) = delete;
+ ThreadLocal& operator=(const ThreadLocal&) = delete;
+
+ ThreadLocal() {
+ // Only support 64-bit platforms, double and void* are the same size
+ static_assert(sizeof(void*) == 8);
+ static_assert(sizeof(double) == 8);
+ _key_valid = pthread_key_create(&_key, nullptr) == 0;
+ // What to do if we can not create a key?
+ assert(_key_valid && "Invalid pthread key");
+ }
+
+ ~ThreadLocal() {
+ if(_key_valid) {
+ pthread_key_delete(_key);
+ } else {
+ assert(_key_valid && "Invalid pthread key");
+ }
+ }
+
+ // double <--> u64 cast, preserve bit format
+ // Can use std::bit_cast after upgrade C++ version to 20
+ void set(double value) {
+ assert(_key_valid && "Invalid pthread key");
+ u64 val;
+ memcpy(&val, &value, sizeof(value));
+ int err = pthread_setspecific(_key, reinterpret_cast<const void*>(val));
+ assert(err == 0);
+ }
+
+ double get() {
+ assert(_key_valid && "Invalid pthread key");
+ void* p = pthread_getspecific(_key);
+ if (p == nullptr) {
+ return 0.0;
+ }
+
+ u64 val = reinterpret_cast<u64>(p);
+ double value;
+ memcpy(&value, &val, sizeof(val));
+ return value;
+ }
+
+ void clear() {
+ assert(_key_valid && "Invalid pthread key");
+ int err = pthread_setspecific(_key, nullptr);
+ assert(err == 0);
+ }
+};
+
+#endif // _THREADLOCAL_H