fix(codecache): make memoryUsage() accurate and live

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

@@ -9,13 +9,14 @@
 #include "os.h"
 #include "safeAccess.h"
 
+#include <cassert>
 #include <stdint.h>
 #include <stdlib.h>
 #include <string.h>
 #include <sys/mman.h>
 
 char *NativeFunc::create(const char *name, short lib_index) {
-  size_t size = align_up(sizeof(NativeFunc) + 1 + strlen(name), sizeof(NativeFunc*));
+  size_t size = allocSize(name);
   NativeFunc *f = (NativeFunc *)aligned_alloc(sizeof(NativeFunc*), size);
   f->_lib_index = lib_index;
   f->_mark = 0;
@@ -69,6 +70,8 @@ CodeCache::CodeCache(const char *name, short lib_index,
   _capacity = INITIAL_CODE_CACHE_CAPACITY;
   _count = 0;
   _blobs = new CodeBlob[_capacity];
+
+  _published.store(false, std::memory_order_relaxed);
 }
 
 void CodeCache::copyFrom(const CodeCache& other) {
@@ -113,6 +116,17 @@ void CodeCache::copyFrom(const CodeCache& other) {
   _count = other._count;
   _blobs = new CodeBlob[_capacity];
   memcpy(_blobs, other._blobs, _count * sizeof(CodeBlob));
+  // The memcpy above copied each CodeBlob's _name *pointer*, so without this the
+  // two caches would share name allocations and both destructors would free
+  // them (double-free / use-after-free). Give this copy its own name strings.
+  for (int i = 0; i < _count; i++) {
+    if (_blobs[i]._name != nullptr) {
+      _blobs[i]._name = NativeFunc::create(other._blobs[i]._name, _lib_index);
+    }
+  }
+
+  // A copy is a fresh, not-yet-registered cache.
+  _published.store(false, std::memory_order_relaxed);
 }
 
 CodeCache::CodeCache(const CodeCache &other) {
@@ -144,7 +158,43 @@ CodeCache::~CodeCache() {
   free(_build_id);  // Free build-id memory
 }
 
+long long CodeCache::memoryUsage() const {
+  // The blob array: _capacity entries of CodeBlob.
+  long long total = (long long)_capacity * sizeof(CodeBlob);
+
+  // This cache's own name, plus each symbol's name string. Each is a
+  // variable-length allocation whose size depends on the name length
+  // (see NativeFunc::allocSize). The lock-free read here is safe only for caches
+  // registered in a CodeCacheArray (Libraries::native_libs): their _blobs array
+  // and name pointers are fixed once published (add()/expand()/setDwarfTable()
+  // run only pre-publish — asserted via _published), which is the same read the
+  // symbolication fast path relies on. It must NOT be called on a continuously
+  // mutated, unpublished cache such as JitCodeCache::_runtime_stubs without
+  // holding JitCodeCache::_stubs_lock (shared), since a concurrent add()/expand()
+  // would free _blobs underneath the reader.
+  total += (long long)NativeFunc::allocSize(_name);
+  for (int i = 0; i < _count; i++) {
+    total += (long long)NativeFunc::allocSize(_blobs[i]._name);
+  }
+
+  // The DWARF unwind table, when present (length only — no pointer deref).
+  total += (long long)_dwarf_table_length * sizeof(FrameDesc);
+
+  // The build-id string is intentionally NOT counted here: the background
+  // library refresher (Libraries::updateBuildIds) frees and replaces _build_id
+  // on already-published caches under _build_id_lock, which dump does not hold,
+  // so dereferencing it here would race. It is negligible (~tens of bytes per
+  // library) next to the symbol tables, so excluding it costs no meaningful
+  // accuracy while keeping this read lock-free.
+
+  return total;
+}
+
 void CodeCache::expand() {
+  // Must not run after publication: memoryUsage() reads _blobs lock-free from
+  // the dump thread and a concurrent realloc would free it underneath.
+  assert(!_published.load(std::memory_order_acquire) &&
+         "expand() on a published CodeCache races memoryUsage()");
   CodeBlob *old_blobs = _blobs;
   CodeBlob *new_blobs = new CodeBlob[_capacity * 2];
 
@@ -157,6 +207,12 @@ void CodeCache::expand() {
 
 void CodeCache::add(const void *start, int length, const char *name,
                     bool update_bounds) {
+  // Symbols are added while parsing a library, before it is registered into a
+  // CodeCacheArray. Adding after publication would race memoryUsage() (which
+  // reads _blobs/_count lock-free at dump time). Unpublished standalone caches
+  // (e.g. _runtime_stubs) are exempt — they are never published.
+  assert(!_published.load(std::memory_order_acquire) &&
+         "add() on a published CodeCache races memoryUsage()");
   char *name_copy = NativeFunc::create(name, _lib_index);
   // Replace non-printable characters
   for (char *s = name_copy; *s != 0; s++) {
@@ -414,6 +470,10 @@ void CodeCache::makeImportsPatchable() {
 }
 
 void CodeCache::setDwarfTable(FrameDesc *table, int length, const FrameDesc &default_frame) {
+  // Set during library parsing, before publication. memoryUsage() reads
+  // _dwarf_table_length lock-free at dump time, so this must not run afterwards.
+  assert(!_published.load(std::memory_order_acquire) &&
+         "setDwarfTable() on a published CodeCache races memoryUsage()");
   _dwarf_table = table;
   _dwarf_table_length = length;
   _default_frame = &default_frame;

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

@@ -12,6 +12,7 @@
 #include "dwarf.h"
 #include "utils.h"
 
+#include <atomic>
 #include <jvmti.h>
 #include <stdlib.h>
 #include <string.h>
@@ -74,6 +75,16 @@ class NativeFunc {
   static char *create(const char *name, short lib_index);
   static void destroy(char *name);
 
+  // Size of the heap allocation backing a name string produced by create().
+  // This is the single source of truth for that size: create() allocates
+  // exactly allocSize(name) bytes, and memoryUsage() accounts for it. 0 if null.
+  static size_t allocSize(const char *name) {
+    if (name == nullptr) {
+      return 0;
+    }
+    return align_up(sizeof(NativeFunc) + 1 + strlen(name), sizeof(NativeFunc *));
+  }
+
   static short libIndex(const char *name) {
     if (name == nullptr) {
       return -1;
@@ -153,6 +164,18 @@ class CodeCache {
   int _count;
   CodeBlob *_blobs;
 
+  // Set once the cache is registered into a CodeCacheArray (see markPublished()).
+  // After that, memoryUsage() may be read lock-free from another thread (dump),
+  // so the mutators that touch the fields it reads (add()/expand()/
+  // setDwarfTable()) must not run — asserted in debug. Standalone caches that
+  // are never registered (e.g. JitCodeCache::_runtime_stubs, which is instead
+  // guarded by its own JitCodeCache::_stubs_lock) stay unpublished, so the
+  // asserts never fire for them. Atomic with release/acquire so the flag is
+  // observably true on any thread that reaches a published cache — the store is
+  // sequenced before CodeCacheArray::add()'s RELEASE publish of the pointer, and
+  // the assert loads it with acquire — without a formal data race.
+  std::atomic<bool> _published;
+
   void expand();
   void makeImportsPatchable();
   void saveImport(ImportId id, void** entry);
@@ -209,6 +232,11 @@ class CodeCache {
   void setBuildId(const char* build_id, size_t build_id_len);
   void setLoadBias(uintptr_t load_bias) { _load_bias = load_bias; }
 
+  // Mark this cache as published into a CodeCacheArray. Call before the array
+  // makes the pointer visible to readers; afterwards add()/expand()/
+  // setDwarfTable() must not be called (see _published).
+  void markPublished() { _published.store(true, std::memory_order_release); }
+
   void add(const void *start, int length, const char *name,
            bool update_bounds = false);
   void updateBounds(const void *start, const void *end);
@@ -251,9 +279,14 @@ class CodeCache {
   void setDwarfTable(FrameDesc *table, int length, const FrameDesc &default_frame = FrameDesc::default_frame);
   FrameDesc findFrameDesc(const void *pc);
 
-  long long memoryUsage() {
-    return _capacity * sizeof(CodeBlob *) + _count * sizeof(NativeFunc);
-  }
+  // Live size of what this CodeCache owns on the heap: the blob array, the
+  // per-symbol name strings (variable length), this cache's own name, and the
+  // DWARF unwind table. Recomputed on demand (called only at dump time), so it
+  // reflects the current contents. The build-id string is deliberately excluded
+  // — it is mutated by the background refresher and negligible in size (see the
+  // definition). Const and lock-free: reads only fields that are stable once the
+  // library is published.
+  long long memoryUsage() const;
 
   int count() { return _count; }
   CodeBlob* blob(int idx) {
@@ -266,11 +299,10 @@ class CodeCacheArray {
   CodeCache *_libs[MAX_NATIVE_LIBS];
   volatile int _reserved;       // next slot to reserve (CAS by writers)
   volatile int _count;          // published count (all indices < _count have non-NULL pointers)
-  volatile size_t _used_memory;
   bool _overflow_reported;
 
 public:
-  CodeCacheArray() : _reserved(0), _count(0), _used_memory(0), _overflow_reported(false) {
+  CodeCacheArray() : _reserved(0), _count(0), _overflow_reported(false) {
     memset(_libs, 0, MAX_NATIVE_LIBS * sizeof(CodeCache *));
   }
 
@@ -296,7 +328,10 @@ class CodeCacheArray {
     } while (!__atomic_compare_exchange_n(&_reserved, &slot, slot + 1,
                                           true, __ATOMIC_RELAXED, __ATOMIC_RELAXED));
     assert(__atomic_load_n(&_libs[slot], __ATOMIC_RELAXED) == nullptr);
-    __atomic_fetch_add(&_used_memory, lib->memoryUsage(), __ATOMIC_RELAXED);
+    // Mark published before the RELEASE store makes the pointer visible, so any
+    // later add()/expand()/setDwarfTable() on this cache trips the assert (its
+    // _blobs would then be read lock-free by memoryUsage() at dump time).
+    lib->markPublished();
     // Store pointer before publishing count. The RELEASE here pairs with
     // the ACQUIRE load in operator[]/at() and count().
     __atomic_store_n(&_libs[slot], lib, __ATOMIC_RELEASE);
@@ -316,8 +351,22 @@ class CodeCacheArray {
     return __atomic_load_n(&_libs[index], __ATOMIC_ACQUIRE);
   }
 
+  // Sum the live memory of all registered libraries. Recomputed on demand
+  // (called only at dump time) so it reflects each library's fully-populated
+  // state at that point — the accurate per-library formula rather than a value
+  // cached at add() time. (Libraries are populated before being registered and
+  // not grown afterwards.) The array is append-only, so iterating the published
+  // prefix is safe alongside concurrent add()s.
   size_t memoryUsage() const {
-    return __atomic_load_n(&_used_memory, __ATOMIC_RELAXED);
+    size_t total = 0;
+    int n = count();
+    for (int i = 0; i < n; i++) {
+      CodeCache *lib = at(i);
+      if (lib != nullptr) {
+        total += (size_t)lib->memoryUsage();
+      }
+    }
+    return total;
   }
 };
 

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

@@ -0,0 +1,101 @@
+/*
+ * Copyright 2026, Datadog, Inc.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include <gtest/gtest.h>
+#include "codeCache.h"
+
+// memoryUsage() must reflect the actual per-symbol name length, not a fixed
+// per-symbol constant: adding one symbol grows the reported usage by exactly
+// the allocation size of its name string.
+TEST(CodeCacheTest, MemoryUsageCountsSymbolNameLength) {
+    CodeCache cc("testlib");
+    long long base = cc.memoryUsage();
+
+    const char *name =
+        "a_long_symbol_name_well_beyond_sizeof_NativeFunc_padding_xxxxxxxxxxxx";
+    cc.add(reinterpret_cast<const void *>(0x1000), 16, name);
+
+    EXPECT_EQ(NativeFunc::allocSize(name), (size_t)(cc.memoryUsage() - base));
+}
+
+// A longer name costs more than a shorter one — the reported size scales with
+// the actual symbol-name length.
+TEST(CodeCacheTest, MemoryUsageGrowsWithNameLength) {
+    CodeCache shortc("lib");
+    CodeCache longc("lib");
+    long long short_base = shortc.memoryUsage();
+    long long long_base = longc.memoryUsage();
+
+    shortc.add(reinterpret_cast<const void *>(0x1000), 8, "f");
+    longc.add(reinterpret_cast<const void *>(0x1000), 8,
+              "an_extremely_long_symbol_name_padded_out_further_and_further_1234");
+
+    EXPECT_GT(longc.memoryUsage() - long_base, shortc.memoryUsage() - short_base);
+}
+
+// A freshly-constructed, empty cache reports exactly the blob array (counted at
+// the full CodeBlob size, not a pointer's worth) plus its own name — nothing
+// more. An exact check so over-counting or an inflated constant would fail.
+TEST(CodeCacheTest, MemoryUsageCountsFullBlobArray) {
+    CodeCache cc("lib");
+    EXPECT_EQ((long long)INITIAL_CODE_CACHE_CAPACITY * sizeof(CodeBlob)
+                  + (long long)NativeFunc::allocSize("lib"),
+              cc.memoryUsage());
+}
+
+// The DWARF unwind table contributes length * sizeof(FrameDesc). Populate it so
+// a regression (e.g. * flipped to /) in that term is caught; the table is
+// zero-length in the other tests.
+TEST(CodeCacheTest, MemoryUsageCountsDwarfTable) {
+    CodeCache cc("lib");
+    long long base = cc.memoryUsage();
+
+    const int entries = 7;
+    // setDwarfTable() takes ownership; ~CodeCache frees it. Contents are
+    // irrelevant here — memoryUsage() reads only the length.
+    FrameDesc *table = (FrameDesc *)malloc(entries * sizeof(FrameDesc));
+    cc.setDwarfTable(table, entries);
+
+    EXPECT_EQ((long long)entries * sizeof(FrameDesc), cc.memoryUsage() - base);
+}
+
+// Copying a CodeCache must deep-copy each blob's name string. With the previous
+// shallow memcpy the copy shared the originals' _name pointers, so destroying
+// both double-freed them. Here both copies are populated, sized equally, and
+// destroyed — a double-free would abort (or trip ASan).
+TEST(CodeCacheTest, CopyIsDeepAndDoesNotDoubleFree) {
+    CodeCache src("lib");
+    src.add(reinterpret_cast<const void *>(0x1000), 16, "sym_one");
+    src.add(reinterpret_cast<const void *>(0x2000), 16, "sym_two_longer_name");
+    long long expected = src.memoryUsage();
+
+    {
+        CodeCache copy(src);
+        EXPECT_EQ(expected, copy.memoryUsage());
+    }  // copy destroyed — must not free src's name strings
+
+    EXPECT_EQ(expected, src.memoryUsage());  // src's names still intact
+}  // src destroyed — no double free
+
+// CodeCacheArray::memoryUsage() is the aggregation production actually consumes
+// (Profiler::dump()). Verify it sums the per-library values — guards the loop
+// bound and the accumulation against regressions the per-instance tests miss.
+TEST(CodeCacheTest, ArrayMemoryUsageSumsLibraries) {
+    CodeCacheArray arr;
+    CodeCache *a = new CodeCache("liba");
+    CodeCache *b = new CodeCache("libb");
+    // Populate before registering: add() must run pre-publication.
+    a->add(reinterpret_cast<const void *>(0x1000), 16, "alpha_symbol");
+    b->add(reinterpret_cast<const void *>(0x2000), 16, "beta_symbol_longer_name");
+
+    long long expected = (long long)a->memoryUsage() + (long long)b->memoryUsage();
+    ASSERT_TRUE(arr.add(a));
+    ASSERT_TRUE(arr.add(b));
+    EXPECT_EQ(expected, (long long)arr.memoryUsage());
+
+    // CodeCacheArray does not own the entries.
+    delete a;
+    delete b;
+}