From 578cca0bc7282701a22d26a2139a3b00d23229fe Mon Sep 17 00:00:00 2001
From: Adrien <adrien@huggingface.co>
Date: Wed, 5 Jul 2023 09:39:56 +0200
Subject: [PATCH] Add more stats

---
 lib/defaultMetrics.js          |  2 ++
 lib/metrics/heapSizeAndUsed.js | 11 +++++++++++
 lib/metrics/heapStats.js       | 30 ++++++++++++++++++++++++++++++
 3 files changed, 43 insertions(+)
 create mode 100644 lib/metrics/heapStats.js

diff --git a/lib/defaultMetrics.js b/lib/defaultMetrics.js
index f285981a..9caa146b 100644
--- a/lib/defaultMetrics.js
+++ b/lib/defaultMetrics.js
@@ -13,6 +13,7 @@ const processHandles = require('./metrics/processHandles');
 const processRequests = require('./metrics/processRequests');
 const processResources = require('./metrics/processResources');
 const heapSizeAndUsed = require('./metrics/heapSizeAndUsed');
+const heapStat = require('./metrics/heapStats');
 const heapSpacesSizeAndUsed = require('./metrics/heapSpacesSizeAndUsed');
 const version = require('./metrics/version');
 const gc = require('./metrics/gc');
@@ -30,6 +31,7 @@ const metrics = {
 	processHandles,
 	processRequests,
 	heapSizeAndUsed,
+	heapStat,
 	heapSpacesSizeAndUsed,
 	version,
 	gc,
diff --git a/lib/metrics/heapSizeAndUsed.js b/lib/metrics/heapSizeAndUsed.js
index ae8a145d..636cb9cc 100644
--- a/lib/metrics/heapSizeAndUsed.js
+++ b/lib/metrics/heapSizeAndUsed.js
@@ -6,6 +6,7 @@ const safeMemoryUsage = require('./helpers/safeMemoryUsage');
 const NODEJS_HEAP_SIZE_TOTAL = 'nodejs_heap_size_total_bytes';
 const NODEJS_HEAP_SIZE_USED = 'nodejs_heap_size_used_bytes';
 const NODEJS_EXTERNAL_MEMORY = 'nodejs_external_memory_bytes';
+const NODEJS_ARRAY_BUFFER = 'nodejs_array_buffers_bytes';
 
 module.exports = (registry, config = {}) => {
 	if (typeof process.memoryUsage !== 'function') {
@@ -24,6 +25,9 @@ module.exports = (registry, config = {}) => {
 			if (memUsage.external !== undefined) {
 				externalMemUsed.set(labels, memUsage.external);
 			}
+			if (memUsage.arrayBuffers !== undefined) {
+				arrayBufferUsed.set(labels, memUsage.arrayBuffers);
+			}
 		}
 	};
 
@@ -47,10 +51,17 @@ module.exports = (registry, config = {}) => {
 		registers,
 		labelNames,
 	});
+	const arrayBufferUsed = new Gauge({
+		name: namePrefix + NODEJS_ARRAY_BUFFER,
+		help: 'Node.js arrayBuffers size in bytes.',
+		registers,
+		labelNames,
+	});
 };
 
 module.exports.metricNames = [
 	NODEJS_HEAP_SIZE_TOTAL,
 	NODEJS_HEAP_SIZE_USED,
 	NODEJS_EXTERNAL_MEMORY,
+	NODEJS_ARRAY_BUFFER,
 ];
diff --git a/lib/metrics/heapStats.js b/lib/metrics/heapStats.js
new file mode 100644
index 00000000..cc4e5c2f
--- /dev/null
+++ b/lib/metrics/heapStats.js
@@ -0,0 +1,30 @@
+'use strict';
+
+const Gauge = require('../gauge');
+const v8 = require('v8');
+
+const NODEJS_HEAP_STATS = 'nodejs_heap_stats';
+
+module.exports = (registry, config = {}) => {
+	const registers = registry ? [registry] : undefined;
+	const namePrefix = config.prefix ? config.prefix : '';
+
+	const labels = config.labels ? config.labels : {};
+	const labelNames = ['stat', ...Object.keys(labels)];
+
+	const gauge = new Gauge({
+		name: namePrefix + NODEJS_HEAP_STATS,
+		help: `v8.getHeapStatistics() stats`,
+		labelNames,
+		registers,
+	});
+
+	// Use this one metric's `collect` to set all metrics' values.
+	gauge.collect = () => {
+		for (const [stat, value] of Object.entries(v8.getHeapStatistics())) {
+			gauge.set({stat: stat, ...labels}, value)
+		}
+	};
+};
+
+module.exports.metricNames = Object.values(NODEJS_HEAP_STATS);