Address review feedback: fix races, remove dead code, clean up tests

Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>

Author: jbachorik

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

@@ -27,56 +27,23 @@
 char* ContextApi::_attribute_keys[MAX_ATTRIBUTE_KEYS] = {};
 int ContextApi::_attribute_key_count = 0;
 
-void ContextApi::initialize() {
-    // No per-session state to initialize — OTEL TLS is always used.
-}
-
 void ContextApi::shutdown() {
     freeAttributeKeys();
 }
 
 /**
  * Initialize OTel TLS for the current thread on first use.
- * Zero-inits the embedded OtelThreadContextRecord and triggers TLS init.
+ * Triggers the first access to custom_labels_current_set_v2 (TLS init).
  * Must be called with signals blocked to prevent musl TLS deadlock.
+ * The OtelThreadContextRecord is already zero-initialized by the ProfiledThread ctor.
  */
-static void initializeOtelTls(ProfiledThread* thrd) {
+void ContextApi::initializeOtelTls(ProfiledThread* thrd) {
     SignalBlocker blocker;
-
-    OtelThreadContextRecord* record = thrd->getOtelContextRecord();
-    record->valid = 0;
-    record->_reserved = 0;
-    record->attrs_data_size = 0;
-
     // First access to custom_labels_current_set_v2 triggers TLS init
     custom_labels_current_set_v2 = nullptr;
-
     thrd->markOtelContextInitialized();
 }
 
-void ContextApi::setFull(u64 local_root_span_id, u64 span_id, u64 trace_id_high, u64 trace_id_low) {
-    ProfiledThread* thrd = ProfiledThread::current();
-    if (thrd == nullptr) return;
-
-    if (!thrd->isOtelContextInitialized()) {
-        initializeOtelTls(thrd);
-    }
-
-    // All-zero IDs = context detachment
-    if (trace_id_high == 0 && trace_id_low == 0 && span_id == 0) {
-        OtelContexts::clear(thrd->getOtelContextRecord());
-        thrd->clearOtelSidecar();
-        return;
-    }
-
-    // Update sidecar BEFORE publishing the OTEP record. OtelContexts::set()
-    // flips valid=1 at the end; a signal handler that sees valid=1 must also
-    // see the updated local_root_span_id.
-    thrd->setOtelLocalRootSpanId(local_root_span_id);
-
-    OtelContexts::set(thrd->getOtelContextRecord(), trace_id_high, trace_id_low, span_id, local_root_span_id);
-}
-
 bool ContextApi::get(u64& span_id, u64& root_span_id) {
     ProfiledThread* thrd = ProfiledThread::currentSignalSafe();
     if (thrd == nullptr || !thrd->isOtelContextInitialized()) {
@@ -100,41 +67,13 @@ Context ContextApi::snapshot() {
     return thrd->snapshotContext(numAttrs);
 }
 
-bool ContextApi::setAttribute(uint8_t key_index, const char* value, uint8_t value_len) {
-    if (key_index >= DD_TAGS_CAPACITY) return false;
-
-    ProfiledThread* thrd = ProfiledThread::current();
-    if (thrd == nullptr) return false;
-
-    if (!thrd->isOtelContextInitialized()) {
-        initializeOtelTls(thrd);
-    }
-
-    // Pre-register the value string in the Dictionary and cache the encoding
-    // in the sidecar for O(1) signal-handler reads (no hash lookup needed).
-    u32 encoding = Profiler::instance()->contextValueMap()->bounded_lookup(
-        value, value_len, 1 << 16);
-    if (encoding == INT_MAX) {
-        thrd->setOtelTagEncoding(key_index, 0);
-        return false;
-    }
-    thrd->setOtelTagEncoding(key_index, encoding);
-
-    // Offset by 1 in the OTEP record: index 0 is reserved for LOCAL_ROOT_SPAN_ATTR_INDEX
-    uint8_t otep_key_index = key_index + 1;
-    return OtelContexts::setAttribute(thrd->getOtelContextRecord(), otep_key_index, value, value_len);
-}
-
 void ContextApi::freeAttributeKeys() {
     int count = _attribute_key_count;
-    char* old_keys[MAX_ATTRIBUTE_KEYS];
-    for (int i = 0; i < count; i++) {
-        old_keys[i] = _attribute_keys[i];
-        _attribute_keys[i] = nullptr;
-    }
     _attribute_key_count = 0;
     for (int i = 0; i < count; i++) {
-        free(old_keys[i]);
+        char* key = _attribute_keys[i];
+        _attribute_keys[i] = nullptr;
+        free(key);
     }
 }
 

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

@@ -21,6 +21,8 @@
 #include "context.h"
 #include <cstdint>
 
+class ProfiledThread;
+
 /**
  * Unified context API for trace/span context storage.
  *
@@ -30,30 +32,17 @@
  */
 class ContextApi {
 public:
-    /**
-     * Initialize context storage.
-     * Must be called once during profiler startup.
-     */
-    static void initialize();
-
     /**
      * Shutdown context storage.
      * Releases resources allocated during initialization.
      */
     static void shutdown();
 
     /**
-     * Write full OTEL context with 128-bit trace ID and local root span ID.
-     *
-     * Writes trace_id + span_id + local_root_span_id to the OTEP record
-     * in a single detach/attach cycle.
-     *
-     * @param local_root_span_id Local root span ID (for endpoint correlation)
-     * @param span_id The span ID
-     * @param trace_id_high Upper 64 bits of 128-bit trace ID
-     * @param trace_id_low Lower 64 bits of 128-bit trace ID
+     * Initialize OTel TLS for the given thread on first use.
+     * Must be called with signals blocked (SignalBlocker).
      */
-    static void setFull(u64 local_root_span_id, u64 span_id, u64 trace_id_high, u64 trace_id_low);
+    static void initializeOtelTls(ProfiledThread* thrd);
 
     /**
      * Read context for the current thread.
@@ -78,19 +67,6 @@ class ContextApi {
      */
     static Context snapshot();
 
-    /**
-     * Set a custom attribute on the current thread's context.
-     *
-     * Encodes into attrs_data of the OTEP record and caches the
-     * encoding in the ProfiledThread sidecar for O(1) signal-handler reads.
-     *
-     * @param key_index Index into the registered attribute key map
-     * @param value UTF-8 string value
-     * @param value_len Length of value in bytes (max 255)
-     * @return true if the attribute was set successfully
-     */
-    static bool setAttribute(uint8_t key_index, const char* value, uint8_t value_len);
-
     /**
      * Register attribute key names and publish them in the process context.
      * Must be called before setAttribute().

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

@@ -549,14 +549,7 @@ Java_com_datadoghq_profiler_JavaProfiler_initializeOtelTls0(JNIEnv* env, jclass
   if (thrd == nullptr) return nullptr;
 
   if (!thrd->isOtelContextInitialized()) {
-    // Block profiling signals during first TLS access
-    SignalBlocker blocker;
-    OtelThreadContextRecord* record = thrd->getOtelContextRecord();
-    record->valid = 0;
-    record->_reserved = 0;
-    record->attrs_data_size = 0;
-    custom_labels_current_set_v2 = nullptr;
-    thrd->markOtelContextInitialized();
+    ContextApi::initializeOtelTls(thrd);
   }
 
   OtelThreadContextRecord* record = thrd->getOtelContextRecord();

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

@@ -35,77 +35,6 @@ u64 OtelContexts::bytesToU64(const uint8_t* in) {
     return val;
 }
 
-void OtelContexts::set(OtelThreadContextRecord* record, u64 trace_id_high, u64 trace_id_low, u64 span_id, u64 local_root_span_id) {
-    if (record == nullptr) {
-        return;
-    }
-
-    // OTEP publication protocol:
-    // 1. Detach — external readers see no record during construction
-    __atomic_store_n(&custom_labels_current_set_v2, (OtelThreadContextRecord*)nullptr, __ATOMIC_RELEASE);
-
-    // 2. Invalidate — internal readers (signal handlers using cached pointer) see invalid record
-    __atomic_store_n(&record->valid, (uint8_t)0, __ATOMIC_RELEASE);
-
-    // 3. Populate record fields
-    u64ToBytes(trace_id_high, record->trace_id);
-    u64ToBytes(trace_id_low, record->trace_id + 8);
-    u64ToBytes(span_id, record->span_id);
-
-    // 4. Rebuild attrs_data: preserve user entries (key_index != 0), replace LRS entry.
-    //    LOCAL_ROOT_SPAN_ATTR_INDEX = 0 is reserved; user attrs use key_index >= 1.
-    //    Use a temp buffer because source and destination regions overlap.
-    uint16_t old_size = record->attrs_data_size;
-    uint8_t user_buf[OTEL_MAX_ATTRS_DATA_SIZE];
-    uint16_t user_size = 0;
-
-    uint16_t pos = 0;
-    while (pos + 2 <= old_size) {
-        uint8_t k = record->attrs_data[pos];
-        uint8_t len = record->attrs_data[pos + 1];
-        if (pos + 2 + len > old_size) break;
-        if (k != 0) {
-            memcpy(user_buf + user_size, record->attrs_data + pos, 2 + len);
-            user_size += 2 + len;
-        }
-        pos += 2 + len;
-    }
-
-    uint16_t new_size = 0;
-
-    // Write local_root_span_id as decimal string attribute at reserved index 0
-    if (local_root_span_id != 0) {
-        // Max digits for u64: 20 ("18446744073709551615"). Entry = 2 + digits.
-        uint8_t digits[20];
-        int ndigits = 0;
-        u64 v = local_root_span_id;
-        do {
-            digits[ndigits++] = '0' + (v % 10);
-            v /= 10;
-        } while (v > 0);
-        record->attrs_data[0] = 0;          // key_index = LOCAL_ROOT_SPAN_ATTR_INDEX
-        record->attrs_data[1] = ndigits;    // length
-        for (int i = 0; i < ndigits; i++) {
-            record->attrs_data[2 + i] = digits[ndigits - 1 - i]; // reverse into big-endian order
-        }
-        new_size = 2 + ndigits;
-    }
-
-    // Re-append preserved user entries
-    if (user_size > 0 && new_size + user_size <= OTEL_MAX_ATTRS_DATA_SIZE) {
-        memcpy(record->attrs_data + new_size, user_buf, user_size);
-        new_size += user_size;
-    }
-
-    record->attrs_data_size = new_size;
-
-    // 5. Mark record valid
-    __atomic_store_n(&record->valid, (uint8_t)1, __ATOMIC_RELEASE);
-
-    // 6. Attach — publish completed record
-    __atomic_store_n(&custom_labels_current_set_v2, record, __ATOMIC_RELEASE);
-}
-
 bool OtelContexts::getSpanId(OtelThreadContextRecord* record, u64& span_id) {
     if (record == nullptr) {
         return false;
@@ -117,67 +46,3 @@ bool OtelContexts::getSpanId(OtelThreadContextRecord* record, u64& span_id) {
     return true;
 }
 
-bool OtelContexts::setAttribute(OtelThreadContextRecord* record, uint8_t key_index, const char* value, uint8_t value_len) {
-    if (record == nullptr) {
-        return false;
-    }
-
-    // Entry size: 1 (key_index) + 1 (length) + value_len
-    uint16_t entry_size = 2 + value_len;
-    uint16_t current_size = record->attrs_data_size;
-
-    // Detach (external readers) and invalidate (signal handler readers)
-    __atomic_store_n(&custom_labels_current_set_v2, (OtelThreadContextRecord*)nullptr, __ATOMIC_RELEASE);
-    __atomic_store_n(&record->valid, (uint8_t)0, __ATOMIC_RELEASE);
-
-    // Scan for duplicate key and compact if found
-    uint16_t read_pos = 0;
-    uint16_t write_pos = 0;
-    bool found = false;
-    while (read_pos + 2 <= current_size) {
-        uint8_t k = record->attrs_data[read_pos];
-        uint8_t len = record->attrs_data[read_pos + 1];
-        if (read_pos + 2 + len > current_size) break;  // corrupt length guard
-        if (k == key_index) {
-            found = true;
-            read_pos += 2 + len;
-        } else {
-            if (found && write_pos < read_pos) {
-                memmove(record->attrs_data + write_pos, record->attrs_data + read_pos, 2 + len);
-            }
-            write_pos += 2 + len;
-            read_pos += 2 + len;
-        }
-    }
-    if (found) {
-        current_size = write_pos;
-    }
-
-    // Re-check fit after compaction
-    if (current_size + entry_size > OTEL_MAX_ATTRS_DATA_SIZE) {
-        record->attrs_data_size = current_size;
-        __atomic_store_n(&record->valid, (uint8_t)1, __ATOMIC_RELEASE);
-        __atomic_store_n(&custom_labels_current_set_v2, record, __ATOMIC_RELEASE);
-        return false;
-    }
-
-    // Append the new entry
-    record->attrs_data[current_size] = key_index;
-    record->attrs_data[current_size + 1] = value_len;
-    memcpy(record->attrs_data + current_size + 2, value, value_len);
-    record->attrs_data_size = current_size + entry_size;
-
-    // Re-publish
-    __atomic_store_n(&record->valid, (uint8_t)1, __ATOMIC_RELEASE);
-    __atomic_store_n(&custom_labels_current_set_v2, record, __ATOMIC_RELEASE);
-    return true;
-}
-
-void OtelContexts::clear(OtelThreadContextRecord* record) {
-    if (record != nullptr) {
-        __atomic_store_n(&record->valid, (uint8_t)0, __ATOMIC_RELEASE);
-        record->attrs_data_size = 0;
-    }
-    // OTEP context detachment: set TLS pointer to nullptr
-    __atomic_store_n(&custom_labels_current_set_v2, (OtelThreadContextRecord*)nullptr, __ATOMIC_RELEASE);
-}

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

@@ -73,38 +73,13 @@ DLLEXPORT extern thread_local OtelThreadContextRecord* custom_labels_current_set
  */
 class OtelContexts {
 public:
-    /**
-     * Write context for the current thread using OTEP publication protocol:
-     * 1. Set TLS pointer to nullptr (detach — external readers)
-     * 2. Set valid = 0 (invalidate — signal handler readers via cached pointer)
-     * 3. Populate record fields (trace_id, span_id, local_root_span_id as hex attr)
-     * 4. Set valid = 1
-     * 5. Set TLS pointer to record (attach)
-     *
-     * local_root_span_id is encoded as a 16-char hex string attribute at index 0.
-     * Existing user attributes (key_index >= 1) are preserved across calls.
-     */
-    static void set(OtelThreadContextRecord* record, u64 trace_id_high, u64 trace_id_low, u64 span_id, u64 local_root_span_id);
-
     /**
      * Read span_id from a record pointer (signal-safe).
-     * Skips trace_id decode for efficiency.
-     * Returns false if the record is null or invalid.
+     * Checks the valid flag with acquire semantics before reading.
+     * Returns false if the record is null or invalid (being mutated).
      */
     static bool getSpanId(OtelThreadContextRecord* record, u64& span_id);
 
-    /**
-     * Write a single attribute into the current thread's OTEP record.
-     * Appends (key_index, length, value) to attrs_data, re-publishes via detach/attach.
-     * Returns false if the attribute doesn't fit in the remaining space.
-     */
-    static bool setAttribute(OtelThreadContextRecord* record, uint8_t key_index, const char* value, uint8_t value_len);
-
-    /**
-     * Clear context for the current thread (set TLS pointer to nullptr).
-     */
-    static void clear(OtelThreadContextRecord* record);
-
     // Byte conversion helpers (big-endian, W3C trace context)
     static void u64ToBytes(u64 val, uint8_t* out);
     static u64 bytesToU64(const uint8_t* in);

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

@@ -1440,9 +1440,6 @@ Error Profiler::start(Arguments &args, bool reset) {
     _libs->updateBuildIds();
   }
 
-  // Initialize OTEL context storage
-  ContextApi::initialize();
-
   enableEngines();
 
   // Always enable library trap to catch wasmtime loading and patch its broken sigaction