feat: report allocation stats in the near-OOM profile

bindings/profilers/near-oom.cc

@@ -16,6 +16,7 @@
 
 #include "near-oom.hh"
 
+#include "allocation-profile.hh"
 #include "defer.hh"
 #include "heap.hh"
 #include "per-isolate-data.hh"
@@ -259,6 +260,48 @@ static size_t ExtendedHeapLimit(size_t current_heap_limit, size_t extension) {
              : current_heap_limit + extension;
 }
 
+static void CaptureProfile(v8::Isolate* isolate,
+                           const std::shared_ptr<HeapProfilerState>& state,
+                           bool deliver_async_callback) {
+  // Release a superseded capture before asking v8 to allocate the next one.
+  state->ResetProfile();
+  std::unique_ptr<v8::AllocationProfile> profile{
+      isolate->GetHeapProfiler()->GetAllocationProfile()};
+  if (!profile) {
+    fprintf(stderr,
+            "NearHeapLimit: heap profiler is not enabled, no allocation "
+            "profile to report\n");
+    return;
+  }
+
+  state->profile = TranslateAllocationProfileToCpp(profile->GetRootNode());
+  if (state->allocations) {
+    state->profile_allocation_stats =
+        BuildAllocationStatsByNodeId(profile->GetSamples());
+  }
+  if (state->dumpProfileOnStderr) {
+    dumpAllocationProfile(stderr, state->profile.get());
+  }
+  if (!state->export_command.empty()) {
+    ExportProfile(*state);
+  }
+
+  if (state->callback.IsEmpty()) {
+    state->ResetProfile();
+    return;
+  }
+  if (state->callbackMode & kInterruptCallback) {
+    isolate->RequestInterrupt(InterruptCallback, nullptr);
+  }
+  if (state->callbackMode & kAsyncCallback) {
+    if (deliver_async_callback) {
+      InterruptCallback(isolate, nullptr);
+    } else {
+      uv_async_send(state->async);
+    }
+  }
+}
+
 size_t NearHeapLimit(void* data,
                      size_t current_heap_limit,
                      size_t initial_heap_limit) {
@@ -337,41 +380,11 @@ size_t NearHeapLimit(void* data,
               stats.object_count());
     }
   }
