worker: add ability to take heap snapshot from parent thread

PR-URL: https://github.com/nodejs/node/pull/31569
Reviewed-By: Denys Otrishko <shishugi@gmail.com>
Reviewed-By: James M Snell <jasnell@gmail.com>
Reviewed-By: Gus Caplan <me@gus.host>
Reviewed-By: Colin Ihrig <cjihrig@gmail.com>
Reviewed-By: Richard Lau <riclau@uk.ibm.com>
Reviewed-By: Rich Trott <rtrott@gmail.com>

doc/api/errors.md

@@ -2053,6 +2053,11 @@ The WASI instance has already started.
 The `execArgv` option passed to the `Worker` constructor contains
 invalid flags.
 
+<a id="ERR_WORKER_NOT_RUNNING"></a>
+### `ERR_WORKER_NOT_RUNNING`
+
+An operation failed because the `Worker` instance is not currently running.
+
 <a id="ERR_WORKER_OUT_OF_MEMORY"></a>
 ### `ERR_WORKER_OUT_OF_MEMORY`
 

doc/api/worker_threads.md

@@ -674,6 +674,21 @@ inside the worker thread. If `stdout: true` was not passed to the
 [`Worker`][] constructor, then data will be piped to the parent thread's
 [`process.stdout`][] stream.
 
+### `worker.takeHeapSnapshot()`
+<!-- YAML
+added: REPLACEME
+-->
+
+* Returns: {Promise} A promise for a Readable Stream containing
+  a V8 heap snapshot
+
+Returns a readable stream for a V8 snapshot of the current state of the Worker.
+See [`v8.getHeapSnapshot()`][] for more details.
+
+If the Worker thread is no longer running, which may occur before the
+[`'exit'` event][] is emitted, the returned `Promise` will be rejected
+immediately with an [`ERR_WORKER_NOT_RUNNING`][] error.
+
 ### `worker.terminate()`
 <!-- YAML
 added: v10.5.0
@@ -716,6 +731,7 @@ active handle in the event system. If the worker is already `unref()`ed calling
 [`'exit'` event]: #worker_threads_event_exit
 [`AsyncResource`]: async_hooks.html#async_hooks_class_asyncresource
 [`Buffer`]: buffer.html
+[`ERR_WORKER_NOT_RUNNING`]: errors.html#ERR_WORKER_NOT_RUNNING
 [`EventEmitter`]: events.html
 [`EventTarget`]: https://developer.mozilla.org/en-US/docs/Web/API/EventTarget
 [`MessagePort`]: #worker_threads_class_messageport
@@ -743,6 +759,7 @@ active handle in the event system. If the worker is already `unref()`ed calling
 [`require('worker_threads').threadId`]: #worker_threads_worker_threadid
 [`require('worker_threads').workerData`]: #worker_threads_worker_workerdata
 [`trace_events`]: tracing.html
+[`v8.getHeapSnapshot()`]: v8.html#v8_v8_getheapsnapshot
 [`vm`]: vm.html
 [`worker.on('message')`]: #worker_threads_event_message_1
 [`worker.postMessage()`]: #worker_threads_worker_postmessage_value_transferlist

lib/internal/errors.js

@@ -1361,6 +1361,7 @@ E('ERR_WASI_ALREADY_STARTED', 'WASI instance has already started', Error);
 E('ERR_WORKER_INVALID_EXEC_ARGV', (errors) =>
   `Initiated Worker with invalid execArgv flags: ${errors.join(', ')}`,
   Error);
