From 20348784da2b5b6e6c5f606730802786616d086b Mon Sep 17 00:00:00 2001
From: Roch Devost <roch.devost@gmail.com>
Date: Mon, 5 Nov 2018 12:54:37 -0500
Subject: [PATCH] add span hooks for web framework integrations (#351)

---
 src/plugins/util/web.js       | 19 ++++++++++++++++---
 test/plugins/util/web.spec.js | 22 +++++++++++++++++++++-
 2 files changed, 37 insertions(+), 4 deletions(-)

diff --git a/src/plugins/util/web.js b/src/plugins/util/web.js
index fde4cd5280e..0816f3069cc 100644
--- a/src/plugins/util/web.js
+++ b/src/plugins/util/web.js
@@ -24,10 +24,12 @@ const web = {
   normalizeConfig (config) {
     const headers = getHeadersToRecord(config)
     const validateStatus = getStatusValidator(config)
+    const hooks = getHooks(config)
 
     return Object.assign({}, config, {
       headers,
-      validateStatus
+      validateStatus,
+      hooks
     })
   },
 
@@ -78,6 +80,7 @@ const web = {
         span: null,
         scope: null,
         paths: [],
+        hooks: [],
         beforeEnd: []
       }
     })
@@ -90,6 +93,8 @@ const web = {
 }
 
 function startSpan (tracer, config, req, res, name) {
+  req._datadog.hooks.push(config.hooks)
+
   if (req._datadog.span) {
     req._datadog.span.context().name = name
     return req._datadog.span
@@ -110,11 +115,12 @@ function startSpan (tracer, config, req, res, name) {
   return span
 }
 
-function finish (req) {
+function finish (req, res) {
   if (req._datadog.finished) return
 
   addResponseTags(req)
 
+  req._datadog.hooks.forEach(hooks => hooks.request(req._datadog.span, req, res))
   req._datadog.span.finish()
   req._datadog.scope && req._datadog.scope.close()
   req._datadog.finished = true
@@ -129,7 +135,7 @@ function wrapEnd (req) {
 
     const returnValue = end.apply(this, arguments)
 
-    finish(req)
+    finish(req, res)
 
     return returnValue
   }
@@ -212,4 +218,11 @@ function getStatusValidator (config) {
   return code => code < 500
 }
 
+function getHooks (config) {
+  const noop = () => {}
+  const request = (config.hooks && config.hooks.request) || noop
+
+  return { request }
+}
+
 module.exports = web
diff --git a/test/plugins/util/web.spec.js b/test/plugins/util/web.spec.js
index 1e7350d543b..cb0b92ae2fe 100644
--- a/test/plugins/util/web.spec.js
+++ b/test/plugins/util/web.spec.js
@@ -37,7 +37,7 @@ describe('plugins/util/web', () => {
     res = {
       end
     }
-    config = {}
+    config = { hooks: {} }
 
     tracer = require('../../..').init({ plugins: false })
     web = require('../../../src/plugins/util/web')
@@ -216,6 +216,26 @@ describe('plugins/util/web', () => {
           [HTTP_ROUTE]: '/custom/route'
         })
       })
+
+      it('should execute the request end hook', () => {
+        config.hooks.request = sinon.spy()
+
+        res.end()
+
+        expect(config.hooks.request).to.have.been.calledWith(span, req, res)
+      })
+
+      it('should execute multiple end hooks', () => {
+        config.hooks = {
+          request: sinon.spy()
+        }
+
+        span = web.instrument(tracer, config, req, res, 'test.request')
+
+        res.end()
+
+        expect(config.hooks.request).to.have.been.calledWith(span, req, res)
+      })
     })
   })