From ad043bac31f4f126d533d6f2b3d73396ee071872 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

PR-URL: https://github.com/nodejs/node/pull/47854
Refs: https://github.com/nodejs/node/issues/47842
Reviewed-By: James M Snell <jasnell@gmail.com>
Reviewed-By: Benjamin Gruenbaum <benjamingr@gmail.com>
Reviewed-By: Joyee Cheung <joyeec9h3@gmail.com>
Reviewed-By: Erick Wendel <erick.workspace@gmail.com>
---
 lib/internal/process/pre_execution.js         | 39 ++++++++++++++-
 .../test-heapdump-flag-custom-dir.js          | 50 +++++++++++++++++++
 2 files changed, 87 insertions(+), 2 deletions(-)
 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 0627b06be89c2b..1f4a08515b5ae9 100644
--- a/lib/internal/process/pre_execution.js
+++ b/lib/internal/process/pre_execution.js
@@ -2,15 +2,23 @@
 
 const {
   ArrayPrototypeForEach,
+  Date,
+  DatePrototypeGetDate,
+  DatePrototypeGetFullYear,
+  DatePrototypeGetHours,
+  DatePrototypeGetMinutes,
+  DatePrototypeGetMonth,
+  DatePrototypeGetSeconds,
   NumberParseInt,
   ObjectDefineProperties,
   ObjectDefineProperty,
   ObjectGetOwnPropertyDescriptor,
   SafeMap,
+  String,
   StringPrototypeStartsWith,
   Symbol,
-  SymbolDispose,
   SymbolAsyncDispose,
+  SymbolDispose,
   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);
+}