fix: propagate context in koa tracing

src/plugins/plugin-koa.ts

@@ -16,26 +16,26 @@
 'use strict';
 
 const shimmer = require('shimmer');
-var urlParse = require('url').parse;
+const urlParse = require('url').parse;
 
 function startSpanForRequest(api, req, res, next) {
-  var originalEnd = res.end;
-  var options = {
+  const originalEnd = res.end;
+  const options = {
     name: urlParse(req.url).pathname,
     url: req.url,
     traceContext: req.headers[api.constants.TRACE_CONTEXT_HEADER_NAME],
     skipFrames: 4
   };
-  api.runInRootSpan(options, function(root) {
+  return api.runInRootSpan(options, function(root) {
     // Set response trace context.
-    var responseTraceContext =
+    const responseTraceContext =
       api.getResponseTraceContext(options.traceContext, !!root);
     if (responseTraceContext) {
       res.setHeader(api.constants.TRACE_CONTEXT_HEADER_NAME, responseTraceContext);
     }
 
     if (!root) {
-      return;
+      return next;
     }
 
     api.wrapEmitter(req);
@@ -74,7 +74,14 @@ function startSpanForRequest(api, req, res, next) {
       root.endSpan();
     });
 
-    api.wrap(next);
+    // In Koa 1, next is a Generator object.
+    // We wrap Generator#next here.
+    if (!next.apply && next.next) {
+      next.next = api.wrap(next.next);
+      return next;
+    } else {
+      return api.wrap(next);
+    }
   });
 }
 
@@ -84,7 +91,7 @@ function createMiddleware(api) {
     const req = this.req;
     const res = this.res;
 
-    startSpanForRequest(api, req, res, next);
+    next = startSpanForRequest(api, req, res, next);
 
     yield next;
   };
@@ -95,7 +102,7 @@ function createMiddleware2x(api) {
     const req = ctx.req;
     const res = ctx.res;
 
-    startSpanForRequest(api, req, res, next);
+    next = startSpanForRequest(api, req, res, next);
 
     return next();
   };

test/plugins/test-trace-koa.ts

@@ -25,8 +25,9 @@ var semver = require('semver');
 var appBuilders: any = {
   koa1: buildKoa1App,
 };
-if (semver.satisfies(process.version, '>4')) {
-  appBuilders.koa2 = buildKoa2App;
+var skipKoa2 = semver.satisfies(process.version, '<=4');
+if (!skipKoa2) {
+  appBuilders.koa2 = buildKoa2App
 }
 
 describe('koa', function() {
@@ -166,6 +167,65 @@ describe('koa', function() {
       });
     });
   });
+
+  describe('execution context propagation', function() {
+    it('should work in koa 1', function(done) {
+      var children: any[] = [];
+      var koa = require('./fixtures/koa1');
+      var app = koa();
+      app.use(function* (next) {
+        children.push(agent.createChildSpan({ name: 'span0' }));
+        yield* next;
+        this.body = '';
+      });
+      app.use(function* () {
+        children.push(agent.createChildSpan({ name: 'span1' }));
+      });
+      server = app.listen(common.serverPort, () => {
+        http.get({ port: common.serverPort, path: '/' }, (res) => {
+          res.on('data', () => {});
+          res.on('end', () => {
+            server.close();
+            assert.strictEqual(children.length, 2);
+            children.forEach((childSpan, index) => {
+              assert.ok(childSpan);
+              assert.strictEqual(childSpan.span.name, `span${index}`);
+            });
+            done();
+          });
+        });
+      });
+    });
+
+    (skipKoa2 ? it.skip : it)('should work in koa 2', function(done) {
+      var children: any[] = [];
+      var Koa = require('./fixtures/koa2');
+      var app = new Koa();
+      app.use(function (ctx, next) {
+        children.push(agent.createChildSpan({ name: 'span0' }));
+        return new Promise(resolve => {
+          ctx.body = '';
+          setTimeout(resolve, 100);
+        }).then(() => {
+          children.push(agent.createChildSpan({ name: 'span1' }));
+        });
+      });
+      server = app.listen(common.serverPort, () => {
+        http.get({ port: common.serverPort, path: '/' }, (res) => {
+          res.on('data', () => {});
+          res.on('end', () => {
+            server.close();
+            assert.strictEqual(children.length, 2);
+            children.forEach((childSpan, index) => {
+              assert.ok(childSpan);
+              assert.strictEqual(childSpan.span.name, `span${index}`);
+            });
+            done();
+          });
+        });
+      });
+    });
+  });
 });
 
 function koaPredicate(span) {