Fix trace preservation across JFR dumps

Author: kaahos

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2025, Datadog, Inc.
+ * Copyright 2025, 2026, Datadog, Inc.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -451,15 +451,34 @@ void CallTraceHashTable::putWithExistingId(CallTrace* source_trace, u64 weight)
 
   u64 hash = calcHash(source_trace->num_frames, source_trace->frames, source_trace->truncated);
 
-  // First check if trace already exists in any table in the chain.
+  // First check if this exact trace ID already exists in any table in the
+  // chain. Different storage generations can assign different IDs to the same
+  // stack. Both entries must survive preservation because recorded events
+  // reference the ID, not the stack hash.
   // Use ACQUIRE to match the RELEASE store in clearTableOnly(); putWithExistingId()
   // is only called on scratch/standby tables with no concurrent writers, so the
   // load is safe, but consistent ordering prevents latent issues if callers change.
   for (LongHashTable *search_table = __atomic_load_n(&_table, __ATOMIC_ACQUIRE);
        search_table != nullptr; search_table = search_table->prev()) {
-    CallTrace *existing_trace = findCallTrace(search_table, hash);
-    if (existing_trace != nullptr) {
-      return;
+    u64 *search_keys = search_table->keys();
+    HashProbe search_probe(hash, search_table->capacity());
+    u32 search_slot = search_probe.slot();
+    while (true) {
+      u64 key = __atomic_load_n(&search_keys[search_slot], __ATOMIC_RELAXED);
+      if (key == 0) {
+        break;
+      }
+      if (key == hash) {
+        CallTrace *existing_trace = search_table->values()[search_slot].acquireTrace();
+        if (existing_trace != nullptr && existing_trace != CallTraceSample::PREPARING &&
+            existing_trace->trace_id == source_trace->trace_id) {
+          return;
+        }
+      }
+      if (!search_probe.hasNext()) {
+        break;
+      }
+      search_slot = search_probe.next();
     }
   }
 

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

@@ -498,7 +498,7 @@ MethodInfo *Lookup::resolveMethod(ASGCT_CallFrame &frame) {
   jmethodID method_id = frame.method_id;
 
   // Resolve native method
-  if (FrameType::isRawPointer(bci)) {
+  if (VM::isHotspot() && FrameType::isRawPointer(bci)) {
     method_id = JVMSupport::resolve(frame.method);
   }
 
@@ -2107,6 +2107,7 @@ bool FlightRecorder::recordTaskBlock(int lock_index, int tid,
     Recording* rec = _rec;
     if (rec != nullptr) {
       Buffer *buf = rec->buffer(lock_index);
+      rec->addThread(lock_index, tid);
       rec->recordTaskBlock(buf, tid, event);
       return true;
     }

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

@@ -426,8 +426,8 @@ void LivenessTracker::onGC() {
 }
 
 void LivenessTracker::getLiveTraceIds(std::unordered_set<u64>& out_buffer) {
-  out_buffer.clear();
-  
+  // Liveness checkers share this buffer and must append their IDs. Clearing it
+  // here would discard references contributed by checkers that ran earlier.
   if (!_enabled || !_initialized) {
     return;
   }

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2021, 2025, Datadog, Inc.
+ * Copyright 2021, 2026, Datadog, Inc.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -94,6 +94,12 @@ class alignas(alignof(SpinLock)) LivenessTracker {
   void track(JNIEnv *env, AllocEvent &event, jint tid, jobject object, u64 call_trace_id);
   void flush(std::set<int> &tracked_thread_ids);
 
+#ifdef UNIT_TEST
+  void getLiveTraceIdsForTest(std::unordered_set<u64>& out_buffer) {
+    getLiveTraceIds(out_buffer);
+  }
+#endif
+
   static void JNICALL GarbageCollectionFinish(jvmtiEnv *jvmti_env);
 
 private:

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

@@ -16,6 +16,7 @@
 
 #include <gtest/gtest.h>
 #include "../../main/cpp/gtest_crash_handler.h"
+#include "../../main/cpp/livenessTracker.h"
 #include <cstdlib>
 #include <cstring>
 
@@ -241,3 +242,12 @@ TEST_F(LivenessTrackerTest, CapacityDoesNotExceedMaxCap) {
     // In the actual code, this would trigger: if (_table_cap != newcap) { ... }
     // which would be false, so no resize would be attempted
 }
+
+TEST_F(LivenessTrackerTest, LiveTraceCollectionPreservesIdsFromOtherCheckers) {
+    LivenessTracker tracker;
+    std::unordered_set<u64> trace_ids = {0x12345678ULL};
+
+    tracker.getLiveTraceIdsForTest(trace_ids);
+
+    EXPECT_EQ(trace_ids.count(0x12345678ULL), 1U);
+}

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2025, Datadog, Inc.
+ * Copyright 2025, 2026, Datadog, Inc.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -727,6 +727,27 @@ TEST_F(CallTraceStorageTest, PutWithExistingIdNoInfiniteLoopWhenFull) {
            "when the scratch table was full";
 }
 
+TEST_F(CallTraceStorageTest, PutWithExistingIdPreservesDuplicateStackIds) {
+    CallTraceHashTable table;
+
+    alignas(alignof(CallTrace)) char first_buf[sizeof(CallTrace)];
+    CallTrace* first = new (first_buf) CallTrace(false, 1, 0x100000001ULL);
+    first->frames[0].bci = 17;
+    first->frames[0].method_id = reinterpret_cast<jmethodID>(0x1234);
+
+    alignas(alignof(CallTrace)) char second_buf[sizeof(CallTrace)];
+    CallTrace* second = new (second_buf) CallTrace(false, 1, 0x200000001ULL);
+    second->frames[0] = first->frames[0];
+
+    table.putWithExistingId(first, 1);
+    table.putWithExistingId(second, 1);
+
+    std::unordered_set<CallTrace*> traces;
+    table.collect(traces);
+    EXPECT_NE(findTraceById(traces, first->trace_id), nullptr);
+    EXPECT_NE(findTraceById(traces, second->trace_id), nullptr);
+}
+
 /**
  * Integration test: processTraces preserves live traces across rotation cycles
  * without hanging.

ddprof-test/src/test/java/com/datadoghq/profiler/wallclock/JavaProfilerTaskBlockApiTest.java

@@ -120,7 +120,9 @@ public void taskBlockSpanningLiveDumpKeepsStackReference() throws Exception {
         assertTrue(recorded.get());
 
         stopProfiler();
-        TaskBlockAssertions.assertContainsStackTrace(verifyEvents("datadog.TaskBlock"));
+        IItemCollection events = verifyEvents("datadog.TaskBlock");
+        TaskBlockAssertions.assertContainsStackTrace(events);
+        TaskBlockAssertions.assertContainsEventThread(events);
     }
 
     @Override

ddprof-test/src/test/java/com/datadoghq/profiler/wallclock/TaskBlockAssertions.java

@@ -142,6 +142,20 @@ static void assertContainsStackTrace(IItemCollection taskBlockEvents) {
         assertTrue(checked > 0, "Expected at least one TaskBlock with a non-empty stackTrace");
     }
 
+    static void assertContainsEventThread(IItemCollection taskBlockEvents) {
+        int checked = 0;
+        for (IItemIterable iterable : taskBlockEvents) {
+            IMemberAccessor<IMCThread, IItem> threadAccessor =
+                    JfrAttributes.EVENT_THREAD.getAccessor(iterable.getType());
+            assertNotNull(threadAccessor, "TaskBlock must expose eventThread");
+            for (IItem item : iterable) {
+                checked++;
+                assertNotNull(threadAccessor.getMember(item), "TaskBlock eventThread must not be null");
+            }
+        }
+        assertTrue(checked > 0, "Expected at least one TaskBlock with an eventThread");
+    }
+
     static void assertContainsCorrelationId(IItemCollection taskBlockEvents) {
         Set<Long> correlationIds = nonZeroValues(taskBlockEvents, CORRELATION_ID);
         assertTrue(correlationIds.size() > 0, "Expected at least one non-zero TaskBlock correlationId");