refactor: prep for reply shenanigans

lumigo-io · Nov 2, 2023 · c971b49 · c971b49
1 parent 3493dcf
commit c971b49
Show file tree

Hide file tree

Showing 3 changed files with 94 additions and 33 deletions.
diff --git a/src/instrumentations/fastify/FastifyInstrumentation.ts b/src/instrumentations/fastify/FastifyInstrumentation.ts
@@ -1,9 +1,6 @@
-import { CommonUtils, ScrubContext } from '@lumigo/node-core';
-import { Span } from '@opentelemetry/api';
-import { FastifyInstrumentation, FastifyRequestInfo } from '@opentelemetry/instrumentation-fastify';
-import { contentType, scrubHttpPayload } from '../../tools/payloads';
-import { getSpanAttributeMaxLength } from '../../utils';
+import { FastifyInstrumentation } from '@opentelemetry/instrumentation-fastify';
 import { Instrumentor } from '../instrumentor';
+import { FastifyHooks } from './fastify';
 
 export default class LumigoFastifyInstrumentation extends Instrumentor<FastifyInstrumentation> {
   getInstrumentedModule(): string {
@@ -12,32 +9,7 @@ export default class LumigoFastifyInstrumentation extends Instrumentor<FastifyIn
 
   getInstrumentation(): FastifyInstrumentation {
     return new FastifyInstrumentation({
-      requestHook: function (span: Span, info: FastifyRequestInfo) {
-        span.setAttribute(
-          'http.request.headers',
-          CommonUtils.payloadStringify(
-            info.request.headers,
-            ScrubContext.HTTP_REQUEST_HEADERS,
-            getSpanAttributeMaxLength()
-          )
-        );
-        span.setAttribute(
-          'http.request.query',
-          CommonUtils.payloadStringify(
-            info.request.query,
-            ScrubContext.HTTP_REQUEST_QUERY,
-            getSpanAttributeMaxLength()
-          )
-        );
-        span.setAttribute(
-          'http.request.body',
-          scrubHttpPayload(
-            info.request.body,
-            contentType(info.request.headers),
-            ScrubContext.HTTP_REQUEST_BODY
-          )
-        );
-      },
+      requestHook: FastifyHooks.requestHook,
     });
   }
 }
diff --git a/src/instrumentations/fastify/fastify.ts b/src/instrumentations/fastify/fastify.ts
@@ -0,0 +1,86 @@
+import { CommonUtils, ScrubContext } from '@lumigo/node-core';
+import { Span } from '@opentelemetry/api';
+
+import { contentType, scrubHttpPayload } from '../../tools/payloads';
+import { getSpanAttributeMaxLength, safeExecute } from '../../utils';
+import { InstrumentationIfc } from '../hooksIfc';
+
+type FastifyRequestType = { request: any; reply?: any };
+
+export const FastifyHooks: InstrumentationIfc<FastifyRequestType, any> = {
+  requestHook(span: Span, { request, reply }: FastifyRequestType): void {
+    if (request.query)
+      span.setAttribute(
+        'http.request.query',
+        CommonUtils.payloadStringify(
+          request.query,
+          ScrubContext.HTTP_REQUEST_QUERY,
+          getSpanAttributeMaxLength()
+        )
+      );
+    if (request.headers)
+      span.setAttribute(
+        'http.request.headers',
+        CommonUtils.payloadStringify(
+          request.headers,
+          ScrubContext.HTTP_REQUEST_HEADERS,
+          getSpanAttributeMaxLength()
+        )
+      );
+    if (request.body) {
+      span.setAttribute(
+        'http.request.body',
+        scrubHttpPayload(request.body, contentType(request.headers), ScrubContext.HTTP_REQUEST_BODY)
+      );
+    }
+
+    let response;
+    if (reply) {
+      const oldReplyEnd = reply.end;
+      const oldReplySend = reply.send;
+      reply.send = function (data: any) {
+        response = data;
+        reply.send = oldReplySend;
+        // eslint-disable-next-line prefer-rest-params
+        return oldReplySend.apply(reply, arguments);
+      };
+      reply.end = function () {
+        return safeExecute(() => {
+          // eslint-disable-next-line prefer-rest-params
+          const origRes = oldReplyEnd.apply(reply, arguments);
+          if (reply.getHeaders())
+            span.setAttribute(
+              'http.response.headers',
+              CommonUtils.payloadStringify(
+                reply.getHeaders(),
+                ScrubContext.HTTP_RESPONSE_HEADERS,
+                getSpanAttributeMaxLength()
+              )
+            ); // TODO This is not compliant with the HTTP semantic conventions
+          if (response)
+            span.setAttribute(
+              'http.response.body',
+              scrubHttpPayload(
+                response,
+                contentType(reply.getHeaders()),
+                ScrubContext.HTTP_RESPONSE_BODY
+              )
+            );
+          if (reply.body)
+            span.setAttribute(
+              'http.request.body',
+              scrubHttpPayload(
+                reply.body,
+                contentType(reply.headers),
+                ScrubContext.HTTP_REQUEST_BODY
+              )
+            );
+          reply.end = oldReplyEnd;
+          return origRes;
+        })();
+      };
+    }
+  },
+  // eslint-disable-next-line @typescript-eslint/no-unused-vars, @typescript-eslint/no-empty-function
+  responseHook(span: Span, response: any): void {},
+};
diff --git a/test/instrumentations/express/app/express_app.js b/test/instrumentations/express/app/express_app.js
@@ -1,10 +1,11 @@
-require('@lumigo/opentelemetry');
+const lumigo = require('@lumigo/opentelemetry');
 const express = require('express');
 const bodyParser = require('body-parser');
 const app = express();
 require('log-timestamp');
 
 let server;
+let tracerProvider;
 
 // parse application/x-www-form-urlencoded
 app.use(bodyParser.urlencoded({ extended: false }));
@@ -29,11 +30,13 @@ app.get('/basic', async (_, res) => {
 
 app.get('/quit', async (_, res) => {
   console.error('Received quit command');
+  await tracerProvider.forceFlush();
   res.status(200).send({});
   server.close();
 });
 
-server = app.listen(0, () => {
+server = app.listen(0, async () => {
+  tracerProvider = (await lumigo.init).tracerProvider;
   const port = server.address().port;
   console.error(`HTTP server listening on port ${port}`);
   if (process.send) {