Properly clean up the thread info maps

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

@@ -992,20 +992,19 @@ void Recording::writeThreads(Buffer* buf) {
     _thread_set.clear();
 
     Profiler* profiler = Profiler::instance();
-    MutexLocker ml(profiler->_thread_names_lock);
-    std::map<int, std::string>& thread_names = profiler->_thread_names;
-    std::map<int, jlong>& thread_ids = profiler->_thread_ids;
+    ThreadInfo t_info = profiler->_thread_info;
+
     char name_buf[32];
 
     buf->putVar64(T_THREAD);
     buf->putVar64(threads.size());
     for (int i = 0; i < threads.size(); i++) {
         const char* thread_name;
         jlong thread_id;
-        std::map<int, std::string>::const_iterator it = thread_names.find(threads[i]);
-        if (it != thread_names.end()) {
-            thread_name = it->second.c_str();
-            thread_id = thread_ids[threads[i]];
+        std::pair<std::shared_ptr<std::string>, u64> info = t_info.get(threads[i]);
+        if (info.first) {
+            thread_name = info.first->c_str();
+            thread_id = info.second;
         } else {
             snprintf(name_buf, sizeof(name_buf), "[tid=%d]", threads[i]);
             thread_name = name_buf;

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

@@ -16,6 +16,7 @@
 
 #include <algorithm>
 #include <fstream>
+#include <memory>
 #include <set>
 #include <dlfcn.h>
 #include <unistd.h>
@@ -888,19 +889,13 @@ void Profiler::setupSignalHandlers() {
     }
 }
 
-void Profiler::setThreadInfo(int tid, const char* name, jlong java_thread_id) {
-    MutexLocker ml(_thread_names_lock);
-    _thread_names[tid] = name;
-    _thread_ids[tid] = java_thread_id;
-}
-
 void Profiler::updateThreadName(jvmtiEnv* jvmti, JNIEnv* jni, jthread thread) {
     JitWriteProtection jit(true);  // workaround for JDK-8262896
     jvmtiThreadInfo thread_info;
     int native_thread_id = VMThread::nativeThreadId(jni, thread);
     if (native_thread_id >= 0 && jvmti->GetThreadInfo(thread, &thread_info) == 0) {
         jlong java_thread_id = VMThread::javaThreadId(jni, thread);
-        setThreadInfo(native_thread_id, thread_info.name, java_thread_id);
+        _thread_info.set(native_thread_id, thread_info.name, java_thread_id);
         jvmti->Deallocate((unsigned char*)thread_info.name);
     }
 }
@@ -923,16 +918,15 @@ void Profiler::updateJavaThreadNames() {
 
 void Profiler::updateNativeThreadNames() {
     ThreadList* thread_list = OS::listThreads();
-    char name_buf[64];
-
     for (int tid; (tid = thread_list->next()) != -1; ) {
-        MutexLocker ml(_thread_names_lock);
-        std::map<int, std::string>::iterator it = _thread_names.lower_bound(tid);
-        if (it == _thread_names.end() || it->first != tid) {
+        _thread_info.updateThreadName(tid, [](int tid) -> std::unique_ptr<char[]> {
+            char *name_buf = new char[64];
             if (OS::threadName(tid, name_buf, sizeof(name_buf))) {
-                _thread_names.insert(it, std::map<int, std::string>::value_type(tid, name_buf));
+                return std::unique_ptr<char[]>(name_buf);
             }
-        }
+            delete[] name_buf;
+            return nullptr;
+        });
     }
 
     delete thread_list;
@@ -1046,9 +1040,7 @@ Error Profiler::start(Arguments& args, bool reset) {
         Counters::reset();
 
         // Reset thread names and IDs
-        MutexLocker ml(_thread_names_lock);
-        _thread_names.clear();
-        _thread_ids.clear();
+        _thread_info.clearAll();
     }
 
     // (Re-)allocate calltrace buffers
@@ -1173,8 +1165,7 @@ Error Profiler::stop() {
     updateNativeThreadNames();
 
     // writing these out before stopping the JFR recording allows to report the correct counts in the recording
-    Counters::set(THREAD_IDS_COUNT, _thread_ids.size());
-    Counters::set(THREAD_NAMES_COUNT, _thread_names.size());
+    _thread_info.reportCounters();
 
     // Acquire all spinlocks to avoid race with remaining signals
     lockAll();
@@ -1265,35 +1256,8 @@ Error Profiler::dump(const char* path, const int length) {
         _class_map.clear();
         _class_map_lock.unlock();
 
-        // // Reset thread names and IDs
-        MutexLocker ml(_thread_names_lock);
-        if (thread_ids.empty()) {
-            // take the fast path
-            _thread_names.clear();
-            _thread_ids.clear();
-        } else {
-            // we need to honor the thread referenced from th liveness tracker
-            std::map<int, std::string>::iterator name_itr = _thread_names.begin();
-            while (name_itr != _thread_names.end()) {
-                if (thread_ids.find(name_itr->first) != thread_ids.end()) {
-                    name_itr = _thread_names.erase(name_itr);
-                } else {
-                    ++name_itr;
-                }
-            }
-            std::map<int, jlong>::iterator id_itr = _thread_ids.begin();
-            while (id_itr != _thread_ids.end()) {
-                if (thread_ids.find(name_itr->first) != thread_ids.end()) {
-                    id_itr = _thread_ids.erase(id_itr);
-                } else {
-                    ++id_itr;
-                }
-            }
-        }
-
-        Counters::set(THREAD_IDS_COUNT, _thread_ids.size());
-        Counters::set(THREAD_NAMES_COUNT, _thread_names.size());
-
+        _thread_info.clearAll(thread_ids);
+        _thread_info.reportCounters();
         return err;
     }
 

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

@@ -32,6 +32,7 @@
 #include "mutex.h"
 #include "spinLock.h"
 #include "threadFilter.h"
+#include "threadInfo.h"
 #include "trap.h"
 #include "vmEntry.h"
 #include "objectSampler.h"
@@ -112,10 +113,7 @@ class Profiler {
     NotifyClassUnloadedFunc _notify_class_unloaded_func;
     // --
 
-    Mutex _thread_names_lock;
-    // TODO: single map?
-    std::map<int, std::string> _thread_names;
-    std::map<int, jlong> _thread_ids;
+    ThreadInfo _thread_info;
     Dictionary _class_map;
     Dictionary _string_label_map;
     Dictionary _context_value_map;
@@ -176,7 +174,6 @@ class Profiler {
     int getJavaTraceInternal(jvmtiFrameInfo* jvmti_frames, ASGCT_CallFrame* frames, int max_depth);
     int convertFrames(jvmtiFrameInfo* jvmti_frames, ASGCT_CallFrame* frames, int num_frames);
     void fillFrameTypes(ASGCT_CallFrame* frames, int num_frames, NMethod* nmethod);
-    void setThreadInfo(int tid, const char* name, jlong java_thread_id);
     void updateThreadName(jvmtiEnv* jvmti, JNIEnv* jni, jthread thread);
     void updateJavaThreadNames();
     void updateNativeThreadNames();

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

@@ -0,0 +1,75 @@
+#include "counters.h"
+#include "mutex.h"
+#include "threadInfo.h"
+
+void ThreadInfo::set(int tid, const char* name, u64 java_thread_id) {
+    MutexLocker ml(_ti_lock);
+    _thread_names[tid] = std::string(name);
+    _thread_ids[tid] = java_thread_id;
+}
+
+std::pair<std::shared_ptr<std::string>, u64> ThreadInfo::get(int threadId) {
+    MutexLocker ml(_ti_lock);
+    auto it = _thread_names.find(threadId);
+    if (it != _thread_names.end()) {
+        return std::make_pair(std::make_shared<std::string>(it->second), _thread_ids[threadId]);
+    }
+    return std::make_pair(nullptr, 0);
+}
+
+void ThreadInfo::clearAll() {
+    MutexLocker ml(_ti_lock);
+    _thread_names.clear();
+    _thread_ids.clear();
+}
+
+void ThreadInfo::clearAll(std::set<int> &live_thread_ids) {
+    // Reset thread names and IDs
+    MutexLocker ml(_ti_lock);
+    if (live_thread_ids.empty()) {
+        // take the fast path
+        _thread_names.clear();
+        _thread_ids.clear();
+    } else {
+        // we need to honor the thread referenced from the liveness tracker
+        std::map<int, std::string>::iterator name_itr = _thread_names.begin();
+        while (name_itr != _thread_names.end()) {
+            if (live_thread_ids.find(name_itr->first) == live_thread_ids.end()) {
+                name_itr = _thread_names.erase(name_itr);
+            } else {
+                ++name_itr;
+            }
+        }
+        std::map<int, u64>::iterator id_itr = _thread_ids.begin();
+        while (id_itr != _thread_ids.end()) {
+            if (live_thread_ids.find(id_itr->first) == live_thread_ids.end()) {
+                id_itr = _thread_ids.erase(id_itr);
+            } else {
+                ++id_itr;
+            }
+        }
+    }
+}
+
+int ThreadInfo::size() {
+    MutexLocker ml(_ti_lock);
+    return _thread_names.size();
+}
+
+void ThreadInfo::updateThreadName(int tid, std::function<std::unique_ptr<char[]>(int)> resolver) {
+    MutexLocker ml(_ti_lock);
+
+    std::map<int, std::string>::iterator it = _thread_names.lower_bound(tid);
+    if (it == _thread_names.end() || it->first != tid) {
+        std::unique_ptr<char[]> namePtr = resolver(tid);
+        if (namePtr.get() != nullptr) {
+            _thread_names.insert(it, std::map<int, std::string>::value_type(tid, static_cast<char*>(namePtr.get())));
+        }
+    }
+}
+
+void ThreadInfo::reportCounters() {
+    MutexLocker ml(_ti_lock);
+    Counters::set(THREAD_IDS_COUNT, _thread_ids.size());
+    Counters::set(THREAD_NAMES_COUNT, _thread_names.size());
+}

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

@@ -0,0 +1,27 @@
+#include "mutex.h"
+#include "os.h"
+#include <functional>
+#include <map>
+#include <memory>
+#include <set>
+
+class ThreadInfo {
+  private:
+    Mutex _ti_lock;
+    std::map<int, std::string> _thread_names;
+    std::map<int, u64> _thread_ids;
+
+  public:
+    void set(int tid, const char* name, u64 java_thread_id);
+    std::pair<std::shared_ptr<std::string>, u64> get(int tid);
+
+    void updateThreadName(int threadId, std::function<std::unique_ptr<char[]>(int)> resolver);
+
+    int size();
+
+    void clearAll(std::set<int> &live_thread_ids);
+    void clearAll();
+
+    void reportCounters();
+};
+

ddprof-lib/src/test/CMakeLists.txt

@@ -51,7 +51,8 @@ file(GLOB_RECURSE SRC_FILES CONFIGURE_DEPENDS
     "${PROJECT_SOURCE_DIR}/../main/cpp/threadFilter.cpp"
     "${PROJECT_SOURCE_DIR}/../main/cpp/dictionary.cpp"
     "${PROJECT_SOURCE_DIR}/../main/cpp/methodCache.cpp"
-    "${PROJECT_SOURCE_DIR}/../main/cpp/counters.cpp"
+    "${PROJECT_SOURCE_DIR}/../main/cpp/mutex.cpp"
+    "${PROJECT_SOURCE_DIR}/../main/cpp/threadInfo.cpp"
 )
 
 add_compile_definitions(DEBUG)

ddprof-lib/src/test/cpp/ddprof_ut.cpp

@@ -6,6 +6,7 @@
     #include "mutex.h"
     #include "os.h"
     #include "threadFilter.h"
+    #include "threadInfo.h"
     #include <vector>
 
     ssize_t callback(char* ptr, int len) {
@@ -140,6 +141,37 @@
         EXPECT_EQ(0, filter.size());
     }
 
+    TEST(Profiler, testThreadInfoCleanupAllDead) {
+        ThreadInfo info;
+        info.set(1, "main", 1);
+        info.set(2, "ephemeral", 2);
+        ASSERT_EQ(2, info.size());
+
+        std::set<int> live_thread_ids;
+        live_thread_ids.insert(1);
+
+        // make sure only the non-live threads are removed
+        info.clearAll(live_thread_ids);
+        ASSERT_EQ(1, info.size());
+        ASSERT_EQ(-1, info.getThreadId(2));
+
+        // sanity check that all threads are removed when no live threads are provided
+        std::set<int> empty_set;
+        info.set(2, "ephemeral-1", 2);
+        info.clearAll(empty_set);
+        ASSERT_EQ(0, info.size());
+    }
+
+    TEST(Profiler, testThreadInfoCleanupAll) {
+        ThreadInfo info;
+        info.set(1, "main", 1);
+        info.set(2, "ephemeral", 2);
+        ASSERT_EQ(2, info.size());
+
+        info.clearAll();
+        ASSERT_EQ(0, info.size());
+    }
+
     int main(int argc, char **argv) {
       ::testing::InitGoogleTest(&argc, argv);
       return RUN_ALL_TESTS();