Support Koa tracing through the plugin API (#380)

PR-URL: #380

config.js

@@ -41,7 +41,8 @@ module.exports = {
     // all.
     // This field is experimental.
     plugins: {
-      'restify': path.join(__dirname, 'src/plugins/plugin-restify.js')
+      'restify': path.join(__dirname, 'src/plugins/plugin-restify.js'),
+      'koa': path.join(__dirname, 'src/plugins/plugin-koa.js')
     },
 
     // Valid entries are:

src/hooks/index.js

@@ -45,8 +45,6 @@ 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',

src/plugins/plugin-koa.js

@@ -0,0 +1,101 @@
+/**
+ * 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 shimmer = require('shimmer');
+var urlParse = require('url').parse;
+
+const SUPPORTED_VERSIONS = '1.x';
+
+var api;
+
+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 req = this.req;
+  const res = this.res;
+  const originalEnd = res.end;
+  var options = {
+    name: urlParse(req.url).pathname,
+    traceContext: this.req.headers[api.constants.TRACE_CONTEXT_HEADER_NAME],
+    skipFrames: 3
+  };
+  api.runInRootSpan(options, function(root) {
+    if (!root) {
+      return;
+    }
+
+    api.wrapEmitter(req);
+    api.wrapEmitter(res);
+
+    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
+    // req.path would be more desirable but is not set at the time our middlewear runs.
+    root.addLabel(api.labels.HTTP_METHOD_LABEL_KEY, req.method);
+    root.addLabel(api.labels.HTTP_URL_LABEL_KEY, url);
+    root.addLabel(api.labels.HTTP_SOURCE_IP, req.connection.remoteAddress);
+
+    var context = root.getTraceContext();
+    if (context) {
+      res.setHeader(api.constants.TRACE_CONTEXT_HEADER_NAME, context);
+    }
+
+    // wrap end
+    res.end = function(chunk, encoding) {
+      res.end = originalEnd;
+      const returned = res.end(chunk, encoding);
+
+      if (req.route && req.route.path) {
+        root.addLabel(
+          'koa/request.route.path', req.route.path);
+      }
+      root.addLabel(
+          api.labels.HTTP_RESPONSE_CODE_LABEL_KEY, res.statusCode);
+      root.endSpan();
+
+      return returned;
+    };
+    api.wrap(next);
+  });
+
+  yield next;
+}
+
+module.exports = [
+  {
+    file: '',
+    versions: SUPPORTED_VERSIONS,
+    patch: function(koa, api_) {
+      api = api_;
+      shimmer.wrap(koa.prototype, 'use', useWrap);
+    },
+    unpatch: function(koa) {
+      shimmer.unwrap(koa.prototype, 'use');
+    }
+  }
+];