From 55cc5aaf04c481c199570311f2861df980a53a4d Mon Sep 17 00:00:00 2001
From: Utsab Chowdhury <utsab.c97@gmail.com>
Date: Thu, 14 Sep 2023 14:25:27 +0530
Subject: [PATCH] fix: pyroscope endpoints

---
 src/controllers/misc.ts | 18 +++++++++++++++
 src/index.ts            |  7 +++---
 src/middleware.js       | 51 ++++++++++-------------------------------
 src/routes/misc.ts      |  2 ++
 src/services/misc.ts    | 11 ++++++++-
 5 files changed, 46 insertions(+), 43 deletions(-)

diff --git a/src/controllers/misc.ts b/src/controllers/misc.ts
index eff0eee244..92ec33f80f 100644
--- a/src/controllers/misc.ts
+++ b/src/controllers/misc.ts
@@ -25,4 +25,22 @@ export default class MiscController {
     ctx.status = 200;
     return ctx;
   }
+
+  public static async getCPUProfile(ctx: Context) {
+    const { seconds } = ctx.query;
+    let secondsData = 10;
+    // if seconds is not null and is not array then parseInt
+    if (seconds && !Array.isArray(seconds)) {
+      secondsData = parseInt(seconds, 10);
+    }
+    ctx.body = await MiscService.getCPUProfile(secondsData);
+    ctx.status = 200;
+    return ctx;
+  }
+
+  public static async getHeapProfile(ctx: Context) {
+    ctx.body = await MiscService.getHeapProfile();
+    ctx.status = 200;
+    return ctx;
+  }
 }
diff --git a/src/index.ts b/src/index.ts
index a928c8c3c5..d1cc95cc36 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -7,7 +7,7 @@ import cluster from './util/cluster';
 import { router } from './legacy/router';
 import { testRouter } from './testRouter';
 import { metricsRouter } from './routes/metricsRouter';
-import { addStatMiddleware, addRequestSizeMiddleware, addPyroscopeMiddleware } from './middleware';
+import { addStatMiddleware, addRequestSizeMiddleware, initPyroscope } from './middleware';
 import { logProcessInfo } from './util/utils';
 import { applicationRoutes, addSwaggerRoutes } from './routes';
 import { RedisDB } from './util/redis/redisConnector';
@@ -15,9 +15,11 @@ import { RedisDB } from './util/redis/redisConnector';
 dotenv.config();
 const clusterEnabled = process.env.CLUSTER_ENABLED !== 'false';
 const useUpdatedRoutes = process.env.ENABLE_NEW_ROUTES !== 'false';
-const port = parseInt(process.env.PORT || '9090', 10);
+const port = parseInt(process.env.PORT ?? '9090', 10);
 const metricsPort = parseInt(process.env.METRICS_PORT || '9091', 10);
 
+initPyroscope();
+
 const app = new Koa();
 addStatMiddleware(app);
 
@@ -30,7 +32,6 @@ app.use(
     jsonLimit: '200mb',
   }),
 );
-addPyroscopeMiddleware(app);
 addRequestSizeMiddleware(app);
 addSwaggerRoutes(app);
 
diff --git a/src/middleware.js b/src/middleware.js
index b8f6a2ac1b..53aabc90e3 100644
--- a/src/middleware.js
+++ b/src/middleware.js
@@ -1,48 +1,19 @@
 const Pyroscope = require('@pyroscope/nodejs');
 const stats = require('./util/stats');
-const logger = require('./logger');
 
-Pyroscope.init({
-  appName: 'rudder-transformer',
-});
-
-async function handlerCpu(ctx) {
-  try {
-    const p = await Pyroscope.collectCpu(Number(ctx.query.seconds));
-    ctx.body = p;
-    ctx.status = 200;
-  } catch (e) {
-    logger.error(e);
-    ctx.status = 500;
-  }
-}
-
-async function handlerHeap(ctx) {
-  try {
-    const p = await Pyroscope.collectHeap();
-    ctx.body = p;
-    ctx.status = 200;
-  } catch (e) {
-    logger.error(e);
-    ctx.status = 500;
-  }
+function initPyroscope() {
+  Pyroscope.init({
+    appName: 'rudder-transformer',
+  });
+  Pyroscope.startHeapCollecting();
 }
 
-function pyroscopeMiddleware() {
-  Pyroscope.startHeapCollecting();
-  return (ctx, next) => {
-    if (ctx.method === 'GET' && ctx.path === '/debug/pprof/profile') {
-      return handlerCpu(ctx).then(() => next());
-    }
-    if (ctx.method === 'GET' && ctx.path === '/debug/pprof/heap') {
-      return handlerHeap(ctx).then(() => next());
-    }
-    return next();
-  };
+function getCPUProfile(seconds) {
+  return Pyroscope.collectCpu(seconds);
 }
 
-function addPyroscopeMiddleware(app) {
-  app.use(pyroscopeMiddleware());
+function getHeapProfile() {
+  return Pyroscope.collectHeap();
 }
 
 function durationMiddleware() {
@@ -90,5 +61,7 @@ function addRequestSizeMiddleware(app) {
 module.exports = {
   addStatMiddleware,
   addRequestSizeMiddleware,
-  addPyroscopeMiddleware,
+  getHeapProfile,
+  getCPUProfile,
+  initPyroscope,
 };
diff --git a/src/routes/misc.ts b/src/routes/misc.ts
index 12ee09b8a9..3e30b9dd39 100644
--- a/src/routes/misc.ts
+++ b/src/routes/misc.ts
@@ -10,5 +10,7 @@ router.get('/transformerBuildVersion', MiscController.buildVersion); // depricia
 router.get('/buildVersion', MiscController.buildVersion);
 router.get('/version', MiscController.version);
 router.get('/features', MiscController.features);
+router.get('/debug/pprof/profile', MiscController.getCPUProfile);
+router.get('/debug/pprof/heap', MiscController.getHeapProfile);
 
 export const miscRoutes = router.routes();
diff --git a/src/services/misc.ts b/src/services/misc.ts
index 6508ced809..2805f58584 100644
--- a/src/services/misc.ts
+++ b/src/services/misc.ts
@@ -3,6 +3,7 @@ import path from 'path';
 import { Context } from 'koa';
 import { DestHandlerMap } from '../constants/destinationCanonicalNames';
 import { Metadata } from '../types';
+import { getCPUProfile, getHeapProfile, } from '../middleware';
 
 export default class MiscService {
   public static getDestHandler(dest: string, version: string) {
@@ -32,7 +33,7 @@ export default class MiscService {
 
   public static getMetaTags(metadata: Metadata) {
     if (!metadata) {
-      return {}
+      return {};
     }
     return {
       sourceType: metadata.sourceType,
@@ -62,4 +63,12 @@ export default class MiscService {
     const obj = JSON.parse(fs.readFileSync(path.resolve(__dirname, '../features.json'), 'utf8'));
     return JSON.stringify(obj);
   }
+
+  public static async getCPUProfile(seconds: number) {
+    return getCPUProfile(seconds);
+  }
+
+  public static async getHeapProfile() {
+    return getHeapProfile()
+  }
 }