fix(ext/node): add v8.queryObjects() and util.queryObjects() (#34159)

Closes denoland/orchid#122

Co-authored-by: divybot <divybot@users.noreply.github.com>
Co-authored-by: Divy Srivastava <me@littledivy.com>

Author: divybot

ext/node/lib.rs

@@ -241,6 +241,7 @@ deno_core::extension!(deno_node,
     ops::v8::op_v8_update_heap_space_statistics,
     ops::v8::op_v8_get_heap_code_statistics,
     ops::v8::op_v8_take_heap_snapshot,
+    ops::v8::op_v8_query_objects_count,
     ops::v8::op_v8_get_wire_format_version,
     ops::v8::op_v8_new_deserializer,
     ops::v8::op_v8_new_serializer,

ext/node/ops/v8.rs

@@ -72,6 +72,103 @@ pub fn op_v8_take_heap_snapshot(scope: &mut v8::PinScope<'_, '_>) -> Vec<u8> {
   buf
 }
 
+// Walks the V8 heap snapshot and counts nodes that look like instances of a
+// class whose constructor name matches `ctor_name`. Used by `util.queryObjects`
+// / `v8.queryObjects` to implement `{ format: 'count' }` without exposing
+// `HeapProfiler::QueryObjects` (which the rusty_v8 crate does not bind).
+//
+// Limitation: matches by the immediate constructor name only, so instances of
+// subclasses of `ctor` won't be counted. This is sufficient for Node's leak
+// tests (which check direct instances of `Channel`, `SourceTextModule`, ...).
+#[op2(nofast)]
+#[smi]
+pub fn op_v8_query_objects_count(
+  scope: &mut v8::PinScope<'_, '_>,
+  #[string] ctor_name: &str,
+) -> u32 {
+  use deno_core::serde_json;
+  use deno_core::serde_json::Value;
+
+  let mut buf = Vec::new();
+  scope.take_heap_snapshot(|chunk| {
+    buf.extend_from_slice(chunk);
+    true
+  });
+  if buf.is_empty() {
+    return 0;
+  }
+
+  let snapshot: Value = match serde_json::from_slice(&buf) {
+    Ok(v) => v,
+    Err(_) => return 0,
+  };
+
+  let meta = match snapshot.get("snapshot").and_then(|s| s.get("meta")) {
+    Some(m) => m,
+    None => return 0,
+  };
+  let node_fields = match meta.get("node_fields").and_then(|f| f.as_array()) {
+    Some(a) => a,
+    None => return 0,
+  };
+  let node_field_count = node_fields.len();
+  if node_field_count == 0 {
+    return 0;
+  }
+  let type_field_index = node_fields.iter().position(|f| f == "type");
+  let name_field_index = node_fields.iter().position(|f| f == "name");
+  let (Some(type_field_index), Some(name_field_index)) =
+    (type_field_index, name_field_index)
+  else {
+    return 0;
+  };
+
+  // `node_types` is an array where the entry at `type_field_index` is the
+  // list of named type variants (the rest are scalars like "string"/"number").
+  let object_type_index = match meta
+    .get("node_types")
+    .and_then(|t| t.as_array())
+    .and_then(|t| t.get(type_field_index))
+    .and_then(|t| t.as_array())
+  {
+    Some(types) => match types.iter().position(|t| t == "object") {
+      Some(i) => i as u64,
+      None => return 0,
+    },
+    None => return 0,
+  };
+
+  let nodes = match snapshot.get("nodes").and_then(|n| n.as_array()) {
+    Some(a) => a,
+    None => return 0,
+  };
+  let strings = match snapshot.get("strings").and_then(|s| s.as_array()) {
+    Some(a) => a,
+    None => return 0,
+  };
+
+  let mut count: u32 = 0;
+  for chunk in nodes.chunks_exact(node_field_count) {
+    let Some(ty) = chunk[type_field_index].as_u64() else {
+      continue;
+    };
+    if ty != object_type_index {
+      continue;
+    }
+    let Some(name_idx) = chunk[name_field_index].as_u64() else {
+      continue;
+    };
+    let Some(name) = strings.get(name_idx as usize).and_then(|s| s.as_str())
+    else {
+      continue;
+    };
+    if name == ctor_name {
+      count = count.saturating_add(1);
+    }
+  }
+  count
+}
+
 #[op2(fast)]
 pub fn op_v8_get_heap_code_statistics(
   scope: &mut v8::PinScope<'_, '_>,

ext/node/polyfills/util.ts

@@ -75,6 +75,7 @@ const { default: binding } = core.loadExtScript(
 const { validateOneOf } = core.loadExtScript(
   "ext:deno_node/internal/validators.mjs",
 );
+const lazyV8 = core.createLazyLoader("node:v8");
 const { os: osConstants } = core.loadExtScript(
   "ext:deno_node/internal_binding/constants.ts",
 );
@@ -344,6 +345,12 @@ function setTraceSigInt(enabled) {
   // facility, but the call should succeed so user code can opt in/out.
 }
 
+// https://nodejs.org/api/util.html#utilqueryobjectsconstructor-options
+// Mirrors `v8.queryObjects` - see ext/node/polyfills/v8.ts for the limitations.
+function queryObjects(ctor, options) {
+  return lazyV8().queryObjects(ctor, options);
+}
+
 function convertProcessSignalToExitCode(signalCode) {
   const { signals } = osConstants;
   validateOneOf(signalCode, "signalCode", ObjectKeys(signals));
@@ -375,6 +382,7 @@ return {
   aborted,
   getCallSites,
   parseEnv,
+  queryObjects,
   setTraceSigInt,
   convertProcessSignalToExitCode,
   getSystemErrorMessage,

ext/node/polyfills/v8.ts

@@ -25,6 +25,7 @@ const {
   op_v8_read_value,
   op_v8_release_buffer,
   op_v8_set_treat_array_buffer_views_as_host_objects,
+  op_v8_query_objects_count,
   op_v8_take_heap_snapshot,
   op_v8_transfer_array_buffer,
   op_v8_transfer_array_buffer_de,
@@ -48,9 +49,10 @@ const { isArrayBufferView } = core.loadExtScript(
 const lazyFsUtils = core.createLazyLoader(
   "ext:deno_node/internal/fs/utils.mjs",
 );
-const { validateObject } = core.loadExtScript(
-  "ext:deno_node/internal/validators.mjs",
-);
+const { validateFunction, validateObject, validateOneOf } = core
+  .loadExtScript(
+    "ext:deno_node/internal/validators.mjs",
+  );
 
 function cachedDataVersionTag() {
   return op_v8_cached_data_version_tag();
@@ -167,6 +169,44 @@ function writeHeapSnapshot(
   return filename;
 }
 
+// https://nodejs.org/api/v8.html#v8queryobjectsctor-options
+//
+// Deno currently only supports `{ format: 'count' }`. Returning live instances
+// would require V8's `HeapProfiler::QueryObjects`, which isn't exposed in the
+// rusty_v8 bindings; the count form is what Node's leak tests rely on.
+function queryObjects(
+  ctor: { name?: string; prototype?: unknown },
+  options:
+    | { format?: "count" | "summary" }
+    | undefined = undefined,
+) {
+  validateFunction(ctor, "constructor");
+  if (options !== undefined) {
+    validateObject(options, "options");
+    if (options.format !== undefined) {
+      validateOneOf(options.format, "options.format", ["count", "summary"]);
+    }
+  }
+  const format = options?.format;
+
+  const name = typeof ctor.name === "string" ? ctor.name : "";
+  if (name === "") {
+    return format === "count" ? 0 : [];
+  }
+  const count = op_v8_query_objects_count(name);
+  if (format === "count") {
+    return count;
+  }
+  if (format === "summary") {
+    if (count === 0) return [];
+    return [`${count} instance(s) of ${name}`];
+  }
+  // Default format returns live object handles, which would require V8's
+  // `HeapProfiler::QueryObjects` (not exposed in rusty_v8). Returning an
+  // empty array keeps the signature sensible.
+  return [];
+}
+
 // deno-lint-ignore no-explicit-any
 function serialize(value: any) {
   const ser = new DefaultSerializer();
@@ -396,6 +436,7 @@ return {
   getHeapSnapshot,
   getHeapSpaceStatistics,
   getHeapStatistics,
+  queryObjects,
   setFlagsFromString,
   stopCoverage,
   takeCoverage,

ext/node/polyfills/v8_esm.ts

@@ -8,6 +8,7 @@ export const {
   getHeapSnapshot,
   getHeapSpaceStatistics,
   getHeapStatistics,
+  queryObjects,
   setFlagsFromString,
   stopCoverage,
   takeCoverage,

tests/node_compat/config.jsonc

@@ -130,6 +130,7 @@
     "es-module/test-typescript-transform.mjs": {},
     "es-module/test-typescript.mjs": {},
     "es-module/test-vm-compile-function-lineoffset.js": {},
+    "es-module/test-vm-source-text-module-leak.js": {},
     "es-module/test-wasm-memory-out-of-bound.js": {},
     "es-module/test-wasm-simple.js": {},
     "internet/test-dns-ipv4.js": {},
@@ -948,6 +949,7 @@
     "parallel/test-diagnostic-channel-http-response-created.js": {},
     "parallel/test-diagnostics-channel-bind-store.js": {},
     "parallel/test-diagnostics-channel-has-subscribers.js": {},
+    "parallel/test-diagnostics-channel-memory-leak.js": {},
     "parallel/test-diagnostics-channel-http-server-start.js": {},
     "parallel/test-diagnostics-channel-http.js": {},
     "parallel/test-diagnostics-channel-http2-client-stream-close-error.js": {},

tests/unit_node/util_test.ts

@@ -244,6 +244,24 @@ Deno.test("[util] stripVTControlCharacters() removes OSC 8 hyperlinks", () => {
   assertEquals(util.stripVTControlCharacters(inputBel), "This is a link hello");
 });
 
+Deno.test("[util] queryObjects() counts instances", () => {
+  class UtilQueryObjectsFixture {}
+  // util.queryObjects is not declared on the bundled @types/node yet, but the
+  // runtime exposes it (mirroring v8.queryObjects).
+  // deno-lint-ignore no-explicit-any
+  const queryObjects = (util as any).queryObjects as (
+    ctor: unknown,
+    options?: { format?: "count" | "summary" },
+  ) => number | string[];
+  const before = queryObjects(UtilQueryObjectsFixture, { format: "count" });
+  const refs = [];
+  for (let i = 0; i < 25; i++) refs.push(new UtilQueryObjectsFixture());
+  const after = queryObjects(UtilQueryObjectsFixture, { format: "count" });
+  assertEquals(typeof before, "number");
+  assertEquals((after as number) - (before as number) >= 25, true);
+  assertEquals(refs.length, 25);
+});
+
 Deno.test("[util] parseEnv()", () => {
   const env =
     "KEY1=VALUE1\nKEY2='VALUE2'\nKEYÄ3=\"VALUE3\"\nKEY4=VALÜE4\nKEY5='VALUE6'INVALID_LINE\nKEY6=A";

tests/unit_node/v8_test.ts

@@ -69,3 +69,59 @@ Deno.test({
     }, Deno.errors.NotCapable);
   },
 });
+
+Deno.test({
+  name: "queryObjects counts instances by constructor",
+  fn() {
+    class QueryObjectsTestFixture {}
+    const before = v8.queryObjects(QueryObjectsTestFixture, {
+      format: "count",
+    });
+    assertEquals(typeof before, "number");
+    const instances = [];
+    for (let i = 0; i < 50; i++) {
+      instances.push(new QueryObjectsTestFixture());
+    }
+    const after = v8.queryObjects(QueryObjectsTestFixture, {
+      format: "count",
+    });
+    assertEquals(after - before >= 50, true);
+
+    const summary = v8.queryObjects(QueryObjectsTestFixture, {
+      format: "summary",
+    });
+    assertEquals(Array.isArray(summary), true);
+    assertEquals((summary as string[]).length, 1);
+    assertEquals(
+      (summary as string[])[0].includes("QueryObjectsTestFixture"),
+      true,
+    );
+
+    // Keep the instances reachable until after the snapshot.
+    assertEquals(instances.length, 50);
+  },
+});
+
+Deno.test({
+  name: "queryObjects validates the constructor argument",
+  fn() {
+    assertThrows(() => {
+      // @ts-expect-error testing invalid input
+      v8.queryObjects("not a function");
+    });
+  },
+});
+
+Deno.test({
+  name: "queryObjects validates the format option",
+  fn() {
+    class Anything {}
+    assertThrows(() => {
+      v8.queryObjects(
+        Anything,
+        // @ts-expect-error testing invalid input
+        { format: "bogus" },
+      );
+    });
+  },
+});