Select rows and symbols in walkVM by attribution address

ddprof-lib/src/main/cpp/hotspot/hotspotSupport.h

@@ -25,9 +25,13 @@ class HotspotSupport {
     friend class HotspotSupportTestAccessor;
 
 private:
+    // pc_is_return_address describes the seed pc: a ucontext pc is the exact
+    // interrupted address, while callerPC() is a real return address on every
+    // architecture except the one where CALLER_PC_IS_RETURN_ADDRESS is false.
     static int walkVM(void* ucontext, ASGCT_CallFrame* frames, int max_depth,
                       StackWalkFeatures features, EventType event_type,
-                      const void* pc, uintptr_t sp, uintptr_t fp, int lock_index, bool* truncated);
+                      const void* pc, bool pc_is_return_address,
+                      uintptr_t sp, uintptr_t fp, int lock_index, bool* truncated);
     static int walkVM(void* ucontext, ASGCT_CallFrame* frames, int max_depth,
                       StackWalkFeatures features, EventType event_type,
                       int lock_index, bool* truncated = nullptr);

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

@@ -82,6 +82,15 @@ class Libraries {
     return _native_libs;
   }
 
+#ifdef UNIT_TEST
+  // Publishes a caller-owned CodeCache into the process-wide set so address
+  // lookups resolve against it, letting a test build a library whose symbol
+  // layout it controls. CodeCacheArray is append-only, so the cache has to
+  // outlive every later lookup. The real population path is updateSymbols()
+  // reading the loaded libraries; compiled only into gtest binaries.
+  bool addLibraryForTest(CodeCache *lib) { return _native_libs.add(lib); }
+#endif
+
   // Delete copy constructor and assignment operator to prevent copies
   Libraries(const Libraries&) = delete;
   Libraries& operator=(const Libraries&) = delete;

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

@@ -35,6 +35,7 @@
 #include "samplerPerf.h"
 #include "stackFrame.h"
 #include "stackWalker.h"
+#include "stackWalker.inline.h"
 #include "symbols.h"
 #include "threadLocalData.inline.h"
 #include "tsc.h"
