Fix reference-chain test-seam races and cross-test candidate pollution

resetSearchStateForTest() and klassPopulationResetForTest() now stop the
BFS thread / take _table_lock respectively before mutating state the
thread otherwise owns exclusively. Also reset LivenessTracker's own
population history before shouldReconstructReferrerChainThroughUnboundedCacheLeak()
so a stale candidate from the prior test can't outrank CachedPayload, and
align its per-round growth rate with the sibling test's.

Author: jbachorik

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

@@ -1168,4 +1168,11 @@ Java_com_datadoghq_profiler_JavaProfiler_drainReferenceChainEventCount0(
   ReferenceChainTracker::instance()->drainPendingChainEvents(&events);
   return (jint)events.size();
 }
+
+extern "C" DLLEXPORT void JNICALL
+Java_com_datadoghq_profiler_JavaProfiler_resetReferenceChainSearchForTest0(
+    JNIEnv *env, jclass unused) {
+  jvmtiEnv *jvmti = VM::jvmti();
+  ReferenceChainTracker::instance()->resetSearchStateForTest(jvmti, env);
+}
 #endif // DEBUG

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

@@ -372,7 +372,17 @@ class alignas(alignof(SpinLock)) LivenessTracker {
       }
     }
   }
-  void klassPopulationResetForTest() { _klass_population_size = 0; }
+  // Unlike the other klassPopulation*ForTest() seams above, this one is
+  // also called from a live-JVM test (not just gtest) while the BFS thread
+  // (ReferenceChainTracker::threadLoop()) may concurrently be inside
+  // cleanup_table()'s epoch-advance pass, which holds _table_lock while
+  // mutating _klass_population_size/_klass_population - so this seam must
+  // take the same lock rather than writing the field unguarded.
+  void klassPopulationResetForTest() {
+    _table_lock.lock();
+    _klass_population_size = 0;
+    _table_lock.unlock();
+  }
 
   // Sets _gc_generations directly, bypassing initialize() (which requires a
   // live JVM - VM::hotspot_version()/VM::jni(), see that method's own code -

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

