Thread filter - add a table with linear probing

Remove usage of set to ensure we are signal safe.

Author: r1viollet

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

@@ -1063,8 +1063,8 @@ void Recording::writeThreads(Buffer *buf) {
   threads.insert(_tid);
 
   for (int i = 0; i < CONCURRENCY_LEVEL; ++i) {
-    // I can not use merge : cpp 17
-    threads.insert(_thread_ids[i][old_index].begin(), _thread_ids[i][old_index].end());
+    // Collect thread IDs from the fixed-size table into the main set
+    _thread_ids[i][old_index].collect(threads);
     _thread_ids[i][old_index].clear();
   }
 

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

@@ -35,6 +35,7 @@
 #include "mutex.h"
 #include "objectSampler.h"
 #include "threadFilter.h"
+#include "threadIdTable.h"
 #include "vmEntry.h"
 
 const u64 MAX_JLONG = 0x7fffffffffffffffULL;
@@ -128,9 +129,9 @@ class Recording {
   static char *_java_command;
 
   RecordingBuffer _buf[CONCURRENCY_LEVEL];
-  // we have several sets to avoid lock contention
-  // we have a second dimension to allow a switch in the active set
-  std::unordered_set<int> _thread_ids[CONCURRENCY_LEVEL][2];
+  // we have several tables to avoid lock contention
+  // we have a second dimension to allow a switch in the active table
+  ThreadIdTable _thread_ids[CONCURRENCY_LEVEL][2];
   std::atomic<int> _active_index{0};  // 0 or 1 globally
 
   int _fd;

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

@@ -125,8 +125,8 @@ Java_com_datadoghq_profiler_JavaProfiler_getSamples(JNIEnv *env,
 
 // some duplication between add and remove, though we want to avoid having an extra branch in the hot path
 extern "C" DLLEXPORT void JNICALL
-Java_com_datadoghq_profiler_JavaProfiler_filterThread_1add(JNIEnv *env,
-                                                       jobject unused) {
+Java_com_datadoghq_profiler_JavaProfiler_filterThreadAdd0(JNIEnv *env,
+                                                          jobject unused) {
   ProfiledThread *current = ProfiledThread::current();
   int tid = current->tid();
   if (unlikely(tid < 0)) {
@@ -141,8 +141,8 @@ Java_com_datadoghq_profiler_JavaProfiler_filterThread_1add(JNIEnv *env,
 }
 
 extern "C" DLLEXPORT void JNICALL
-Java_com_datadoghq_profiler_JavaProfiler_filterThread_1remove(JNIEnv *env,
-                                                       jobject unused) {
+Java_com_datadoghq_profiler_JavaProfiler_filterThreadRemove0(JNIEnv *env,
+                                                             jobject unused) {
   ProfiledThread *current = ProfiledThread::current();
   int tid = current->tid();
   if (unlikely(tid < 0)) {
@@ -156,6 +156,28 @@ Java_com_datadoghq_profiler_JavaProfiler_filterThread_1remove(JNIEnv *env,
   thread_filter->remove(slot_id);
 }
 
+// Backward compatibility for existing code
+extern "C" DLLEXPORT void JNICALL
+Java_com_datadoghq_profiler_JavaProfiler_filterThread0(JNIEnv *env,
+                                                       jobject unused,
+                                                       jboolean enable) {
+  ProfiledThread *current = ProfiledThread::current();
+  int tid = current->tid();
+  if (unlikely(tid < 0)) {
+    return;
+  }
+  ThreadFilter *thread_filter = Profiler::instance()->threadFilter();
+  int slot_id = current->filterSlotId();
+  if (unlikely(slot_id == -1)) {
+    return;
+  }
+  
+  if (enable) {
+    thread_filter->add(tid, slot_id);
+  } else {
+    thread_filter->remove(slot_id);
+  }
+}
 
 extern "C" DLLEXPORT jobject JNICALL
 Java_com_datadoghq_profiler_JavaProfiler_getContextPage0(JNIEnv *env,

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

@@ -106,8 +106,9 @@ void Profiler::onThreadStart(jvmtiEnv *jvmti, JNIEnv *jni, jthread thread) {
   ProfiledThread *current = ProfiledThread::current();
   int tid = current->tid();
   if (_thread_filter.enabled()) {
-    current->setFilterSlotId(_thread_filter.registerThread());
-    _thread_filter.remove(tid);
+    int slot_id = _thread_filter.registerThread();
+    current->setFilterSlotId(slot_id);
+    _thread_filter.remove(slot_id);  // Remove from filtering initially
   }
   updateThreadName(jvmti, jni, thread, true);
 

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

@@ -0,0 +1,70 @@
+/*
+ * Copyright The async-profiler authors
+ * SPDX-License-Identifier: Apache-2.0
+ * Copyright 2021, 2025 Datadog, Inc
+ */
+
+#ifndef _THREADIDTABLE_H
+#define _THREADIDTABLE_H
+
+#include <atomic>
+#include <unordered_set>
+
+// Signal-safe thread ID table with fixed size
+class ThreadIdTable {
+private:
+    static const int TABLE_SIZE = 256;  // Should handle most realistic thread counts
+    std::atomic<int> table[TABLE_SIZE];
+    
+    int hash(int tid) const {
+        // Simple hash function - could be improved if needed
+        return tid % TABLE_SIZE;
+    }
+    
+public:
+    ThreadIdTable() {
+        clear();
+    }
+    
+    // Signal-safe insertion using atomic operations only
+    void insert(int tid) {
+        if (tid == 0) return; // Invalid thread ID, 0 is reserved for empty slots
+        
+        int start_slot = hash(tid);
+        for (int probe = 0; probe < TABLE_SIZE; probe++) {
+            int slot = (start_slot + probe) % TABLE_SIZE;
+            int expected = 0;
+            
+            // Try to claim empty slot
+            if (table[slot].compare_exchange_strong(expected, tid, std::memory_order_relaxed)) {
+                return; // Successfully inserted
+            }
+            
+            // Check if already present
+            if (table[slot].load(std::memory_order_relaxed) == tid) {
+                return; // Already exists
+            }
+        }
+        // Table full - thread ID will be lost, but this is rare and non-critical
+        // Could increment a counter here for diagnostics if needed
+    }
+    
+    // Clear the table (not signal-safe, called during buffer switch)
+    void clear() {
+        for (int i = 0; i < TABLE_SIZE; i++) {
+            table[i].store(0, std::memory_order_relaxed);
+        }
+    }
+    
+    // Collect all thread IDs into a set (not signal-safe, called during buffer switch)
+    void collect(std::unordered_set<int>& result) {
+        for (int i = 0; i < TABLE_SIZE; i++) {
+            int tid = table[i].load(std::memory_order_relaxed);
+            if (tid != 0) {
+                result.insert(tid);
+            }
+        }
+    }
+};
+
+#endif // _THREADIDTABLE_H 

ddprof-lib/src/main/java/com/datadoghq/profiler/JavaProfiler.java

@@ -207,15 +207,15 @@ public boolean recordTraceRoot(long rootSpanId, String endpoint, int sizeLimit)
      * 'filter' option must be enabled to use this method.
      */
     public void addThread() {
-        filterThread_add();
+        filterThreadAdd0();
     }
 
     /**
      * Remove the given thread to the set of profiled threads.
      * 'filter' option must be enabled to use this method.
      */
     public void removeThread() {
-        filterThread_remove();
+        filterThreadRemove0();
     }
 
     /**
@@ -446,8 +446,10 @@ public Map<String, Long> getDebugCounters() {
     private native void stop0() throws IllegalStateException;
     private native String execute0(String command) throws IllegalArgumentException, IllegalStateException, IOException;
 
-    private native void filterThread_add();
-    private native void filterThread_remove();
+    private native void filterThreadAdd0();
+    private native void filterThreadRemove0();
+    // Backward compatibility for existing code
+    private native void filterThread0(boolean enable);
 
     private static native int getTid0();
     private static native ByteBuffer getContextPage0(int tid);