heap: auto near oom headroom extension

bindings/profilers/heap.cc

@@ -23,8 +23,10 @@
 #include "translate-heap-profile.hh"
 
 #include <chrono>
+#include <limits>
 #include <memory>
 #include <mutex>
+#include <optional>
 #include <unordered_set>
 #include <vector>
 
@@ -132,6 +134,10 @@ struct HeapProfilerState {
 
   v8::Isolate* isolate = nullptr;
   uint32_t heap_extension_size = 0;
+  // When true, heap_extension_size is ignored in favour of one maximum-sized
+  // young generation, sampled once into automatic_heap_extension_size.
+  bool automatic_heap_extension = false;
+  std::optional<size_t> automatic_heap_extension_size;
   uint32_t max_heap_extension_count = 0;
   uint32_t current_heap_extension_count = 0;
   uv_async_t* async = nullptr;
@@ -364,6 +370,15 @@ static void ExportProfile(HeapProfilerState& state) {
   uv_fs_req_cleanup(&fs_req);
 }
 
+// V8 only raises the limit when the returned value is strictly greater than
+// current_heap_limit, and clamps it to its own allocator maximum, so
+// saturating is enough to stay well-defined in the extreme case.
+static size_t ExtendedHeapLimit(size_t current_heap_limit, size_t extension) {
+  return extension > std::numeric_limits<size_t>::max() - current_heap_limit
+             ? std::numeric_limits<size_t>::max()
+             : current_heap_limit + extension;
+}
+
 size_t NearHeapLimit(void* data,
                      size_t current_heap_limit,
                      size_t initial_heap_limit) {
@@ -385,14 +400,35 @@ size_t NearHeapLimit(void* data,
     return current_heap_limit;
   }
 
+  size_t extension = state->heap_extension_size;
+  if (state->automatic_heap_extension) {
+    if (!state->automatic_heap_extension_size.has_value()) {
+      // Grant at most one young generation, as Node.js does for its near-OOM
+      // heap snapshot callback. In current V8, heap_size_limit() is the
+      // old-generation limit this callback was handed plus the maximum
+      // young-generation size, so the delta is that young generation. It is
+      // fixed for the isolate, so sample it once and reuse it.
+      v8::HeapStatistics heap_statistics;
+      isolate->GetHeapStatistics(&heap_statistics);
+      const size_t total_heap_limit = heap_statistics.heap_size_limit();
+      // Only cache a usable sample: a degenerate one must not disable
+      // automatic sizing for the rest of the isolate's lifetime.
+      if (total_heap_limit > current_heap_limit) {
+        state->automatic_heap_extension_size =
+            total_heap_limit - current_heap_limit;
+      }
+    }
+    extension = state->automatic_heap_extension_size.value_or(0);
+  }
+
   if (state->insideCallback) {
-    // Reentrant call detected, try to increase heap limit a bit so that
-    // previous callback can proceed
-    const uint32_t default_heap_extension_size = 10 * 1024 * 1024;
-    auto extension_size = state->heap_extension_size
-                              ? state->heap_extension_size
-                              : default_heap_extension_size;
-    return current_heap_limit + extension_size;
+    // Reentrant call: GetAllocationProfile() allocated its way back into us.
+    // The in-progress capture still needs room to finish, so rescue it even
+    // when the caller asked for no top-level extension at all.
+    constexpr size_t kReentrantRescueExtension = 10 * 1024 * 1024;
+    return ExtendedHeapLimit(
+        current_heap_limit,
+        extension != 0 ? extension : kReentrantRescueExtension);
   }
   state->insideCallback = true;
   defer {
@@ -472,18 +508,15 @@ size_t NearHeapLimit(void* data,
     return current_heap_limit + kExtraHeapAllowance + 1;
   }
 
-  size_t new_heap_limit =
-      current_heap_limit +
-      ((state->current_heap_extension_count <= state->max_heap_extension_count)
-           ? state->heap_extension_size
-           : 0);
   if (state->current_heap_extension_count >= state->max_heap_extension_count) {
     // On Node 14, NearLimitCallback is sometimes called many times, without the
     // process aborting, even when returned limit is not increased. Disable
     // callback until next call to GetAllocationProfile()
     state->UninstallNearHeapLimitCallback();
   }
-  return new_heap_limit;
+  return state->current_heap_extension_count <= state->max_heap_extension_count
+             ? ExtendedHeapLimit(current_heap_limit, extension)
+             : current_heap_limit;
 }
 
 NAN_METHOD(HeapProfiler::StartSamplingHeapProfiler) {
@@ -646,8 +679,8 @@ NAN_METHOD(HeapProfiler::MapAllocationProfile) {
 }
 
 NAN_METHOD(HeapProfiler::MonitorOutOfMemory) {
-  if (info.Length() != 7) {
-    return Nan::ThrowTypeError("MonitorOOMCondition must have 7 arguments.");
+  if (info.Length() != 8) {
+    return Nan::ThrowTypeError("MonitorOOMCondition must have 8 arguments.");
   }
   if (!info[0]->IsUint32()) {
     return Nan::ThrowTypeError("Heap limit extension size must be a uint32.");
@@ -671,6 +704,10 @@ NAN_METHOD(HeapProfiler::MonitorOutOfMemory) {
   if (!info[6]->IsBoolean()) {
     return Nan::ThrowTypeError("IsMainThread must be a boolean.");
   }
+  if (!info[7]->IsBoolean()) {
+    return Nan::ThrowTypeError(
+        "AutomaticHeapLimitExtension must be a boolean.");
+  }
 
   auto isolate = v8::Isolate::GetCurrent();
 
@@ -684,6 +721,7 @@ NAN_METHOD(HeapProfiler::MonitorOutOfMemory) {
   }
 
   state->current_heap_extension_count = 0;
+  state->automatic_heap_extension_size.reset();
   state->profile.reset();
   state->export_command.clear();
   state->callback.Reset();
@@ -693,6 +731,7 @@ NAN_METHOD(HeapProfiler::MonitorOutOfMemory) {
   state->dumpProfileOnStderr = info[2].As<v8::Boolean>()->Value();
   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());

ts/src/heap-profiler-bindings.ts

@@ -64,6 +64,7 @@ export function monitorOutOfMemory(
   callback: NearHeapLimitCallback | undefined,
   callbackMode: number,
   isMainThread: boolean,
+  automaticHeapLimitExtension: boolean,
 ) {
   profiler.heapProfiler.monitorOutOfMemory(
     heapLimitExtensionSize,
@@ -73,5 +74,6 @@ export function monitorOutOfMemory(
     callback,
     callbackMode,
     isMainThread,
+    automaticHeapLimitExtension,
   );
 }

ts/src/heap-profiler.ts

@@ -242,25 +242,39 @@ export const CallbackMode = {
   Both: 3,
 };
 
+/**
+ * How much the heap limit is raised when v8 signals it is near the limit.
+ *
+ * A number is an exact byte count, and 0 means "grant no extension and let v8
+ * run its normal OOM handling". `'auto'` instead sizes the extension to one
+ * maximum young generation - the same budget Node.js grants its own near-OOM
+ * heap snapshot callback - which is what v8 actually needs to finish one more
+ * GC while the profile is captured.
+ */
+export type HeapLimitExtensionSize = number | 'auto';
+
 /**
  * Add monitoring for v8 heap, heap profiler must already be started.
  * When an out of heap memory event occurs:
- *  - an extension of heap memory of |heapLimitExtensionSize| bytes is
- *    requested to v8. This extension can occur |maxHeapLimitExtensionCount|
- *    number of times. If the extension amount is not enough to satisfy
- *    memory allocation that triggers GC and OOM, process will abort.
+ *  - the heap limit is extended by |heapLimitExtensionSize| so a profile can
+ *    be captured before the process dies. If the extension amount is not
+ *    enough to satisfy the memory allocation that triggers GC and OOM, the
+ *    process will abort, so prefer 'auto' over a hand-picked constant. This
+ *    top-level extension can occur |maxHeapLimitExtensionCount| times.
+ *    Reentrant rescue extensions used to finish an in-progress capture are
+ *    additional and are not included in that count.
  *  - heap profile is dumped as folded stacks on stderr if
  *    |dumpHeapProfileOnSdterr| is true
  *  - heap profile is dumped in temporary file and a new process is spawned
  *    with |exportCommand| arguments and profile path appended at the end.
- *  - |callback| is called. Callback can be invoked only if
- *    heapLimitExtensionSize is enough for the process to continue. Invocation
- *    will be done by a RequestInterrupt if |callbackMode| is Interrupt or Both,
- *    this might be unsafe since Isolate should not be reentered
- *    from RequestInterrupt, but this allows to interrupt synchronous code.
- *    Otherwise the callback is scheduled to be called asynchronously.
+ *  - |callback| is called. Callback can be invoked only if the extension is
+ *    enough for the process to continue. Invocation will be done by a
+ *    RequestInterrupt if |callbackMode| is Interrupt or Both, this might be
+ *    unsafe since Isolate should not be reentered from RequestInterrupt, but
+ *    this allows to interrupt synchronous code. Otherwise the callback is
+ *    scheduled to be called asynchronously.
  * @param heapLimitExtensionSize - amount of bytes heap should be expanded
- *  with upon OOM
+ *  with upon OOM, or 'auto' to size it to one maximum young generation
  * @param maxHeapLimitExtensionCount - maximum number of times heap size
  *  extension can occur
  * @param dumpHeapProfileOnSdterr - dump heap profile on stderr upon OOM
@@ -270,7 +284,7 @@ export const CallbackMode = {
  * @param callbackMode
  */
 export function monitorOutOfMemory(
-  heapLimitExtensionSize: number,
+  heapLimitExtensionSize: HeapLimitExtensionSize,
   maxHeapLimitExtensionCount: number,
   dumpHeapProfileOnSdterr: boolean,
   exportCommand?: Array<string>,
@@ -288,13 +302,15 @@ export function monitorOutOfMemory(
       callback(convertProfile(profile));
     };
   }
+  const automatic = heapLimitExtensionSize === 'auto';
   monitorOutOfMemoryImported(
-    heapLimitExtensionSize,
+    automatic ? 0 : heapLimitExtensionSize,
     maxHeapLimitExtensionCount,
     dumpHeapProfileOnSdterr,
     exportCommand || [],
     newCallback,
     typeof callbackMode !== 'undefined' ? callbackMode : CallbackMode.Async,
     isMainThread,
+    automatic,
   );
 }

ts/src/index.ts

@@ -26,6 +26,8 @@ export {
   LabelSet,
 } from './v8-types';
 
+export {HeapLimitExtensionSize} from './heap-profiler';
+
 export {encode, encodeSync} from './profile-encoder';
 export {SourceMapper} from './sourcemapper/sourcemapper';
 export {setLogger} from './logger';

ts/test/oom-heap-limit-extension.ts

@@ -0,0 +1,65 @@
+/*
+ * Copyright 2026 Datadog, Inc
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+'use strict';
+
+import * as v8 from 'v8';
+
+import {heap, HeapLimitExtensionSize} from '../src/index';
+
+const MB = 1024 * 1024;
+const CHUNK_SIZE = 4 * MB;
+const MAX_CHUNKS = 64;
+const heapLimitExtensionSize: HeapLimitExtensionSize =
+  process.argv[2] === 'auto' ? 'auto' : Number(process.argv[2] || 0);
+
+function heapLimit() {
+  return v8.getHeapStatistics().heap_size_limit;
+}
+
+heap.start(MB, 64);
+heap.monitorOutOfMemory(heapLimitExtensionSize, 1, false);
+
+const initialLimit = heapLimit();
+const retained: number[][] = [];
+let chunks = 0;
+
+// Report every heap limit the process observes so the parent can tell whether
+// a top-level extension was granted even if v8 aborts us mid-leak. A near-heap
+// limit event is not necessarily fatal - v8 may free enough and carry on - so
+// the limit is the only reliable signal here, not survival.
+console.log(`limit ${initialLimit}`);
+
+function leak() {
+  const limit = heapLimit();
+  if (limit !== initialLimit) {
+    console.log(`limit ${limit}`);
+    process.exit(0);
+  }
+  if (chunks >= MAX_CHUNKS) {
+    process.exit(0);
+  }
+
+  const chunk = new Array<number>(CHUNK_SIZE / 8);
+  for (let i = 0; i < chunk.length; i++) {
+    chunk[i] = i + 0.1;
+  }
+  retained.push(chunk);
+  chunks++;
+  setTimeout(leak, 5);
+}
+
+leak();

ts/test/oom-restore-heap-limit.ts

@@ -18,11 +18,13 @@
 
 import * as v8 from 'v8';
 
-import {heap} from '../src/index';
+import {heap, HeapLimitExtensionSize} from '../src/index';
 
 const MB = 1024 * 1024;
 const LIMIT_TOLERANCE = 16 * MB;
 const CHUNK_SIZE = 4 * MB;
+const heapLimitExtensionSize: HeapLimitExtensionSize =
+  process.argv[2] === 'auto' ? 'auto' : Number(process.argv[2] || 0);
 const gc = (global as typeof globalThis & {gc?: () => void}).gc;
 
 function heapLimit() {
@@ -44,7 +46,7 @@ async function main() {
 
   heap.start(MB, 64);
   try {
-    heap.monitorOutOfMemory(64 * MB, 1, false);
+    heap.monitorOutOfMemory(heapLimitExtensionSize, 1, false);
 
     const initialLimit = heapLimit();
     const retained: number[][] = [];