From abe6a2e3d1c268f5bebbf05864f2588c6f4858b5 Mon Sep 17 00:00:00 2001
From: Anna Henningsen <anna@addaleax.net>
Date: Tue, 28 Jan 2020 14:25:10 +0000
Subject: [PATCH] process: report ArrayBuffer memory in `memoryUsage()`

Report memory allocations performed by the `ArrayBuffer::Allocator`.

PR-URL: https://github.com/nodejs/node/pull/31550
Reviewed-By: Richard Lau <riclau@uk.ibm.com>
Reviewed-By: Colin Ihrig <cjihrig@gmail.com>
Reviewed-By: Gus Caplan <me@gus.host>
Reviewed-By: Ben Noordhuis <info@bnoordhuis.nl>
Reviewed-By: James M Snell <jasnell@gmail.com>
Reviewed-By: Rich Trott <rtrott@gmail.com>
Reviewed-By: Gireesh Punathil <gpunathi@in.ibm.com>
---
 doc/api/process.md                 | 28 ++++++++++++++++++----------
 lib/internal/process/per_thread.js |  5 +++--
 src/api/environment.cc             | 30 ++++++++++++++++++++++++++++--
 src/node_internals.h               | 20 ++++++++++++--------
 src/node_process_methods.cc        |  7 ++++++-
 test/parallel/test-memory-usage.js | 13 +++++++++++++
 6 files changed, 80 insertions(+), 23 deletions(-)

diff --git a/doc/api/process.md b/doc/api/process.md
index fde2316302a2de..3ccd9fc776def6 100644
--- a/doc/api/process.md
+++ b/doc/api/process.md
@@ -1510,6 +1510,9 @@ is no entry script.
 <!-- YAML
 added: v0.1.16
 changes:
+  - version: REPLACEME
+    pr-url: https://github.com/nodejs/node/pull/31550
+    description: Added `arrayBuffers` to the returned object.
   - version: v7.2.0
     pr-url: https://github.com/nodejs/node/pull/9587
     description: Added `external` to the returned object.
@@ -1520,6 +1523,7 @@ changes:
   * `heapTotal` {integer}
   * `heapUsed` {integer}
   * `external` {integer}
+  * `arrayBuffers` {integer}
 
 The `process.memoryUsage()` method returns an object describing the memory usage
 of the Node.js process measured in bytes.
@@ -1538,19 +1542,22 @@ Will generate:
   rss: 4935680,
   heapTotal: 1826816,
   heapUsed: 650472,
