feat: aggregate metrics in a worker thread (#3458)

src/routes/metricsRouter.js

@@ -16,6 +16,16 @@ if (enableStats) {
       ctx.body = error.message;
     }
   });
+
+  metricsRouter.get('/resetMetrics', async (ctx) => {
+    try {
+      await stats.resetMetricsController(ctx);
+    } catch (error) {
+      logger.error(error);
+      ctx.status = 400;
+      ctx.body = error.message;
+    }
+  });
 }
 
 module.exports = { metricsRouter };

src/util/cluster.js

@@ -3,6 +3,7 @@ const gracefulShutdown = require('http-graceful-shutdown');
 const logger = require('../logger');
 const { logProcessInfo } = require('./utils');
 const { RedisDB } = require('./redis/redisConnector');
+const { shutdownMetricsClient } = require('./stats');
 
 const numWorkers = parseInt(process.env.NUM_PROCS || '1', 10);
 const metricsPort = parseInt(process.env.METRICS_PORT || '9091', 10);
@@ -18,11 +19,12 @@ function finalFunction() {
 
 // This function works only in master.
 // It sends SIGTERM to all the workers
-function shutdownWorkers() {
+async function shutdownWorkers() {
   Object.values(cluster.workers).forEach((worker) => {
     process.kill(worker.process.pid);
     logger.error(`Sent kill signal to worker ${worker.id} (pid: ${worker.process.pid})`);
   });
+  await shutdownMetricsClient();
 }
 
 function start(port, app, metricsApp) {
@@ -51,35 +53,35 @@ function start(port, app, metricsApp) {
     });
 
     let isShuttingDown = false;
-    cluster.on('exit', (worker) => {
+    cluster.on('exit', async (worker) => {
       if (!isShuttingDown) {
         logger.error(`Worker (pid: ${worker.process.pid}) died`);
         logger.error(`Killing other workers to avoid any side effects of the dead worker`);
         logProcessInfo();
         isShuttingDown = true;
-        shutdownWorkers();
+        await shutdownWorkers();
       }
     });
 
-    process.on('SIGTERM', () => {
+    process.on('SIGTERM', async () => {
       logger.error('SIGTERM signal received. Closing workers...');
       logProcessInfo();
       isShuttingDown = true;
-      shutdownWorkers();
+      await shutdownWorkers();
     });
 
-    process.on('SIGINT', () => {
+    process.on('SIGINT', async () => {
       logger.error('SIGINT signal received. Closing workers...');
       logProcessInfo();
       isShuttingDown = true;
-      shutdownWorkers();
+      await shutdownWorkers();
     });
 
-    process.on('SIGSEGV', () => {
+    process.on('SIGSEGV', async () => {
       logger.error('SIGSEGV - JavaScript memory error occurred. Closing workers...');
       logProcessInfo();
       isShuttingDown = true;
-      shutdownWorkers();
+      await shutdownWorkers();
     });
   } else {
     const server = app.listen(port);
@@ -90,16 +92,19 @@ function start(port, app, metricsApp) {
       finally: finalFunction,
     });
 
-    process.on('SIGTERM', () => {
+    process.on('SIGTERM', async () => {
       logger.error(`SIGTERM signal received in the worker`);
+      await shutdownMetricsClient();
     });
 
-    process.on('SIGINT', () => {
+    process.on('SIGINT', async () => {
       logger.error(`SIGINT signal received in the worker`);
+      await shutdownMetricsClient();
     });
 
-    process.on('SIGSEGV', () => {
+    process.on('SIGSEGV', async () => {
       logger.error(`SIGSEGV - JavaScript memory error occurred in the worker`);
+      await shutdownMetricsClient();
     });
 
     logger.info(`Worker (pid: ${process.pid}) has started`);

src/util/metricsAggregator.js

@@ -0,0 +1,181 @@
+/* eslint-disable */
+const cluster = require('cluster');
+const logger = require('../logger');
+const { Worker, isMainThread } = require('worker_threads');
+
+const MESSAGE_TYPES = {
+  GET_METRICS_REQ: 'rudder-transformer:getMetricsReq',
+  GET_METRICS_RES: 'rudder-transformer:getMetricsRes',
+  AGGREGATE_METRICS_REQ: 'rudder-transformer:aggregateMetricsReq',
+  AGGREGATE_METRICS_RES: 'rudder-transformer:aggregateMetricsRes',
+  RESET_METRICS_REQ: 'rudder-transformer:resetMetricsReq',
+};
+
+const config = {
+  isPeriodicResetEnabled: process.env.METRICS_AGGREGATOR_PERIODIC_RESET_ENABLED === 'true',
+  periodicResetInterval: process.env.METRICS_AGGREGATOR_PERIODIC_RESET_INTERVAL_SECONDS
+    ? parseInt(process.env.METRICS_AGGREGATOR_PERIODIC_RESET_INTERVAL_SECONDS, 10)
+    : 30 * 60,
+};
+
+class MetricsAggregator {
+  constructor(prometheusInstance) {
+    this.metricsBuffer = [];
+    this.pendingMetricRequests = 0;
+    this.resolveFunc = null;
+    this.rejectFunc = null;
+    this.prometheusInstance = prometheusInstance;
+    this.createWorkerThread();
+    this.registerCallbacks();
+  }
+
+  // onWorkerMessage is called when the master receives a message from a worker
+  async onWorkerMessage(worker, message) {
+    if (message.type === MESSAGE_TYPES.GET_METRICS_RES) {
+      logger.debug(`[MetricsAggregator] Master received metrics from worker ${worker.id}`);
+      await this.handleMetricsResponse(message);
+    }
+  }
+
+  // onMasterMessage is called when a worker receives a message from the master
+  async onMasterMessage(message) {
+    if (message.type === MESSAGE_TYPES.GET_METRICS_REQ) {
+      logger.debug(`[MetricsAggregator] Worker ${cluster.worker.id} received metrics request`);
+      try {
+        const metrics = await this.prometheusInstance.prometheusRegistry.getMetricsAsJSON();
+        cluster.worker.send({ type: MESSAGE_TYPES.GET_METRICS_RES, metrics });
+      } catch (error) {
+        cluster.worker.send({ type: MESSAGE_TYPES.GET_METRICS_RES, error: error.message });
+      }
+    } else if (message.type === MESSAGE_TYPES.RESET_METRICS_REQ) {
+      logger.info(`[MetricsAggregator] Worker ${cluster.worker.id} received reset metrics request`);
+      this.prometheusInstance.prometheusRegistry.resetMetrics();
+      logger.info(`[MetricsAggregator] Worker ${cluster.worker.id} reset metrics successfully`);
+    }
+  }
+
+  registerCallbacks() {
+    if (cluster.isPrimary) {
+      // register callback for master process
+      cluster.on('message', this.onWorkerMessage.bind(this));
+      if (config.isPeriodicResetEnabled) {
+        // register callback to reset metrics if enabled
+        this.registerCallbackForPeriodicReset(config.periodicResetInterval);
+      }
+      return;
+    }
+    // register callback for worker process
+    cluster.worker.on('message', this.onMasterMessage.bind(this));
+  }
+
+  registerCallbackForPeriodicReset(intervalSeconds) {
+    // store the timer in the aggregator for future operations like shutdown
+    this.periodicResetTimer = setInterval(() => {
+      logger.info(
+        `[MetricsAggregator] Periodic reset interval of ${intervalSeconds} seconds expired, reseting metrics`,
+      );
+      this.resetMetrics();
+    }, intervalSeconds * 1000);
+  }
+
+  createWorkerThread() {
+    if (cluster.isPrimary && isMainThread) {
+      this.workerThread = new Worker('./src/util/worker.js');
+      logger.info(
+        `[MetricsAggregator] Worker thread created with threadId ${this.workerThread.threadId}`,
+      );
+
+      this.workerThread.on('message', (message) => {
+        if (message.type === MESSAGE_TYPES.AGGREGATE_METRICS_RES) {
+          if (message.error) {
+            this.rejectFunc(new Error(message.error));
+            this.resetAggregator();
+            return;
+          }
+          this.resolveFunc(message.metrics);
+          this.resetAggregator();
+        }
+      });
+    }
+  }
+
+  resetAggregator() {
+    this.metricsBuffer = [];
+    this.pendingMetricRequests = 0;
+    this.resolveFunc = null;
+    this.rejectFunc = null;
+  }
+
+  async aggregateMetrics() {
+    // If a request is already being processed, reject the new request
+    // Use resolveFunc to check if a request is already being processed
+    // we dont support concurrent /metrics requests for now - we would need to implement a requestId mechanism and then handle all message calls accoridng to this requestId
+    // this is how it is implemented in prom-client [https://github.com/siimon/prom-client/blob/564e46724e258704df52ab329a7be833aaed4b69/lib/cluster.js#L43]
+    // we are not implementing it for now to keep things simple, once we validate the solution we can implement it
+    if (this.resolveFunc !== null) {
+      logger.error(
+        '[MetricsAggregator] Failed to serve /metrics request, a request is already being processed.',
+      );
+      throw new Error(
+        '[MetricsAggregator] Currently processing a request, please try again later.',
+      );
+    }
+    return new Promise((resolve, reject) => {
+      this.resolveFunc = resolve;
+      this.rejectFunc = reject;
+      for (const id in cluster.workers) {
+        this.pendingMetricRequests++;
+        logger.debug(`[MetricsAggregator] Requesting metrics from worker ${id}`);
+        cluster.workers[id].send({ type: MESSAGE_TYPES.GET_METRICS_REQ });
+      }
+    });
+  }
+
+  async aggregateMetricsInWorkerThread() {
+    this.workerThread.postMessage({
+      type: MESSAGE_TYPES.AGGREGATE_METRICS_REQ,
+      metrics: this.metricsBuffer,
+    });
+  }
+
+  async handleMetricsResponse(message) {
+    if (message.error) {
+      this.rejectFunc(new Error(message.error));
+      this.resetAggregator();
+      return;
+    }
+    this.metricsBuffer.push(message.metrics);
+    this.pendingMetricRequests--;
+    if (this.pendingMetricRequests === 0) {
+      this.aggregateMetricsInWorkerThread();
+    }
+  }
+
+  async terminateWorkerThread() {
+    logger.info(
+      `[MetricsAggregator] Worker thread terminated with exit code ${await this.workerThread.terminate()}`,
+    );
+  }
+
+  resetMetrics() {
+    for (const id in cluster.workers) {
+      logger.info(`[MetricsAggregator] Resetting metrics for worker ${id}`);
+      cluster.workers[id].send({ type: MESSAGE_TYPES.RESET_METRICS_REQ });
+    }
+  }
+
+  async shutdown() {
+    // terminate worker thread if the current process is the master
+    if (cluster.isPrimary) {
+      await this.terminateWorkerThread();
+    }
+    if (this.periodicResetTimer) {
+      clearInterval(this.periodicResetTimer);
+    }
+  }
+}
+
+module.exports = {
+  MetricsAggregator,
+  MESSAGE_TYPES,
+};

src/util/prometheus.js

@@ -1,7 +1,9 @@
 const prometheusClient = require('prom-client');
 const logger = require('../logger');
+const { MetricsAggregator } = require('./metricsAggregator');
 
 const clusterEnabled = process.env.CLUSTER_ENABLED !== 'false';
+const useMetricsAggregator = process.env.USE_METRICS_AGGREGATOR === 'true';
 const instanceID = process.env.INSTANCE_ID || 'localhost';
 const prefix = 'transformer';
 const defaultLabels = { instanceName: instanceID };
@@ -12,6 +14,9 @@ function appendPrefix(name) {
 
 class Prometheus {
   constructor(enableSummaryMetrics = true) {
+    if (clusterEnabled && useMetricsAggregator) {
+      this.metricsAggregator = new MetricsAggregator(this);
+    }
     this.prometheusRegistry = new prometheusClient.Registry();
     this.prometheusRegistry.setDefaultLabels(defaultLabels);
     prometheusClient.collectDefaultMetrics({
@@ -27,15 +32,29 @@ class Prometheus {
   async metricsController(ctx) {
     ctx.status = 200;
     if (clusterEnabled) {
-      ctx.type = this.aggregatorRegistry.contentType;
-      ctx.body = await this.aggregatorRegistry.clusterMetrics();
+      if (useMetricsAggregator) {
+        ctx.type = this.prometheusRegistry.contentType;
+        ctx.body = await this.metricsAggregator.aggregateMetrics();
+      } else {
+        ctx.type = this.aggregatorRegistry.contentType;
+        ctx.body = await this.aggregatorRegistry.clusterMetrics();
+      }
     } else {
       ctx.type = this.prometheusRegistry.contentType;
       ctx.body = await this.prometheusRegistry.metrics();
     }
     return ctx.body;
   }
 
+  async resetMetricsController(ctx) {
+    ctx.status = 200;
+    if (clusterEnabled && useMetricsAggregator) {
+      this.metricsAggregator.resetMetrics();
+    }
+    ctx.body = 'Metrics reset';
+    return ctx.body;
+  }
+
   newCounterStat(name, help, labelNames) {
     const counter = new prometheusClient.Counter({
       name,
@@ -192,6 +211,12 @@ class Prometheus {
     }
   }
 
+  async shutdown() {
+    if (this.metricsAggregator) {
+      await this.metricsAggregator.shutdown();
+    }
+  }
+
   createMetrics(enableSummaryMetrics) {
     const metrics = [
       // Counters

src/util/stats.js

@@ -105,6 +105,32 @@ async function metricsController(ctx) {
   ctx.body = `Not supported`;
 }
 
+async function resetMetricsController(ctx) {
+  if (!enableStats || !statsClient) {
+    ctx.status = 501;
+    ctx.body = `Not supported`;
+    return;
+  }
+
+  if (statsClientType === 'prometheus') {
+    await statsClient.resetMetricsController(ctx);
+    return;
+  }
+
+  ctx.status = 501;
+  ctx.body = `Not supported`;
+}
+
+async function shutdownMetricsClient() {
+  if (!enableStats || !statsClient) {
+    return;
+  }
+
+  if (statsClientType === 'prometheus') {
+    await statsClient.shutdown();
+  }
+}
+
 init();
 
 module.exports = {
@@ -117,4 +143,6 @@ module.exports = {
   gauge,
   histogram,
   metricsController,
+  resetMetricsController,
+  shutdownMetricsClient,
 };