v8: add setHeapSnapshotNearHeapLimit

Author: theanarkh

doc/api/v8.md

@@ -356,6 +356,21 @@ if (isMainThread) {
 }
 ```
 
+## `v8.setHeapSnapshotNearHeapLimit(limit)`
+
+<!-- YAML
+added: REPLACEME
+-->
+
+> Stability: 1 - Experimental
+
+The API is a no-op if `--heapsnapshot-near-heap-limit` is already set from the
+command line. And if `--heapsnapshot-near-heap-limit` is not set, the API can
+be called with the value greater than 0. After the limit is set, you can call
+this API with 0 to tell Node.js stop writing heap snapshot to disk.
+Other case is also a no-op. See [`--heapsnapshot-near-heap-limit`][] for
+more information.
+
 ## Serialization API
 
 The serialization API provides means of serializing JavaScript values in a way
@@ -1020,6 +1035,7 @@ Returns true if the Node.js instance is run to build a snapshot.
 [HTML structured clone algorithm]: https://developer.mozilla.org/en-US/docs/Web/API/Web_Workers_API/Structured_clone_algorithm
 [Hook Callbacks]: #hook-callbacks
 [V8]: https://developers.google.com/v8/
+[`--heapsnapshot-near-heap-limit`]: cli.md#--heapsnapshot-near-heap-limitmax_count
 [`AsyncLocalStorage`]: async_context.md#class-asynclocalstorage
 [`Buffer`]: buffer.md
 [`DefaultDeserializer`]: #class-v8defaultdeserializer

lib/v8.js

@@ -33,7 +33,7 @@ const {
 } = primordials;
 
 const { Buffer } = require('buffer');
-const { validateString } = require('internal/validators');
+const { validateString, validateNumber } = require('internal/validators');
 const {
   Serializer,
   Deserializer
@@ -59,7 +59,6 @@ const {
 } = internalBinding('heap_utils');
 const { HeapSnapshotStream } = require('internal/heap_utils');
 const promiseHooks = require('internal/promise_hooks');
-
 /**
  * Generates a snapshot of the current V8 heap
  * and writes it to a JSON file.
@@ -95,6 +94,7 @@ const {
   updateHeapStatisticsBuffer,
   updateHeapSpaceStatisticsBuffer,
   updateHeapCodeStatisticsBuffer,
+  setHeapSnapshotNearHeapLimit: _setHeapSnapshotNearHeapLimit,
 
   // Properties for heap statistics buffer extraction.
   kTotalHeapSizeIndex,
@@ -226,6 +226,11 @@ function getHeapCodeStatistics() {
   };
 }
 
+function setHeapSnapshotNearHeapLimit(limit) {
+  validateNumber(limit, 'limit', 0);
+  _setHeapSnapshotNearHeapLimit(limit);
+}
+
 /* V8 serialization API */
 
 /* JS methods for the base objects */
@@ -387,5 +392,6 @@ module.exports = {
   serialize,
   writeHeapSnapshot,
   promiseHooks,
-  startupSnapshot
+  startupSnapshot,
+  setHeapSnapshotNearHeapLimit,
 };

src/env.cc

@@ -720,6 +720,8 @@ Environment::Environment(IsolateData* isolate_data,
   inspector_host_port_ = std::make_shared<ExclusiveAccess<HostPort>>(
       options_->debug_options().host_port);
 
+  set_heap_snapshot_near_heap_limit(options_->heap_snapshot_near_heap_limit);
+
   if (!(flags_ & EnvironmentFlags::kOwnsProcessState)) {
     set_abort_on_uncaught_exception(false);
   }
@@ -834,7 +836,8 @@ Environment::~Environment() {
   // FreeEnvironment() should have set this.
   CHECK(is_stopping());
 
-  if (options_->heap_snapshot_near_heap_limit > heap_limit_snapshot_taken_) {
+  if (heapsnapshot_near_heap_limit_callback_added()) {
+    set_heapsnapshot_near_heap_limit_callback_added(false);
     isolate_->RemoveNearHeapLimitCallback(Environment::NearHeapLimitCallback,
                                           0);
   }
@@ -1952,6 +1955,7 @@ size_t Environment::NearHeapLimitCallback(void* data,
     Debug(env,
           DebugCategory::DIAGNOSTICS,
           "Not generating snapshots because it's too risky.\n");
+    env->set_heapsnapshot_near_heap_limit_callback_added(false);
     env->isolate()->RemoveNearHeapLimitCallback(NearHeapLimitCallback,
                                                 initial_heap_limit);
     // The new limit must be higher than current_heap_limit or V8 might
@@ -1973,16 +1977,17 @@ size_t Environment::NearHeapLimitCallback(void* data,
 
   // Remove the callback first in case it's triggered when generating
   // the snapshot.
+  env->set_heapsnapshot_near_heap_limit_callback_added(false);
   env->isolate()->RemoveNearHeapLimitCallback(NearHeapLimitCallback,
                                               initial_heap_limit);
 
   heap::WriteSnapshot(env, filename.c_str());
-  env->heap_limit_snapshot_taken_ += 1;
+  env->set_heap_limit_snapshot_taken(env->heap_limit_snapshot_taken() + 1);
 
   // Don't take more snapshots than the number specified by
   // --heapsnapshot-near-heap-limit.
-  if (env->heap_limit_snapshot_taken_ <
-      env->options_->heap_snapshot_near_heap_limit) {
+  if (env->heap_limit_snapshot_taken() < env->heap_snapshot_near_heap_limit()) {
+    env->set_heapsnapshot_near_heap_limit_callback_added(true);
     env->isolate()->AddNearHeapLimitCallback(NearHeapLimitCallback, env);
   }
 

src/env.h

@@ -1489,6 +1489,30 @@ class Environment : public MemoryRetainer {
   template <typename T>
   void ForEachBaseObject(T&& iterator);
 
+  inline int64_t heap_limit_snapshot_taken() {
+    return heap_limit_snapshot_taken_;
+  }
+
+  inline int64_t set_heap_limit_snapshot_taken(int64_t count) {
+    return heap_limit_snapshot_taken_ = count;
+  }
+
+  inline int64_t heap_snapshot_near_heap_limit() {
+    return heap_snapshot_near_heap_limit_;
+  }
+
+  inline void set_heap_snapshot_near_heap_limit(int64_t limit) {
+    heap_snapshot_near_heap_limit_ = limit;
+  }
+
+  inline bool heapsnapshot_near_heap_limit_callback_added() {
+    return near_heap_callback_is_added_;
+  }
+
+  inline void set_heapsnapshot_near_heap_limit_callback_added(bool added) {
+    near_heap_callback_is_added_ = added;
+  }
+
  private:
   inline void ThrowError(v8::Local<v8::Value> (*fun)(v8::Local<v8::String>),
                          const char* errmsg);
@@ -1547,6 +1571,8 @@ class Environment : public MemoryRetainer {
 
   bool is_processing_heap_limit_callback_ = false;
   int64_t heap_limit_snapshot_taken_ = 0;
+  int64_t heap_snapshot_near_heap_limit_ = 0;
+  bool near_heap_callback_is_added_ = false;
 
   uint32_t module_id_counter_ = 0;
   uint32_t script_id_counter_ = 0;

src/node.cc

@@ -273,7 +273,8 @@ static void AtomicsWaitCallback(Isolate::AtomicsWaitEvent event,
 void Environment::InitializeDiagnostics() {
   isolate_->GetHeapProfiler()->AddBuildEmbedderGraphCallback(
       Environment::BuildEmbedderGraph, this);
-  if (options_->heap_snapshot_near_heap_limit > 0) {
+  if (heap_snapshot_near_heap_limit() > 0) {
+    set_heapsnapshot_near_heap_limit_callback_added(true);
     isolate_->AddNearHeapLimitCallback(Environment::NearHeapLimitCallback,
                                        this);
   }

src/node_v8.cc

@@ -157,6 +157,30 @@ void CachedDataVersionTag(const FunctionCallbackInfo<Value>& args) {
   args.GetReturnValue().Set(result);
 }
 
+void SetHeapSnapshotNearHeapLimit(const FunctionCallbackInfo<Value>& args) {
+  CHECK(args[0]->IsNumber());
+  Environment* env = Environment::GetCurrent(args);
+  Isolate* const isolate = args.GetIsolate();
+  int64_t limit = args[0].As<v8::Number>()->Value();
+  CHECK_GE(limit, 0);
+  // Do not set twice which will make the process crash.
+  if (limit == 0) {
+    if (env->heapsnapshot_near_heap_limit_callback_added()) {
+      env->set_heapsnapshot_near_heap_limit_callback_added(false);
+      isolate->RemoveNearHeapLimitCallback(
+        Environment::NearHeapLimitCallback, 0);
+      env->set_heap_snapshot_near_heap_limit(limit);
+    }
+  } else {
+    if (!env->heapsnapshot_near_heap_limit_callback_added()) {
+      env->set_heapsnapshot_near_heap_limit_callback_added(true);
+      isolate->AddNearHeapLimitCallback(
+        Environment::NearHeapLimitCallback, env);
+      env->set_heap_snapshot_near_heap_limit(limit);
+    }
+  }
+}
+
 void UpdateHeapStatisticsBuffer(const FunctionCallbackInfo<Value>& args) {
   BindingData* data = Environment::GetBindingData<BindingData>(args);
   HeapStatistics s;
@@ -212,6 +236,10 @@ void Initialize(Local<Object> target,
 
   SetMethodNoSideEffect(
       context, target, "cachedDataVersionTag", CachedDataVersionTag);
+  SetMethodNoSideEffect(context,
+                        target,
+                        "setHeapSnapshotNearHeapLimit",
+                        SetHeapSnapshotNearHeapLimit);
   SetMethod(context,
             target,
             "updateHeapStatisticsBuffer",
@@ -267,6 +295,7 @@ void RegisterExternalReferences(ExternalReferenceRegistry* registry) {
   registry->Register(UpdateHeapCodeStatisticsBuffer);
   registry->Register(UpdateHeapSpaceStatisticsBuffer);
   registry->Register(SetFlagsFromString);
+  registry->Register(SetHeapSnapshotNearHeapLimit);
 }
 
 }  // namespace v8_utils

test/fixtures/workload/grow-and-set-near-heap-limit.js

@@ -0,0 +1,9 @@
+'use strict';
+const path = require('path');
+const v8 = require('v8');
+
+v8.setHeapSnapshotNearHeapLimit(+process.env.limit);
+if (process.env.limit2) {
+  v8.setHeapSnapshotNearHeapLimit(+process.env.limit2);
+}
+require(path.resolve(__dirname, 'grow.js'));

test/fixtures/workload/grow-worker-and-set-near-heap-limit.js

@@ -0,0 +1,15 @@
+'use strict';
+const path = require('path');
+const { Worker } = require('worker_threads');
+const max_snapshots = parseInt(process.env.TEST_SNAPSHOTS) || 1;
+new Worker(path.join(__dirname, 'grow-and-set-near-heap-limit.js'), {
+  env: {
+    ...process.env,
+    limit: max_snapshots,
+  },
+  resourceLimits: {
+    maxOldGenerationSizeMb:
+      parseInt(process.env.TEST_OLD_SPACE_SIZE) || 20
+  }
+});
+

test/parallel/test-heapsnapshot-near-heap-limit-by-api-in-worker.js

@@ -0,0 +1,41 @@
+// Copy from test-heapsnapshot-near-heap-limit-worker.js
+'use strict';
+
+require('../common');
+const tmpdir = require('../common/tmpdir');
+const assert = require('assert');
+const { spawnSync } = require('child_process');
+const fixtures = require('../common/fixtures');
+const fs = require('fs');
+
+const env = {
+  ...process.env,
+  NODE_DEBUG_NATIVE: 'diagnostics'
+};
+
+{
+  tmpdir.refresh();
+  const child = spawnSync(process.execPath, [
+    fixtures.path('workload', 'grow-worker-and-set-near-heap-limit.js'),
+  ], {
+    cwd: tmpdir.path,
+    env: {
+      TEST_SNAPSHOTS: 1,
+      TEST_OLD_SPACE_SIZE: 50,
+      ...env
+    }
+  });
+  console.log(child.stdout.toString());
+  const stderr = child.stderr.toString();
+  console.log(stderr);
+  const risky = /Not generating snapshots because it's too risky/.test(stderr);
+  if (!risky) {
+    // There should be one snapshot taken and then after the
+    // snapshot heap limit callback is popped, the OOM callback
+    // becomes effective.
+    assert(stderr.includes('ERR_WORKER_OUT_OF_MEMORY'));
+    const list = fs.readdirSync(tmpdir.path)
+      .filter((file) => file.endsWith('.heapsnapshot'));
+    assert.strictEqual(list.length, 1);
+  }
+}