feat: use inspector for heap profiles

psx95 · psx95 · commit 23eab92d983c · 2022-11-26T04:07:05.000Z
diff --git a/ts/src/heap-profiler-inspector.ts b/ts/src/heap-profiler-inspector.ts
@@ -0,0 +1,180 @@
+import * as inspector from 'node:inspector';
+import {AllocationProfileNode, Allocation} from './v8-types';
+
+const session = new inspector.Session();
+
+export interface SamplingHeapProfileSample {
+  size: number;
+  nodeId: number;
+  ordinal: number;
+}
+
+export interface SamplingHeapProfileNode {
+  callFrame: inspector.Runtime.CallFrame;
+  selfSize: number;
+  id: number;
+  children: SamplingHeapProfileNode[];
+}
+
+/**
+ * Need to create this interface since the type definitions file for node inspector
+ * at https://github.com/DefinitelyTyped/DefinitelyTyped/blob/master/types/node/inspector.d.ts
+ * has not been updated with the latest changes yet.
+ *
+ * The types defined through this interface are in sync with the documentation found at -
+ * https://chromedevtools.github.io/devtools-protocol/v8/HeapProfiler/
+ */
+export interface CompatibleSamplingHeapProfile {
+  head: SamplingHeapProfileNode;
+  samples: SamplingHeapProfileSample[];
+}
+
+export function startSamplingHeapProfiler(
+  heapIntervalBytes: number,
+  stackDepth: number
+): Promise<void> {
+  session.connect();
+  return new Promise<void>((resolve, reject) => {
+    session.post(
+      'HeapProfiler.startSampling',
+      {heapIntervalBytes},
+      (err: Error | null): void => {
+        if (err !== null) {
+          console.error(`Error starting heap sampling: ${err}`);
+          reject(err);
+          return;
+        }
+        console.log(
+          `Started Heap Sampling with interval bytes ${heapIntervalBytes}`
+        );
+        resolve();
+      }
+    );
+  });
+}
+
+/**
+ * Stops the sampling heap profile and discards the current profile.
+ */
+export function stopSamplingHeapProfiler(): Promise<void> {
+  return new Promise<void>((resolve, reject) => {
+    session.post(
+      'HeapProfiler.stopSampling',
+      (
+        err: Error | null,
+        profile: inspector.HeapProfiler.StopSamplingReturnType
+      ) => {
+        if (err !== null) {
+          console.error(`Error stopping heap sampling ${err}`);
+          reject(err);
+          return;
+        }
+        console.log(
+          `Stopped sampling heap, discarding current profile: ${profile}`
+        );
+        session.disconnect();
+        console.log('Disconnected from current profiling session');
+        resolve();
+      }
+    );
+  });
+}
+
+export async function getAllocationProfile(): Promise<AllocationProfileNode> {
+  return new Promise<AllocationProfileNode>((resolve, reject) => {
+    session.post(
+      'HeapProfiler.getSamplingProfile',
+      (
+        err: Error | null,
+        result: inspector.HeapProfiler.GetSamplingProfileReturnType
+      ) => {
+        if (err !== null) {
+          console.error(`Error getting sampling profile ${err}`);
+          reject(err);
+          return;
+        }
+        const compatibleHeapProfile =
+          result.profile as CompatibleSamplingHeapProfile;
+        resolve(
+          translateToAllocationProfileNode(
+            compatibleHeapProfile.head,
+            compatibleHeapProfile.samples
+          )
+        );
+      }
+    );
+  });
+}
+
+function translateToAllocationProfileNode(
+  node: SamplingHeapProfileNode,
+  samples: SamplingHeapProfileSample[]
+): AllocationProfileNode {
+  const allocationProfileNode: AllocationProfileNode = {
+    allocations: [],
+    name: node.callFrame.functionName,
+    scriptName: node.callFrame.url,
+    scriptId: Number(node.callFrame.scriptId),
+    lineNumber: node.callFrame.lineNumber,
+    columnNumber: node.callFrame.columnNumber,
+    children: [],
+  };
+
+  const children: AllocationProfileNode[] = new Array<AllocationProfileNode>(
+    node.children.length
+  );
+  for (let i = 0; i < node.children.length; i++) {
+    children.splice(
+      i,
+      1,
+      translateToAllocationProfileNode(node.children[i], samples)
+    );
+  }
+  allocationProfileNode.children = children;
+
+  // find all samples belonging to this node Id
+  const samplesForCurrentNodeId: SamplingHeapProfileSample[] =
+    filterSamplesBasedOnNodeId(node.id, samples);
+  const mappedAllocationsForNodeId: Allocation[] =
+    createAllocationsFromSamplesForNode(samplesForCurrentNodeId);
+
+  allocationProfileNode.allocations = mappedAllocationsForNodeId;
+  return allocationProfileNode;
+}
+
+function filterSamplesBasedOnNodeId(
+  nodeId: number,
+  samples: SamplingHeapProfileSample[]
+): SamplingHeapProfileSample[] {
+  const filtered = samples.filter((sample: SamplingHeapProfileSample) => {
+    return sample.nodeId === nodeId;
+  });
+  return filtered;
+}
+
+function createAllocationsFromSamplesForNode(
+  samplesForNode: SamplingHeapProfileSample[]
+): Allocation[] {
+  const sampleSizeToCountMap = new Map<number, number>();
+  samplesForNode.forEach((sample: SamplingHeapProfileSample) => {
+    const currentCountForSize: number | undefined = sampleSizeToCountMap.get(
+      sample.size
+    );
+    if (currentCountForSize !== undefined) {
+      sampleSizeToCountMap.set(sample.size, currentCountForSize + 1);
+    } else {
+      sampleSizeToCountMap.set(sample.size, 1);
+    }
+  });
+
+  const mappedAllocations: Allocation[] = [];
+  sampleSizeToCountMap.forEach((size: number, count: number) => {
+    const mappedAllocation: Allocation = {
+      sizeBytes: size,
+      count: count,
+    };
+    mappedAllocations.push(mappedAllocation);
+  });
+
+  return mappedAllocations;
+}
diff --git a/ts/src/heap-profiler.ts b/ts/src/heap-profiler.ts
@@ -20,26 +20,25 @@ import {
   getAllocationProfile,
   startSamplingHeapProfiler,
   stopSamplingHeapProfiler,
-} from './heap-profiler-bindings';
+} from './heap-profiler-inspector';
 import {serializeHeapProfile} from './profile-serializer';
 import {SourceMapper} from './sourcemapper/sourcemapper';
 import {AllocationProfileNode} from './v8-types';
 
 let enabled = false;
 let heapIntervalBytes = 0;
