lib: add heap profile labels API to v8 module

Add JS API for labeling heap profiler samples via AsyncLocalStorage.
Labels are pre-flattened at set time to avoid V8 property access
during resolution.

Signed-off-by: Rudolf Meijering <[email protected]>

Author: rudolf

lib/v8.js

@@ -26,7 +26,9 @@ const {
   Int32Array,
   Int8Array,
   JSONParse,
+  ObjectKeys,
   ObjectPrototypeToString,
+  ReflectGet,
   SymbolDispose,
   Uint16Array,
   Uint32Array,
@@ -39,6 +41,8 @@ const {
 
 const { Buffer } = require('buffer');
 const {
+  validateFunction,
+  validateObject,
   validateString,
   validateUint32,
   validateOneOf,
@@ -156,6 +160,11 @@ const {
   heapSpaceStatisticsBuffer,
   getCppHeapStatistics: _getCppHeapStatistics,
   detailLevel,
+
+  startSamplingHeapProfiler: _startSamplingHeapProfiler,
+  stopSamplingHeapProfiler: _stopSamplingHeapProfiler,
+  getAllocationProfile: _getAllocationProfile,
+  setHeapProfileLabelsStore: _setHeapProfileLabelsStore,
 } = binding;
 
 const kNumberOfHeapSpaces = kHeapSpaces.length;
@@ -494,6 +503,109 @@ class GCProfiler {
   }
 }
 
+// --- Heap profile labels API ---
+// Internal AsyncLocalStorage for propagating labels through async context.
+// Requires --experimental-async-context-frame (Node 22) or Node 24+.
+// Lazily initialized on first use so processes that never use heap profiling
+// pay zero cost (no ALS instance, no async_hooks require).
+let _heapProfileLabelsALS;
+
+function ensureHeapProfileLabelsALS() {
+  if (_heapProfileLabelsALS === undefined) {
+    // When V8_HEAP_PROFILER_SAMPLE_LABELS is compiled out, the C++ binding
+    // for _setHeapProfileLabelsStore is not registered — labels are a no-op.
+    if (typeof _setHeapProfileLabelsStore !== 'function') return;
+    const { AsyncLocalStorage } = require('async_hooks');
+    _heapProfileLabelsALS = new AsyncLocalStorage();
+    // The ALS instance is passed to C++ as the key used to look up labels in
+    // the AsyncContextFrame map (via ContinuationPreservedEmbedderData).
+    // This relies on the async-context-frame implementation storing ALS
+    // instances as Map keys — if that internal representation changes, the
+    // C++ label-resolution code in node_v8.cc must be updated to match.
+    _setHeapProfileLabelsStore(_heapProfileLabelsALS);
+  }
+  return _heapProfileLabelsALS;
+}
+
+/**
+ * Convert a labels object to a flat array [key1, val1, key2, val2, ...].
+ * Pre-flattened at label-set time (not per allocation/sample) because the
+ * C++ callback runs in GetAllocationProfile() BEFORE BuildSamples() where
+ * it resolves labels from CPED captured at allocation time.
+ * @param {Record<string, string>} labels
+ * @returns {string[]}
+ */
+function labelsToFlat(labels) {
+  const keys = ObjectKeys(labels);
+  const len = keys.length;
+  const flat = new Array(len * 2);
+  for (let i = 0; i < len; i++) {
+    const key = keys[i];
+    const val = ReflectGet(labels, key);
+    validateString(val, `labels.${key}`);
+    flat[i * 2] = key;
+    flat[i * 2 + 1] = val;
+  }
+  return flat;
+}
+
+/**
+ * Starts the V8 sampling heap profiler.
+ * @param {number} [sampleInterval] - Average bytes between samples (default 512 KB).
+ * @param {number} [stackDepth] - Maximum stack depth for samples (default 16).
+ * @param {object} [options] - Options object.
+ * @param {boolean} [options.includeCollectedObjects] - If true, retain
+ *   samples for objects collected by GC (allocation-rate mode).
+ */
+function startSamplingHeapProfiler(sampleInterval, stackDepth, options) {
+  if (sampleInterval !== undefined) validateUint32(sampleInterval, 'sampleInterval', true);
+  if (stackDepth !== undefined) validateUint32(stackDepth, 'stackDepth');
+  if (options !== undefined) validateObject(options, 'options');
+  return _startSamplingHeapProfiler(sampleInterval, stackDepth, options);
+}
+
+/**
+ * Runs `fn` with the given heap profile labels active. Labels propagate
+ * across `await` boundaries via AsyncLocalStorage. If `fn` returns a
+ * Promise, labels remain active until the Promise settles.
+ *
+ * @param {Record<string, string>} labels
+ * @param {Function} fn
+ * @returns {*} The return value of `fn`.
+ */
+function withHeapProfileLabels(labels, fn) {
+  validateObject(labels, 'labels');
+  validateFunction(fn, 'fn');
+  // Store the flat [key1, val1, key2, val2, ...] array in ALS.
+  // Conversion happens once at label-set time (not per allocation).
+  // The C++ callback resolves labels from CPED in GetAllocationProfile()
+  // before BuildSamples() — pre-flattening avoids V8 Object property
+  // access during label resolution.
+  const flat = labelsToFlat(labels);
+  const als = ensureHeapProfileLabelsALS();
+  // When labels are compiled out, still run the callback — just without
+  // label tracking.
+  if (als === undefined) return fn();
+  return als.run(flat, fn);
+}
+
+/**
+ * Sets heap profile labels for the current async scope using
+ * `enterWith` semantics. Labels persist until overwritten or the
+ * async scope ends. Useful for frameworks (e.g. Hapi) where the
+ * handler runs after the extension returns.
+ *
+ * @param {Record<string, string>} labels
+ */
+function setHeapProfileLabels(labels) {
+  validateObject(labels, 'labels');
+  const flat = labelsToFlat(labels);
+  const als = ensureHeapProfileLabelsALS();
+  // When labels are compiled out, setHeapProfileLabels is a no-op.
+  if (als === undefined) return;
+  als.enterWith(flat);
+}
+
 module.exports = {
   cachedDataVersionTag,
   getHeapSnapshot,
@@ -518,4 +630,9 @@ module.exports = {
   GCProfiler,
   isStringOneByteRepresentation,
   startCpuProfile,
+  startSamplingHeapProfiler,
+  stopSamplingHeapProfiler: _stopSamplingHeapProfiler,
+  getAllocationProfile: _getAllocationProfile,
+  withHeapProfileLabels,
+  setHeapProfileLabels,
 };