feat: add profiling middleware fastify

package.json

@@ -62,6 +62,7 @@
     "eslint-config-prettier": "^10.1.8",
     "eslint-plugin-prettier": "^5.5.4",
     "express": "^5.1.0",
+    "fastify": "^5.6.1",
     "husky": "^9.1.7",
     "pinst": "^3.0.0",
     "pprof-format": "^2.2.1",
@@ -70,6 +71,18 @@
     "typescript": "^5.9.3",
     "vitest": "^4.0.6"
   },
+  "peerDependencies": {
+    "express": "^4.0.0 || ^5.0.0",
+    "fastify": "^4.0.0 || ^5.0.0"
+  },
+  "peerDependenciesMeta": {
+    "express": {
+      "optional": true
+    },
+    "fastify": {
+      "optional": true
+    }
+  },
   "engines": {
     "node": ">=v18"
   }

src/index.ts

@@ -1,7 +1,8 @@
 import 'regenerator-runtime/runtime.js';
 
 import { setLogger as datadogSetLogger } from '@datadog/pprof';
-import expressMiddleware from './express/middleware.js';
+import expressMiddleware from './middleware/express.js';
+import fastifyMiddleware from './middleware/fastify.js';
 import { Logger, setLogger as ourSetLogger } from './logger.js';
 import { PyroscopeProfiler } from './profilers/pyroscope-profiler.js';
 import {
@@ -94,6 +95,7 @@ function setLogger(logger: Logger): void {
 export default {
   SourceMapper,
   expressMiddleware,
+  fastifyMiddleware,
   init,
   getWallLabels,
   setWallLabels,

src/middleware/fastify.ts

@@ -0,0 +1,111 @@
+import debug from 'debug';
+import type {
+  FastifyRequest,
+  FastifyReply,
+  FastifyPluginCallback,
+} from 'fastify';
+import { Profile } from 'pprof-format';
+import { Profiler } from '../profilers/profiler.js';
+import { PyroscopeProfiler } from '../profilers/pyroscope-profiler.js';
+import { WallProfilerStartArgs } from '../profilers/wall-profiler.js';
+import { getProfiler } from '../utils/pyroscope-profiler.js';
+import { encode } from '@datadog/pprof';
+import { HeapProfilerStartArgs } from '../profilers/heap-profiler.js';
+
+const log = debug('pyroscope');
+
+async function collectProfile<TStartArgs>(
+  profiler: Profiler<TStartArgs>
+): Promise<Buffer> {
+  const profile: Profile = profiler.profile().profile;
+
+  profiler.stop();
+
+  return encode(profile);
+}
+
+async function collectProfileAfterMs<TStartArgs>(
+  profiler: Profiler<TStartArgs>,
+  args: TStartArgs,
+  delayMs: number
+): Promise<Buffer> {
+  profiler.start(args);
+
+  if (delayMs === 0) {
+    return collectProfile(profiler);
+  }
+
+  return new Promise(
+    (resolve: (buffer: Buffer | PromiseLike<Buffer>) => void) => {
+      setTimeout(() => {
+        resolve(collectProfile(profiler));
+      }, delayMs);
+    }
+  );
+}
+
+function collectHeap(): Promise<Buffer> {
+  const profiler: PyroscopeProfiler = getProfiler();
+
+  const heapProfilerArgs: HeapProfilerStartArgs =
+    profiler.heapProfiler.startArgs;
+  const heapProfiler: Profiler<HeapProfilerStartArgs> =
+    profiler.heapProfiler.profiler;
+
+  return collectProfileAfterMs(heapProfiler, heapProfilerArgs, 0);
+}
+
+function collectWall(ms: number): Promise<Buffer> {
+  const profiler: PyroscopeProfiler = getProfiler();
+
+  const wallProfilerArgs: WallProfilerStartArgs =
+    profiler.wallProfiler.startArgs;
+  const wallProfiler: Profiler<WallProfilerStartArgs> =
+    profiler.wallProfiler.profiler;
+
+  return collectProfileAfterMs(wallProfiler, wallProfilerArgs, ms);
+}
+
+async function heapHandler(
+  request: FastifyRequest,
+  reply: FastifyReply
+): Promise<void> {
+  log('Fetching Heap Profile');
+  try {
+    const profileBuffer = await collectHeap();
+    reply.status(200).type('application/octet-stream').send(profileBuffer);
+  } catch (error: unknown) {
+    log('Error collecting Heap', error);
+    reply.status(500).send({ error: 'Internal Server Error' });
+  }
+}
+
+async function wallHandler(
+  request: FastifyRequest<{ Querystring: { seconds?: string } }>,
+  reply: FastifyReply
+): Promise<void> {
+  log('Fetching Wall Profile');
+  try {
+    const seconds = Number(request.query.seconds || 1);
+    const profileBuffer = await collectWall(1000 * seconds);
+    reply.status(200).type('application/octet-stream').send(profileBuffer);
+  } catch (error: unknown) {
+    log('Error collecting Wall', error);
+    reply.status(500).send({ error: 'Internal Server Error' });
+  }
+}
+
+const fastifyMiddleware = (): FastifyPluginCallback => {
+  const plugin: FastifyPluginCallback = (fastify, _options, done) => {
+    // Register route for heap profiling
+    fastify.get('/debug/pprof/heap', heapHandler);
+
+    // Register route for wall/CPU profiling
+    fastify.get('/debug/pprof/profile', wallHandler);
+
+    done();
+  };
+  return plugin;
+};
+
+export default fastifyMiddleware;

test/fastify.test.ts

@@ -0,0 +1,104 @@
+import { describe, it, expect } from 'vitest';
+
+import Pyroscope from '../src/index.js';
+import Fastify from 'fastify';
+
+// You only need appName for the pull mode
+Pyroscope.init();
+
+describe('fastify middleware', () => {
+  it('should be a function', () => {
+    expect(typeof Pyroscope.fastifyMiddleware).toBe('function');
+  });
+  it('should respond to cpu calls', async () => {
+    const app = Fastify();
+    await app.register(Pyroscope.fastifyMiddleware());
+    const response = await app.inject({
+      method: 'GET',
+      url: '/debug/pprof/profile?seconds=1',
+    });
+    expect(response.statusCode).toBe(200);
+  });
+  it('should respond to repetitive cpu calls', async () => {
+    const app = Fastify();
+    await app.register(Pyroscope.fastifyMiddleware());
+    const response = await app.inject({
+      method: 'GET',
+      url: '/debug/pprof/profile?seconds=1',
+    });
+    expect(response.statusCode).toBe(200);
+  });
+
+  // it('should respond to simultaneous cpu calls', async () => {
+  //   const app = Fastify()
+  //   await app.register(Pyroscope.fastifyMiddleware())
+  //   console.log('0', Date.now()/1000);
+  //   const [response1, response2] = await Promise.all([
+  //     app.inject({
+  //       method: 'GET',
+  //       url: '/debug/pprof/profile?seconds=1'
+  //     }),
+  //     app.inject({
+  //       method: 'GET',
+  //       url: '/debug/pprof/profile?seconds=1'
+  //     }),
+  //   ])
+  //   expect(response1.statusCode).toBe(200)
+  //   expect(response2.statusCode).toBe(200)
+  // })
+  it('should respond to heap profiling calls', async () => {
+    const app = Fastify();
+    await app.register(Pyroscope.fastifyMiddleware());
+    const response = await app.inject({
+      method: 'GET',
+      url: '/debug/pprof/heap',
+    });
+    expect(response.statusCode).toBe(200);
+  });
+  it('should respond to repetitive heap profiling calls', async () => {
+    const app = Fastify();
+    await app.register(Pyroscope.fastifyMiddleware());
+    const response = await app.inject({
+      method: 'GET',
+      url: '/debug/pprof/heap',
+    });
+    expect(response.statusCode).toBe(200);
+  });
+
+  it('should respond to simultaneous heap profiling calls', async () => {
+    const app = Fastify();
+    await app.register(Pyroscope.fastifyMiddleware());
+    const [response1, response2] = await Promise.all([
+      app.inject({
+        method: 'GET',
+        url: '/debug/pprof/heap?seconds=1',
+      }),
+      app.inject({
+        method: 'GET',
+        url: '/debug/pprof/heap?seconds=1',
+      }),
+    ]);
+    expect(response1.statusCode).toBe(200);
+    expect(response2.statusCode).toBe(200);
+  });
+
+  it('should be fine using two middlewares at the same time', async () => {
+    const app = Fastify();
+    await app.register(Pyroscope.fastifyMiddleware());
+
+    const app2 = Fastify();
+    await app2.register(Pyroscope.fastifyMiddleware());
+
+    const response1 = await app.inject({
+      method: 'GET',
+      url: '/debug/pprof/heap',
+    });
+    expect(response1.statusCode).toBe(200);
+
+    const response2 = await app2.inject({
+      method: 'GET',
+      url: '/debug/pprof/heap',
+    });
+    expect(response2.statusCode).toBe(200);
+  });
+});