-let heapStackDepth = 0;
 
 /*
  * Collects a heap profile when heapProfiler is enabled. Otherwise throws
  * an error.
  *
  * Data is returned in V8 allocation profile format.
  */
-export function v8Profile(): AllocationProfileNode {
+export async function v8Profile(): Promise<AllocationProfileNode> {
   if (!enabled) {
     throw new Error('Heap profiler is not enabled.');
   }
-  return getAllocationProfile();
+  return await getAllocationProfile();
 }
 
 /**
@@ -49,12 +48,12 @@ export function v8Profile(): AllocationProfileNode {
  * @param ignoreSamplePath
  * @param sourceMapper
  */
-export function profile(
+export async function profile(
   ignoreSamplePath?: string,
   sourceMapper?: SourceMapper
-): perftools.profiles.IProfile {
+): Promise<perftools.profiles.IProfile> {
   const startTimeNanos = Date.now() * 1000 * 1000;
-  const result = v8Profile();
+  const result = await v8Profile();
   // Add node for external memory usage.
   // Current type definitions do not have external.
   // TODO: remove any once type definition is updated to include external.
@@ -84,17 +83,17 @@ export function profile(
  * started with different parameters, this throws an error.
  *
  * @param intervalBytes - average number of bytes between samples.
- * @param stackDepth - maximum stack depth for samples collected.
+ * @param stackDepth - maximum stack depth for samples collected. This is currently no-op.
+ *                     Default stack depth of 128 will be used. Kept to avoid making breaking change.
  */
 export function start(intervalBytes: number, stackDepth: number) {
   if (enabled) {
     throw new Error(
-      `Heap profiler is already started  with intervalBytes ${heapIntervalBytes} and stackDepth ${stackDepth}`
+      `Heap profiler is already started  with intervalBytes ${heapIntervalBytes} and stackDepth 128`
     );
   }
   heapIntervalBytes = intervalBytes;
-  heapStackDepth = stackDepth;
-  startSamplingHeapProfiler(heapIntervalBytes, heapStackDepth);
+  startSamplingHeapProfiler(heapIntervalBytes, stackDepth);
   enabled = true;
 }
 
diff --git a/ts/test/test-heap-profiler.ts b/ts/test/test-heap-profiler.ts
@@ -17,7 +17,7 @@
 import * as sinon from 'sinon';
 
 import * as heapProfiler from '../src/heap-profiler';
-import * as v8HeapProfiler from '../src/heap-profiler-bindings';
+import * as inspectorHeapProfiler from '../src/heap-profiler-inspector';
 import {AllocationProfileNode} from '../src/v8-types';
 
 import {
@@ -32,14 +32,14 @@ const copy = require('deep-copy');
 const assert = require('assert');
 
 describe('HeapProfiler', () => {
-  let startStub: sinon.SinonStub<[number, number], void>;
-  let stopStub: sinon.SinonStub<[], void>;
-  let profileStub: sinon.SinonStub<[], AllocationProfileNode>;
+  let startStub: sinon.SinonStub<[number, number], Promise<void>>;
+  let stopStub: sinon.SinonStub<[], Promise<void>>;
+  let profileStub: sinon.SinonStub<[], Promise<AllocationProfileNode>>;
   let dateStub: sinon.SinonStub<[], number>;
   let memoryUsageStub: sinon.SinonStub<[], NodeJS.MemoryUsage>;
   beforeEach(() => {
-    startStub = sinon.stub(v8HeapProfiler, 'startSamplingHeapProfiler');
-    stopStub = sinon.stub(v8HeapProfiler, 'stopSamplingHeapProfiler');
+    startStub = sinon.stub(inspectorHeapProfiler, 'startSamplingHeapProfiler');
+    stopStub = sinon.stub(inspectorHeapProfiler, 'stopSamplingHeapProfiler');
     dateStub = sinon.stub(Date, 'now').returns(0);
   });
 
@@ -54,7 +54,7 @@ describe('HeapProfiler', () => {
   describe('profile', () => {
     it('should return a profile equal to the expected profile when external memory is allocated', async () => {
       profileStub = sinon
-        .stub(v8HeapProfiler, 'getAllocationProfile')
+        .stub(inspectorHeapProfiler, 'getAllocationProfile')
         .returns(copy(v8HeapProfile));
       memoryUsageStub = sinon.stub(process, 'memoryUsage').returns({
         external: 1024,
@@ -66,13 +66,13 @@ describe('HeapProfiler', () => {
       const intervalBytes = 1024 * 512;
       const stackDepth = 32;
       heapProfiler.start(intervalBytes, stackDepth);
-      const profile = heapProfiler.profile();
+      const profile = await heapProfiler.profile();
       assert.deepEqual(heapProfileWithExternal, profile);
     });
 
     it('should return a profile equal to the expected profile when including all samples', async () => {
       profileStub = sinon
-        .stub(v8HeapProfiler, 'getAllocationProfile')
+        .stub(inspectorHeapProfiler, 'getAllocationProfile')
         .returns(copy(v8HeapWithPathProfile));
       memoryUsageStub = sinon.stub(process, 'memoryUsage').returns({
         external: 0,
@@ -84,13 +84,13 @@ describe('HeapProfiler', () => {
       const intervalBytes = 1024 * 512;
       const stackDepth = 32;
       heapProfiler.start(intervalBytes, stackDepth);
-      const profile = heapProfiler.profile();
+      const profile = await heapProfiler.profile();
       assert.deepEqual(heapProfileIncludePath, profile);
     });
 
     it('should return a profile equal to the expected profile when excluding profiler samples', async () => {
       profileStub = sinon
-        .stub(v8HeapProfiler, 'getAllocationProfile')
+        .stub(inspectorHeapProfiler, 'getAllocationProfile')
         .returns(copy(v8HeapWithPathProfile));
       memoryUsageStub = sinon.stub(process, 'memoryUsage').returns({
         external: 0,
@@ -102,14 +102,14 @@ describe('HeapProfiler', () => {
       const intervalBytes = 1024 * 512;
       const stackDepth = 32;
       heapProfiler.start(intervalBytes, stackDepth);
-      const profile = heapProfiler.profile('@google-cloud/profiler');
+      const profile = await heapProfiler.profile('@google-cloud/profiler');
       assert.deepEqual(heapProfileExcludePath, profile);
     });
 
     it('should throw error when not started', () => {
-      assert.throws(
-        () => {
-          heapProfiler.profile();
+      assert.rejects(
+        async () => {
+          await heapProfiler.profile();
         },
         (err: Error) => {
           return err.message === 'Heap profiler is not enabled.';
@@ -122,9 +122,9 @@ describe('HeapProfiler', () => {
       const stackDepth = 32;
       heapProfiler.start(intervalBytes, stackDepth);
       heapProfiler.stop();
-      assert.throws(
-        () => {
-          heapProfiler.profile();
+      assert.rejects(
+        async () => {
+          await heapProfiler.profile();
         },
         (err: Error) => {
           return err.message === 'Heap profiler is not enabled.';
@@ -160,7 +160,7 @@ describe('HeapProfiler', () => {
         assert.strictEqual(
           e.message,
           'Heap profiler is already started  with intervalBytes 524288 and' +
-            ' stackDepth 64'
+            ' stackDepth 128'
         );
       }
       assert.ok(