@@ -388,14 +389,21 @@ void Profiler::populateRemoteFrame(ASGCT_CallFrame* frame, uintptr_t pc, CodeCac
  * - Checks marks after symbol resolution (same O(log n) + O(1) cost)
  * - If no symbol found but PC is in a known library, packs as
  *   BCI_NATIVE_FRAME_REMOTE for library-relative rendering ([lib+0xoffset])
+ *
+ * Range-based lookups (findLibraryByAddress, binarySearch) key off the
+ * attribution address, so a call that is the last instruction of its caller
+ * still selects the caller rather than whatever follows it. The emitted
+ * pc_offset keeps using the raw pc: it is the remote-symbolication wire value
+ * and its meaning is a cross-team contract, not a local lookup detail.
  */
-Profiler::NativeFrameResolution Profiler::resolveNativeFrameForWalkVM(uintptr_t pc, int lock_index) {
-  CodeCache* lib = _libs->findLibraryByAddress((void*)pc);
+Profiler::NativeFrameResolution Profiler::resolveNativeFrameForWalkVM(uintptr_t pc, bool pc_is_return_address, int lock_index) {
+  const void* lookup_pc = attributionPC((const void*)pc, pc_is_return_address);
+  CodeCache* lib = _libs->findLibraryByAddress(lookup_pc);
 
   if (_remote_symbolication && lib != nullptr && lib->hasBuildId()) {
     // Get symbol name and check mark
     const char *method_name = nullptr;
-    lib->binarySearch((void*)pc, &method_name);
+    lib->binarySearch(lookup_pc, &method_name);
     char mark = (method_name != nullptr) ? NativeFunc::read_mark(method_name) : 0;
 
     if (mark != 0) {
@@ -413,7 +421,7 @@ Profiler::NativeFrameResolution Profiler::resolveNativeFrameForWalkVM(uintptr_t
   // Traditional symbol resolution
   const char *method_name = nullptr;
   if (lib != nullptr) {
-    lib->binarySearch((void*)pc, &method_name);
+    lib->binarySearch(lookup_pc, &method_name);
   }
   if (method_name != nullptr) {
     char mark = NativeFunc::read_mark(method_name);

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

@@ -437,7 +437,7 @@ class alignas(alignof(SpinLock)) Profiler {
   };
 
   void populateRemoteFrame(ASGCT_CallFrame* frame, uintptr_t pc, CodeCache* lib, char mark);
-  NativeFrameResolution resolveNativeFrameForWalkVM(uintptr_t pc, int lock_index);
+  NativeFrameResolution resolveNativeFrameForWalkVM(uintptr_t pc, bool pc_is_return_address, int lock_index);
   int convertNativeTrace(int native_frames, const void **callchain,
                          ASGCT_CallFrame *frames, int lock_index,
                          bool skip_hook_prefix);

ddprof-lib/src/main/cpp/stackWalker.inline.h

@@ -82,6 +82,11 @@ class WalkPc {
     // The address to symbolize with or to select an unwind row with.
     const void* attribution() const { return attributionPC(_pc, _is_return_address); }
 
+    // For handing the pair to a callee that derives the attribution address on
+    // its own. Reading the flag is fine; it is assigning the pc without
+    // restating its nature that the mutators below exist to prevent.
+    bool isReturnAddress() const { return _is_return_address; }
+
     // A pc read out of a return-address slot, a link register, or a
     // DW_CFA_val_expression on the return-address register column.
     void setReturnAddress(const void* pc) {

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

@@ -0,0 +1,126 @@
+/*
+ * Copyright 2026, Datadog, Inc.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+// Gates the address split in Profiler::resolveNativeFrameForWalkVM, which
+// HotspotSupport::walkVM relies on for every non-Java frame: range lookups
+// (findLibraryByAddress/binarySearch) key off the attribution address so a
+// call that is the last instruction of its caller still resolves to the
+// caller, while the remote-symbolication pc_offset keeps deriving from the
+// raw pc because its meaning is a cross-team wire contract.
+//
+// These use synthetic CodeCaches rather than the test binary's own symbols,
+// so unlike returnAddressAttribution_ut.cpp they need no GNU-as/ELF-CFI asm
+// and no reliance on updateSymbols() parsing the main executable -- they run
+// on every platform this repo builds.
+
+#include <gtest/gtest.h>
+#include <cstdint>
+#include "codeCache.h"
+#include "libraries.h"
+#include "profiler.h"
+#include "vmEntry.h"
+
+namespace {
+
+// Far above anything the loader maps, so these never overlap a real library
+// in the process-wide set they are published into.
+const char* const kSymbolLibBase = (const char*)0x5a5a00000000ULL;
+const char* const kBareLibBase   = (const char*)0x5a5b00000000ULL;
+const size_t kLibSpan = 0x1000;
+
+// Zero-gap pair: `second` starts on the byte immediately after `first` ends,
+// which is exactly where a return address lands when the call is the last
+// instruction of `first`.
+const int kFirstOff  = 0x100;
+const int kFuncLen   = 0x10;
+const int kBoundary  = kFirstOff + kFuncLen;
+
+// No symbol covers this, so resolution falls through to the library-relative
+// packing path that carries pc_offset.
+const int kUnnamedOff = 0x800;
+
+class WalkVmAttributionTest : public ::testing::Test {
+  protected:
+    static void SetUpTestSuite() {
+        // CodeCacheArray is append-only and the set is process-wide, so both
+        // caches are published once and outlive every test in this binary.
+        static CodeCache symbol_lib("walkvm_attr_symbols", /*lib_index=*/-1,
+                                    kSymbolLibBase, kSymbolLibBase + kLibSpan,
+                                    /*image_base=*/kSymbolLibBase);
+        symbol_lib.add(kSymbolLibBase + kFirstOff, kFuncLen, "walkvm_attr_first");
+        symbol_lib.add(kSymbolLibBase + kBoundary, kFuncLen, "walkvm_attr_second");
+        symbol_lib.sort();
+
+        static CodeCache bare_lib("walkvm_attr_nosymbols", /*lib_index=*/-1,
+                                  kBareLibBase, kBareLibBase + kLibSpan,
+                                  /*image_base=*/kBareLibBase);
+
+        ASSERT_TRUE(Libraries::instance()->addLibraryForTest(&symbol_lib));
+        ASSERT_TRUE(Libraries::instance()->addLibraryForTest(&bare_lib));
+    }
+
+    static Profiler::NativeFrameResolution resolve(const char* pc, bool pc_is_ra) {
+        return Profiler::instance()->resolveNativeFrameForWalkVM(
+            (uintptr_t)pc, pc_is_ra, /*lock_index=*/0);
+    }
+};
+
+// The defect this fix addresses: at a zero-gap boundary the raw return
+// address names the following function, and only the attribution address
+// still names the caller that actually made the call.
+TEST_F(WalkVmAttributionTest, ReturnAddressAtZeroGapBoundaryResolvesToTheCaller) {
+    const char* boundary = kSymbolLibBase + kBoundary;
+
+    Profiler::NativeFrameResolution exact = resolve(boundary, /*pc_is_ra=*/false);
+    ASSERT_EQ(BCI_NATIVE_FRAME, exact.bci);
+    ASSERT_NE(nullptr, exact.method_name);
+    EXPECT_STREQ("walkvm_attr_second", exact.method_name)
+        << "an exact pc on the first byte of a function resolves to that function";
+
+    Profiler::NativeFrameResolution as_ra = resolve(boundary, /*pc_is_ra=*/true);
+    ASSERT_EQ(BCI_NATIVE_FRAME, as_ra.bci);
+    ASSERT_NE(nullptr, as_ra.method_name);
+    EXPECT_STREQ("walkvm_attr_first", as_ra.method_name)
+        << "the same address flagged as a return address must resolve to the caller, "
+        << "not to whatever happens to follow it";
+}
+
+// Inside a function the adjustment must be invisible -- it only ever moves the
+// lookup by one byte, so it may not reclassify a pc that is nowhere near a
+// boundary.
+TEST_F(WalkVmAttributionTest, AddressInsideAFunctionResolvesTheSameEitherWay) {
+    const char* inside = kSymbolLibBase + kFirstOff + kFuncLen / 2;
+
+    Profiler::NativeFrameResolution exact = resolve(inside, /*pc_is_ra=*/false);
+    Profiler::NativeFrameResolution as_ra = resolve(inside, /*pc_is_ra=*/true);
+
+    ASSERT_NE(nullptr, exact.method_name);
+    ASSERT_NE(nullptr, as_ra.method_name);
+    EXPECT_STREQ("walkvm_attr_first", exact.method_name);
+    EXPECT_STREQ("walkvm_attr_first", as_ra.method_name);
+}
+
+// The wire contract: the lookup may move, the emitted offset may not. This is
+// what keeps walkVM out of the unresolved cross-team question about what
+// pc_offset means to the backend symbolizer.
+TEST_F(WalkVmAttributionTest, PcOffsetStaysDerivedFromTheRawPc) {
+    const char* pc = kBareLibBase + kUnnamedOff;
+
+    Profiler::NativeFrameResolution exact = resolve(pc, /*pc_is_ra=*/false);
+    Profiler::NativeFrameResolution as_ra = resolve(pc, /*pc_is_ra=*/true);
+
+    ASSERT_EQ(BCI_NATIVE_FRAME_REMOTE, exact.bci);
+    ASSERT_EQ(BCI_NATIVE_FRAME_REMOTE, as_ra.bci);
+
+    uintptr_t exact_off = Profiler::RemoteFramePacker::unpackPcOffset(exact.packed_remote_frame);
+    uintptr_t as_ra_off = Profiler::RemoteFramePacker::unpackPcOffset(as_ra.packed_remote_frame);
+
+    EXPECT_EQ((uintptr_t)kUnnamedOff, exact_off);
+    EXPECT_EQ((uintptr_t)kUnnamedOff, as_ra_off)
+        << "flagging the pc as a return address must not shift the emitted offset -- "
+        << "the attribution address is a lookup detail and must not reach the wire";
+}
+
+}  // namespace

doc/reference/RemoteSymbolication.md

@@ -50,7 +50,7 @@ Modified frame collection to support dual modes:
 **Key Functions**:
 - `populateRemoteFrame()`: Packs pc_offset, mark, and lib_index into jmethodID field
 - `resolveNativeFrameForWalkVM()`: Resolves native frames for walkVM/walkVMX modes
-  - Performs binarySearch() to get symbol name
+  - Performs binarySearch() to get symbol name, keyed off the attribution address
   - Extracts mark via NativeFunc::read_mark() (O(1))
   - Packs data using RemoteFramePacker::pack()
 - `convertNativeTrace()`: Converts raw PCs to frames for walkFP/walkDwarf modes
@@ -65,7 +65,7 @@ Modified frame collection to support dual modes:
 
 **Stack Walker Integration**:
 - **walkFP/walkDwarf**: Return raw PCs → `convertNativeTrace()` → `populateRemoteFrame()`
-- **walkVM/walkVMX**: Directly call `resolveNativeFrameForWalkVM(pc, lock_index)` during stack walk (patched via gradle/patching.gradle)
+- **walkVM/walkVMX**: Directly call `resolveNativeFrameForWalkVM(pc, pc_is_return_address, lock_index)` during stack walk. `pc_is_return_address` says whether the walker took this pc out of a return-address slot: the symbol and library lookups then key off the attribution address (`pc - 1`), while the emitted `pc_offset` keeps deriving from the raw pc, so the wire value is unchanged.
 
 ### 5. **JFR Serialization** (`flightRecorder.cpp/h`)
 
@@ -109,7 +109,7 @@ Patches async-profiler's `stackWalker.h` and `stackWalker.cpp` to integrate remo
 
 **Implementation Patches (stackWalker.cpp)**:
 - Updates all `walkVM` signatures to accept and propagate `lock_index`
-- **Critical patch at line 454**: Replaces `profiler->findNativeMethod(pc)` with `profiler->resolveNativeFrameForWalkVM(pc, lock_index)`
+- **Critical patch at line 454**: Replaces `profiler->findNativeMethod(pc)` with `profiler->resolveNativeFrameForWalkVM(pc, pc_is_return_address, lock_index)`
 - Adds dynamic BCI selection (BCI_NATIVE_FRAME vs BCI_NATIVE_FRAME_REMOTE)
 - Adds `fillFrame()` overload for void* method_id to support both symbol names and RemoteFrameInfo pointers
 - Handles marked C++ interpreter frames (terminates scan if detected)