From ae3929e958dfa72bcf9f9efeb53c89f825c228ce Mon Sep 17 00:00:00 2001
From: Joyee Cheung <joyeec9h3@gmail.com>
Date: Sat, 8 Feb 2020 01:49:43 +0800
Subject: [PATCH] vm: implement vm.measureMemory() for per-context memory
 measurement

This patch implements `vm.measureMemory()` with the new
`v8::Isolate::MeasureMemory()` API to measure per-context memory
usage. This should be experimental, since detailed memory
measurement requires further integration with the V8 API
that should be available in a future V8 update.

PR-URL: https://github.com/nodejs/node/pull/31824
Refs: https://github.com/ulan/performance-measure-memory
Reviewed-By: Ben Noordhuis <info@bnoordhuis.nl>
Reviewed-By: Anna Henningsen <anna@addaleax.net>
Reviewed-By: James M Snell <jasnell@gmail.com>
Reviewed-By: Denys Otrishko <shishugi@gmail.com>
Reviewed-By: Colin Ihrig <cjihrig@gmail.com>
---
 doc/api/errors.md                       |  8 +++
 doc/api/vm.md                           | 50 ++++++++++++++++++
 lib/internal/errors.js                  |  1 +
 lib/vm.js                               | 37 ++++++++++++-
 src/node_contextify.cc                  | 44 ++++++++++++++++
 test/parallel/test-vm-measure-memory.js | 70 +++++++++++++++++++++++++
 6 files changed, 208 insertions(+), 2 deletions(-)
 create mode 100644 test/parallel/test-vm-measure-memory.js

diff --git a/doc/api/errors.md b/doc/api/errors.md
index a7726cd03019cd..15be1fd2061618 100644
--- a/doc/api/errors.md
+++ b/doc/api/errors.md
@@ -710,6 +710,14 @@ STDERR/STDOUT, and the data's length is longer than the `maxBuffer` option.
 `Console` was instantiated without `stdout` stream, or `Console` has a
 non-writable `stdout` or `stderr` stream.
 
+<a id="ERR_CONTEXT_NOT_INITIALIZED"></a>
+### `ERR_CONTEXT_NOT_INITIALIZED`
+
+The vm context passed into the API is not yet initialized. This could happen
+when an error occurs (and is caught) during the creation of the
+context, for example, when the allocation fails or the maximum call stack
+size is reached when the context is created.
+
 <a id="ERR_CONSTRUCT_CALL_REQUIRED"></a>
 ### `ERR_CONSTRUCT_CALL_REQUIRED`
 