-  // GetAllocationProfile returns null when V8's sampling heap profiler isn't
-  // running, and that can happen while this callback is still installed:
-  // HeapProfilerCleanupHook stops V8's sampler without touching our state, so
-  // between that hook and the isolate actually going away we stay registered
-  // with nothing to sample. The heap-limit bookkeeping below still has to run,
-  // so skip only the profile-dependent work.
-  std::unique_ptr<v8::AllocationProfile> profile{
-      isolate->GetHeapProfiler()->GetAllocationProfile()};
-  if (profile) {
-    state->profile = TranslateAllocationProfileToCpp(profile->GetRootNode());
-    if (state->dumpProfileOnStderr) {
-      dumpAllocationProfile(stderr, state->profile.get());
-    }
-
-    if (!state->export_command.empty()) {
-      ExportProfile(*state);
-    }
-
-    if (!state->callback.IsEmpty()) {
-      if (state->callbackMode & kInterruptCallback) {
-        isolate->RequestInterrupt(InterruptCallback, nullptr);
-      }
-      if (state->callbackMode & kAsyncCallback) {
-        uv_async_send(state->async);
-      }
-    } else {
-      state->profile.reset();
-    }
+  // kSamplingForceGC needs the extension returned below to already be active.
+  if (state->allocations) {
+    uv_async_send(state->async);
   } else {
-    // Drop any profile retained from an earlier invocation: it is stale, and
-    // nothing below is going to consume or replace it.
-    state->profile.reset();
-    fprintf(stderr,
-            "NearHeapLimit: heap profiler is not enabled, no allocation "
-            "profile to report\n");
+    CaptureProfile(isolate, state, false);
   }
 
   if (!state->isMainThread) {
@@ -443,7 +456,7 @@ NAN_METHOD(HeapProfiler::MonitorOutOfMemory) {
 
   state->current_heap_extension_count = 0;
   state->automatic_heap_extension_size.reset();
-  state->profile.reset();
+  state->ResetProfile();
   state->export_command.clear();
   state->callback.Reset();
 
@@ -453,7 +466,6 @@ NAN_METHOD(HeapProfiler::MonitorOutOfMemory) {
   state->callbackMode = info[5].As<v8::Integer>()->Value();
   state->isMainThread = info[6].As<v8::Boolean>()->Value();
   state->automatic_heap_extension = info[7].As<v8::Boolean>()->Value();
-  state->InstallNearHeapLimitCallback();
   if (!info[4]->IsNullOrUndefined() && state->callbackMode != kNoCallback) {
     state->callback.Reset(Nan::To<v8::Function>(info[4]).ToLocalChecked());
   }
@@ -467,9 +479,11 @@ NAN_METHOD(HeapProfiler::MonitorOutOfMemory) {
     }
   }
 
-  if (!state->callback.IsEmpty() && (state->callbackMode & kAsyncCallback)) {
+  if (state->allocations ||
+      (!state->callback.IsEmpty() && (state->callbackMode & kAsyncCallback))) {
     state->RegisterAsyncCallback();
   }
+  state->InstallNearHeapLimitCallback();
 }
 
 void InterruptCallback(v8::Isolate* isolate, void* data) {
@@ -480,16 +494,28 @@ void InterruptCallback(v8::Isolate* isolate, void* data) {
   if (!state || !state->profile) {
     return;
   }
-  v8::Local<v8::Value> argv[1] = {
-      dd::TranslateAllocationProfile(state->profile.get())};
+  // Own the capture locally before translating: translation and the JS
+  // callback both allocate in the v8 heap and can re-enter NearHeapLimit,
+  // which overwrites these fields and would dangle or lose that capture.
+  auto profile = std::move(state->profile);
+  auto allocation_stats = std::move(state->profile_allocation_stats);
+  state->ResetProfile();
+
+  v8::Local<v8::Value> argv[1] = {dd::TranslateAllocationProfile(
+      profile.get(), allocation_stats ? &*allocation_stats : nullptr)};
   Nan::AsyncResource resource("NearHeapLimit");
   state->callback.Call(1, argv, &resource);
-  // Release the retained native profile once the callback has been invoked.
-  state->profile.reset();
 }
 
 void AsyncCallback(uv_async_t* handle) {
-  InterruptCallback(v8::Isolate::GetCurrent(), nullptr);
+  auto isolate = v8::Isolate::GetCurrent();
+  v8::HandleScope scope(isolate);
+  auto state = PerIsolateData::For(isolate)->GetHeapProfilerState();
+  if (state && state->allocations) {
+    CaptureProfile(isolate, state, true);
+    return;
+  }
+  InterruptCallback(isolate, nullptr);
 }
 
 }  // namespace dd

bindings/profilers/near-oom.hh

@@ -100,8 +100,13 @@ struct HeapProfilerState {
     uv_unref(reinterpret_cast<uv_handle_t*>(async));
   }
 
-  void OnNewProfile() {
+  void ResetProfile() {
     profile.reset();
+    profile_allocation_stats.reset();
+  }
+
+  void OnNewProfile() {
+    ResetProfile();
     // Only (re)install the NearHeapLimit callback when OOM monitoring is
     // configured. Otherwise a plain start()+profile() flow would silently
     // register a callback that the user never asked for.
@@ -120,6 +125,9 @@ struct HeapProfilerState {
   uint32_t current_heap_extension_count = 0;
   uv_async_t* async = nullptr;
   std::shared_ptr<Node> profile;
+  // Engaged iff |profile| was captured in allocation mode, so delivery reads
+  // the capture's own mode rather than |allocations|, which may have moved on.
+  std::optional<AllocationProfileNodeStatsMap> profile_allocation_stats;
   std::vector<std::string> export_command;
   bool allocations = false;
   bool dumpProfileOnStderr = false;

bindings/translate-heap-profile.cc

@@ -42,18 +42,13 @@ class HeapProfileTranslator : ProfileTranslator {
 
  public:
   v8::Local<v8::Value> TranslateAllocationProfile(
-      v8::AllocationProfile::Node* node) {
+      v8::AllocationProfile::Node* node,
+      const AllocationProfileNodeStatsMap* allocation_stats) {
     v8::Local<v8::Array> children = NewArray(node->children.size());
     for (size_t i = 0; i < node->children.size(); i++) {
-      Set(children, i, TranslateAllocationProfile(node->children[i]));
-    }
-
-    v8::Local<v8::Array> allocations = NewArray(node->allocations.size());
-    for (size_t i = 0; i < node->allocations.size(); i++) {
-      auto alloc = node->allocations[i];
-      Set(allocations,
+      Set(children,
           i,
-          CreateAllocation(NewNumber(alloc.count), NewNumber(alloc.size)));
+          TranslateAllocationProfile(node->children[i], allocation_stats));
     }
 
     return CreateNode(node->name,
@@ -62,61 +57,53 @@ class HeapProfileTranslator : ProfileTranslator {
                       NewInteger(node->line_number),
                       NewInteger(node->column_number),
                       children,
-                      allocations);
+                      TranslateAllocations(
+                          node->node_id, node->allocations, allocation_stats));
   }
 
   v8::Local<v8::Value> TranslateAllocationProfile(
-      v8::AllocationProfile::Node* node,
-      const AllocationProfileNodeStatsMap* allocation_stats) {
-    if (!allocation_stats) {
-      return TranslateAllocationProfile(node);
-    }
-
+      Node* node, const AllocationProfileNodeStatsMap* allocation_stats) {
     v8::Local<v8::Array> children = NewArray(node->children.size());
     for (size_t i = 0; i < node->children.size(); i++) {
       Set(children,
           i,
-          TranslateAllocationProfile(node->children[i], allocation_stats));
+          TranslateAllocationProfile(node->children[i].get(),
+                                     allocation_stats));
     }
 
-    auto node_stats = allocation_stats->find(node->node_id);
-    v8::Local<v8::Array> allocations = TranslateAllocationStats(
-        isolate,
-        node_stats == allocation_stats->end() ? nullptr : &node_stats->second);
-
-    return CreateNode(node->name,
-                      node->script_name,
+    return CreateNode(NewString(node->name.c_str()),
+                      NewString(node->script_name.c_str()),
                       NewInteger(node->script_id),
                       NewInteger(node->line_number),
                       NewInteger(node->column_number),
                       children,
-                      allocations);
+                      TranslateAllocations(
+                          node->node_id, node->allocations, allocation_stats));
   }
 
-  v8::Local<v8::Value> TranslateAllocationProfile(Node* node) {
-    v8::Local<v8::Array> children = NewArray(node->children.size());
-    for (size_t i = 0; i < node->children.size(); i++) {
-      Set(children, i, TranslateAllocationProfile(node->children[i].get()));
+ private:
+  v8::Local<v8::Array> TranslateAllocations(
+      uint32_t node_id,
+      const std::vector<v8::AllocationProfile::Allocation>& node_allocations,
+      const AllocationProfileNodeStatsMap* allocation_stats) {
+    if (allocation_stats) {
+      auto node_stats = allocation_stats->find(node_id);
+      return TranslateAllocationStats(isolate,
+                                      node_stats == allocation_stats->end()
+                                          ? nullptr
+                                          : &node_stats->second);
     }
 
-    v8::Local<v8::Array> allocations = NewArray(node->allocations.size());
-    for (size_t i = 0; i < node->allocations.size(); i++) {
-      auto alloc = node->allocations[i];
+    v8::Local<v8::Array> allocations = NewArray(node_allocations.size());
+    for (size_t i = 0; i < node_allocations.size(); i++) {
+      auto alloc = node_allocations[i];
       Set(allocations,
           i,
           CreateAllocation(NewNumber(alloc.count), NewNumber(alloc.size)));
     }
-
-    return CreateNode(NewString(node->name.c_str()),
-                      NewString(node->script_name.c_str()),
-                      NewInteger(node->script_id),
-                      NewInteger(node->line_number),
-                      NewInteger(node->column_number),
-                      children,
-                      allocations);
+    return allocations;
   }
 
- private:
   v8::Local<v8::Object> CreateNode(v8::Local<v8::String> name,
                                    v8::Local<v8::String> scriptName,
                                    v8::Local<v8::Integer> scriptId,
@@ -153,6 +140,7 @@ std::shared_ptr<Node> TranslateAllocationProfileToCpp(
   new_node->line_number = node->line_number;
   new_node->column_number = node->column_number;
   new_node->script_id = node->script_id;
+  new_node->node_id = node->node_id;
   Nan::Utf8String name(node->name);
   new_node->name.assign(*name, name.length());
   Nan::Utf8String script_name(node->script_name);
@@ -170,20 +158,17 @@ std::shared_ptr<Node> TranslateAllocationProfileToCpp(
   return new_node;
 }
 
-v8::Local<v8::Value> TranslateAllocationProfile(
-    v8::AllocationProfile::Node* node) {
-  return HeapProfileTranslator().TranslateAllocationProfile(node);
-}
-
 v8::Local<v8::Value> TranslateAllocationProfile(
     v8::AllocationProfile::Node* node,
     const AllocationProfileNodeStatsMap* allocation_stats) {
   return HeapProfileTranslator().TranslateAllocationProfile(node,
                                                             allocation_stats);
 }
 
-v8::Local<v8::Value> TranslateAllocationProfile(Node* node) {
-  return HeapProfileTranslator().TranslateAllocationProfile(node);
+v8::Local<v8::Value> TranslateAllocationProfile(
+    Node* node, const AllocationProfileNodeStatsMap* allocation_stats) {
+  return HeapProfileTranslator().TranslateAllocationProfile(node,
+                                                            allocation_stats);
 }
 
 }  // namespace dd

bindings/translate-heap-profile.hh

@@ -20,6 +20,7 @@
 
 #include <v8-profiler.h>
 #include <v8.h>
+#include <cstdint>
 #include <memory>
 #include <string>
 #include <vector>
@@ -33,18 +34,22 @@ struct Node {
   int line_number;
   int column_number;
   int script_id;
+  // Joins the retained AllocationProfileNodeStatsMap onto this tree.
+  uint32_t node_id = 0;
   std::vector<std::shared_ptr<Node>> children;
+  // v8 only decrements this on GC without the include-collected-objects flags:
+  // the live set in legacy mode, the cumulative allocated total with them.
   std::vector<Allocation> allocations;
 };
 
 std::shared_ptr<Node> TranslateAllocationProfileToCpp(
     v8::AllocationProfile::Node* node);
 
-v8::Local<v8::Value> TranslateAllocationProfile(Node* node);
 v8::Local<v8::Value> TranslateAllocationProfile(
-    v8::AllocationProfile::Node* node);
+    Node* node,
+    const AllocationProfileNodeStatsMap* allocation_stats = nullptr);
 v8::Local<v8::Value> TranslateAllocationProfile(
     v8::AllocationProfile::Node* node,
-    const AllocationProfileNodeStatsMap* allocation_stats);
+    const AllocationProfileNodeStatsMap* allocation_stats = nullptr);
 
 }  // namespace dd

ts/src/heap-profiler-bindings.ts

@@ -54,7 +54,9 @@ export function mapAllocationProfile<T>(
   return profiler.heapProfiler.mapAllocationProfile(callback);
 }
 
-export type NearHeapLimitCallback = (profile: AllocationProfileNode) => void;
+export type NearHeapLimitCallback = (
+  profile: AllocationProfileNode | AllocationProfileNodeWithStats,
+) => void;
 
 export function monitorOutOfMemory(
   heapLimitExtensionSize: number,

ts/src/heap-profiler.ts

@@ -298,7 +298,9 @@ export function monitorOutOfMemory(
   }
   let newCallback;
   if (typeof callback !== 'undefined') {
-    newCallback = (profile: AllocationProfileNode) => {
+    newCallback = (
+      profile: AllocationProfileNode | AllocationProfileNodeWithStats,
+    ) => {
       callback(convertProfile(profile));
     };
   }