Tracing hooks for koa

Fixes #191

bin/run-test.sh

@@ -42,7 +42,10 @@ function run {
 run test test/hooks
 for test in test/standalone/test-*.js ;
 do
-  run "${test}"
+  if [[ ! $(node --version) =~ v0\.12\..* || ! "${test}" =~ .*trace\-koa\.js ]]
+  then
+    run "${test}"
+  fi
 done
 
 # Conditionally publish coverage

lib/hooks/index.js

@@ -43,6 +43,8 @@ var toInstrument = Object.create(null, {
       patches: {} } },
   'http': { enumerable: true, value: { file: './core/hook-http.js',
       patches: {} } },
+  'koa': { enumerable: true, value: { file: './userspace/hook-koa.js',
+      patches: {} } },
   'mongodb-core': { enumerable: true, value: { file: './userspace/hook-mongodb-core.js',
       patches: {} } },
   'mysql': { enumerable: true, value: { file: './userspace/hook-mysql.js',

lib/hooks/userspace/hook-koa.js

@@ -0,0 +1,131 @@
+/**
+ * Copyright 2015 Google Inc. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+'use strict';
+
+const cls = require('../../cls.js');
+const TraceLabels = require('../../trace-labels.js');
+const shimmer = require('shimmer');
+const semver = require('semver');
+const constants = require('../../constants.js');
+var agent;
+
+const SUPPORTED_VERSIONS = '1.x';
+
+function useWrap(use) {
+  return function useTrace() {
+    if (!this._google_trace_patched) {
+      this._google_trace_patched = true;
+      this.use(middleware);
+    }
+    return use.apply(this, arguments);
+  };
+}
+
+function* middleware(next) {
+  /* jshint validthis:true */
+  const namespace = cls.getNamespace();
+  if (!namespace) {
+    agent.logger.info('Koa: no namespace found, ignoring request');
+    return;
+  }
+  if (!agent.shouldTrace(this.req.url)) {
+    return;
+  }
+  const req = this.req;
+  const res = this.res;
+
+  namespace.bindEmitter(req);
+  namespace.bindEmitter(res);
+
+  const originalEnd = res.end;
+
+  namespace.run(function() {
+    const rootContext = startRootSpanForRequest(req);
+
+    // wrap end
+    res.end = function(chunk, encoding) {
+      res.end = originalEnd;
+      const returned = res.end(chunk, encoding);
+
+      endRootSpanForRequest(rootContext, req, res);
+      return returned;
+    };
+    namespace.bind(next);
+  });
+  yield next;
+}
+
+/**
+ * Creates and sets up a new root span for the given request.
+ * @param {Object} req The request being processed.
+ * @returns {!SpanData} The new initialized trace span data instance.
+ */
+function startRootSpanForRequest(req) {
+  const result = agent.parseContextFromHeader(
+    req.headers[constants.TRACE_CONTEXT_HEADER_NAME]) || {};
+
+  const traceId = result.traceId;
+  const parentSpanId = result.spanId;
+  const url = (req.headers['X-Forwarded-Proto'] || 'http') +
+    '://' + req.headers.host + req.url;
+
+  // we use the path part of the url as the span name and add the full
+  // url as a label
+  const rootContext = agent.createRootSpanData(req.url, traceId,
+    parentSpanId);
+  rootContext.addLabel(TraceLabels.HTTP_METHOD_LABEL_KEY, req.method);
+  rootContext.addLabel(TraceLabels.HTTP_URL_LABEL_KEY, url);
+  rootContext.addLabel(TraceLabels.HTTP_SOURCE_IP, req.connection.remoteAddress);
+  return rootContext;
+}
+
+
+/**
+ * Ends the root span for the given request.
+ * @param {!SpanData} rootContext The span to close out.
+ * @param {Object} req The request being processed.
+ * @param {Object} res The response being processed.
+ */
+function endRootSpanForRequest(rootContext, req, res) {
+  if (req.route && req.route.path) {
+    rootContext.addLabel(
+      'koa/request.route.path', req.route.path);
+  }
+  rootContext.addLabel(
+      TraceLabels.HTTP_RESPONSE_CODE_LABEL_KEY, res.statusCode);
+  rootContext.close();
+}
+
+module.exports = function(version_, agent_) {
+  if (!semver.satisfies(version_, SUPPORTED_VERSIONS)) {
+    agent_.logger.info('Koa: unsupported version ' + version_ + ' loaded');
+    return {};
+  }
+  return {
+    // An empty relative path here matches the root module being loaded.
+    '': {
+      patch: function(koa) {
+        agent = agent_;
+        shimmer.wrap(koa.prototype, 'use', useWrap);
+      },
+      unpatch: function(koa) {
+        shimmer.unwrap(koa.prototype, 'use');
+        agent_.logger.info('Koa: unpatched');
+      }
+    }
+  };
+};

test/hooks/fixtures/koa1/index.js

@@ -0,0 +1 @@
+module.exports = require('koa');

test/hooks/fixtures/koa1/package.json

@@ -0,0 +1,8 @@
+{
+  "name": "koa1",
+  "version": "1.0.0",
+  "main": "index.js",
+  "dependencies": {
+    "koa": "^1.1.2"
+  }
+}

test/standalone/test-trace-koa.js

@@ -0,0 +1,80 @@
+/**
+ * Copyright 2015 Google Inc. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+'use strict';
+
+var common = require('../hooks/common.js');
+var koa = require('../hooks/fixtures/koa1');
+var http = require('http');
+var assert = require('assert');
+var TraceLabels = require('../../lib/trace-labels.js');
+
+var server;
+
+describe('test-trace-koa', function() {
+  afterEach(function() {
+    common.cleanTraces();
+    server.close();
+  });
+
+  it('should accurately measure get time, get', function(done) {
+    var app = koa();
+    app.use(function* () {
+      this.body = yield function(cb) {
+        setTimeout(function() {
+          cb(null, common.serverRes);
+        }, common.serverWait);
+      };
+    });
+    server = app.listen(common.serverPort, function() {
+      common.doRequest('GET', done, koaPredicate);
+    });
+  });
+
+  it('should have required labels', function(done) {
+    var app = koa();
+    app.use(function* () {
+      this.body = yield function(cb) {
+        setTimeout(function() {
+          cb(null, common.serverRes);
+        }, common.serverWait);
+      };
+    });
+    server = app.listen(common.serverPort, function() {
+      http.get({port: common.serverPort}, function(res) {
+        var result = '';
+        res.on('data', function(data) { result += data; });
+        res.on('end', function() {
+          assert.equal(common.serverRes, result);
+          var expectedKeys = [
+            TraceLabels.HTTP_METHOD_LABEL_KEY,
+            TraceLabels.HTTP_URL_LABEL_KEY,
+            TraceLabels.HTTP_SOURCE_IP,
+            TraceLabels.HTTP_RESPONSE_CODE_LABEL_KEY
+          ];
+          var span = common.getMatchingSpan(koaPredicate);
+          expectedKeys.forEach(function(key) {
+            assert(span.labels[key]);
+          });
+          done();
+        });
+      });
+    });
+  });
+});
+
+function koaPredicate(span) {
+  return span.name === '/';
+}