v8: add setHeapSnapshotNearHeapLimit

PR-URL: https://github.com/nodejs/node/pull/44420
Refs: https://github.com/nodejs/node/pull/33010
Reviewed-By: Joyee Cheung <joyeec9h3@gmail.com>
Reviewed-By: Chengzhong Wu <legendecas@gmail.com>

5 files changed, 52 insertions, 14 deletions

diff --git a/src/env-inl.h b/src/env-inl.h
index b78fd1dfd3d..58e4540d7d4 100644
--- a/src/env-inl.h
+++ b/src/env-inl.h
@@ -892,6 +892,24 @@ Realm* Environment::principal_realm() const {
   return principal_realm_.get();
 }
 
+inline void Environment::set_heap_snapshot_near_heap_limit(uint32_t limit) {
+  heap_snapshot_near_heap_limit_ = limit;
+}
+
+inline void Environment::AddHeapSnapshotNearHeapLimitCallback() {
+  DCHECK(!heapsnapshot_near_heap_limit_callback_added_);
+  heapsnapshot_near_heap_limit_callback_added_ = true;
+  isolate_->AddNearHeapLimitCallback(Environment::NearHeapLimitCallback, this);
+}
+
+inline void Environment::RemoveHeapSnapshotNearHeapLimitCallback(
+    size_t heap_limit) {
+  DCHECK(heapsnapshot_near_heap_limit_callback_added_);
+  heapsnapshot_near_heap_limit_callback_added_ = false;
+  isolate_->RemoveNearHeapLimitCallback(Environment::NearHeapLimitCallback,
+                                        heap_limit);
+}
+
 }  // namespace node
 
 // These two files depend on each other. Including base_object-inl.h after this
