From 58491eb8097b0b521a1a6e3d8e3d30d0409605df Mon Sep 17 00:00:00 2001
From: Joyee Cheung <joyeec9h3@gmail.com>
Date: Fri, 3 Nov 2023 22:53:21 +0000
Subject: [PATCH] src: implement countObjectsWithPrototype

This implements an internal utility for counting objects
in the heap with a specified prototype. In addition this
adds a checkIfCollectableByCounting() test helper.

PR-URL: https://github.com/nodejs/node/pull/50572
Refs: https://github.com/v8/v8/commit/0fd478bcdabd3400d9d74c47c4883c085ef37d18
Reviewed-By: Geoffrey Booth <webadmin@geoffreybooth.com>
Reviewed-By: Stephen Belanger <admin@stephenbelanger.com>
---
 src/heap_utils.cc | 36 ++++++++++++++++++++++++++++++++++++
 test/common/gc.js | 47 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 83 insertions(+)

diff --git a/src/heap_utils.cc b/src/heap_utils.cc
index e385955a5d5fce..75f7e7e1bdc1b4 100644
--- a/src/heap_utils.cc
+++ b/src/heap_utils.cc
@@ -474,6 +474,39 @@ void TriggerHeapSnapshot(const FunctionCallbackInfo<Value>& args) {
   return args.GetReturnValue().Set(filename_v);
 }
 
+class PrototypeChainHas : public v8::QueryObjectPredicate {
+ public:
+  PrototypeChainHas(Local<Context> context, Local<Object> search)
+      : context_(context), search_(search) {}
+
+  // What we can do in the filter can be quite limited, but looking up
+  // the prototype chain is something that the inspector console API
+  // queryObject() does so it is supported.
+  bool Filter(Local<Object> object) override {
+    for (Local<Value> proto = object->GetPrototype(); proto->IsObject();
+         proto = proto.As<Object>()->GetPrototype()) {
+      if (search_ == proto) return true;
+    }
+    return false;
+  }
+
+ private:
+  Local<Context> context_;
+  Local<Object> search_;
+};
+
+void CountObjectsWithPrototype(const FunctionCallbackInfo<Value>& args) {
+  CHECK_EQ(args.Length(), 1);
+  CHECK(args[0]->IsObject());
+  Local<Object> proto = args[0].As<Object>();
+  Isolate* isolate = args.GetIsolate();
+  Local<Context> context = isolate->GetCurrentContext();
+  PrototypeChainHas prototype_chain_has(context, proto);
+  std::vector<Global<Object>> out;
+  isolate->GetHeapProfiler()->QueryObjects(context, &prototype_chain_has, &out);
+  args.GetReturnValue().Set(static_cast<uint32_t>(out.size()));
+}
+
 void Initialize(Local<Object> target,
                 Local<Value> unused,
                 Local<Context> context,
@@ -482,12 +515,15 @@ void Initialize(Local<Object> target,
   SetMethod(context, target, "triggerHeapSnapshot", TriggerHeapSnapshot);
   SetMethod(
       context, target, "createHeapSnapshotStream", CreateHeapSnapshotStream);
+  SetMethod(
+      context, target, "countObjectsWithPrototype", CountObjectsWithPrototype);
 }
 
 void RegisterExternalReferences(ExternalReferenceRegistry* registry) {
   registry->Register(BuildEmbedderGraph);
   registry->Register(TriggerHeapSnapshot);
   registry->Register(CreateHeapSnapshotStream);
+  registry->Register(CountObjectsWithPrototype);
 }
 
 }  // namespace heap
diff --git a/test/common/gc.js b/test/common/gc.js
index 9860bb191ee072..4671a5422641ca 100644
--- a/test/common/gc.js
+++ b/test/common/gc.js
@@ -75,7 +75,54 @@ async function runAndBreathe(fn, repeat, waitTime = 20) {
   }
 }
 
+/**
+ * This requires --expose-internals.
+ * This function can be used to check if an object factory leaks or not by
+ * iterating over the heap and count objects with the specified class
+ * (which is checked by looking up the prototype chain).
+ * @param {(i: number) => number} fn The factory receiving iteration count
+ *   and returning number of objects created. The return value should be
+ *   precise otherwise false negatives can be produced.
+ * @param {Function} klass The class whose object is used to count the objects
+ * @param {number} count Number of iterations that this check should be done
+ * @param {number} waitTime Optional breathing time for GC.
+ */
+async function checkIfCollectableByCounting(fn, klass, count, waitTime = 20) {
+  const { internalBinding } = require('internal/test/binding');
+  const { countObjectsWithPrototype } = internalBinding('heap_utils');
+  const { prototype, name } = klass;
+  const initialCount = countObjectsWithPrototype(prototype);
+  console.log(`Initial count of ${name}: ${initialCount}`);
+  let totalCreated = 0;
+  for (let i = 0; i < count; ++i) {
+    const created = await fn(i);
+    totalCreated += created;
+    console.log(`#${i}: created ${created} ${name}, total ${totalCreated}`);
+    await wait(waitTime);  // give GC some breathing room.
+    const currentCount = countObjectsWithPrototype(prototype);
+    const collected = totalCreated - (currentCount - initialCount);
+    console.log(`#${i}: counted ${currentCount} ${name}, collected ${collected}`);
+    if (collected > 0) {
+      console.log(`Detected ${collected} collected ${name}, finish early`);
+      return;
+    }
+  }
+
+  await wait(waitTime);  // give GC some breathing room.
+  const currentCount = countObjectsWithPrototype(prototype);
+  const collected = totalCreated - (currentCount - initialCount);
+  console.log(`Last count: counted ${currentCount} ${name}, collected ${collected}`);
+  // Some objects with the prototype can be collected.
+  if (collected > 0) {
+    console.log(`Detected ${collected} collected ${name}`);
+    return;
+  }
+
+  throw new Error(`${name} cannot be collected`);
+}
+
 module.exports = {
   checkIfCollectable,
   runAndBreathe,
+  checkIfCollectableByCounting,
 };