diff --git a/doc/api/vm.md b/doc/api/vm.md
index bd64b23484eb24..ed676414471b8e 100644
--- a/doc/api/vm.md
+++ b/doc/api/vm.md
@@ -295,6 +295,56 @@ console.log(globalVar);
 // 1000
 ```
 
+## `vm.measureMemory([options])`
+
+<!-- YAML
+added: REPLACEME
+-->
+
+> Stability: 1 - Experimental
+
+Measure the memory known to V8 and used by the current execution context
+or a specified context.
+
+* `options` {Object} Optional.
+  * `mode` {string} Either `'summary'` or `'detailed'`.
+    **Default:** `'summary'`
+  * `context` {Object} Optional. A [contextified][] object returned
+    by `vm.createContext()`. If not specified, measure the memory
+    usage of the current context where `vm.measureMemory()` is invoked.
+* Returns: {Promise} If the memory is successfully measured the promise will
+  resolve with an object containing information about the memory usage.
+
+The format of the object that the returned Promise may resolve with is
+specific to the V8 engine and may change from one version of V8 to the next.
+
+The returned result is different from the statistics returned by
+`v8.getHeapSpaceStatistics()` in that `vm.measureMemory()` measures
+the memory reachable by V8 from a specific context, while
+`v8.getHeapSpaceStatistics()` measures the memory used by an instance
+of V8 engine, which can switch among multiple contexts that reference
+objects in the heap of one engine.
+
+```js
+const vm = require('vm');
+// Measure the memory used by the current context and return the result
+// in summary.
+vm.measureMemory({ mode: 'summary' })
+  // Is the same as vm.measureMemory()
+  .then((result) => {
+    // The current format is:
+    // { total: { jsMemoryEstimate: 2211728, jsMemoryRange: [ 0, 2211728 ] } }
+    console.log(result);
+  });
+
+const context = vm.createContext({});
+vm.measureMemory({ mode: 'detailed' }, context)
+  .then((result) => {
+    // At the moment the detailed format is the same as the summary one.
+    console.log(result);
+  });
+```
+
 ## Class: `vm.Module`
 <!-- YAML
 added: v13.0.0
diff --git a/lib/internal/errors.js b/lib/internal/errors.js
index b67034d4558834..cfab59f2b721d3 100644
--- a/lib/internal/errors.js
+++ b/lib/internal/errors.js
@@ -758,6 +758,7 @@ E('ERR_CHILD_PROCESS_STDIO_MAXBUFFER', '%s maxBuffer length exceeded',
   RangeError);
 E('ERR_CONSOLE_WRITABLE_STREAM',
   'Console expects a writable stream instance for %s', TypeError);
+E('ERR_CONTEXT_NOT_INITIALIZED', 'context used is not initialized', Error);
 E('ERR_CPU_USAGE', 'Unable to obtain cpu usage %s', Error);
 E('ERR_CRYPTO_CUSTOM_ENGINE_NOT_SUPPORTED',
   'Custom engines not supported by this OpenSSL', Error);
diff --git a/lib/vm.js b/lib/vm.js
index 211755b9728566..a864f94fef9100 100644
--- a/lib/vm.js
+++ b/lib/vm.js
@@ -25,16 +25,21 @@ const {
   ArrayIsArray,
   ArrayPrototypeForEach,
   Symbol,
+  PromiseReject
 } = primordials;
 
 const {
   ContextifyScript,
   makeContext,
   isContext: _isContext,
-  compileFunction: _compileFunction
+  constants,
+  compileFunction: _compileFunction,
+  measureMemory: _measureMemory,
 } = internalBinding('contextify');
 const {
+  ERR_CONTEXT_NOT_INITIALIZED,
   ERR_INVALID_ARG_TYPE,
+  ERR_INVALID_ARG_VALUE,
 } = require('internal/errors').codes;
 const {
   isArrayBufferView,
@@ -44,7 +49,10 @@ const {
   validateUint32,
   validateString
 } = require('internal/validators');
-const { kVmBreakFirstLineSymbol } = require('internal/util');
+const {
+  kVmBreakFirstLineSymbol,
+  emitExperimentalWarning,
+} = require('internal/util');
 const kParsingContext = Symbol('script parsing context');
 
 class Script extends ContextifyScript {
@@ -401,6 +409,30 @@ function compileFunction(code, params, options = {}) {
   return result.function;
 }
 
+const measureMemoryModes = {
+  summary: constants.measureMemory.mode.SUMMARY,
+  detailed: constants.measureMemory.mode.DETAILED,
+};
+
+function measureMemory(options = {}) {
+  emitExperimentalWarning('vm.measureMemory');
+  validateObject(options, 'options');
+  const { mode = 'summary', context } = options;
+  if (mode !== 'summary' && mode !== 'detailed') {
+    throw new ERR_INVALID_ARG_VALUE(
+      'options.mode', options.mode,
+      'must be either \'summary\' or \'detailed\'');
+  }
+  if (context !== undefined &&
+    (typeof context !== 'object' || context === null || !_isContext(context))) {
+    throw new ERR_INVALID_ARG_TYPE('options.context', 'vm.Context', context);
+  }
+  const result = _measureMemory(measureMemoryModes[mode], context);
+  if (result === undefined) {
+    return PromiseReject(new ERR_CONTEXT_NOT_INITIALIZED());
+  }
+  return result;
+}
 
 module.exports = {
   Script,
@@ -411,6 +443,7 @@ module.exports = {
   runInThisContext,
   isContext,
   compileFunction,
+  measureMemory,
 };
 
 if (require('internal/options').getOptionValue('--experimental-vm-modules')) {
diff --git a/src/node_contextify.cc b/src/node_contextify.cc
index 5f289aecb3414c..6de27188be6c7f 100644
--- a/src/node_contextify.cc
+++ b/src/node_contextify.cc
@@ -47,17 +47,20 @@ using v8::FunctionCallbackInfo;
 using v8::FunctionTemplate;
 using v8::HandleScope;
 using v8::IndexedPropertyHandlerConfiguration;
+using v8::Int32;
 using v8::Integer;
 using v8::Isolate;
 using v8::Local;
 using v8::Maybe;
 using v8::MaybeLocal;
+using v8::MeasureMemoryMode;
 using v8::Name;
 using v8::NamedPropertyHandlerConfiguration;
 using v8::Number;
 using v8::Object;
 using v8::ObjectTemplate;
 using v8::PrimitiveArray;
+using v8::Promise;
 using v8::PropertyAttribute;
 using v8::PropertyCallbackInfo;
 using v8::PropertyDescriptor;
@@ -1203,11 +1206,39 @@ static void WatchdogHasPendingSigint(const FunctionCallbackInfo<Value>& args) {
   args.GetReturnValue().Set(ret);
 }
 
+static void MeasureMemory(const FunctionCallbackInfo<Value>& args) {
+  CHECK(args[0]->IsInt32());
+  int32_t mode = args[0].As<v8::Int32>()->Value();
+  Isolate* isolate = args.GetIsolate();
+  Environment* env = Environment::GetCurrent(args);
+  Local<Context> context;
+  if (args[1]->IsUndefined()) {
+    context = isolate->GetCurrentContext();
+  } else {
+    CHECK(args[1]->IsObject());
+    ContextifyContext* sandbox =
+        ContextifyContext::ContextFromContextifiedSandbox(env,
+                                                          args[1].As<Object>());
+    CHECK_NOT_NULL(sandbox);
+    context = sandbox->context();
+    if (context.IsEmpty()) {  // Not yet fully initilaized
+      return;
+    }
+  }
+  v8::Local<v8::Promise> promise;
+  if (!isolate->MeasureMemory(context, static_cast<v8::MeasureMemoryMode>(mode))
+           .ToLocal(&promise)) {
+    return;
+  }
+  args.GetReturnValue().Set(promise);
+}
+
 void Initialize(Local<Object> target,
                 Local<Value> unused,
                 Local<Context> context,
                 void* priv) {
   Environment* env = Environment::GetCurrent(context);
+  Isolate* isolate = env->isolate();
   ContextifyContext::Init(env, target);
   ContextifyScript::Init(env, target);
 
@@ -1224,6 +1255,19 @@ void Initialize(Local<Object> target,
 
     env->set_compiled_fn_entry_template(tpl->InstanceTemplate());
   }
+
+  Local<Object> constants = Object::New(env->isolate());
+  Local<Object> measure_memory = Object::New(env->isolate());
+  Local<Object> memory_mode = Object::New(env->isolate());
+  MeasureMemoryMode SUMMARY = MeasureMemoryMode::kSummary;
+  MeasureMemoryMode DETAILED = MeasureMemoryMode::kDetailed;
+  NODE_DEFINE_CONSTANT(memory_mode, SUMMARY);
+  NODE_DEFINE_CONSTANT(memory_mode, DETAILED);
+  READONLY_PROPERTY(measure_memory, "mode", memory_mode);
+  READONLY_PROPERTY(constants, "measureMemory", measure_memory);
+  target->Set(context, env->constants_string(), constants).Check();
+
+  env->SetMethod(target, "measureMemory", MeasureMemory);
 }
 
 }  // namespace contextify
diff --git a/test/parallel/test-vm-measure-memory.js b/test/parallel/test-vm-measure-memory.js
new file mode 100644
index 00000000000000..7e620304e0a276
--- /dev/null
+++ b/test/parallel/test-vm-measure-memory.js
@@ -0,0 +1,70 @@
+'use strict';
+const common = require('../common');
+const assert = require('assert');
+const vm = require('vm');
+
+common.expectWarning('ExperimentalWarning',
+                     'vm.measureMemory is an experimental feature. ' +
+                     'This feature could change at any time');
+
+// The formats could change when V8 is updated, then the tests should be
+// updated accordingly.
+function assertSummaryShape(result) {
+  assert.strictEqual(typeof result, 'object');
+  assert.strictEqual(typeof result.total, 'object');
+  assert.strictEqual(typeof result.total.jsMemoryEstimate, 'number');
+  assert(Array.isArray(result.total.jsMemoryRange));
+  assert.strictEqual(typeof result.total.jsMemoryRange[0], 'number');
+  assert.strictEqual(typeof result.total.jsMemoryRange[1], 'number');
+}
+
+function assertDetailedShape(result) {
+  // For now, the detailed shape is the same as the summary shape. This
+  // should change in future versions of V8.
+  return assertSummaryShape(result);
+}
+
+// Test measuring memory of the current context
+{
+  vm.measureMemory()
+    .then(assertSummaryShape);
+
+  vm.measureMemory({})
+    .then(assertSummaryShape);
+
+  vm.measureMemory({ mode: 'summary' })
+    .then(assertSummaryShape);
+
+  vm.measureMemory({ mode: 'detailed' })
+    .then(assertDetailedShape);
+
+  assert.throws(() => vm.measureMemory(null), {
+    code: 'ERR_INVALID_ARG_TYPE'
+  });
+  assert.throws(() => vm.measureMemory('summary'), {
+    code: 'ERR_INVALID_ARG_TYPE'
+  });
+  assert.throws(() => vm.measureMemory({ mode: 'random' }), {
+    code: 'ERR_INVALID_ARG_VALUE'
+  });
+}
+
+// Test measuring memory of the sandbox
+{
+  const context = vm.createContext();
+  vm.measureMemory({ context })
+    .then(assertSummaryShape);
+
+  vm.measureMemory({ mode: 'summary', context },)
+    .then(assertSummaryShape);
+
+  vm.measureMemory({ mode: 'detailed', context })
+    .then(assertDetailedShape);
+
+  assert.throws(() => vm.measureMemory({ mode: 'summary', context: null }), {
+    code: 'ERR_INVALID_ARG_TYPE'
+  });
+  assert.throws(() => vm.measureMemory({ mode: 'summary', context: {} }), {
+    code: 'ERR_INVALID_ARG_TYPE'
+  });
+}