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 (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 (void*)(new std::uniform_real_distribution<>(0, 1.0));
+}
+
+static void free_mt19937(void* p) {
+  std::mt19937* mt = (std::mt19937*)p;
+  delete mt;
+}
+
+static void free_uniform_real_distribution(void* p) {
+  std::uniform_real_distribution<>* urd = (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::mt19937 gen(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,164 @@
+/*
+ * 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);
+        }
+    }
+
+    /**
+     * 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);
+        int err = pthread_setspecific(_key, reinterpret_cast<const void*>(value));
+        assert(err == 0);
+    }
+
+    T get() {
+        assert(_key_valid);
+        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);
+        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);
+    }
+
+    ~ThreadLocal() {
+        if(_key_valid) {
+            pthread_key_delete(_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);
+        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);
+        void* p = pthread_getspecific(_key);
+        if (p == nullptr) {
+            return 0.0;
+        }
+
+        u64 val = (u64)p;
+        double value;
+        memcpy(&value, &val, sizeof(val));
+        return value;
+    }
+
+    void clear() {
+        assert(_key_valid);
+        int err = pthread_setspecific(_key, nullptr);
+        assert(err == 0);
+    }
+};
+
+#endif // _THREADLOCAL_H