From fed97533a24ba722539409e8ecc107dc019bbfb9 Mon Sep 17 00:00:00 2001
From: Ben Noordhuis <info@bnoordhuis.nl>
Date: Tue, 23 Jan 2018 02:02:29 +0100
Subject: [PATCH] fix: Fix `takeHeapSnapshot()` truncation bug

Fix a logic bug in `takeHeapSnapshot()` where it interpreted the item
count as the byte count, resulting in truncated snapshots.

This change hinges on the observation that the completion callback
always comes after any and all `'addHeapSnapshotChunk'` events.

I'm 95% sure after digging into the V8 inspector and its integration
with Node.js that the assumption above is correct.

If it turns out I'm mistaken, then we will most likely treat that as
a bug in Node.js, not node-inspect.

Fixes: https://github.com/nodejs/node-inspect/issues/56
---
 lib/internal/inspect_repl.js   | 20 ++++++++++----------
 test/cli/heap-profiler.test.js | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 44 insertions(+), 10 deletions(-)
 create mode 100644 test/cli/heap-profiler.test.js

diff --git a/lib/internal/inspect_repl.js b/lib/internal/inspect_repl.js
index 937c184..38fe468 100644
--- a/lib/internal/inspect_repl.js
+++ b/lib/internal/inspect_repl.js
@@ -900,10 +900,8 @@ function createRepl(inspector) {
         return new Promise((resolve, reject) => {
           const absoluteFile = Path.resolve(filename);
           const writer = FS.createWriteStream(absoluteFile);
-          let totalSize;
           let sizeWritten = 0;
           function onProgress({ done, total, finished }) {
-            totalSize = total;
             if (finished) {
               print('Heap snaphost prepared.');
             } else {
@@ -913,13 +911,18 @@ function createRepl(inspector) {
           function onChunk({ chunk }) {
             sizeWritten += chunk.length;
             writer.write(chunk);
-            print(`Writing snapshot: ${sizeWritten}/${totalSize}`, true);
-            if (sizeWritten >= totalSize) {
-              writer.end();
+            print(`Writing snapshot: ${sizeWritten}`, true);
+          }
+          function onResolve() {
+            writer.end(() => {
               teardown();
               print(`Wrote snapshot: ${absoluteFile}`);
               resolve();
-            }
+            });
+          }
+          function onReject(error) {
+            teardown();
+            reject(error);
           }
           function teardown() {
             HeapProfiler.removeListener(
@@ -932,10 +935,7 @@ function createRepl(inspector) {
 
           print('Heap snapshot: 0/0', true);
           HeapProfiler.takeHeapSnapshot({ reportProgress: true })
-            .then(null, (error) => {
-              teardown();
-              reject(error);
-            });
+            .then(onResolve, onReject);
         });
       },
 
diff --git a/test/cli/heap-profiler.test.js b/test/cli/heap-profiler.test.js
new file mode 100644
index 0000000..ebd734e
--- /dev/null
+++ b/test/cli/heap-profiler.test.js
@@ -0,0 +1,34 @@
+'use strict';
+const { test } = require('tap');
+const { readFileSync, unlinkSync } = require('fs');
+
+const startCLI = require('./start-cli');
+const filename = 'node.heapsnapshot';
+
+function cleanup() {
+  try {
+    unlinkSync(filename);
+  } catch (_) {
+    // Ignore.
+  }
+}
+
+cleanup();
+
+test('Heap profiler take snapshot', (t) => {
+  const cli = startCLI(['examples/empty.js']);
+
+  function onFatal(error) {
+    cli.quit();
+    throw error;
+  }
+
+  // Check that the snapshot is valid JSON.
+  return cli.waitForInitialBreak()
+    .then(() => cli.waitForPrompt())
+    .then(() => cli.command('takeHeapSnapshot()'))
+    .then(() => JSON.parse(readFileSync(filename, 'utf8')))
+    .then(() => cleanup())
+    .then(() => cli.quit())
+    .then(null, onFatal);
+});