+E('ERR_WORKER_NOT_RUNNING', 'Worker instance not running', Error);
 E('ERR_WORKER_OUT_OF_MEMORY', 'Worker terminated due to reaching memory limit',
   Error);
 E('ERR_WORKER_PATH',

lib/internal/heap_utils.js

@@ -0,0 +1,41 @@
+'use strict';
+const {
+  Symbol
+} = primordials;
+const {
+  kUpdateTimer,
+  onStreamRead,
+} = require('internal/stream_base_commons');
+const { owner_symbol } = require('internal/async_hooks').symbols;
+const { Readable } = require('stream');
+
+const kHandle = Symbol('kHandle');
+
+class HeapSnapshotStream extends Readable {
+  constructor(handle) {
+    super({ autoDestroy: true });
+    this[kHandle] = handle;
+    handle[owner_symbol] = this;
+    handle.onread = onStreamRead;
+  }
+
+  _read() {
+    if (this[kHandle])
+      this[kHandle].readStart();
+  }
+
+  _destroy() {
+    // Release the references on the handle so that
+    // it can be garbage collected.
+    this[kHandle][owner_symbol] = undefined;
+    this[kHandle] = undefined;
+  }
+
+  [kUpdateTimer]() {
+    // Does nothing
+  }
+}
+
+module.exports = {
+  HeapSnapshotStream
+};

lib/internal/worker.js

@@ -19,6 +19,7 @@ const path = require('path');
 
 const errorCodes = require('internal/errors').codes;
 const {
+  ERR_WORKER_NOT_RUNNING,
   ERR_WORKER_PATH,
   ERR_WORKER_UNSERIALIZABLE_ERROR,
   ERR_WORKER_UNSUPPORTED_EXTENSION,
@@ -314,6 +315,17 @@ class Worker extends EventEmitter {
 
     return makeResourceLimits(this[kHandle].getResourceLimits());
   }
+
+  getHeapSnapshot() {
+    const heapSnapshotTaker = this[kHandle] && this[kHandle].takeHeapSnapshot();
+    return new Promise((resolve, reject) => {
+      if (!heapSnapshotTaker) return reject(new ERR_WORKER_NOT_RUNNING());
+      heapSnapshotTaker.ondone = (handle) => {
+        const { HeapSnapshotStream } = require('internal/heap_utils');
+        resolve(new HeapSnapshotStream(handle));
+      };
+    });
+  }
 }
 
 function pipeWithoutWarning(source, dest) {

lib/v8.js

@@ -25,7 +25,6 @@ const {
   Int8Array,
   Map,
   ObjectPrototypeToString,
-  Symbol,
   Uint16Array,
   Uint32Array,
   Uint8Array,
@@ -48,14 +47,7 @@ const {
   createHeapSnapshotStream,
   triggerHeapSnapshot
 } = internalBinding('heap_utils');
-const { Readable } = require('stream');
-const { owner_symbol } = require('internal/async_hooks').symbols;
-const {
-  kUpdateTimer,
-  onStreamRead,
-} = require('internal/stream_base_commons');
-const kHandle = Symbol('kHandle');
-
+const { HeapSnapshotStream } = require('internal/heap_utils');
 
 function writeHeapSnapshot(filename) {
   if (filename !== undefined) {
@@ -65,31 +57,6 @@ function writeHeapSnapshot(filename) {
   return triggerHeapSnapshot(filename);
 }
 
-class HeapSnapshotStream extends Readable {
-  constructor(handle) {
-    super({ autoDestroy: true });
-    this[kHandle] = handle;
-    handle[owner_symbol] = this;
-    handle.onread = onStreamRead;
-  }
-
-  _read() {
-    if (this[kHandle])
-      this[kHandle].readStart();
-  }
-
-  _destroy() {
-    // Release the references on the handle so that
-    // it can be garbage collected.
-    this[kHandle][owner_symbol] = undefined;
-    this[kHandle] = undefined;
-  }
-
-  [kUpdateTimer]() {
-    // Does nothing
-  }
-}
-
 function getHeapSnapshot() {
   const handle = createHeapSnapshotStream();
   assert(handle);
@@ -321,5 +288,5 @@ module.exports = {
   DefaultDeserializer,
   deserialize,
   serialize,
-  writeHeapSnapshot
+  writeHeapSnapshot,
 };

node.gyp

@@ -138,6 +138,7 @@
       'lib/internal/fs/utils.js',
       'lib/internal/fs/watchers.js',
       'lib/internal/http.js',
+      'lib/internal/heap_utils.js',
       'lib/internal/idna.js',
       'lib/internal/inspector_async_hook.js',
       'lib/internal/js_stream_socket.js',

src/async_wrap.h

@@ -70,6 +70,7 @@ namespace node {
   V(UDPWRAP)                                                                  \
   V(SIGINTWATCHDOG)                                                           \
   V(WORKER)                                                                   \
+  V(WORKERHEAPSNAPSHOT)                                                       \
   V(WRITEWRAP)                                                                \
   V(ZLIB)
 

src/env.h

@@ -414,7 +414,8 @@ constexpr size_t kFsStatsBufferLength =
   V(streambaseoutputstream_constructor_template, v8::ObjectTemplate)           \
   V(tcp_constructor_template, v8::FunctionTemplate)                            \
   V(tty_constructor_template, v8::FunctionTemplate)                            \
-  V(write_wrap_template, v8::ObjectTemplate)
+  V(write_wrap_template, v8::ObjectTemplate)                                   \
+  V(worker_heap_snapshot_taker_template, v8::ObjectTemplate)
 
 #define ENVIRONMENT_STRONG_PERSISTENT_VALUES(V)                                \
   V(as_callback_data, v8::Object)                                              \

src/heap_utils.cc

@@ -236,18 +236,16 @@ class HeapSnapshotStream : public AsyncWrap,
  public:
   HeapSnapshotStream(
       Environment* env,
-      const HeapSnapshot* snapshot,
+      HeapSnapshotPointer&& snapshot,
       v8::Local<v8::Object> obj) :
       AsyncWrap(env, obj, AsyncWrap::PROVIDER_HEAPSNAPSHOT),
       StreamBase(env),
-      snapshot_(snapshot) {
+      snapshot_(std::move(snapshot)) {
     MakeWeak();
     StreamBase::AttachToObject(GetObject());
   }
 
-  ~HeapSnapshotStream() override {
-    Cleanup();
-  }
+  ~HeapSnapshotStream() override {}
 
   int GetChunkSize() override {
     return 65536;  // big chunks == faster
@@ -255,7 +253,7 @@ class HeapSnapshotStream : public AsyncWrap,
 
   void EndOfStream() override {
     EmitRead(UV_EOF);
-    Cleanup();
+    snapshot_.reset();
   }
 
   WriteResult WriteAsciiChunk(char* data, int size) override {
@@ -309,22 +307,13 @@ class HeapSnapshotStream : public AsyncWrap,
   SET_SELF_SIZE(HeapSnapshotStream)
 
  private:
-  void Cleanup() {
-    if (snapshot_ != nullptr) {
-      const_cast<HeapSnapshot*>(snapshot_)->Delete();
-      snapshot_ = nullptr;
-    }
-  }
-
-
-  const HeapSnapshot* snapshot_;
+  HeapSnapshotPointer snapshot_;
 };
 
 inline void TakeSnapshot(Isolate* isolate, v8::OutputStream* out) {
-  const HeapSnapshot* const snapshot =
-      isolate->GetHeapProfiler()->TakeHeapSnapshot();
+  HeapSnapshotPointer snapshot {
+      isolate->GetHeapProfiler()->TakeHeapSnapshot() };
   snapshot->Serialize(out, HeapSnapshot::kJSON);
-  const_cast<HeapSnapshot*>(snapshot)->Delete();
 }
 
 inline bool WriteSnapshot(Isolate* isolate, const char* filename) {
@@ -339,20 +328,44 @@ inline bool WriteSnapshot(Isolate* isolate, const char* filename) {
 
 }  // namespace
 
-void CreateHeapSnapshotStream(const FunctionCallbackInfo<Value>& args) {
-  Environment* env = Environment::GetCurrent(args);
+void DeleteHeapSnapshot(const v8::HeapSnapshot* snapshot) {
+  const_cast<HeapSnapshot*>(snapshot)->Delete();
+}
+
+BaseObjectPtr<AsyncWrap> CreateHeapSnapshotStream(
+    Environment* env, HeapSnapshotPointer&& snapshot) {
   HandleScope scope(env->isolate());
-  const HeapSnapshot* const snapshot =
-      env->isolate()->GetHeapProfiler()->TakeHeapSnapshot();
-  CHECK_NOT_NULL(snapshot);
+
+  if (env->streambaseoutputstream_constructor_template().IsEmpty()) {
+    // Create FunctionTemplate for HeapSnapshotStream
+    Local<FunctionTemplate> os = FunctionTemplate::New(env->isolate());
+    os->Inherit(AsyncWrap::GetConstructorTemplate(env));
+    Local<ObjectTemplate> ost = os->InstanceTemplate();
+    ost->SetInternalFieldCount(StreamBase::kStreamBaseFieldCount);
+    os->SetClassName(
+        FIXED_ONE_BYTE_STRING(env->isolate(), "HeapSnapshotStream"));
+    StreamBase::AddMethods(env, os);
+    env->set_streambaseoutputstream_constructor_template(ost);
+  }
+
   Local<Object> obj;
   if (!env->streambaseoutputstream_constructor_template()
            ->NewInstance(env->context())
            .ToLocal(&obj)) {
-    return;
+    return {};
   }
-  HeapSnapshotStream* out = new HeapSnapshotStream(env, snapshot, obj);
-  args.GetReturnValue().Set(out->object());
+  return MakeBaseObject<HeapSnapshotStream>(env, std::move(snapshot), obj);
+}
+
+void CreateHeapSnapshotStream(const FunctionCallbackInfo<Value>& args) {
+  Environment* env = Environment::GetCurrent(args);
+  HeapSnapshotPointer snapshot {
+      env->isolate()->GetHeapProfiler()->TakeHeapSnapshot() };
+  CHECK(snapshot);
+  BaseObjectPtr<AsyncWrap> stream =
+      CreateHeapSnapshotStream(env, std::move(snapshot));
+  if (stream)
+    args.GetReturnValue().Set(stream->object());
 }
 
 void TriggerHeapSnapshot(const FunctionCallbackInfo<Value>& args) {
@@ -388,15 +401,6 @@ void Initialize(Local<Object> target,
   env->SetMethod(target, "buildEmbedderGraph", BuildEmbedderGraph);
   env->SetMethod(target, "triggerHeapSnapshot", TriggerHeapSnapshot);
   env->SetMethod(target, "createHeapSnapshotStream", CreateHeapSnapshotStream);
-
-  // Create FunctionTemplate for HeapSnapshotStream
-  Local<FunctionTemplate> os = FunctionTemplate::New(env->isolate());
-  os->Inherit(AsyncWrap::GetConstructorTemplate(env));
-  Local<ObjectTemplate> ost = os->InstanceTemplate();
-  ost->SetInternalFieldCount(StreamBase::kStreamBaseFieldCount);
-  os->SetClassName(FIXED_ONE_BYTE_STRING(env->isolate(), "HeapSnapshotStream"));
-  StreamBase::AddMethods(env, os);
-  env->set_streambaseoutputstream_constructor_template(ost);
 }
 
 }  // namespace heap

src/node_internals.h

@@ -384,6 +384,16 @@ class TraceEventScope {
   void* id_;
 };
 
+namespace heap {
+
+void DeleteHeapSnapshot(const v8::HeapSnapshot* snapshot);
+using HeapSnapshotPointer =
+  DeleteFnPtr<const v8::HeapSnapshot, DeleteHeapSnapshot>;
+
+BaseObjectPtr<AsyncWrap> CreateHeapSnapshotStream(
+    Environment* env, HeapSnapshotPointer&& snapshot);
+}  // namespace heap
+
 }  // namespace node
 
 #endif  // defined(NODE_WANT_INTERNALS) && NODE_WANT_INTERNALS

src/node_worker.cc

@@ -657,6 +657,53 @@ void Worker::MemoryInfo(MemoryTracker* tracker) const {
   tracker->TrackField("parent_port", parent_port_);
 }
 
+class WorkerHeapSnapshotTaker : public AsyncWrap {
+ public:
+  WorkerHeapSnapshotTaker(Environment* env, Local<Object> obj)
+    : AsyncWrap(env, obj, AsyncWrap::PROVIDER_WORKERHEAPSNAPSHOT) {}
+
+  SET_NO_MEMORY_INFO()
+  SET_MEMORY_INFO_NAME(WorkerHeapSnapshotTaker)
+  SET_SELF_SIZE(WorkerHeapSnapshotTaker)
+};
+
+void Worker::TakeHeapSnapshot(const FunctionCallbackInfo<Value>& args) {
+  Worker* w;
+  ASSIGN_OR_RETURN_UNWRAP(&w, args.This());
+
+  Debug(w, "Worker %llu taking heap snapshot", w->thread_id_);
+
+  Environment* env = w->env();
+  AsyncHooks::DefaultTriggerAsyncIdScope trigger_id_scope(w);
+  Local<Object> wrap;
+  if (!env->worker_heap_snapshot_taker_template()
+      ->NewInstance(env->context()).ToLocal(&wrap)) {
+    return;
+  }
+  BaseObjectPtr<WorkerHeapSnapshotTaker> taker =
+      MakeDetachedBaseObject<WorkerHeapSnapshotTaker>(env, wrap);
+
+  // Interrupt the worker thread and take a snapshot, then schedule a call
+  // on the parent thread that turns that snapshot into a readable stream.
+  bool scheduled = w->RequestInterrupt([taker, env](Environment* worker_env) {
+    heap::HeapSnapshotPointer snapshot {
+        worker_env->isolate()->GetHeapProfiler()->TakeHeapSnapshot() };
+    CHECK(snapshot);
+    env->SetImmediateThreadsafe(
+        [taker, snapshot = std::move(snapshot)](Environment* env) mutable {
+          HandleScope handle_scope(env->isolate());
+          Context::Scope context_scope(env->context());
+
+          AsyncHooks::DefaultTriggerAsyncIdScope trigger_id_scope(taker.get());
+          BaseObjectPtr<AsyncWrap> stream = heap::CreateHeapSnapshotStream(
+              env, std::move(snapshot));
+          Local<Value> args[] = { stream->object() };
+          taker->MakeCallback(env->ondone_string(), arraysize(args), args);
+        });
+  });
+  args.GetReturnValue().Set(scheduled ? taker->object() : Local<Object>());
+}
+
 namespace {
 
 // Return the MessagePort that is global for this Environment and communicates
@@ -689,6 +736,7 @@ void InitWorker(Local<Object> target,
     env->SetProtoMethod(w, "ref", Worker::Ref);
     env->SetProtoMethod(w, "unref", Worker::Unref);
     env->SetProtoMethod(w, "getResourceLimits", Worker::GetResourceLimits);
+    env->SetProtoMethod(w, "takeHeapSnapshot", Worker::TakeHeapSnapshot);
 
     Local<String> workerString =
         FIXED_ONE_BYTE_STRING(env->isolate(), "Worker");
@@ -698,6 +746,18 @@ void InitWorker(Local<Object> target,
                 w->GetFunction(env->context()).ToLocalChecked()).Check();
   }
 
+  {
+    Local<FunctionTemplate> wst = FunctionTemplate::New(env->isolate());
+
+    wst->InstanceTemplate()->SetInternalFieldCount(1);
+    wst->Inherit(AsyncWrap::GetConstructorTemplate(env));
+
+    Local<String> wst_string =
+        FIXED_ONE_BYTE_STRING(env->isolate(), "WorkerHeapSnapshotTaker");
+    wst->SetClassName(wst_string);
+    env->set_worker_heap_snapshot_taker_template(wst->InstanceTemplate());
+  }
+
   env->SetMethod(target, "getEnvMessagePort", GetEnvMessagePort);
 
   target

src/node_worker.h

@@ -59,6 +59,7 @@ class Worker : public AsyncWrap {
   static void GetResourceLimits(
       const v8::FunctionCallbackInfo<v8::Value>& args);
   v8::Local<v8::Float64Array> GetResourceLimits(v8::Isolate* isolate) const;
+  static void TakeHeapSnapshot(const v8::FunctionCallbackInfo<v8::Value>& args);
 
  private:
   void CreateEnvMessagePort(Environment* env);

test/common/heap.js

@@ -14,8 +14,7 @@ try {
 const { buildEmbedderGraph } = internalBinding('heap_utils');
 const { getHeapSnapshot } = require('v8');
 
-function createJSHeapSnapshot() {
-  const stream = getHeapSnapshot();
+function createJSHeapSnapshot(stream = getHeapSnapshot()) {
   stream.pause();
   const dump = JSON.parse(stream.read());
   const meta = dump.snapshot.meta;
@@ -106,8 +105,8 @@ function isEdge(edge, { node_name, edge_name }) {
 }
 
 class State {
-  constructor() {
-    this.snapshot = createJSHeapSnapshot();
+  constructor(stream) {
+    this.snapshot = createJSHeapSnapshot(stream);
     this.embedderGraph = buildEmbedderGraph();
   }
 
@@ -189,8 +188,8 @@ class State {
   }
 }
 
-function recordState() {
-  return new State();
+function recordState(stream = undefined) {
+  return new State(stream);
 }
 
 function validateSnapshotNodes(...args) {

test/parallel/test-worker-heapdump-failure.js

@@ -0,0 +1,15 @@
+'use strict';
+const common = require('../common');
+const assert = require('assert');
+const { Worker } = require('worker_threads');
+const { once } = require('events');
+
+(async function() {
+  const w = new Worker('', { eval: true });
+
+  await once(w, 'exit');
+  await assert.rejects(() => w.getHeapSnapshot(), {
+    name: 'Error',
+    code: 'ERR_WORKER_NOT_RUNNING'
+  });
+})().then(common.mustCall());

test/pummel/test-worker-take-heapsnapshot.js

@@ -0,0 +1,22 @@
+// Flags: --expose-internals
+'use strict';
+const common = require('../common');
+const { recordState } = require('../common/heap');
+const { Worker } = require('worker_threads');
+const { once } = require('events');
+
+(async function() {
+  const w = new Worker('setInterval(() => {}, 100)', { eval: true });
+
+  await once(w, 'online');
+  const stream = await w.getHeapSnapshot();
+  const snapshot = recordState(stream);
+  snapshot.validateSnapshot('Node / MessagePort', [
+    {
+      children: [
+        { node_name: 'Node / MessagePortData', edge_name: 'data' }
+      ]
+    }
+  ], { loose: true });
+  await w.terminate();
+})().then(common.mustCall());

test/sequential/test-async-wrap-getasyncid.js

@@ -52,6 +52,7 @@ const { getSystemErrorName } = require('util');
     delete providers.HTTPINCOMINGMESSAGE;
     delete providers.ELDHISTOGRAM;
     delete providers.SIGINTWATCHDOG;
+    delete providers.WORKERHEAPSNAPSHOT;
 
     const objKeys = Object.keys(providers);
     if (objKeys.length > 0)