Potential memory leak and race with the JVMTI wallclock sampler

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

@@ -607,6 +607,7 @@ u32 Profiler::recordJVMTISample(u64 counter, int tid, jthread thread, jint event
 
     int num_frames = 0;
 
+    // This is a racy call, because the thread may not be alive at this moment.
     if (VM::jvmti()->GetStackTrace(thread, 0, _max_stack_depth, jvmti_frames, &num_frames) == JVMTI_ERROR_NONE && num_frames > 0) {
       // Convert to AsyncGetCallTrace format.
       // Note: jvmti_frames and frames may overlap.
@@ -618,18 +619,17 @@ u32 Profiler::recordJVMTISample(u64 counter, int tid, jthread thread, jint event
         // see https://github.com/async-profiler/async-profiler/pull/1090
         LP64_ONLY(frames[i].padding = 0;)
       }
-    }
 
-    call_trace_id = _call_trace_storage.put(num_frames, frames, false, counter);
-    u64 duration = TSC::ticks() - startTime;
-    if (duration > 0) {
-      Counters::increment(UNWINDING_TIME_JVMTI, duration); // increment the JVMTI specific counter
-    }
-  }
-  if (!deferred) {
-    _jfr.recordEvent(lock_index, tid, call_trace_id, event_type, event);
+        call_trace_id = _call_trace_storage.put(num_frames, frames, false, counter);
+        u64 duration = TSC::ticks() - startTime;
+        if (duration > 0) {
+          Counters::increment(UNWINDING_TIME_JVMTI, duration); // increment the JVMTI specific counter
+        }
+      }
+      if (!deferred) {
+        _jfr.recordEvent(lock_index, tid, call_trace_id, event_type, event);
+      }
   }
-
   _locks[lock_index].unlock();
   return call_trace_id;
 }

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

@@ -183,7 +183,7 @@ void WallClockJVMTI::timerLoop() {
             }
             int tid = nThread->osThreadId();
             if (tid != self && (!do_filter || Profiler::instance()->threadFilter()->accept(tid))) {
-              threads.push_back({nThread, thread});
+              threads.push_back({nThread, thread, tid});
             }
           }
         }
@@ -200,25 +200,32 @@ void WallClockJVMTI::timerLoop() {
                           raw_thread_state < ddprof::JVMJavaThreadState::_thread_max_state;
     OSThreadState state = OSThreadState::UNKNOWN;
     ExecutionMode mode = ExecutionMode::UNKNOWN;
-    if (vm_thread && is_initialized) {
-      OSThreadState os_state = vm_thread->osThreadState();
-      if (os_state != OSThreadState::UNKNOWN) {
-        state = os_state;
-      }
-      mode = convertJvmExecutionState(raw_thread_state);
+    if (vm_thread == nullptr || !is_initialized) {
+        return false;
     }
-    if (state == OSThreadState::UNKNOWN) {
+    OSThreadState os_state = vm_thread->osThreadState();
+    if (state == OSThreadState::TERMINATED) {
+      return false;
+    } else if (state == OSThreadState::UNKNOWN) {
       state = OSThreadState::RUNNABLE;
+    } else {
+      state = os_state;
     }
+    mode = convertJvmExecutionState(raw_thread_state);
+
     event._thread_state = state;
     event._execution_mode = mode;
     event._weight =  1;
 
-    Profiler::instance()->recordJVMTISample(1, thread_entry.native->osThreadId(), thread_entry.java, BCI_WALL, &event, false);
+    Profiler::instance()->recordJVMTISample(1, thread_entry.tid, thread_entry.java, BCI_WALL, &event, false);
     return true;
   };
 
-  timerLoopCommon<ThreadEntry>(collectThreads, sampleThreads, _reservoir_size, _interval);
+  auto cleanThreads = [](ThreadEntry& thread_entry) {
+    VM::jni()->DeleteLocalRef(thread_entry.java);
+  };
+
+  timerLoopCommon<ThreadEntry>(collectThreads, sampleThreads, cleanThreads, _reservoir_size, _interval);
   // Don't forget to detach the thread
   VM::detachThread();
 }
@@ -257,5 +264,8 @@ void WallClockASGCT::timerLoop() {
       return true;
     };
 
-    timerLoopCommon<int>(collectThreads, sampleThreads, _reservoir_size, _interval);
+    auto doNothing = [](int tid) {
+    };
+
+    timerLoopCommon<int>(collectThreads, sampleThreads, doNothing, _reservoir_size, _interval);
 }

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

@@ -50,8 +50,8 @@ class BaseWallClock : public Engine {
 
     bool isEnabled() const;
 
-    template <typename ThreadType, typename CollectThreadsFunc, typename SampleThreadsFunc>
-    void timerLoopCommon(CollectThreadsFunc collectThreads, SampleThreadsFunc sampleThreads, int reservoirSize, u64 interval) {
+    template <typename ThreadType, typename CollectThreadsFunc, typename SampleThreadsFunc, typename CleanThreadFunc>
+    void timerLoopCommon(CollectThreadsFunc collectThreads, SampleThreadsFunc sampleThreads, CleanThreadFunc cleanThreads, int reservoirSize, u64 interval) {
       if (!_enabled.load(std::memory_order_acquire)) {
         return;
       }
@@ -103,6 +103,10 @@ class BaseWallClock : public Engine {
           epoch.clean();
         }
 
+        for (ThreadType thread : threads) {
+          cleanThreads(thread);
+        }
+
         threads.clear();
         // Get a random sleep duration
         // clamp the random interval to <1,2N-1>
@@ -161,6 +165,7 @@ class WallClockJVMTI : public BaseWallClock {
     struct ThreadEntry {
         ddprof::VMThread* native;
         jthread java;
+        int tid;
     };
     WallClockJVMTI() : BaseWallClock() {}
     const char* name() override {