feat(nativemem): instrument remaining allocation sites

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

@@ -470,7 +470,7 @@ u64 CallTraceHashTable::put(int num_frames, ASGCT_CallFrame *frames,
   }
 }
 
-void CallTraceHashTable::collect(std::unordered_set<CallTrace *> &traces, std::function<void(CallTrace*)> trace_hook) {
+void CallTraceHashTable::collect(CallTraceSet &traces, std::function<void(CallTrace*)> trace_hook) {
   // Lock-free collection for read-only tables.
   // Use ACQUIRE to pair with the ACQ_REL CAS in put()'s expansion path and the
   // RELEASE store in clearTableOnly(); ensures we see the fully-initialised table

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

@@ -7,7 +7,9 @@
 #define _CALLTRACEHASHTABLE_H
 
 #include "arch.h"
+#include "countingAllocator.h"
 #include "linearAllocator.h"
+#include "nativeMem.h"
 #include "vmEntry.h"
 #include <unordered_set>
 #include <atomic>
@@ -21,11 +23,21 @@ struct CallTrace {
   u64 trace_id;  // 64-bit for JFR constant pool compatibility
   ASGCT_CallFrame frames[1];
 
-  CallTrace(bool truncated, int num_frames, u64 trace_id) 
+  CallTrace(bool truncated, int num_frames, u64 trace_id)
     : truncated(truncated), num_frames(num_frames), trace_id(trace_id) {
   }
 };
 
+// Traces collected for JFR/liveness processing. Uses CountingAllocator so the
+// real per-node heap cost is attributed to NM_CALLTRACE, the same category the
+// rest of the call-trace arena already accounts into.
+using CallTraceSet =
+    std::unordered_set<CallTrace *, std::hash<CallTrace *>, std::equal_to<CallTrace *>,
+                        CountingAllocator<CallTrace *, NM_CALLTRACE>>;
+using CallTraceIdSet =
+    std::unordered_set<u64, std::hash<u64>, std::equal_to<u64>,
+                        CountingAllocator<u64, NM_CALLTRACE>>;
+
 struct CallTraceSample {
   CallTrace *trace;
 
@@ -122,7 +134,7 @@ class CallTraceHashTable {
    */
   ChunkList clearTableOnly();
 
-  void collect(std::unordered_set<CallTrace *> &traces, std::function<void(CallTrace*)> trace_hook = nullptr);
+  void collect(CallTraceSet &traces, std::function<void(CallTrace*)> trace_hook = nullptr);
 
   u64 put(int num_frames, ASGCT_CallFrame *frames, bool truncated, u64 weight);
   void putWithExistingId(CallTrace* trace, u64 weight);  // For standby tables with no contention

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

@@ -166,7 +166,7 @@ u64 CallTraceStorage::put(int num_frames, ASGCT_CallFrame* frames, bool truncate
  * This function is safe to call concurrently with put() operations.
  * It is not designed to be called concurrently with itself.
  */
-void CallTraceStorage::processTraces(std::function<void(const std::unordered_set<CallTrace*>&)> processor) {
+void CallTraceStorage::processTraces(std::function<void(const CallTraceSet&)> processor) {
     // PHASE 1: Collect liveness information with simple lock (rare operation)
     {
         SharedLockGuard lock(&_liveness_lock);

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

@@ -26,7 +26,7 @@ class CallTraceHashTable;
 // Liveness checker function type
 // Fills the provided set with 64-bit call_trace_id values that should be preserved
 // Using reference parameter avoids malloc() for vector creation and copying
-typedef std::function<void(std::unordered_set<u64>&)> LivenessChecker;
+typedef std::function<void(CallTraceIdSet&)> LivenessChecker;
 
 class CallTraceStorage {
 public:
@@ -65,8 +65,8 @@ class CallTraceStorage {
 
     // Pre-allocated collections for processTraces (single-threaded operation)
     // These collections are reused to eliminate malloc/free cycles
-    std::unordered_set<CallTrace*> _traces_buffer;           // All traces for JFR processing
-    std::unordered_set<u64> _preserve_set_buffer;           // Preserve set for current cycle
+    CallTraceSet _traces_buffer;           // All traces for JFR processing
+    CallTraceIdSet _preserve_set_buffer;           // Preserve set for current cycle
 
 public:
     CallTraceStorage();
@@ -85,7 +85,7 @@ class CallTraceStorage {
     // Lock-free trace processing with RefCountGuard protection
     // The callback receives traces that are guaranteed to be valid during execution
     // Uses atomic table swapping with grace period for safe memory reclamation
-    void processTraces(std::function<void(const std::unordered_set<CallTrace*>&)> processor);
+    void processTraces(std::function<void(const CallTraceSet&)> processor);
 
     // Enhanced clear with liveness preservation (rarely called - uses atomic operations)
     void clear();

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

@@ -474,6 +474,16 @@ void CodeCache::setDwarfTable(FrameDesc *table, int length, const FrameDesc &def
   // _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()");
+  // The parser (SFrameParser/DwarfParser) builds this table with capacity-doubling
+  // malloc/realloc, so the incoming allocation is typically larger than length
+  // entries. Shrink it to the exact size so memoryUsage()'s length-based formula
+  // isn't an undercount, and so the doubling slack doesn't linger in real RSS.
+  if (length > 0) {
+    FrameDesc *trimmed = (FrameDesc *)realloc(table, length * sizeof(FrameDesc));
+    if (trimmed != NULL) {
+      table = trimmed;
+    }
+  }
   _dwarf_table = table;
   _dwarf_table_length = length;
   _default_frame = &default_frame;

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

@@ -0,0 +1,55 @@
+/*
+ * Copyright 2026, Datadog, Inc.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+#ifndef _COUNTINGALLOCATOR_H
+#define _COUNTINGALLOCATOR_H
+
+#include "nativeMem.h"
+#include <cstddef>
+#include <new>
+
+// A stateless, C++11-Allocator-conformant wrapper around ::operator new /
+// ::operator delete that records every allocation/deallocation into the given
+// NativeMem category. STL containers rebind the supplied Allocator<value_type>
+// to their actual node type before calling allocate(), so this yields the
+// exact real per-node byte count the implementation uses -- not an estimate.
+template <typename T, NativeMemCategory Cat>
+class CountingAllocator {
+public:
+  using value_type = T;
+
+  CountingAllocator() noexcept = default;
+  template <typename U>
+  CountingAllocator(const CountingAllocator<U, Cat> &) noexcept {}
+
+  T *allocate(std::size_t n) {
+    T *p = static_cast<T *>(::operator new(n * sizeof(T)));
+    NativeMem::record(Cat, (long long)(n * sizeof(T)));
+    return p;
+  }
+
+  void deallocate(T *p, std::size_t n) noexcept {
+    NativeMem::record(Cat, -(long long)(n * sizeof(T)));
+    ::operator delete(p);
+  }
+
+  template <typename U>
+  struct rebind {
+    using other = CountingAllocator<U, Cat>;
+  };
+};
+
+template <typename T, NativeMemCategory Cat>
+inline bool operator==(const CountingAllocator<T, Cat> &,
+                        const CountingAllocator<T, Cat> &) {
+  return true;
+}
+
+template <typename T, NativeMemCategory Cat>
+inline bool operator!=(const CountingAllocator<T, Cat> &,
+                        const CountingAllocator<T, Cat> &) {
+  return false;
+}
+
+#endif // _COUNTINGALLOCATOR_H

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

@@ -1615,15 +1615,14 @@ int Recording::writeStackTraces(Buffer *buf, Lookup *lookup) {
   // via processCallTraces, but no T_STACK_TRACE section is emitted in that case.
   int trace_count = 0;
   // Use safe trace processing with guaranteed lifetime during callback execution
-  Profiler::instance()->processCallTraces([this, buf, lookup, &trace_count](const std::unordered_set<CallTrace*>& traces) {
+  Profiler::instance()->processCallTraces([this, buf, lookup, &trace_count](const CallTraceSet& traces) {
     if (traces.empty()) {
       return;
     }
     trace_count = (int)traces.size();
     buf->putVar64(T_STACK_TRACE);
     buf->putVar64(traces.size());
-    for (std::unordered_set<CallTrace *>::const_iterator it = traces.begin();
-         it != traces.end(); ++it) {
+    for (auto it = traces.begin(); it != traces.end(); ++it) {
       CallTrace *trace = *it;
       buf->putVar64(trace->trace_id);
       if (trace->num_frames > 0) {

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

@@ -7,6 +7,7 @@
 #ifndef _FLIGHTRECORDER_H
 #define _FLIGHTRECORDER_H
 
+#include <functional>
 #include <map>
 #include <unordered_map>
 #include <unordered_set>
@@ -18,6 +19,7 @@
 #include "arch.h"
 #include "arguments.h"
 #include "buffers.h"
+#include "countingAllocator.h"
 #include "counters.h"
 #include "dictionary.h"
 #include "stringDictionary.h"
@@ -106,7 +108,10 @@ class MethodInfo {
 //   10 - void* address (native frame names)
 //   01 - RemoteFrameInfo (packed remote symbolication)
 //   11 - vtable_receiver class_id (BCI_VTABLE_RECEIVER frames)
-class MethodMap : public std::map<unsigned long, MethodInfo> {
+class MethodMap
+    : public std::map<unsigned long, MethodInfo, std::less<unsigned long>,
+                       CountingAllocator<std::pair<const unsigned long, MethodInfo>,
+                                          NM_METHOD_MAP>> {
 public:
   static constexpr unsigned long ADDRESS_MARK = 0x8000000000000000ULL;
   static constexpr unsigned long REMOTE_FRAME_MARK = 0x4000000000000000ULL;
@@ -352,7 +357,9 @@ class Lookup {
   // as the MethodMap key, so distinct Symbol* addresses for the same
   // class name (class unload/reload mid-chunk) collapse to a single
   // MethodInfo row.
-  std::unordered_map<void*, u32> _vtable_receiver_cache;
+  std::unordered_map<void*, u32, std::hash<void*>, std::equal_to<void*>,
+                      CountingAllocator<std::pair<void* const, u32>, NM_JFR_BUFFERS>>
+      _vtable_receiver_cache;
   Dictionary _packages;
   Dictionary _symbols;
 

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

@@ -16,6 +16,7 @@
 #include "jniHelper.h"
 #include "livenessTracker.h"
 #include "log.h"
+#include "nativeMem.h"
 #include "os.h"
 #include "profiler.h"
 #include "threadLocalData.h"
@@ -193,7 +194,7 @@ Error LivenessTracker::start(Arguments &args) {
   }
 
   // Self-register with the profiler for liveness checking
-  Profiler::instance()->registerLivenessChecker([this](std::unordered_set<u64>& buffer) {
+  Profiler::instance()->registerLivenessChecker([this](CallTraceIdSet& buffer) {
     this->getLiveTraceIds(buffer);
   });
 
@@ -277,6 +278,9 @@ Error LivenessTracker::initialize(Arguments &args) {
       std::min(2048, _table_max_cap); // with default 512k sampling interval, it's
                                    // enough for 1G of heap
   _table = (TrackingEntry *)malloc(sizeof(TrackingEntry) * _table_cap);
+  if (_table != NULL) {
+    NativeMem::record(NM_LIVENESS, (long long)sizeof(TrackingEntry) * _table_cap);
+  }
 
   _gc_epoch = 0;
   _last_gc_epoch = 0;
@@ -406,6 +410,8 @@ void LivenessTracker::track(JNIEnv *env, AllocEvent &event, jint tid,
           TrackingEntry *tmp = (TrackingEntry *)realloc(
                 _table, sizeof(TrackingEntry) * newcap);
           if (tmp != nullptr) {
+              NativeMem::record(NM_LIVENESS,
+                  (long long)sizeof(TrackingEntry) * (newcap - _table_cap));
               _table = tmp;
               _table_cap = newcap;
               Log::debug(
@@ -450,7 +456,7 @@ void LivenessTracker::onGC() {
   }
 }
 
-void LivenessTracker::getLiveTraceIds(std::unordered_set<u64>& out_buffer) {
+void LivenessTracker::getLiveTraceIds(CallTraceIdSet& out_buffer) {
   out_buffer.clear();
 
   if (!_enabled || !_initialized) {

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

@@ -7,6 +7,7 @@
 #define _LIVENESSTRACKER_H
 
 #include "arch.h"
+#include "callTraceHashTable.h"
 #include "context.h"
 #include "engine.h"
 #include "event.h"
@@ -104,7 +105,7 @@ class alignas(alignof(SpinLock)) LivenessTracker {
   static void JNICALL GarbageCollectionFinish(jvmtiEnv *jvmti_env);
 
 private:
-  void getLiveTraceIds(std::unordered_set<u64>& out_buffer);
+  void getLiveTraceIds(CallTraceIdSet& out_buffer);
 };
 
 #endif // _LIVENESSTRACKER_H

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

@@ -26,6 +26,8 @@
   X(PERF, "perf")                                                              \
   X(THREAD_LOCAL, "thread_local")                                              \
   X(JFR_BUFFERS, "jfr_buffers")                                                \
+  X(METHOD_MAP, "method_map")                                                  \
+  X(LIVENESS, "liveness")                                                      \
   X(MISC, "misc")
 
 #define X_NM_ENUM(a, b) NM_##a,

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

@@ -304,12 +304,12 @@ class alignas(alignof(SpinLock)) Profiler {
 
   const char* cstack() const;
   int lookupClass(const char *key, size_t length);
-  void processCallTraces(std::function<void(const std::unordered_set<CallTrace*>&)> processor) {
+  void processCallTraces(std::function<void(const CallTraceSet&)> processor) {
     if (!_omit_stacktraces) {
       _call_trace_storage.processTraces(processor);
     } else {
       // If stack traces are omitted, call processor with empty set
-      static std::unordered_set<CallTrace*> empty_traces;
+      static CallTraceSet empty_traces;
       processor(empty_traces);
     }
   }

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

@@ -42,15 +42,15 @@ void ThreadInfo::clearAll(std::set<int> &live_thread_ids) {
     _thread_ids.clear();
   } else {
     // we need to honor the thread referenced from the liveness tracker
-    std::map<int, std::string>::iterator name_itr = _thread_names.begin();
+    auto name_itr = _thread_names.begin();
     while (name_itr != _thread_names.end()) {
       if (live_thread_ids.find(name_itr->first) == live_thread_ids.end()) {
         name_itr = _thread_names.erase(name_itr);
       } else {
         ++name_itr;
       }
     }
-    std::map<int, u64>::iterator id_itr = _thread_ids.begin();
+    auto id_itr = _thread_ids.begin();
     while (id_itr != _thread_ids.end()) {
       if (live_thread_ids.find(id_itr->first) == live_thread_ids.end()) {
         id_itr = _thread_ids.erase(id_itr);

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

@@ -1,4 +1,6 @@
+#include "countingAllocator.h"
 #include "mutex.h"
+#include "nativeMem.h"
 #include "os.h"
 #include <functional>
 #include <map>
@@ -8,9 +10,16 @@
 
 class ThreadInfo {
 private:
+  using ThreadNamesMap =
+      std::map<int, std::string, std::less<int>,
+               CountingAllocator<std::pair<const int, std::string>, NM_THREAD_LOCAL>>;
+  using ThreadIdsMap =
+      std::map<int, u64, std::less<int>,
+               CountingAllocator<std::pair<const int, u64>, NM_THREAD_LOCAL>>;
+
   Mutex _ti_lock;
-  std::map<int, std::string> _thread_names;
-  std::map<int, u64> _thread_ids;
+  ThreadNamesMap _thread_names;
+  ThreadIdsMap _thread_ids;
 
 public:
   // disallow copy and assign to avoid issues with the mutex

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

@@ -2,6 +2,7 @@
 #define STUB_UNWIND_STATS_H
 
 #include "common.h"
+#include "nativeMem.h"
 #include "spinLock.h"
 
 #include <stddef.h>
@@ -30,9 +31,15 @@ class UnwindFailures {
      _counters = new u64[MAX_UNWIND_FAILURE_NAMES][UNWIND_FAILURE_ANY + 1];
      memset((void*)_names, 0, MAX_UNWIND_FAILURE_NAMES * MAX_NAME_LENGTH);
      memset((void*)_counters, 0, MAX_UNWIND_FAILURE_NAMES * (UNWIND_FAILURE_ANY + 1) * sizeof(u64));
+     NativeMem::record(NM_THREAD_LOCAL,
+         (long long)(MAX_UNWIND_FAILURE_NAMES * MAX_NAME_LENGTH) +
+         (long long)(MAX_UNWIND_FAILURE_NAMES * (UNWIND_FAILURE_ANY + 1) * sizeof(u64)));
    }
 
    ~UnwindFailures() {
+     NativeMem::record(NM_THREAD_LOCAL,
+         -((long long)(MAX_UNWIND_FAILURE_NAMES * MAX_NAME_LENGTH) +
+           (long long)(MAX_UNWIND_FAILURE_NAMES * (UNWIND_FAILURE_ANY + 1) * sizeof(u64))));
      delete[] _names;
      delete[] _counters;
    }

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

@@ -9,6 +9,7 @@
 
 #include <cassert>
 #include "engine.h"
+#include "nativeMem.h"
 #include "os.h"
 #include "profiler.h"
 #include "reservoirSampler.h"
@@ -79,6 +80,7 @@ class BaseWallClock : public Engine {
 
       while (_running.load(std::memory_order_relaxed)) {
         collectThreads(threads);
+        NativeMem::setLive(NM_MISC, (long long)threads.capacity() * sizeof(ThreadType));
 
         int num_failures = 0;
         int threads_already_exited = 0;