From 4d62c9289f41424106ebc64a4d51da686cb990b8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ketil=20=C3=98vreb=C3=B8?= <ketil@aptoma.com>
Date: Wed, 10 Aug 2022 15:33:26 +0200
Subject: [PATCH] fix(hapi): ensure route wrapper starts a new context (#1094)

Co-authored-by: Daniel Dyla <dyladan@users.noreply.github.com>
Co-authored-by: Valentin Marchaud <contact@vmarchaud.fr>
---
 .../src/instrumentation.ts                    | 11 ++--
 .../test/hapi.test.ts                         | 51 +++++++++++++++++++
 2 files changed, 57 insertions(+), 5 deletions(-)

diff --git a/plugins/node/opentelemetry-instrumentation-hapi/src/instrumentation.ts b/plugins/node/opentelemetry-instrumentation-hapi/src/instrumentation.ts
index 344350018a..b3a53cc9d4 100644
--- a/plugins/node/opentelemetry-instrumentation-hapi/src/instrumentation.ts
+++ b/plugins/node/opentelemetry-instrumentation-hapi/src/instrumentation.ts
@@ -380,12 +380,10 @@ export class HapiInstrumentation extends InstrumentationBase {
     const oldHandler = route.options?.handler ?? route.handler;
     if (typeof oldHandler === 'function') {
       const newHandler: Hapi.Lifecycle.Method = async function (
-        request: Hapi.Request,
-        h: Hapi.ResponseToolkit,
-        err?: Error
+        ...params: Parameters<Hapi.Lifecycle.Method>
       ) {
         if (api.trace.getSpan(api.context.active()) === undefined) {
-          return await oldHandler(request, h, err);
+          return await oldHandler(...params);
         }
         const rpcMetadata = getRPCMetadata(api.context.active());
         if (rpcMetadata?.type === RPCType.HTTP) {
@@ -398,7 +396,10 @@ export class HapiInstrumentation extends InstrumentationBase {
           attributes: metadata.attributes,
         });
         try {
-          return await oldHandler(request, h, err);
+          return await api.context.with(
+            api.trace.setSpan(api.context.active(), span),
+            () => oldHandler(...params)
+          );
         } catch (err) {
           span.recordException(err);
           span.setStatus({
diff --git a/plugins/node/opentelemetry-instrumentation-hapi/test/hapi.test.ts b/plugins/node/opentelemetry-instrumentation-hapi/test/hapi.test.ts
index b2fcabaf6c..ba3d87614c 100644
--- a/plugins/node/opentelemetry-instrumentation-hapi/test/hapi.test.ts
+++ b/plugins/node/opentelemetry-instrumentation-hapi/test/hapi.test.ts
@@ -267,6 +267,57 @@ describe('Hapi Instrumentation - Core Tests', () => {
       );
     });
 
+    it('should start a new context for the handler', async () => {
+      const rootSpan = tracer.startSpan('rootSpan');
+      server.route([
+        {
+          method: 'GET',
+          path: '/route',
+          handler: (request, h) => {
+            const span = tracer.startSpan('handler');
+            span.end();
+            return 'ok';
+          },
+        },
+      ]);
+
+      await server.start();
+      assert.strictEqual(memoryExporter.getFinishedSpans().length, 0);
+
+      await context.with(
+        trace.setSpan(context.active(), rootSpan),
+        async () => {
+          const res = await server.inject({
+            method: 'GET',
+            url: '/route',
+          });
+
+          assert.strictEqual(res.statusCode, 200);
+
+          rootSpan.end();
+          assert.deepStrictEqual(memoryExporter.getFinishedSpans().length, 3);
+
+          const routeSpan = memoryExporter
+            .getFinishedSpans()
+            .find(span => span.name === 'route - /route');
+          assert.notStrictEqual(routeSpan, undefined);
+          assert.strictEqual(
+            routeSpan?.attributes[AttributeNames.HAPI_TYPE],
+            HapiLayerType.ROUTER
+          );
+
+          const handlerSpan = memoryExporter
+            .getFinishedSpans()
+            .find(span => span.name === 'handler');
+          assert.notStrictEqual(routeSpan, undefined);
+          assert.strictEqual(
+            handlerSpan?.parentSpanId,
+            routeSpan?.spanContext().spanId
+          );
+        }
+      );
+    });
+
     it('should access route parameters and add to span', async () => {
       const rootSpan = tracer.startSpan('rootSpan');
       server.route({