diff --git a/src/env.cc b/src/env.cc
index 62d3e66a483..6d129c3e932 100644
--- a/src/env.cc
+++ b/src/env.cc
@@ -682,6 +682,9 @@ Environment::Environment(IsolateData* isolate_data,
   inspector_host_port_ = std::make_shared<ExclusiveAccess<HostPort>>(
       options_->debug_options().host_port);
 
+  heap_snapshot_near_heap_limit_ =
+      static_cast<uint32_t>(options_->heap_snapshot_near_heap_limit);
+
   if (!(flags_ & EnvironmentFlags::kOwnsProcessState)) {
     set_abort_on_uncaught_exception(false);
   }
@@ -797,9 +800,8 @@ Environment::~Environment() {
   // FreeEnvironment() should have set this.
   CHECK(is_stopping());
 
-  if (options_->heap_snapshot_near_heap_limit > heap_limit_snapshot_taken_) {
-    isolate_->RemoveNearHeapLimitCallback(Environment::NearHeapLimitCallback,
-                                          0);
+  if (heapsnapshot_near_heap_limit_callback_added_) {
+    RemoveHeapSnapshotNearHeapLimitCallback(0);
   }
 
   isolate()->GetHeapProfiler()->RemoveBuildEmbedderGraphCallback(
@@ -1788,8 +1790,7 @@ size_t Environment::NearHeapLimitCallback(void* data,
     Debug(env,
           DebugCategory::DIAGNOSTICS,
           "Not generating snapshots because it's too risky.\n");
-    env->isolate()->RemoveNearHeapLimitCallback(NearHeapLimitCallback,
-                                                initial_heap_limit);
+    env->RemoveHeapSnapshotNearHeapLimitCallback(initial_heap_limit);
     // The new limit must be higher than current_heap_limit or V8 might
     // crash.
     return current_heap_limit + 1;
@@ -1809,17 +1810,15 @@ size_t Environment::NearHeapLimitCallback(void* data,
 
   // Remove the callback first in case it's triggered when generating
   // the snapshot.
-  env->isolate()->RemoveNearHeapLimitCallback(NearHeapLimitCallback,
-                                              initial_heap_limit);
+  env->RemoveHeapSnapshotNearHeapLimitCallback(initial_heap_limit);
 
   heap::WriteSnapshot(env, filename.c_str());
   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) {
-    env->isolate()->AddNearHeapLimitCallback(NearHeapLimitCallback, env);
+  if (env->heap_limit_snapshot_taken_ < env->heap_snapshot_near_heap_limit_) {
+    env->AddHeapSnapshotNearHeapLimitCallback();
   }
 
   FPrintF(stderr, "Wrote snapshot to %s\n", filename.c_str());
diff --git a/src/env.h b/src/env.h
index faced47d195..d2f1adf5cf1 100644
--- a/src/env.h
+++ b/src/env.h
@@ -1040,6 +1040,12 @@ class Environment : public MemoryRetainer {
   template <typename T>
   void ForEachBaseObject(T&& iterator);
 
+  inline void set_heap_snapshot_near_heap_limit(uint32_t limit);
+
+  inline void AddHeapSnapshotNearHeapLimitCallback();
+
+  inline void RemoveHeapSnapshotNearHeapLimitCallback(size_t heap_limit);
+
  private:
   inline void ThrowError(v8::Local<v8::Value> (*fun)(v8::Local<v8::String>),
                          const char* errmsg);
@@ -1097,7 +1103,9 @@ class Environment : public MemoryRetainer {
   std::string exec_path_;
 
   bool is_processing_heap_limit_callback_ = false;
-  int64_t heap_limit_snapshot_taken_ = 0;
+  uint32_t heap_limit_snapshot_taken_ = 0;
+  uint32_t heap_snapshot_near_heap_limit_ = 0;
+  bool heapsnapshot_near_heap_limit_callback_added_ = false;
 
   uint32_t module_id_counter_ = 0;
   uint32_t script_id_counter_ = 0;
diff --git a/src/node.cc b/src/node.cc
index 494c95046a0..d20d3c4c400 100644
--- a/src/node.cc
+++ b/src/node.cc
@@ -241,9 +241,8 @@ static void AtomicsWaitCallback(Isolate::AtomicsWaitEvent event,
 void Environment::InitializeDiagnostics() {
   isolate_->GetHeapProfiler()->AddBuildEmbedderGraphCallback(
       Environment::BuildEmbedderGraph, this);
-  if (options_->heap_snapshot_near_heap_limit > 0) {
-    isolate_->AddNearHeapLimitCallback(Environment::NearHeapLimitCallback,
-                                       this);
+  if (heap_snapshot_near_heap_limit_ > 0) {
+    AddHeapSnapshotNearHeapLimitCallback();
   }
   if (options_->trace_uncaught)
     isolate_->SetCaptureStackTraceForUncaughtExceptions(true);
diff --git a/src/node_v8.cc b/src/node_v8.cc
index 001d3464ec0..9ac82f25880 100644
--- a/src/node_v8.cc
+++ b/src/node_v8.cc
@@ -157,6 +157,15 @@ void CachedDataVersionTag(const FunctionCallbackInfo<Value>& args) {
   args.GetReturnValue().Set(result);
 }
 
+void SetHeapSnapshotNearHeapLimit(const FunctionCallbackInfo<Value>& args) {
+  CHECK(args[0]->IsUint32());
+  Environment* env = Environment::GetCurrent(args);
+  uint32_t limit = args[0].As<v8::Uint32>()->Value();
+  CHECK_GT(limit, 0);
+  env->AddHeapSnapshotNearHeapLimitCallback();
+  env->set_heap_snapshot_near_heap_limit(limit);
+}
+
 void UpdateHeapStatisticsBuffer(const FunctionCallbackInfo<Value>& args) {
   BindingData* data = Environment::GetBindingData<BindingData>(args);
   HeapStatistics s;
@@ -212,6 +221,10 @@ void Initialize(Local<Object> target,
 
   SetMethodNoSideEffect(
       context, target, "cachedDataVersionTag", CachedDataVersionTag);
+  SetMethodNoSideEffect(context,
+                        target,
+                        "setHeapSnapshotNearHeapLimit",
+                        SetHeapSnapshotNearHeapLimit);
   SetMethod(context,
             target,
             "updateHeapStatisticsBuffer",
@@ -267,6 +280,7 @@ void RegisterExternalReferences(ExternalReferenceRegistry* registry) {
   registry->Register(UpdateHeapCodeStatisticsBuffer);
   registry->Register(UpdateHeapSpaceStatisticsBuffer);
   registry->Register(SetFlagsFromString);
+  registry->Register(SetHeapSnapshotNearHeapLimit);
 }
 
 }  // namespace v8_utils