@@ -829,11 +829,15 @@ void Profiler::writeReferenceChainAbandoned(ReferenceChainAbandonedEvent *event)
 void Profiler::writeReferenceChain(ReferenceChainEvent *event, u64 deadline_ns) {
   int tid = ProfiledThread::currentTid();
   if (tid < 0) {
+    TEST_LOG("Profiler::writeReferenceChain drop: currentTid() < 0");
     return;
   }
   u32 lock_index;
   bool locked = false;
+  int sweeps = 0;
+  u64 start_ns = OS::nanotime();
   for (;;) {
+    sweeps++;
     lock_index = getLockIndex(tid);
     if (_locks[lock_index].tryLock() ||
         _locks[lock_index = (lock_index + 1) % CONCURRENCY_LEVEL].tryLock() ||
@@ -857,8 +861,14 @@ void Profiler::writeReferenceChain(ReferenceChainEvent *event, u64 deadline_ns)
     // (REFERENCE_CHAIN_EVENTS_DROPPED's own comment) rather than dropping it
     // silently.
     Counters::increment(REFERENCE_CHAIN_WRITE_DROPPED);
+    TEST_LOG("Profiler::writeReferenceChain drop: lock contention exhausted shared "
+             "deadline after sweeps=%d waited_us=%llu",
+             sweeps, (unsigned long long)((OS::nanotime() - start_ns) / 1000));
     return;
   }
+  TEST_LOG("Profiler::writeReferenceChain locked lock_index=%u after sweeps=%d "
+           "waited_us=%llu",
+           lock_index, sweeps, (unsigned long long)((OS::nanotime() - start_ns) / 1000));
   _jfr.recordReferenceChain(lock_index, event);
   _locks[lock_index].unlock();
 }
@@ -1811,9 +1821,13 @@ Error Profiler::dump(const char *path, const int length) {
     // _locks[] contention.
     const u64 kChainDrainBudgetNs = 50 * 1000000ULL;
     u64 chain_drain_deadline_ns = OS::nanotime() + kChainDrainBudgetNs;
+    long long write_dropped_before = Counters::getCounter(REFERENCE_CHAIN_WRITE_DROPPED);
     for (auto &rc_event : pending_chain_events) {
       writeReferenceChain(&rc_event, chain_drain_deadline_ns);
     }
+    TEST_LOG("Profiler::dump reference-chain batch=%d write_dropped=%lld",
+             (int)pending_chain_events.size(),
+             Counters::getCounter(REFERENCE_CHAIN_WRITE_DROPPED) - write_dropped_before);
 
     Libraries::instance()->refresh();
     updateJavaThreadNames();

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

@@ -352,6 +352,13 @@ void ReferenceChainTracker::startThread() {
   // previously call pthread_kill() on a still value-initialized (0) or
   // stale/joined pthread_t, which is undefined behavior. Publishing _thread
   // before _running closes that window.
+  // Reset from any previous stopThread() call - a dynamic-attach profiler
+  // can go through multiple start()/stop() cycles in one JVM lifetime (this
+  // class's own start()/stop() header comments), and a stale abort request
+  // left set from the prior cycle would make heapReferenceCallback() abort
+  // this new cycle's very first pass instantly.
+  _abort_pass_requested.store(false, std::memory_order_relaxed);
+
   pthread_t thread;
   if (pthread_create(&thread, NULL, threadEntry, this) != 0) {
     Log::warn("Unable to create ReferenceChains BFS thread");
@@ -366,6 +373,11 @@ void ReferenceChainTracker::stopThread() {
     return;
   }
   _running = false;
+  // Ask any in-flight JVMTI FollowReferences walk (heapReferenceCallback())
+  // to abort at its next callback invocation - set before pthread_kill()
+  // below, since that signal alone cannot interrupt a call already inside
+  // the JVM/JVMTI implementation.
+  _abort_pass_requested.store(true, std::memory_order_relaxed);
   // Same wake-then-join shape as BaseWallClock::stop() (wallClock.cpp:324-333):
   // pthread_kill(WAKEUP_SIGNAL) interrupts threadLoop()'s OS::sleep() early
   // (WAKEUP_SIGNAL/SIGIO is installed with a no-op handler unconditionally
@@ -601,6 +613,63 @@ void ReferenceChainTracker::restartSearch() {
   // _search_start_ns again on this restarted search's own first pass.
 }
 
+void ReferenceChainTracker::resetSearchStateForTest(jvmtiEnv *jvmti,
+                                                     JNIEnv *jni) {
+  // Every field touched below is otherwise only ever mutated by the BFS
+  // thread itself (threadLoop()/runPass()/pollWatchedTargets()) - without
+  // stopping it first, a pass already in flight on that thread can observe
+  // this reset only partially, or overwrite it right back (e.g. finish a
+  // pass that was already headed for SearchState::ABANDONED after this
+  // method has just forced SearchState::RUNNING below), a race found in
+  // practice, not just in theory. stopThread() (now that it can abort an
+  // in-flight JVMTI walk promptly - see its own comment) makes this a cheap,
+  // clean stop/reset/restart rather than an indefinite wait.
+  stopThread();
+
+  // Clear every live tag this search still holds before resetting - the
+  // same ordering restartSearch() itself requires (its own assert), so a
+  // stale tag from whatever search a previous test left running cannot
+  // collide with the fresh search's own tags once _next_tag is rewound
+  // below.
+  if (jvmti != nullptr && jni != nullptr) {
+    releaseSearchTags(jvmti, jni);
+  }
+  _tags_released = true;
+
+  _pain_budget.spend(_search_pain_ms);
+  _search_pain_ms = 0;
+
+  if (_frontier != nullptr) {
+    _frontier->resetForRestart();
+  }
+  _next_tag = 1;
+
+  _search_started = false;
+  store(_search_state, (u8)SearchState::RUNNING);
+  store(_abandon_reason, (u8)SearchAbandonReason::NONE);
+  store(_search_start_ns, (u64)0);
+  _expand_cursor = 1;
+  _last_pass_gc_finish_epoch = 0;
+  store(_last_pass_ns, (u64)0);
+  store(_passes_run, 0);
+
+  // Unlike restartSearch(), which leaves these for pollWatchedTargets() to
+  // notice the _search_start_ns change and clear naturally on its own next
+  // call, this reset must take effect immediately - there is no next real
+  // pass here to trigger that self-clear.
+  _emitted_target_tags.clear();
+  _emitted_search_start_ns = 0;
+
+  _pending_chain_events_lock.lock();
+  _pending_chain_events.clear();
+  _pending_chain_events_lock.unlock();
+
+  // Restart the BFS thread against this freshly reset state - startThread()
+  // itself clears _abort_pass_requested, so the new thread's very first
+  // pass is not instantly aborted by the flag stopThread() just set above.
+  startThread();
+}
+
 jlong ReferenceChainTracker::tagObject(jvmtiEnv *jvmti, jobject obj) {
   assert(!t_inGCCallback &&
          "SetTag is a JVMTI Heap-category call and must not be made from "
@@ -797,6 +866,21 @@ jint JNICALL ReferenceChainTracker::heapReferenceCallback(
     jlong *referrer_tag_ptr, jint length, void *user_data) {
   PassContext *ctx = (PassContext *)user_data;
 
+  if (ctx->tracker->_abort_pass_requested.load(std::memory_order_relaxed)) {
+    // stopThread() has set this right before pthread_kill()/pthread_join() -
+    // see that method's own comment. pthread_kill(WAKEUP_SIGNAL) only
+    // interrupts threadLoop()'s OS::sleep(); it cannot interrupt an
+    // in-flight JVMTI FollowReferences call, so without this check
+    // pthread_join() would block until this pass's walk finishes on its own
+    // - potentially the whole reachable graph, well past any caller's
+    // shutdown timeout. Treat it exactly like an ordinary budget exhaustion
+    // (ctx->truncated = true): this pass ends early and the search stays
+    // non-terminal - fine, since the tracker is shutting down and simply
+    // never resumes it.
+    ctx->truncated = true;
+    return JVMTI_VISIT_ABORT;
+  }
+
   if (*tag_ptr < 0 || reference_kind == JVMTI_HEAP_REFERENCE_CLASS ||
       reference_kind == JVMTI_HEAP_REFERENCE_SYSTEM_CLASS) {
     // Referee is a class object - either already tagged negative by
@@ -1133,6 +1217,10 @@ bool ReferenceChainTracker::runPass(jvmtiEnv *jvmti, JNIEnv *jni,
 
   resolveLoadedClasses(jvmti, jni);
 
+  TEST_LOG("ReferenceChainTracker::runPass starting JVMTI walk: "
+           "search_started=%d frontierSize=%zu",
+           _search_started, _frontier != nullptr ? _frontier->size() : (size_t)0);
+
   int edges_admitted = 0;
   bool truncated = false;
   bool frontier_cap_hit = false;
@@ -1462,8 +1550,13 @@ void ReferenceChainTracker::enqueueChainEvent(ReferenceChainEvent &&event) {
     // "dropped-event-without-counter" review lens).
     _pending_chain_events.erase(_pending_chain_events.begin());
     Counters::increment(REFERENCE_CHAIN_EVENTS_DROPPED);
+    TEST_LOG("ReferenceChainTracker::enqueueChainEvent dropped oldest queued event, "
+             "queue_size=%d (at MAX_PENDING_CHAIN_EVENTS=%d)",
+             (int)_pending_chain_events.size(), MAX_PENDING_CHAIN_EVENTS);
   }
   _pending_chain_events.push_back(std::move(event));
+  TEST_LOG("ReferenceChainTracker::enqueueChainEvent queued, queue_size=%d",
+           (int)_pending_chain_events.size());
   _pending_chain_events_lock.unlock();
 }
 
@@ -1480,4 +1573,6 @@ void ReferenceChainTracker::drainPendingChainEvents(
     _pending_chain_events.clear();
   }
   _pending_chain_events_lock.unlock();
+  TEST_LOG("ReferenceChainTracker::drainPendingChainEvents drained=%d",
+           (int)out->size());
 }

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

@@ -692,6 +692,22 @@ class ReferenceChainTracker {
   pthread_t _thread;
   volatile bool _running;
 
+  // Cooperative-cancellation flag for an in-flight JVMTI FollowReferences
+  // walk: stopThread() sets this before pthread_kill()/pthread_join()
+  // (that signal alone cannot interrupt a call already inside the JVM/JVMTI
+  // implementation), and heapReferenceCallback() checks it on every
+  // invocation, aborting the walk within one callback rather than letting
+  // pthread_join() block until the walk finishes on its own - see both
+  // methods' own comments. startThread() resets it back to false, since a
+  // dynamic-attach profiler can cycle through multiple start()/stop() calls
+  // in one JVM lifetime and a stale abort request would instantly kill the
+  // next cycle's very first pass. std::atomic<bool>: written from
+  // stopThread()/startThread() on the calling (shutdown) thread and read
+  // from heapReferenceCallback() on the BFS thread - the same cross-thread
+  // shape _running above already has, just made explicit via atomic rather
+  // than a plain volatile bool.
+  std::atomic<bool> _abort_pass_requested;
+
   ReferenceChainTracker()
       : _enabled(false), _frontier(nullptr), _last_resolved_class_count(0),
         _gc_start_epoch(0),
@@ -703,7 +719,8 @@ class ReferenceChainTracker {
         _abandon_reason(SearchAbandonReason::NONE), _search_start_ns(0),
         _expand_cursor(1), _last_pass_gc_finish_epoch(0), _last_pass_ns(0),
         _passes_run(0), _emitted_search_start_ns(0), _pain_budget(0.0),
-        _search_pain_ms(0), _thread(), _running(false) {}
+        _search_pain_ms(0), _thread(), _running(false),
+        _abort_pass_requested(false) {}
 
   void onGCStart();
   void onGCFinish();
@@ -1099,6 +1116,22 @@ class ReferenceChainTracker {
   // failure (obj/jvmti/jni null, SetTag failed, or the frontier table is at
   // capacity).
   jlong tagAsRootForTest(jvmtiEnv *jvmti, JNIEnv *jni, jobject obj);
+
+  // Test seam - not part of the production API. Since ReferenceChainTracker
+  // is a process-wide singleton (ExternalProcessReferenceChainTest's own
+  // class javadoc explains why that matters: only the *first* test to ever
+  // call runPass() in a shared JVM gets a real root-seeded walk, since
+  // runPass() only re-walks from the roots once per search's whole
+  // lifetime), an in-process test that needs its own genuine first-ever
+  // root walk calls this at the start of its test body to force exactly
+  // that - releasing any tags a previous test's search still held, then
+  // resetting search/frontier state to the same "brand-new tracker" state
+  // restartSearch() (referenceChains.cpp) produces, plus the target-
+  // dedup/pending-event state restartSearch() itself intentionally leaves
+  // for pollWatchedTargets()/drainPendingChainEvents() to self-clear (this
+  // is an immediate, out-of-band reset - there is no next real pass here to
+  // observe the change and clear them the ordinary way).
+  void resetSearchStateForTest(jvmtiEnv *jvmti, JNIEnv *jni);
 };
 
 #endif // _REFERENCECHAINS_H

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

@@ -771,6 +771,16 @@ public ContextStorageMode contextStorageMode() {
      */
     public static native int drainReferenceChainEventCount0();
 
+    /**
+     * Test seam (debug native builds only): resets ReferenceChainTracker's search/frontier state
+     * back to a brand-new tracker's, releasing any tags a previous search still held. Since the
+     * tracker is a process-wide singleton, an in-process test that needs its own genuine first
+     * root-seeded walk (runPass() only re-walks from the roots once per search's whole lifetime)
+     * calls this at the start of its test body to force one, rather than depending on being the
+     * first reference-chain test to run in a shared test JVM.
+     */
+    public static native void resetReferenceChainSearchForTest0();
+
     /**
      * Resets the cached ThreadContext for the current storage slot — the calling thread in
      * {@link ContextStorageMode#THREAD}, or its current carrier in

ddprof-test/src/test/java/com/datadoghq/profiler/referencechains/ReferenceChainTrackingTest.java

@@ -6,6 +6,7 @@
 package com.datadoghq.profiler.referencechains;
 
 import com.datadoghq.profiler.AbstractProfilerTest;
+import com.datadoghq.profiler.JavaProfiler;
 import com.datadoghq.profiler.Platform;
 import org.junit.jupiter.api.MethodOrderer;
 import org.junit.jupiter.api.Order;
@@ -239,6 +240,18 @@ public void shouldExposeReferenceChainsSettingWhenEnabled() {
   @Test
   @Order(1)
   public void shouldReconstructReferrerChainToGcRoot() throws Exception {
+    if ("debug".equals(System.getProperty("ddprof_test.config"))) {
+      // Being pinned to run first *within this class* (this class's own header comment, "Why
+      // @TestMethodOrder/@Order(1)") does not guarantee this is the first reference-chain test
+      // to ever call into the shared, process-wide singleton ReferenceChainTracker in this whole
+      // test JVM - no forkEvery is configured (same header comment), so another test class can
+      // run first and leave the singleton search already non-RUNNING/already-tagged, in which
+      // case runPass() never takes its real root-seeded-walk branch again. Force a genuine fresh
+      // search so this test's own ChainLink population is guaranteed reachable by a real walk,
+      // regardless of what ran earlier in this JVM. Debug-only: this native seam does not exist
+      // in a release build.
+      JavaProfiler.resetReferenceChainSearchForTest0();
+    }
     List<Object> gcRootHolder = new ArrayList<>();
     Path scratchDumpPath = Paths.get("referencechains-population-scratch.jfr");
     try {
@@ -385,6 +398,26 @@ public void shouldReconstructReferrerChainToGcRoot() throws Exception {
   @Test
   @Order(2)
   public void shouldReconstructReferrerChainThroughUnboundedCacheLeak() throws Exception {
+    if ("debug".equals(System.getProperty("ddprof_test.config"))) {
+      // Mirrors shouldReconstructReferrerChainToGcRoot()'s own reset (see that method's comment):
+      // the natural restartSearch() cycle (shouldRunPass(), referenceChains.cpp) that would
+      // otherwise give this method its own fresh root walk once shouldReconstructReferrerChainToGcRoot()'s
+      // own ChainLink candidate is found and its tags released is gated on cadence/pacing budget
+      // (canAffordNewSearch()) - not guaranteed to fire again before this method's own round/retry
+      // budget runs out. Force a genuine fresh search here too, rather than depend on that timing,
+      // so cache (below) is guaranteed reachable by a real walk regardless of it. Debug-only: this
+      // native seam does not exist in a release build.
+      //
+      // resetReferenceChainSearchForTest0() only resets ReferenceChainTracker's own search state
+      // (frontier/tags/search-progress fields); it does not touch LivenessTracker's per-klass
+      // population-history rings, which selectLeakCandidates() consults independently to decide
+      // which klass is "trending". Without also resetting those, a klass ChainLink already
+      // accumulated a positive slope for during shouldReconstructReferrerChainToGcRoot() can still
+      // outrank CachedPayload as the leak candidate here, so the fresh search below ends up
+      // reconstructing ChainLink's chain again instead of CachedPayload's.
+      JavaProfiler.resetKlassPopulationForTest0();
+      JavaProfiler.resetReferenceChainSearchForTest0();
+    }
     Map<String, CachedPayload> cache = new HashMap<>();
     Path scratchDumpPath = Paths.get("referencechains-cache-leak-scratch.jfr");
     try {
@@ -397,7 +430,13 @@ public void shouldReconstructReferrerChainThroughUnboundedCacheLeak() throws Exc
       ReferenceChainAssertions.ChainMatch match = null;
       int totalRounds = 16;
       for (int round = 1; round <= totalRounds && match == null; round++) {
-        int newEntries = Math.min(round, 10) * 300;
+        // Matches shouldReconstructReferrerChainToGcRoot()'s own *600 growth rate (not *300, this
+        // method's previous value) - CachedPayload and ChainLink are near-identical in size (same
+        // padding fields; CachedPayload is actually slightly smaller, missing ChainLink's `next`
+        // reference), so halving the growth rate here bought no headroom and instead just let
+        // CachedPayload's own population ring stall short of KLASS_POPULATION_MIN_FILL_FOR_TREND
+        // within the same totalRounds budget both methods share.
+        int newEntries = Math.min(round, 10) * 600;
         int roundNumber = round;
         Thread allocator = new Thread(() -> {
           for (int i = 0; i < newEntries; i++) {