-  external: 49879
+  external: 49879,
+  arrayBuffers: 9386
 }
 ```
 
-`heapTotal` and `heapUsed` refer to V8's memory usage.
-`external` refers to the memory usage of C++ objects bound to JavaScript
-objects managed by V8. `rss`, Resident Set Size, is the amount of space
-occupied in the main memory device (that is a subset of the total allocated
-memory) for the process, which includes the _heap_, _code segment_ and _stack_.
-
-The _heap_ is where objects, strings, and closures are stored. Variables are
-stored in the _stack_ and the actual JavaScript code resides in the
-_code segment_.
+* `heapTotal` and `heapUsed` refer to V8's memory usage.
+* `external` refers to the memory usage of C++ objects bound to JavaScript
+  objects managed by V8.
+* `rss`, Resident Set Size, is the amount of space occupied in the main
+  memory device (that is a subset of the total allocated memory) for the
+  process, including all C++ and JavaScript objects and code.
+* `arrayBuffers` refers to memory allocated for `ArrayBuffer`s and
+  `SharedArrayBuffer`s, including all Node.js [`Buffer`][]s.
+  This is also included in the `external` value. When Node.js is used as an
+  embedded library, this value may be `0` because allocations for `ArrayBuffer`s
+  may not be tracked in that case.
 
 When using [`Worker`][] threads, `rss` will be a value that is valid for the
 entire process, while the other fields will only refer to the current thread.
@@ -2518,6 +2525,7 @@ cases:
 [`'exit'`]: #process_event_exit
 [`'message'`]: child_process.html#child_process_event_message
 [`'uncaughtException'`]: #process_event_uncaughtexception
+[`Buffer`]: buffer.html
 [`ChildProcess.disconnect()`]: child_process.html#child_process_subprocess_disconnect
 [`ChildProcess.send()`]: child_process.html#child_process_subprocess_send_message_sendhandle_options_callback
 [`ChildProcess`]: child_process.html#child_process_class_childprocess
diff --git a/lib/internal/process/per_thread.js b/lib/internal/process/per_thread.js
index 5cb55d74939302..d219cd7298fc15 100644
--- a/lib/internal/process/per_thread.js
+++ b/lib/internal/process/per_thread.js
@@ -146,14 +146,15 @@ function wrapProcessMethods(binding) {
     return hrBigintValues[0];
   }
 
-  const memValues = new Float64Array(4);
+  const memValues = new Float64Array(5);
   function memoryUsage() {
     _memoryUsage(memValues);
     return {
       rss: memValues[0],
       heapTotal: memValues[1],
       heapUsed: memValues[2],
-      external: memValues[3]
+      external: memValues[3],
+      arrayBuffers: memValues[4]
     };
   }
 
diff --git a/src/api/environment.cc b/src/api/environment.cc
index 4999349d282de1..fea02128af2670 100644
--- a/src/api/environment.cc
+++ b/src/api/environment.cc
@@ -87,10 +87,34 @@ static void HostCleanupFinalizationGroupCallback(
 }
 
 void* NodeArrayBufferAllocator::Allocate(size_t size) {
+  void* ret;
   if (zero_fill_field_ || per_process::cli_options->zero_fill_all_buffers)
-    return UncheckedCalloc(size);
+    ret = UncheckedCalloc(size);
   else
-    return UncheckedMalloc(size);
+    ret = UncheckedMalloc(size);
+  if (LIKELY(ret != nullptr))
+    total_mem_usage_.fetch_add(size, std::memory_order_relaxed);
+  return ret;
+}
+
+void* NodeArrayBufferAllocator::AllocateUninitialized(size_t size) {
+  void* ret = node::UncheckedMalloc(size);
+  if (LIKELY(ret != nullptr))
+    total_mem_usage_.fetch_add(size, std::memory_order_relaxed);
+  return ret;
+}
+
+void* NodeArrayBufferAllocator::Reallocate(
+    void* data, size_t old_size, size_t size) {
+  void* ret = UncheckedRealloc<char>(static_cast<char*>(data), size);
+  if (LIKELY(ret != nullptr) || UNLIKELY(size == 0))
+    total_mem_usage_.fetch_add(size - old_size, std::memory_order_relaxed);
+  return ret;
+}
+
+void NodeArrayBufferAllocator::Free(void* data, size_t size) {
+  total_mem_usage_.fetch_sub(size, std::memory_order_relaxed);
+  free(data);
 }
 
 DebuggingArrayBufferAllocator::~DebuggingArrayBufferAllocator() {
@@ -140,11 +164,13 @@ void* DebuggingArrayBufferAllocator::Reallocate(void* data,
 
 void DebuggingArrayBufferAllocator::RegisterPointer(void* data, size_t size) {
   Mutex::ScopedLock lock(mutex_);
+  NodeArrayBufferAllocator::RegisterPointer(data, size);
   RegisterPointerInternal(data, size);
 }
 
 void DebuggingArrayBufferAllocator::UnregisterPointer(void* data, size_t size) {
   Mutex::ScopedLock lock(mutex_);
+  NodeArrayBufferAllocator::UnregisterPointer(data, size);
   UnregisterPointerInternal(data, size);
 }
 
diff --git a/src/node_internals.h b/src/node_internals.h
index 2bd0ff78da7c33..6e50b22af2f915 100644
--- a/src/node_internals.h
+++ b/src/node_internals.h
@@ -109,20 +109,24 @@ class NodeArrayBufferAllocator : public ArrayBufferAllocator {
   inline uint32_t* zero_fill_field() { return &zero_fill_field_; }
 
   void* Allocate(size_t size) override;  // Defined in src/node.cc
-  void* AllocateUninitialized(size_t size) override
-    { return node::UncheckedMalloc(size); }
-  void Free(void* data, size_t) override { free(data); }
-  virtual void* Reallocate(void* data, size_t old_size, size_t size) {
-    return static_cast<void*>(
-        UncheckedRealloc<char>(static_cast<char*>(data), size));
+  void* AllocateUninitialized(size_t size) override;
+  void Free(void* data, size_t size) override;
+  virtual void* Reallocate(void* data, size_t old_size, size_t size);
+  virtual void RegisterPointer(void* data, size_t size) {
+    total_mem_usage_.fetch_add(size, std::memory_order_relaxed);
+  }
+  virtual void UnregisterPointer(void* data, size_t size) {
+    total_mem_usage_.fetch_sub(size, std::memory_order_relaxed);
   }
-  virtual void RegisterPointer(void* data, size_t size) {}
-  virtual void UnregisterPointer(void* data, size_t size) {}
 
   NodeArrayBufferAllocator* GetImpl() final { return this; }
+  inline uint64_t total_mem_usage() const {
+    return total_mem_usage_.load(std::memory_order_relaxed);
+  }
 
  private:
   uint32_t zero_fill_field_ = 1;  // Boolean but exposed as uint32 to JS land.
+  std::atomic<size_t> total_mem_usage_ {0};
 };
 
 class DebuggingArrayBufferAllocator final : public NodeArrayBufferAllocator {
diff --git a/src/node_process_methods.cc b/src/node_process_methods.cc
index 7b91f89e79159a..4d53749f0989f5 100644
--- a/src/node_process_methods.cc
+++ b/src/node_process_methods.cc
@@ -200,10 +200,13 @@ static void MemoryUsage(const FunctionCallbackInfo<Value>& args) {
   HeapStatistics v8_heap_stats;
   isolate->GetHeapStatistics(&v8_heap_stats);
 
+  NodeArrayBufferAllocator* array_buffer_allocator =
+      env->isolate_data()->node_allocator();
+
   // Get the double array pointer from the Float64Array argument.
   CHECK(args[0]->IsFloat64Array());
   Local<Float64Array> array = args[0].As<Float64Array>();
-  CHECK_EQ(array->Length(), 4);
+  CHECK_EQ(array->Length(), 5);
   Local<ArrayBuffer> ab = array->Buffer();
   double* fields = static_cast<double*>(ab->GetBackingStore()->Data());
 
@@ -211,6 +214,8 @@ static void MemoryUsage(const FunctionCallbackInfo<Value>& args) {
   fields[1] = v8_heap_stats.total_heap_size();
   fields[2] = v8_heap_stats.used_heap_size();
   fields[3] = v8_heap_stats.external_memory();
+  fields[4] = array_buffer_allocator == nullptr ?
+      0 : array_buffer_allocator->total_mem_usage();
 }
 
 void RawDebug(const FunctionCallbackInfo<Value>& args) {
diff --git a/test/parallel/test-memory-usage.js b/test/parallel/test-memory-usage.js
index 08ef07a34627ca..a1d07c64776068 100644
--- a/test/parallel/test-memory-usage.js
+++ b/test/parallel/test-memory-usage.js
@@ -30,3 +30,16 @@ if (!common.isIBMi)
 assert.ok(r.heapTotal > 0);
 assert.ok(r.heapUsed > 0);
 assert.ok(r.external > 0);
+
+assert.strictEqual(typeof r.arrayBuffers, 'number');
+if (r.arrayBuffers > 0) {
+  const size = 10 * 1024 * 1024;
+  // eslint-disable-next-line no-unused-vars
+  const ab = new ArrayBuffer(size);
+
+  const after = process.memoryUsage();
+  assert(after.external - r.external >= size,
+         `${after.external} - ${r.external} >= ${size}`);
+  assert.strictEqual(after.arrayBuffers - r.arrayBuffers, size,
+                     `${after.arrayBuffers} - ${r.arrayBuffers} >= ${size}`);
+}