Fix reference-chain class-tag reuse after StringDictionary reset

StringDictionary::clearAll() (on profiler restart) wipes dict ids but
JVMTI class tags survive, so ReferenceChainTracker::_class_tags held
ids into a namespace that no longer existed. Add a generation counter
to StringDictionary and have resolveLoadedClasses() detect a mismatch
and force re-resolution, reusing existing JVMTI tags instead of always
minting new ones. Also adds a firstpassbudget override for the
root-seeded first BFS pass.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

Author: jbachorik

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

@@ -82,7 +82,7 @@ static const Multiplier UNIVERSAL[] = {
 //                          and keep the liveness track of 10% of the allocation
 //                          samples
 //     generations        - track surviving generations
-//     referencechains[=BOOL[:hops=N][:budget=N][:ttl=N][:framecap=N][:pausetarget=N][:painbudget=N]]
+//     referencechains[=BOOL[:hops=N][:budget=N][:ttl=N][:framecap=N][:pausetarget=N][:painbudget=N][:firstpassbudget=N]]
 //                        - (PROF-15341, off by default) tag/BFS-walk live-heap
 //                          samples' referrer chains back toward a GC root.
 //                          pausetarget=N (ms) is the pause-time-SLO ceiling
@@ -92,6 +92,12 @@ static const Multiplier UNIVERSAL[] = {
 //                          wall-clock time a *restarted* search (one begun
 //                          after a prior search already completed/abandoned)
 //                          may spend on average - see PainBudget (painBudget.h).
+//                          firstpassbudget=N overrides just the search's
+//                          one-shot, root-seeded first pass's edge budget
+//                          (default 0 - falls back to budget=N) since that
+//                          pass alone decides which GC roots ever enter the
+//                          frontier at all, unlike every later pass's cheap,
+//                          incremental per-node expansion.
 //                          Sub-options are placeholders pending future tuning;
 //                          see doc/architecture/LiveHeapReferenceChains*.md
 //     lightweight[=BOOL] - enable lightweight profiling - events without
@@ -510,6 +516,8 @@ Error Arguments::parse(const char *args) {
             } else if (strcasecmp(cursor, "painbudget") == 0) {
               _reference_chains_pain_budget_percent =
                   std::min(std::max(atoi(eq), 0), 100);
+            } else if (strcasecmp(cursor, "firstpassbudget") == 0) {
+              _reference_chains_first_pass_budget = std::max(atoi(eq), 0);
             }
           }
           cursor = next;

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

@@ -89,6 +89,19 @@ const long DEFAULT_REFERENCE_CHAINS_PAUSE_TARGET_MS = 5; // ms per pass
 // choice of how much background cost is acceptable. Round, provisional
 // default like every other constant in this block.
 const int DEFAULT_REFERENCE_CHAINS_PAIN_BUDGET_PERCENT = 1;
+// First-pass edge budget override: the search's one-and-only root-seeded
+// FollowReferences(0, nullptr, nullptr, ...) call (ReferenceChainTracker::runPass()'s
+// !_search_started branch) enumerates every GC root in one JVMTI-controlled
+// traversal order and stops admitting once this budget is spent - any root
+// FollowReferences had not yet reached is excluded from the frontier for the
+// rest of that search (every later pass only expands forward from already-
+// admitted frontier entries, see expandFrontier()'s own comment). Unlike
+// DEFAULT_REFERENCE_CHAINS_BUDGET, which bounds every pass including the many
+// cheap, per-node expansion passes that follow, this only ever spends once
+// per search, so a much larger one-time ceiling is affordable. 0 (the
+// default) means "no override - use the same budget as every other pass",
+// preserving prior behavior for anyone not setting this explicitly.
+const int DEFAULT_REFERENCE_CHAINS_FIRST_PASS_BUDGET = 0;
 
 const char *const EVENT_NOOP = "noop";
 const char *const EVENT_CPU = "cpu";
@@ -251,6 +264,7 @@ class Arguments {
   int _reference_chains_frontier_cap;
   long _reference_chains_pause_target_ms;
   int _reference_chains_pain_budget_percent;
+  int _reference_chains_first_pass_budget;
   long _nativemem;
   int  _jstackdepth;
   int _safe_mode;
@@ -299,6 +313,7 @@ class Arguments {
         _reference_chains_frontier_cap(DEFAULT_REFERENCE_CHAINS_FRONTIER_CAP),
         _reference_chains_pause_target_ms(DEFAULT_REFERENCE_CHAINS_PAUSE_TARGET_MS),
         _reference_chains_pain_budget_percent(DEFAULT_REFERENCE_CHAINS_PAIN_BUDGET_PERCENT),
+        _reference_chains_first_pass_budget(DEFAULT_REFERENCE_CHAINS_FIRST_PASS_BUDGET),
         _nativemem(-1),
         _jstackdepth(DEFAULT_JSTACKDEPTH),
         _safe_mode(0),

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

@@ -243,6 +243,13 @@ Error ReferenceChainTracker::start(Arguments &args) {
 
   _hop_cap = args._reference_chains_hop_cap;
   _budget = args._reference_chains_budget;
+  // 0 (unset) falls back to _budget - see this field's own comment
+  // (referenceChains.h) for why the first pass needs its own, larger
+  // ceiling rather than sharing _budget/_effective_budget with every later
+  // pass.
+  _first_pass_budget = args._reference_chains_first_pass_budget > 0
+                            ? args._reference_chains_first_pass_budget
+                            : _budget;
   _ttl_ms = args._reference_chains_ttl_ms;
 
   // Pause-time pacing controller: (re)seed the controller's ceiling and the
@@ -603,7 +610,7 @@ void ReferenceChainTracker::restartSearch() {
   store(_search_state, (u8)SearchState::RUNNING);
   store(_abandon_reason, (u8)SearchAbandonReason::NONE);
   store(_search_start_ns, (u64)0);
-  _expand_cursor = 1;
+  _pending_expand.clear();
   _last_pass_gc_finish_epoch = 0;
   store(_last_pass_ns, (u64)0);
   store(_passes_run, 0);
@@ -648,7 +655,7 @@ void ReferenceChainTracker::resetSearchStateForTest(jvmtiEnv *jvmti,
   store(_search_state, (u8)SearchState::RUNNING);
   store(_abandon_reason, (u8)SearchAbandonReason::NONE);
   store(_search_start_ns, (u64)0);
-  _expand_cursor = 1;
+  _pending_expand.clear();
   _last_pass_gc_finish_epoch = 0;
   store(_last_pass_ns, (u64)0);
   store(_passes_run, 0);
@@ -760,6 +767,27 @@ jlong ReferenceChainTracker::tagAsRootForTest(jvmtiEnv *jvmti, JNIEnv *jni,
 
 void ReferenceChainTracker::resolveLoadedClasses(jvmtiEnv *jvmti,
                                                   JNIEnv *jni) {
+  // Profiler::start() resets the class-name StringDictionary
+  // (_class_map.clearAll(), profiler.cpp) whenever `reset || _start_time ==
+  // 0` - which restarts its id namespace at 1, but does NOT touch any
+  // class's JVMTI-level class-object tag (JVM-level state, unrelated to our
+  // dictionary). Detect that reset via the dictionary's own generation
+  // counter and drop every id this table cached from the now-gone
+  // generation before the scan below - see _last_class_map_generation's own
+  // comment (referenceChains.h) for why leaving them in place would keep
+  // resolving heap references to the wrong (or nonexistent) class name.
+  u64 current_generation = Profiler::instance()->classMap()->generation();
+  bool class_map_reset = current_generation != _last_class_map_generation;
+  if (class_map_reset) {
+    _class_tags.clear();
+    // Force the scan below to run even if GetLoadedClasses()'s count happens
+    // to match the last-seen count - -1 can never equal `class_count`
+    // (always >= 0), unlike 0 which is a legitimate "no classes loaded yet"
+    // starting value.
+    _last_resolved_class_count = -1;
+    _last_class_map_generation = current_generation;
+  }
+
   jclass *classes = nullptr;
   jint class_count = 0;
   if (jvmti->GetLoadedClasses(&class_count, &classes) != JVMTI_ERROR_NONE ||
@@ -793,12 +821,18 @@ void ReferenceChainTracker::resolveLoadedClasses(jvmtiEnv *jvmti,
     for (jint i = 0; i < class_count; i++) {
       jclass klass = classes[i];
       jlong tag = 0;
-      if (jvmti->GetTag(klass, &tag) == JVMTI_ERROR_NONE && tag == 0) {
-        // Not yet tagged (by us or a prior pass) - resolve its name now, via
-        // the same GetClassSignature + normalizeClassSignature +
-        // Profiler::lookupClass sequence ObjectSampler::recordAllocation()
-        // already uses (objectSampler.cpp:76-90), reused rather than
-        // re-derived.
+      // Resolve if not yet tagged (ordinary case: a newly-loaded class), or
+      // unconditionally on a class-map reset (class_map_reset above) - a
+      // class already tagged from a prior generation still carries that same
+      // JVMTI tag (untouched by clearAll()), but the dictionary id it used to
+      // map to is gone, so its name must be re-resolved into the new
+      // generation too.
+      if (jvmti->GetTag(klass, &tag) == JVMTI_ERROR_NONE &&
+          (tag == 0 || class_map_reset)) {
+        // Resolve its name now, via the same GetClassSignature +
+        // normalizeClassSignature + Profiler::lookupClass sequence
+        // ObjectSampler::recordAllocation() already uses
+        // (objectSampler.cpp:76-90), reused rather than re-derived.
         char *class_name = nullptr;
         if (jvmti->GetClassSignature(klass, &class_name, nullptr) ==
                 JVMTI_ERROR_NONE &&
@@ -809,8 +843,12 @@ void ReferenceChainTracker::resolveLoadedClasses(jvmtiEnv *jvmti,
                                                       &name_len)) {
             int id = Profiler::instance()->lookupClass(name_slice, name_len);
             if (id != -1) {
-              jlong class_tag = nextClassTag();
-              if (jvmti->SetTag(klass, class_tag) == JVMTI_ERROR_NONE) {
+              // Reuse the existing tag if this class was already tagged by a
+              // prior generation - only the resolved id needs refreshing,
+              // not the tag identity heapReferenceCallback() keys off of.
+              jlong class_tag = tag != 0 ? tag : nextClassTag();
+              if (tag != 0 ||
+                  jvmti->SetTag(klass, class_tag) == JVMTI_ERROR_NONE) {
                 _class_tags.insert(class_tag, (u32)id);
               }
             }
@@ -965,6 +1003,10 @@ jint JNICALL ReferenceChainTracker::heapReferenceCallback(
     }
     *tag_ptr = tag;
     ctx->edges_admitted++;
+    // Queue for expandFrontier()/markAllFrontierExpanded() - see
+    // _pending_expand's own declaration comment for why this replaces a
+    // scan over the admitted range.
+    ctx->tracker->_pending_expand.push_back(tag);
   }
 
   return JVMTI_VISIT_OBJECTS;
@@ -975,15 +1017,10 @@ jint JNICALL ReferenceChainTracker::heapReferenceCallback(
 // ---------------------------------------------------------------------------
 
 void ReferenceChainTracker::markAllFrontierExpanded() {
-  jlong scan_limit = _frontier->size();
-  for (jlong tag = 1; tag <= scan_limit; tag++) {
-    FrontierEntry entry{};
-    if (_frontier->lookup(tag, &entry) &&
-        entry.state == FrontierEntryState::FRONTIER) {
-      _frontier->markExpanded(tag);
-    }
+  while (!_pending_expand.empty()) {
+    _frontier->markExpanded(_pending_expand.front());
+    _pending_expand.pop_front();
   }
-  _expand_cursor = scan_limit + 1;
 }
 
 void ReferenceChainTracker::expandFrontier(jvmtiEnv *jvmti, JNIEnv *jni,
@@ -1012,29 +1049,26 @@ void ReferenceChainTracker::expandFrontier(jvmtiEnv *jvmti, JNIEnv *jni,
   while (!ctx.truncated && progress) {
     progress = false;
 
-    // Re-read size() every outer iteration: expanding one entry can insert
-    // new FRONTIER children with higher tags, extending the range this loop
-    // needs to cover before this pass's budget runs out (design doc
-    // Algorithm step 3/4 - admitting them immediately, rather than waiting
-    // for a whole extra pass, is a strictly tighter use of the budget this
-    // call already has).
-    jlong scan_limit = _frontier->size();
-    if (_expand_cursor > scan_limit) {
+    if (_pending_expand.empty()) {
       break; // nothing pending
     }
 
-    std::vector<jlong> candidate_tags;
-    for (jlong tag = _expand_cursor; tag <= scan_limit; tag++) {
-      FrontierEntry entry{};
-      if (_frontier->lookup(tag, &entry) &&
-          entry.state == FrontierEntryState::FRONTIER) {
-        candidate_tags.push_back(tag);
-      }
-    }
-    if (candidate_tags.empty()) {
-      _expand_cursor = scan_limit + 1;
-      break; // caught up; nothing pending
-    }
+    // Snapshot at most `budget` pending tags per iteration, not the whole
+    // queue: _pending_expand can hold far more than this call could ever
+    // admit (e.g. every hop_cap-limited entry from a huge one-shot first
+    // pass), and GetObjectsWithTags()'s own cost scales with the batch size
+    // handed to it - resolving the entire backlog up front just to abandon
+    // most of it once `budget` is hit would reintroduce the same
+    // admitted-range-sized cost this queue was meant to avoid. Entries
+    // expanded below can push new FRONTIER children onto the back of
+    // _pending_expand (design doc Algorithm step 3/4 - admitting them
+    // immediately, rather than waiting for a whole extra pass, is a
+    // strictly tighter use of the budget this call already has); those wait
+    // for a later outer-loop iteration.
+    size_t batch_size =
+        std::min(_pending_expand.size(), (size_t)std::max(budget, 1));
+    std::vector<jlong> candidate_tags(_pending_expand.begin(),
+                                       _pending_expand.begin() + batch_size);
 
     // Design doc Algorithm step 2: "resolve currently-live tagged frontier
     // objects" - batched via one GetObjectsWithTags call rather than one
@@ -1064,14 +1098,18 @@ void ReferenceChainTracker::expandFrontier(jvmtiEnv *jvmti, JNIEnv *jni,
       if (ctx.truncated) {
         break; // budget/frontier-cap already hit by a sibling in this batch
       }
+      // `tag` is always _pending_expand's current front here: candidate_tags
+      // is a front-to-back snapshot, and nothing removes from the front
+      // except this loop itself - concurrent admissions (below) only append
+      // to the back.
       auto it = live.find(tag);
       if (it == live.end()) {
         // Design doc Algorithm step 2: "objects that fail to resolve are
         // dropped (dead - free pruning)". The object died between passes;
         // JVMTI already forgot its tag along with it, so there is nothing
         // left to expand and no live object to release a tag on.
         _frontier->clear(tag);
-        _expand_cursor = tag + 1;
+        _pending_expand.pop_front();
         progress = true;
         continue;
       }
@@ -1087,14 +1125,15 @@ void ReferenceChainTracker::expandFrontier(jvmtiEnv *jvmti, JNIEnv *jni,
       if (follow_err != JVMTI_ERROR_NONE || ctx.truncated) {
         // Either this call failed outright, or ran into the budget/
         // frontier-cap mid-way through this entry's own children. Leave
-        // `tag` FRONTIER (not EXPANDED) so a later pass retries it -
-        // already-admitted children stay tagged/in the table, and
-        // heapReferenceCallback()'s *tag_ptr == 0 check makes re-expanding
-        // this entry idempotent, so no work is lost, only deferred.
+        // `tag` at the front of _pending_expand (still FRONTIER, not
+        // EXPANDED) so a later pass retries it - already-admitted children
+        // stay tagged/in the table, and heapReferenceCallback()'s
+        // *tag_ptr == 0 check makes re-expanding this entry idempotent, so
+        // no work is lost, only deferred.
         break;
       }
       _frontier->markExpanded(tag);
-      _expand_cursor = tag + 1;
+      _pending_expand.pop_front();
       progress = true;
     }
 
@@ -1237,6 +1276,10 @@ bool ReferenceChainTracker::runPass(jvmtiEnv *jvmti, JNIEnv *jni,
   // unavailable/disabled, so this is a strict upgrade with no behavior change
   // on hosts without a usable timestamp counter.
   u64 pass_wall_ticks = 0;
+  // Captured before _search_started flips below - see its use at
+  // updatePacing() further down for why the first pass's own duration must
+  // not feed that controller the same way every later pass's does.
+  bool was_first_pass = !_search_started;
 
   if (!_search_started) {
     _search_started = true;
@@ -1246,7 +1289,7 @@ bool ReferenceChainTracker::runPass(jvmtiEnv *jvmti, JNIEnv *jni,
     ctx.tracker = this;
     ctx.frontier = _frontier;
     ctx.hop_cap = _hop_cap;
-    ctx.budget = _effective_budget;
+    ctx.budget = _first_pass_budget;
     ctx.edges_admitted = 0;
     ctx.truncated = false;
     ctx.frontier_cap_hit = false;
@@ -1320,7 +1363,15 @@ bool ReferenceChainTracker::runPass(jvmtiEnv *jvmti, JNIEnv *jni,
   store(_passes_run, load(_passes_run) + 1);
   _last_pass_gc_finish_epoch = gcFinishEpoch();
   store(_last_pass_ns, OS::nanotime());
-  updatePacing(pass_wall_ticks);
+  if (!was_first_pass) {
+    // The first pass spends _first_pass_budget, not _effective_budget - its
+    // duration is not a signal about the per-pass cost updatePacing() is
+    // trying to regulate (expandFrontier()'s cheap, per-node expansion
+    // calls), so feeding it in here would throttle _effective_budget down
+    // for every one of those unrelated later passes based on a single,
+    // deliberately oversized outlier.
+    updatePacing(pass_wall_ticks);
+  }
   // Search restart (this class's own header comment): accumulate this
   // pass's own cost toward the running total restartSearch() will spend into
   // _pain_budget once the search reaches a terminal state - same