From 87b92a7a624c3d8401d98c9218a00c03d4ddcac4 Mon Sep 17 00:00:00 2001
From: hyperdefined <contact@hyper.lol>
Date: Wed, 4 Jun 2025 00:43:34 -0400
Subject: [PATCH] metrics: support clusters

---
 api/src/core/api.js     | 35 ++++++++++++++++++++---------------
 api/src/misc/metrics.js | 33 +++++++++++++++++++--------------
 2 files changed, 39 insertions(+), 29 deletions(-)

diff --git a/api/src/core/api.js b/api/src/core/api.js
index eb93563a..aeb856ee 100644
--- a/api/src/core/api.js
+++ b/api/src/core/api.js
@@ -3,8 +3,7 @@ import http from "node:http";
 import rateLimit from "express-rate-limit";
 import { setGlobalDispatcher, ProxyAgent } from "undici";
 import { getCommit, getBranch, getRemote, getVersion } from "@imput/version-info";
-import registry from "../misc/metrics.js"
-import { httpRequests, httpRequestDuration } from "../misc/metrics.js"
+import { httpRequests, httpRequestDuration, WORKER_ID, aggregatorRegistry } from "../misc/metrics.js";
 
 import jwt from "../security/jwt.js";
 import stream from "../stream/stream.js";
@@ -118,11 +117,11 @@ export const runAPI = async (express, app, __dirname, isPrimary = true) => {
 
     if (metrics) {
         app.use((req, res, next) => {
-            const end = httpRequestDuration.startTimer({ method: req.method });
-          
+            const end = httpRequestDuration.startTimer({ method: req.method, worker_id: WORKER_ID });
+
             res.on('finish', () => {
-              httpRequests.labels(req.method, res.statusCode.toString()).inc();
-              end();
+                httpRequests.labels(req.method, res.statusCode.toString(), WORKER_ID).inc();
+                end();
             });
           
             next();
@@ -389,22 +388,28 @@ export const runAPI = async (express, app, __dirname, isPrimary = true) => {
             YouTubeSession.setup();
         }
 
-        if (metrics) {
+        if (metrics && isPrimary) {
             const metricsApp = express();
-
+        
             metricsApp.get('/metrics', async (req, res) => {
-                res.set('Content-Type', registry.contentType);
-                res.send(await registry.metrics());
-            });
-            metricsApp.listen(env.metricsPort, '127.0.0.1', () => {
-                console.log(`${Green('[✓]')} prometheus metrics running on 127.0.0.1:${env.metricsPort}/metrics`);
+                try {
+                    const data = await aggregatorRegistry.clusterMetrics();
+                    res.set('Content-Type', 'text/plain');
+                    res.end(data);
+                } catch (err) {
+                    res.status(500).end(err.message);
+                }
             });
 
             metricsApp.get('/*', (req, res) => {
                 res.redirect('/metrics');
-            })
+            });
+        
+            metricsApp.listen(env.metricsPort, '127.0.0.1', () => {
+                console.log(`${Green('[✓]')} prometheus metrics running on 127.0.0.1:${env.metricsPort}/metrics`);
+            });
         }
     });
 
     setupTunnelHandler();
-}
+}
\ No newline at end of file
diff --git a/api/src/misc/metrics.js b/api/src/misc/metrics.js
index 29af31e6..02a87278 100644
--- a/api/src/misc/metrics.js
+++ b/api/src/misc/metrics.js
@@ -1,33 +1,38 @@
-import { Registry } from 'prom-client';
-import { collectDefaultMetrics, Counter, Histogram } from "prom-client";
+import { collectDefaultMetrics, Counter, Histogram, Registry, AggregatorRegistry } from 'prom-client';
+  
+import cluster from 'node:cluster';
 
-const registry = new Registry();
-export default registry;
+export const WORKER_ID = `worker_${cluster.worker?.id ?? process.pid}`;
 
-collectDefaultMetrics({ register: registry });
+export const registry = new Registry();
+export const aggregatorRegistry = new AggregatorRegistry();
+
+collectDefaultMetrics({
+    register: registry
+});
 
 export const failedRequests = new Counter({
     name: 'cobalt_fail_request_count',
     help: 'Total number of failed requests',
-    labelNames: ['service'],
+    labelNames: ['service', 'worker_id'],
 });
 
 export const successfulRequests = new Counter({
     name: 'cobalt_success_request_count',
     help: 'Total number of successful requests',
-    labelNames: ['service'],
+    labelNames: ['service', 'worker_id'],
 });
 
 export const httpRequests = new Counter({
     name: 'http_requests_total',
     help: 'Total number of HTTP requests',
-    labelNames: ['method', 'status'],
+    labelNames: ['method', 'status', 'worker_id'],
 });
-  
+
 export const httpRequestDuration = new Histogram({
     name: 'http_request_duration_seconds',
     help: 'Duration of HTTP requests in seconds',
-    labelNames: ['method'],
+    labelNames: ['method', 'worker_id'],
     buckets: [0.1, 0.5, 1, 1.5, 2, 5],
 });
 
@@ -36,10 +41,10 @@ registry.registerMetric(successfulRequests);
 registry.registerMetric(httpRequests);
 registry.registerMetric(httpRequestDuration);
 
-export function incrementFailed(type) {
-    failedRequests.labels(type).inc();
+export function incrementFailed(service) {
+    failedRequests.labels(service, WORKER_ID).inc();
 }
 
-export function incrementSuccessful(type) {
-    successfulRequests.labels(type).inc();
+export function incrementSuccessful(service) {
+    successfulRequests.labels(service, WORKER_ID).inc();
 }
\ No newline at end of file