From a5002765e76e52f79b205cb9a6de490793364464 Mon Sep 17 00:00:00 2001
From: Thomas Hunter II <tlhunter@datadog.com>
Date: Mon, 26 Aug 2024 13:03:50 -0700
Subject: [PATCH] hapi: migrate from AsyncResource to TracingChannel, see #4597
 (#4622)

* fix: remove async scope in hapi dd-instrumentation
* fix linting
* replace asyncResource for tracingChannel
* move tracing channels to helpers
* add test
* remove tracingChannel map

---------

Co-authored-by: Cris Naranjo <cris@yulife.com>
---
 packages/datadog-instrumentations/src/hapi.js | 11 +++-----
 packages/datadog-plugin-hapi/src/index.js     |  4 +--
 .../datadog-plugin-hapi/test/index.spec.js    | 25 +++++++++++++++++++
 3 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/packages/datadog-instrumentations/src/hapi.js b/packages/datadog-instrumentations/src/hapi.js
index 2f0024c741..a435328b03 100644
--- a/packages/datadog-instrumentations/src/hapi.js
+++ b/packages/datadog-instrumentations/src/hapi.js
@@ -1,12 +1,13 @@
 'use strict'
 
+const tracingChannel = require('dc-polyfill').tracingChannel
 const shimmer = require('../../datadog-shimmer')
-const { addHook, channel, AsyncResource } = require('./helpers/instrument')
+const { addHook, channel } = require('./helpers/instrument')
 
 const handleChannel = channel('apm:hapi:request:handle')
 const routeChannel = channel('apm:hapi:request:route')
 const errorChannel = channel('apm:hapi:request:error')
-const enterChannel = channel('apm:hapi:extension:enter')
+const hapiTracingChannel = tracingChannel('apm:hapi:extension')
 
 function wrapServer (server) {
   return function (options) {
@@ -96,11 +97,7 @@ function wrapHandler (handler) {
 
     if (!req) return handler.apply(this, arguments)
 
-    const asyncResource = new AsyncResource('bound-anonymous-fn')
-
-    return asyncResource.runInAsyncScope(() => {
-      enterChannel.publish({ req })
-
+    return hapiTracingChannel.traceSync(() => {
       return handler.apply(this, arguments)
     })
   }
diff --git a/packages/datadog-plugin-hapi/src/index.js b/packages/datadog-plugin-hapi/src/index.js
index e4954acd79..b72df6951b 100644
--- a/packages/datadog-plugin-hapi/src/index.js
+++ b/packages/datadog-plugin-hapi/src/index.js
@@ -32,8 +32,8 @@ class HapiPlugin extends RouterPlugin {
       }
     })
 
-    this.addSub('apm:hapi:extension:enter', ({ req }) => {
-      this.enter(this._requestSpans.get(req))
+    this.addBind('apm:hapi:extension:start', ({ req }) => {
+      return this._requestSpans.get(req)
     })
   }
 }
diff --git a/packages/datadog-plugin-hapi/test/index.spec.js b/packages/datadog-plugin-hapi/test/index.spec.js
index 2e67022f49..48093e2904 100644
--- a/packages/datadog-plugin-hapi/test/index.spec.js
+++ b/packages/datadog-plugin-hapi/test/index.spec.js
@@ -4,6 +4,7 @@ const axios = require('axios')
 const semver = require('semver')
 const agent = require('../../dd-trace/test/plugins/agent')
 const { ERROR_MESSAGE, ERROR_TYPE, ERROR_STACK } = require('../../dd-trace/src/constants')
+const { AsyncLocalStorage } = require('async_hooks')
 
 const versionRange = parseInt(process.versions.node.split('.')[0]) > 14
   ? '<17 || >18'
@@ -348,6 +349,30 @@ describe('Plugin', () => {
           .get(`http://localhost:${port}/user/123`)
           .catch(() => {})
       })
+
+      it('should persist AsyncLocalStorage context', (done) => {
+        const als = new AsyncLocalStorage()
+        const path = '/path'
+
+        server.ext('onRequest', (request, h) => {
+          als.enterWith({ path: request.path })
+          return reply(request, h)
+        })
+
+        server.route({
+          method: 'GET',
+          path,
+          handler: async (request, h) => {
+            expect(als.getStore()).to.deep.equal({ path })
+            done()
+            return h.response ? h.response() : h()
+          }
+        })
+
+        axios
+          .get(`http://localhost:${port}${path}`)
+          .catch(() => {})
+      })
     })
   })
 })