From 4e39087890149cfe98e20552cbd50e8c6e861b9b Mon Sep 17 00:00:00 2001
From: Jithil P Ponnan <jithil@outlook.com>
Date: Sat, 20 May 2023 16:01:39 +1000
Subject: [PATCH] process: add custom dir support for heapsnapshot-signal

---
 lib/internal/process/pre_execution.js         | 37 +++++++++++++-
 .../test-heapdump-flag-custom-dir.js          | 50 +++++++++++++++++++
 2 files changed, 86 insertions(+), 1 deletion(-)
 create mode 100644 test/sequential/test-heapdump-flag-custom-dir.js

diff --git a/lib/internal/process/pre_execution.js b/lib/internal/process/pre_execution.js
index 89b916a6ee66ae..c5714e23145c02 100644
--- a/lib/internal/process/pre_execution.js
+++ b/lib/internal/process/pre_execution.js
@@ -11,6 +11,14 @@ const {
   Symbol,
   SymbolDispose,
   SymbolAsyncDispose,
+  Date,
+  DatePrototypeGetFullYear,
+  DatePrototypeGetMonth,
+  DatePrototypeGetDate,
+  DatePrototypeGetHours,
+  DatePrototypeGetMinutes,
+  DatePrototypeGetSeconds,
+  String,
   globalThis,
 } = primordials;
 
@@ -400,6 +408,7 @@ function initializeReportSignalHandlers() {
 
 function initializeHeapSnapshotSignalHandlers() {
   const signal = getOptionValue('--heapsnapshot-signal');
+  const diagnosticDir = getOptionValue('--diagnostic-dir');
 
   if (!signal)
     return;
@@ -408,7 +417,8 @@ function initializeHeapSnapshotSignalHandlers() {
   const { writeHeapSnapshot } = require('v8');
 
   function doWriteHeapSnapshot() {
-    writeHeapSnapshot();
+    const heapSnapshotFilename = getHeapSnapshotFilename(diagnosticDir);
+    writeHeapSnapshot(heapSnapshotFilename);
   }
   process.on(signal, doWriteHeapSnapshot);
 
@@ -700,6 +710,31 @@ function markBootstrapComplete() {
   internalBinding('performance').markBootstrapComplete();
 }
 
+// Sequence number for diagnostic filenames
+let sequenceNumOfheapSnapshot = 0;
+
+// To generate the HeapSnapshotFilename while using custom diagnosticDir
+function getHeapSnapshotFilename(diagnosticDir) {
+  if (!diagnosticDir) return undefined;
+
+  const date = new Date();
+
+  const year = DatePrototypeGetFullYear(date);
+  const month = String(DatePrototypeGetMonth(date) + 1).padStart(2, '0');
+  const day = String(DatePrototypeGetDate(date)).padStart(2, '0');
+  const hours = String(DatePrototypeGetHours(date)).padStart(2, '0');
+  const minutes = String(DatePrototypeGetMinutes(date)).padStart(2, '0');
+  const seconds = String(DatePrototypeGetSeconds(date)).padStart(2, '0');
+
+  const dateString = `${year}${month}${day}`;
+  const timeString = `${hours}${minutes}${seconds}`;
+  const pid = process.pid;
+  const threadId = internalBinding('worker').threadId;
+  const fileSequence = (++sequenceNumOfheapSnapshot).toString().padStart(3, '0');
+
+  return `${diagnosticDir}/Heap.${dateString}.${timeString}.${pid}.${threadId}.${fileSequence}.heapsnapshot`;
+}
+
 module.exports = {
   setupUserModules,
   prepareMainThreadExecution,
diff --git a/test/sequential/test-heapdump-flag-custom-dir.js b/test/sequential/test-heapdump-flag-custom-dir.js
new file mode 100644
index 00000000000000..ea53cb660441d3
--- /dev/null
+++ b/test/sequential/test-heapdump-flag-custom-dir.js
@@ -0,0 +1,50 @@
+'use strict';
+const common = require('../common');
+
+
+if (common.isWindows)
+  common.skip('test not supported on Windows');
+
+const assert = require('assert');
+
+const validateHeapSnapshotFile = () => {
+  const fs = require('fs');
+
+  assert.strictEqual(process.listenerCount('SIGUSR2'), 1);
+  process.kill(process.pid, 'SIGUSR2');
+  process.kill(process.pid, 'SIGUSR2');
+
+  // Asynchronously wait for the snapshot. Use an async loop to be a bit more
+  // robust in case platform or machine differences throw off the timing.
+  (function validate() {
+    const files = fs.readdirSync(process.cwd());
+
+    if (files.length === 0)
+      return setImmediate(validate);
+
+    assert.strictEqual(files.length, 2);
+
+    for (let i = 0; i < files.length; i++) {
+      assert.match(files[i], /^Heap\..+\.heapsnapshot$/);
+
+      // Check the file is parsable as JSON
+      JSON.parse(fs.readFileSync(files[i]));
+    }
+  })();
+};
+
+if (process.argv[2] === 'child') {
+  validateHeapSnapshotFile();
+} else {
+  // Modify the timezone. So we can check the file date string still returning correctly.
+  process.env.TZ = 'America/New_York';
+  const { spawnSync } = require('child_process');
+  const tmpdir = require('../common/tmpdir');
+
+  tmpdir.refresh();
+  const args = ['--heapsnapshot-signal', 'SIGUSR2', '--diagnostic-dir', tmpdir.path, __filename, 'child'];
+  const child = spawnSync(process.execPath, args, { cwd: tmpdir.path });
+
+  assert.strictEqual(child.status, 0);
+  assert.strictEqual(child.signal, null);
+}