Reference chains: recycled leak-tag identity and Java integration tests

ddprof-lib/build.gradle.kts

@@ -72,6 +72,38 @@ gtest {
   failFast.set(true)
 }
 
+// referenceChainJfrRoundtrip_ut writes a .jfr that :ddprof-test's ReferenceChainJfrParserTest
+// reads back. Both sides get the path from this env var (set here and in ddprof-test's
+// build.gradle.kts). The path is scoped per build configuration (derived from the gtest task
+// name, which encodes it) so a multi-configuration invocation cannot delete or overwrite
+// another configuration's recording, and the producer deletes the file before it runs, so
+// the consumer never parses a recording left behind by an older binary.
+val roundtripSuiteSuffix = "_referenceChainJfrRoundtrip_ut"
+val roundtripProducers = tasks.withType<Exec>().matching {
+  it.name.startsWith("gtest") && it.name.endsWith(roundtripSuiteSuffix)
+}
+// The gtest plugin registers one producer per build configuration; zero matches can only
+// mean the suite id or the gtest task naming changed. Fail the build instead of letting
+// ReferenceChainJfrParserTest silently degrade to a skip.
+gradle.projectsEvaluated {
+  check(roundtripProducers.isNotEmpty()) {
+    "No gtest task matching 'gtest*$roundtripSuiteSuffix' was registered - the " +
+      "reference-chain JFR roundtrip wiring is stale (suite renamed?)"
+  }
+}
+roundtripProducers.configureEach {
+  val config = name.removePrefix("gtest").removeSuffix(roundtripSuiteSuffix).lowercase()
+  val jfr = rootProject.layout.buildDirectory
+    .file("tmp/reference-chain-roundtrip/$config/datadog_reference_chain_roundtrip.jfr")
+    .get()
+    .asFile
+  environment("DDPROF_REFERENCE_CHAIN_ROUNDTRIP_JFR", jfr.absolutePath)
+  doFirst {
+    jfr.parentFile.mkdirs()
+    jfr.delete()
+  }
+}
+
 java {
   sourceCompatibility = JavaVersion.VERSION_1_8
   targetCompatibility = JavaVersion.VERSION_1_8

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

@@ -94,6 +94,11 @@ class ObjectLivenessEvent : public Event {
   // flush_table() overwrites it from the entry's own tag (track() zeroes it
   // at insert).
   int64_t leak_tag = 0;
+  // LivenessTracker::leakTagIncarnation(leak_tag), read by flush_table() under the exclusive
+  // _table_lock (so it is the owner's stamp); 0 when untagged. Serialized as
+  // HeapLiveObject.leakTagIncarnation: leak tags are recycled, so (leakTag, leakTagIncarnation)
+  // identifies the object across chunks.
+  u32 leak_tag_incarnation = 0;
   Context _ctx;
 };
 
@@ -139,10 +144,17 @@ class ReferenceChainEvent : public Event {
   u64 _target_tag;
   u32 _depth;
   u8 _root_kind;
+  // LivenessTracker's ownership-period stamp of _target_tag when it is a leak tag, captured when the
+  // tag was bound to the target's frontier entry; 0 otherwise. Serialized as
+  // ReferenceChain.leakTagIncarnation. Leak tags are recycled, so consumers join
+  // ReferenceChain(targetTag, leakTagIncarnation) to HeapLiveObject(leakTag, leakTagIncarnation):
+  // the tag alone is only unambiguous within one chunk.
+  u32 _leak_tag_incarnation;
   std::vector<ReferenceChainHop> _hops;
 
   ReferenceChainEvent()
-      : Event(), _start_time(0), _target_tag(0), _depth(0), _root_kind(0) {}
+      : Event(), _start_time(0), _target_tag(0), _depth(0), _root_kind(0),
+        _leak_tag_incarnation(0) {}
 };
 
 // Search-level abandonment signal: reports why the reference-chain search

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

@@ -2423,12 +2423,14 @@ void Recording::recordHeapLiveObject(Buffer *buf, int tid, u64 call_trace_id,
           ? ((event->_alloc._weight * event->_alloc._size) + event->_skipped) /
                 event->_alloc._size
           : 0);
-  // leak_tag precedes the context snapshot to match the metadata field
-  // order (jfrMetadata.cpp: leakTag is declared between weight and spanId,
-  // before || contextAttributes) - a field written on the opposite side of
-  // writeContextSnapshot() than where the metadata declares it makes every
-  // parser read it from the first context-attribute byte.
+  // leak_tag and leak_tag_incarnation precede the context snapshot to match
+  // the metadata field order (jfrMetadata.cpp: leakTag, leakTagIncarnation
+  // are declared between weight and spanId, before || contextAttributes) - a
+  // field written on the opposite side of writeContextSnapshot() than where
+  // the metadata declares it makes every parser read it from the first
+  // context-attribute byte.
   buf->putVar64(event->leak_tag);
+  buf->putVar64((u64)event->leak_tag_incarnation);
   writeContextSnapshot(buf, event->_ctx);
   writeEventSizePrefix(buf, start);
   flushIfNeeded(buf);
@@ -2524,6 +2526,7 @@ void Recording::recordReferenceChain(Buffer *buf, ReferenceChainEvent *event) {
   buf->putVar64(T_REFERENCE_CHAIN);
   buf->putVar64(event->_start_time);
   buf->putVar64(event->_target_tag);
+  buf->putVar64((u64)event->_leak_tag_incarnation);
   buf->putVar32(event->_depth);
   buf->putUtf8(root_kind_name);
   // Original (pre-truncation) chain length, so a consumer can tell a full

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

@@ -68,7 +68,8 @@ const int MAX_VAR32_LENGTH = 5;
 // a release-mode write past the buffer).
 const int REFERENCE_CHAIN_EVENT_FIXED_BYTES =
     MAX_VAR32_LENGTH /* multi-byte event size prefix */ +
-    3 * MAX_VAR64_LENGTH /* type id, start_time, target_tag */ +
+    4 * MAX_VAR64_LENGTH /* type id, start_time, target_tag,
+                            leak_tag_incarnation */ +
     3 * MAX_VAR32_LENGTH /* depth, totalHops, chain count */ +
     32 /* rootKind string, generous */ +
     MAX_VAR32_LENGTH /* edges count */;

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

@@ -1109,6 +1109,7 @@ Java_com_datadoghq_profiler_JavaProfiler_dumpContext(JNIEnv* env, jclass unused)
 // no-op stubs below.
 #ifdef DEBUG
 #include "livenessTracker.h"
+#include "objectSampler.h"
 #include "referenceChains.h"
 #include <vector>
 
@@ -1173,6 +1174,20 @@ Java_com_datadoghq_profiler_JavaProfiler_resetKlassPopulationForTest0(
     JNIEnv *env, jclass unused) {
   ProfiledThread::initCurrentThreadSignalSafe();
   LivenessTracker::instance()->klassPopulationResetForTest();
+  // Re-baseline the class-map generation: selectLeakCandidates() returns
+  // nothing while classMap()->generation() differs from the tracker's
+  // baseline, and that baseline is only re-synced by cleanup_table()'s
+  // GC-driven sweep - but this JVM has typically already run several
+  // recordings (each Profiler::start() clearAll() bumps the generation)
+  // before this test's seeding, and no GC may land between start() and
+  // the seeded selectLeakCandidates() call. After the wipe above there is
+  // no stale-namespace population data left, and
+  // classMapGenerationSetForTest() itself invalidates every tracking entry's
+  // cached_klass_id when the generation changes (the same thing
+  // cleanup_table()'s reset branch does), so adopting the current
+  // generation is safe.
+  LivenessTracker::instance()->classMapGenerationSetForTest(
+      Profiler::instance()->classMap()->generation());
 }
 
 extern "C" DLLEXPORT jintArray JNICALL
@@ -1227,13 +1242,48 @@ Java_com_datadoghq_profiler_JavaProfiler_pollReferenceChainTargets0(
   ReferenceChainTracker::instance()->finishLoopIterationSerialized(jvmti, env);
 }
 
+// Counts only the drained chains whose leaf hop (hops are stored leaf-first) is `leaf`'s class,
+// so a chain left cached by an earlier test in the same JVM cannot satisfy the caller. The class
+// id is resolved with the same GetClassSignature + normalizeClassSignature() + lookupClass()
+// sequence the chain collector uses (LivenessTracker::resolveKlassId()).
 extern "C" DLLEXPORT jint JNICALL
-Java_com_datadoghq_profiler_JavaProfiler_drainReferenceChainEventCount0(
-    JNIEnv *env, jclass unused) {
+Java_com_datadoghq_profiler_JavaProfiler_drainReferenceChainEventCountForLeaf0(
+    JNIEnv *env, jclass unused, jobject leaf) {
   ProfiledThread::initCurrentThreadSignalSafe();
+  jvmtiEnv *jvmti = VM::jvmti();
+  if (leaf == nullptr || jvmti == nullptr) {
+    return 0;
+  }
+  u32 leaf_id = 0;
+  jclass clz = env->GetObjectClass(leaf);
+  if (clz != nullptr) {
+    char *sig = nullptr;
+    if (jvmti->GetClassSignature(clz, &sig, nullptr) == JVMTI_ERROR_NONE &&
+        sig != nullptr) {
+      const char *name = nullptr;
+      size_t name_len = 0;
+      if (ObjectSampler::normalizeClassSignature(sig, &name, &name_len)) {
+        int id = Profiler::instance()->lookupClass(name, name_len);
+        if (id > 0) {
+          leaf_id = (u32)id;
+        }
+      }
+      jvmti->Deallocate((unsigned char *)sig);
+    }
+    env->DeleteLocalRef(clz);
+  }
+  if (leaf_id == 0) {
+    return 0;
+  }
   std::vector<ReferenceChainEvent> events;
   ReferenceChainTracker::instance()->drainPendingChainEvents(&events);
-  return (jint)events.size();
+  jint count = 0;
+  for (const ReferenceChainEvent &event : events) {
+    if (!event._hops.empty() && event._hops[0].klass_id == leaf_id) {
+      count++;
+    }
+  }
+  return count;
 }
 
 extern "C" DLLEXPORT void JNICALL
@@ -1249,6 +1299,13 @@ Java_com_datadoghq_profiler_JavaProfiler_resetReferenceChainSearchForTest0(
 
 // Reads target's existing tag without tagging it. Return values are described
 // at pendingExpandPositionForTest().
+//
+// The tag read and the lookup both run under _engine_lock: the BFS thread
+// releases and reissues frontier tags (releaseSearchTags()/restartSearch())
+// only while holding it, so the tag cannot be reassigned to another object
+// between GetTag() and the lookup. The lock is recursive (mutex.cpp), so the
+// tracker seam's own nested acquisition is fine. The same applies to the two
+// seams below.
 extern "C" DLLEXPORT jlong JNICALL
 Java_com_datadoghq_profiler_JavaProfiler_getReferenceChainPendingPositionForTest0(
     JNIEnv *env, jclass unused, jobject target) {
@@ -1257,13 +1314,14 @@ Java_com_datadoghq_profiler_JavaProfiler_getReferenceChainPendingPositionForTest
   if (jvmti == nullptr || target == nullptr) {
     return -2;
   }
+  ReferenceChainTracker *tracker = ReferenceChainTracker::instance();
+  MutexLocker engine_guard(tracker->_engine_lock);
   jlong tag = 0;
   jvmtiError err = jvmti->GetTag(target, &tag);
   if (err != JVMTI_ERROR_NONE) {
     return -2;
   }
-  return (jlong)ReferenceChainTracker::instance()->pendingExpandPositionForTest(
-      tag);
+  return (jlong)tracker->pendingExpandPositionForTest(tag);
 }
 
 extern "C" DLLEXPORT jlong JNICALL
@@ -1273,6 +1331,45 @@ Java_com_datadoghq_profiler_JavaProfiler_getReferenceChainPendingSizeForTest0(
   return (jlong)ReferenceChainTracker::instance()->pendingExpandSizeForTest();
 }
 
+// Reads target's existing tag without tagging it, under _engine_lock (see
+// getReferenceChainPendingPositionForTest0()); see frontierStateForTest().
+extern "C" DLLEXPORT jint JNICALL
+Java_com_datadoghq_profiler_JavaProfiler_getReferenceChainFrontierStateForTest0(
+    JNIEnv *env, jclass unused, jobject target) {
+  ProfiledThread::initCurrentThreadSignalSafe();
+  jvmtiEnv *jvmti = VM::jvmti();
+  if (jvmti == nullptr || target == nullptr) {
+    return -1;
+  }
+  ReferenceChainTracker *tracker = ReferenceChainTracker::instance();
+  MutexLocker engine_guard(tracker->_engine_lock);
+  jlong tag = 0;
+  if (jvmti->GetTag(target, &tag) != JVMTI_ERROR_NONE) {
+    return -1;
+  }
+  return (jint)tracker->frontierStateForTest(tag);
+}
+
+// Reads target's existing tag without tagging it, under _engine_lock (see
+// getReferenceChainPendingPositionForTest0()); see
+// resolvedChainTargetTagForTest().
+extern "C" DLLEXPORT jlong JNICALL
+Java_com_datadoghq_profiler_JavaProfiler_getReferenceChainResolvedTargetTagForTest0(
+    JNIEnv *env, jclass unused, jobject target) {
+  ProfiledThread::initCurrentThreadSignalSafe();
+  jvmtiEnv *jvmti = VM::jvmti();
+  if (jvmti == nullptr || target == nullptr) {
+    return 0;
+  }
+  ReferenceChainTracker *tracker = ReferenceChainTracker::instance();
+  MutexLocker engine_guard(tracker->_engine_lock);
+  jlong tag = 0;
+  if (jvmti->GetTag(target, &tag) != JVMTI_ERROR_NONE) {
+    return 0;
+  }
+  return tracker->resolvedChainTargetTagForTest(tag);
+}
+
 // Seeds one heap-floor sample for secondsToOOM() without waiting for a GC.
 // Timestamps are only compared with each other, so any increasing sequence works.
 extern "C" DLLEXPORT void JNICALL
@@ -1299,6 +1396,24 @@ Java_com_datadoghq_profiler_JavaProfiler_setHeapFloorRecordingForTest0(
   LivenessTracker::instance()->setHeapFloorRecordingForTest(enabled == JNI_TRUE);
 }
 
+// Clears the heap-floor rings seeded by heapFloorRecordForTest0() so the
+// next test in the same JVM starts from an empty ring.
+extern "C" DLLEXPORT void JNICALL
+Java_com_datadoghq_profiler_JavaProfiler_heapFloorResetForTest0(
+    JNIEnv *env, jclass unused) {
+  ProfiledThread::initCurrentThreadSignalSafe();
+  LivenessTracker::instance()->heapFloorResetForTest();
+}
+
+// Reports whether the seeded heap-floor ring currently shows a sustained rise
+// (heapFloorRising()) - the direct read over the ring the reset seam clears.
+extern "C" DLLEXPORT jboolean JNICALL
+Java_com_datadoghq_profiler_JavaProfiler_heapFloorRisingForTest0(
+    JNIEnv *env, jclass unused) {
+  ProfiledThread::initCurrentThreadSignalSafe();
+  return LivenessTracker::instance()->heapFloorRisingForTest() ? JNI_TRUE : JNI_FALSE;
+}
+
 // Evaluates the real scheduling gate without executing a BFS pass. This is not read-only: under
 // the engine lock it may account completed-search pain and restart/reset terminal search state.
 extern "C" DLLEXPORT jboolean JNICALL
@@ -1368,8 +1483,8 @@ Java_com_datadoghq_profiler_JavaProfiler_pollReferenceChainTargets0(
     JNIEnv *env, jclass unused) {}
 
 extern "C" DLLEXPORT jint JNICALL
-Java_com_datadoghq_profiler_JavaProfiler_drainReferenceChainEventCount0(
-    JNIEnv *env, jclass unused) {
+Java_com_datadoghq_profiler_JavaProfiler_drainReferenceChainEventCountForLeaf0(
+    JNIEnv *env, jclass unused, jobject leaf) {
   return 0;
 }
 
@@ -1389,6 +1504,18 @@ Java_com_datadoghq_profiler_JavaProfiler_getReferenceChainPendingSizeForTest0(
   return 0;
 }
 
+extern "C" DLLEXPORT jint JNICALL
+Java_com_datadoghq_profiler_JavaProfiler_getReferenceChainFrontierStateForTest0(
+    JNIEnv *env, jclass unused, jobject target) {
+  return -1;
+}
+
+extern "C" DLLEXPORT jlong JNICALL
+Java_com_datadoghq_profiler_JavaProfiler_getReferenceChainResolvedTargetTagForTest0(
+    JNIEnv *env, jclass unused, jobject target) {
+  return 0;
+}
+
 extern "C" DLLEXPORT void JNICALL
 Java_com_datadoghq_profiler_JavaProfiler_heapFloorRecordForTest0(
     JNIEnv *env, jclass unused, jlong usedBytes, jlong timestampNs) {}
@@ -1401,6 +1528,16 @@ extern "C" DLLEXPORT void JNICALL
 Java_com_datadoghq_profiler_JavaProfiler_setHeapFloorRecordingForTest0(
     JNIEnv *env, jclass unused, jboolean enabled) {}
 
+extern "C" DLLEXPORT void JNICALL
+Java_com_datadoghq_profiler_JavaProfiler_heapFloorResetForTest0(
+    JNIEnv *env, jclass unused) {}
+
+extern "C" DLLEXPORT jboolean JNICALL
+Java_com_datadoghq_profiler_JavaProfiler_heapFloorRisingForTest0(
+    JNIEnv *env, jclass unused) {
+  return JNI_FALSE;
+}
+
 extern "C" DLLEXPORT jboolean JNICALL
 Java_com_datadoghq_profiler_JavaProfiler_shouldRunPassForTest0(JNIEnv *env,
                                                                 jclass unused) {

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

@@ -202,6 +202,8 @@ void JfrMetadata::initialize(
                   << field("size", T_LONG, "Original Size", F_BYTES)
                   << field("weight", T_FLOAT, "Sample weight")
                   << field("leakTag", T_LONG, "Leak Tag", F_UNSIGNED)
+                  << field("leakTagIncarnation", T_LONG,
+                           "Leak Tag Incarnation", F_UNSIGNED)
                   << field("spanId", T_LONG, "Span ID")
                   << field("localRootSpanId", T_LONG, "Local Root Span ID") ||
               contextAttributes)
@@ -211,6 +213,8 @@ void JfrMetadata::initialize(
               << category("Datadog", "Profiling")
               << field("startTime", T_LONG, "Start Time", F_TIME_TICKS)
               << field("targetTag", T_LONG, "Frontier Tag", F_UNSIGNED)
+              << field("leakTagIncarnation", T_LONG, "Leak Tag Incarnation",
+                       F_UNSIGNED)
               << field("depth", T_INT, "Depth")
               << field("rootKind", T_STRING, "GC Root Kind")
               << field("totalHops", T_INT,