fix context propagation for route handlers in fastify

packages/datadog-plugin-fastify/src/fastify.js

@@ -1,5 +1,6 @@
 'use strict'
 
+const methods = require('methods').concat('all')
 const web = require('../../dd-trace/src/plugins/util/web')
 
 function createWrapFastify (tracer, config) {
@@ -15,11 +16,19 @@ function createWrapFastify (tracer, config) {
     fastify._datadog_wrapper = function () {
       const app = fastify.apply(this, arguments)
 
-      if (app && typeof app.addHook === 'function') {
+      if (!app) return app
+
+      if (typeof app.addHook === 'function') {
         app.addHook('onRequest', createOnRequest(tracer, config))
         app.addHook('preHandler', preHandler)
       }
 
+      methods.forEach(method => {
+        app[method] = wrapMethod(app[method])
+      })
+
+      app.route = wrapRoute(app.route)
+
       return app
     }
 
@@ -58,6 +67,36 @@ function wrapSend (send) {
   }
 }
 
+function wrapRoute (route) {
+  if (typeof route !== 'function') return route
+
+  return function routeWithTrace (opts) {
+    opts.handler = wrapHandler(opts.handler)
+
+    return route.apply(this, arguments)
+  }
+}
+
+function wrapMethod (method) {
+  if (typeof method !== 'function') return method
+
+  return function methodWithTrace (url, opts, handler) {
+    const lastIndex = arguments.length - 1
+
+    arguments[lastIndex] = wrapHandler(arguments[lastIndex])
+
+    return method.apply(this, arguments)
+  }
+}
+
+function wrapHandler (handler) {
+  return function handlerWithTrace (request, reply) {
+    const req = getReq(request)
+
+    return web.reactivate(req, () => handler.apply(this, arguments))
+  }
+}
+
 function unwrapFastify (fastify) {
   fastify._datadog_wrapper = fastify
 }

packages/datadog-plugin-fastify/test/index.spec.js

@@ -141,6 +141,49 @@ describe('Plugin', () => {
           })
         })
 
+        it('should run POST handlers in the request scope', done => {
+          if (process.env.DD_CONTEXT_PROPAGATION === 'false') return done()
+
+          app.post('/user', (request, reply) => {
+            expect(tracer.scope().active()).to.not.be.null
+            reply.send()
+          })
+
+          getPort().then(port => {
+            app.listen(port, 'localhost', () => {
+              axios.post(`http://localhost:${port}/user`, { foo: 'bar' })
+                .then(() => done())
+                .catch(done)
+            })
+          })
+        })
+
+        it('should run routes in the request scope', done => {
+          if (process.env.DD_CONTEXT_PROPAGATION === 'false') return done()
+
+          app.use((req, res, next) => {
+            expect(tracer.scope().active()).to.not.be.null
+            next()
+          })
+
+          app.route({
+            method: 'POST',
+            url: '/user',
+            handler: (request, reply) => {
+              expect(tracer.scope().active()).to.not.be.null
+              reply.send()
+            }
+          })
+
+          getPort().then(port => {
+            app.listen(port, 'localhost', () => {
+              axios.post(`http://localhost:${port}/user`, { foo: 'bar' })
+                .then(() => done())
+                .catch(done)
+            })
+          })
+        })
+
         it('should run hooks in the request scope', done => {
           if (process.env.DD_CONTEXT_PROPAGATION === 'false') return done()