From 8d6662feb77bff5bf85a31418b91903c4ea9dde5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rapha=C3=ABl=20Th=C3=A9riault?=
 <raphael.theriault@solarwinds.com>
Date: Fri, 10 Jan 2025 14:44:34 -0500
Subject: [PATCH] feat(instrumentation-express): propagate context and measure
 full handler spans

---
 .../src/instrumentation.ts                       | 16 +++++++++++-----
 .../test/express.test.ts                         |  4 ++--
 2 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/plugins/node/opentelemetry-instrumentation-express/src/instrumentation.ts b/plugins/node/opentelemetry-instrumentation-express/src/instrumentation.ts
index 2e283158b8..2c64ef79f8 100644
--- a/plugins/node/opentelemetry-instrumentation-express/src/instrumentation.ts
+++ b/plugins/node/opentelemetry-instrumentation-express/src/instrumentation.ts
@@ -215,6 +215,9 @@ export class ExpressInstrumentation extends InstrumentationBase<ExpressInstrumen
           attributes: Object.assign(attributes, metadata.attributes),
         });
 
+        const parentContext = context.active();
+        let currentContext = trace.setSpan(parentContext, span);
+
         const { requestHook } = instrumentation.getConfig();
         if (requestHook) {
           safeExecuteInTheMiddle(
@@ -234,12 +237,15 @@ export class ExpressInstrumentation extends InstrumentationBase<ExpressInstrumen
         }
 
         let spanHasEnded = false;
+        // TODO: Fix router spans (getRouterPath does not work properly) to
+        // have useful names before removing this branch
         if (
-          metadata.attributes[AttributeNames.EXPRESS_TYPE] !==
-          ExpressLayerType.MIDDLEWARE
+          metadata.attributes[AttributeNames.EXPRESS_TYPE] ===
+          ExpressLayerType.ROUTER
         ) {
           span.end();
           spanHasEnded = true;
+          currentContext = parentContext;
         }
         // listener for response.on('finish')
         const onResponseFinish = () => {
@@ -278,12 +284,12 @@ export class ExpressInstrumentation extends InstrumentationBase<ExpressInstrumen
               (req[_LAYERS_STORE_PROPERTY] as string[]).pop();
             }
             const callback = args[callbackIdx] as Function;
-            return callback.apply(this, arguments);
+            return context.bind(parentContext, callback).apply(this, arguments);
           };
         }
 
         try {
-          return original.apply(this, arguments);
+          return context.bind(currentContext, original).apply(this, arguments);
         } catch (anyError) {
           const [error, message] = asErrorAndMessage(anyError);
           span.recordException(error);
@@ -296,7 +302,7 @@ export class ExpressInstrumentation extends InstrumentationBase<ExpressInstrumen
           /**
            * At this point if the callback wasn't called, that means either the
            * layer is asynchronous (so it will call the callback later on) or that
-           * the layer directly end the http response, so we'll hook into the "finish"
+           * the layer directly ends the http response, so we'll hook into the "finish"
            * event to handle the later case.
            */
           if (!spanHasEnded) {
diff --git a/plugins/node/opentelemetry-instrumentation-express/test/express.test.ts b/plugins/node/opentelemetry-instrumentation-express/test/express.test.ts
index 5991688ffd..9fce3707f8 100644
--- a/plugins/node/opentelemetry-instrumentation-express/test/express.test.ts
+++ b/plugins/node/opentelemetry-instrumentation-express/test/express.test.ts
@@ -96,7 +96,7 @@ describe('ExpressInstrumentation', () => {
           );
           assert.strictEqual(response, 'tata');
           rootSpan.end();
-          assert.strictEqual(finishListenerCount, 2);
+          assert.strictEqual(finishListenerCount, 3);
           assert.notStrictEqual(
             memoryExporter
               .getFinishedSpans()
@@ -201,7 +201,7 @@ describe('ExpressInstrumentation', () => {
           );
           assert.strictEqual(response, 'tata');
           rootSpan.end();
-          assert.strictEqual(finishListenerCount, 2);
+          assert.strictEqual(finishListenerCount, 3);
           assert.notStrictEqual(
             memoryExporter
               .getFinishedSpans()