feat(mustCollect): move mustCollect to cls session

RisingStack · Dec 2, 2015 · e6ceda2 · e6ceda2
1 parent 5afdb48
commit e6ceda2
Showing 1 changed file with 8 additions and 6 deletions.
diff --git a/lib/providers/httpTransaction/wraps/http.Server.prototype.js b/lib/providers/httpTransaction/wraps/http.Server.prototype.js
@@ -14,7 +14,7 @@ function wrapListener (listener, collector, config) {
   var ignoreHeaders = config.ignoreHeaders
 
   return function (request, response) {
-    var mustCollect
+    var session = getNamespace('trace')
     var serverRecieveTime
 
     var headers = request.headers
@@ -36,7 +36,10 @@ function wrapListener (listener, collector, config) {
 
     var requestId = headers['request-id'] || requestQuery
 
-    mustCollect = headers['X-Must-Collect']
+    var mustCollectHeader = headers['X-Must-Collect']
+    if (mustCollectHeader) {
+      session.set('mustCollect', mustCollectHeader)
+    }
 
     if (!requestId) {
       requestId = uuid.v1()
@@ -54,8 +57,6 @@ function wrapListener (listener, collector, config) {
       headers: headers
     }
 
-    var session = getNamespace('trace')
-
     // Collect request start
     process.nextTick(function () {
       collector.emit(Collector.SERVER_RECV, collectorDataBag)
@@ -69,7 +70,7 @@ function wrapListener (listener, collector, config) {
       var responseTime = serverSendTime - serverRecieveTime
 
       var collectorDataBag = {
-        mustCollect: mustCollect,
+        mustCollect: session.get('mustCollect'),
         id: requestId,
         host: headers.host,
         url: requestUrl.pathname,
@@ -89,7 +90,8 @@ function wrapListener (listener, collector, config) {
 
     response.writeHead = function () {
       serverSendTime = microtime.now()
-      mustCollect = mustCollect || response.statusCode > 399
+      var mustCollect = session.get('mustCollect') || response.statusCode > 399
+      session.set('mustCollect', mustCollect)
 
       /* Service name may be unavailable due to uninitialized reporter */
       var serviceName = collector.getService()