fix(koa): fix instrumentation of ESM-imported koa (#1736)

Co-authored-by: Trent Mick <trentm@gmail.com>

plugins/node/opentelemetry-instrumentation-koa/package.json

@@ -49,6 +49,8 @@
     "@koa/router": "12.0.0",
     "@opentelemetry/api": "^1.3.0",
     "@opentelemetry/context-async-hooks": "^1.8.0",
+    "@opentelemetry/contrib-test-utils": "^0.35.0",
+    "@opentelemetry/instrumentation-http": "^0.45.1",
     "@opentelemetry/sdk-trace-base": "^1.8.0",
     "@opentelemetry/sdk-trace-node": "^1.8.0",
     "@types/mocha": "7.0.2",

plugins/node/opentelemetry-instrumentation-koa/src/instrumentation.ts

@@ -55,7 +55,11 @@ export class KoaInstrumentation extends InstrumentationBase<typeof koa> {
     return new InstrumentationNodeModuleDefinition<typeof koa>(
       'koa',
       ['^2.0.0'],
-      moduleExports => {
+      (module: any) => {
+        const moduleExports: typeof koa =
+          module[Symbol.toStringTag] === 'Module'
+            ? module.default // ESM
+            : module; // CommonJS
         if (moduleExports == null) {
           return moduleExports;
         }
@@ -68,9 +72,13 @@ export class KoaInstrumentation extends InstrumentationBase<typeof koa> {
           'use',
           this._getKoaUsePatch.bind(this)
         );
-        return moduleExports;
+        return module;
       },
-      moduleExports => {
+      (module: any) => {
+        const moduleExports: typeof koa =
+          module[Symbol.toStringTag] === 'Module'
+            ? module.default // ESM
+            : module; // CommonJS
         api.diag.debug('Unpatching Koa');
         if (isWrapped(moduleExports.prototype.use)) {
           this._unwrap(moduleExports.prototype, 'use');

plugins/node/opentelemetry-instrumentation-koa/test/fixtures/use-koa.mjs

@@ -0,0 +1,71 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ *
+ * 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
+ *
+ *      https://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 koa from an ES module:
+//    node --experimental-loader=@opentelemetry/instrumentation/hook.mjs use-koa.mjs
+
+import { promisify } from 'util';
+import { createTestNodeSdk } from '@opentelemetry/contrib-test-utils';
+
+import { HttpInstrumentation } from '@opentelemetry/instrumentation-http';
+import { KoaInstrumentation } from '../../build/src/index.js';
+
+const sdk = createTestNodeSdk({
+  serviceName: 'use-koa',
+  instrumentations: [
+    new KoaInstrumentation(),
+    new HttpInstrumentation()
+  ]
+})
+sdk.start();
+
+import Koa from 'koa';
+import KoaRouter from '@koa/router';
+import * as http from 'http';
+
+const app = new Koa();
+
+app.use(async function simpleMiddleware(ctx, next) {
+  // Wait a short delay to ensure this "middleware - ..." span clearly starts
+  // before the "router - ..." span. The test rely on a start-time-based sort
+  // of the produced spans. If they start in the same millisecond, then tests
+  // can be flaky.
+  await promisify(setTimeout)(10);
+  await next();
+});
+
+const router = new KoaRouter();
+router.get('/post/:id', ctx => {
+  ctx.body = `Post id: ${ctx.params.id}`;
+});
+
+app.use(router.routes());
+
+const server = http.createServer(app.callback());
+await new Promise(resolve => server.listen(0, resolve));
+const port = server.address().port;
+
+await new Promise(resolve => {
+  http.get(`http://localhost:${port}/post/0`, (res) => {
+    res.resume();
+    res.on('end', () => {
+      resolve();
+    });
+  })
+});
+
+await new Promise(resolve => server.close(resolve));
+await sdk.shutdown();

plugins/node/opentelemetry-instrumentation-koa/test/koa.test.ts

@@ -15,9 +15,10 @@
  */
 
 import * as KoaRouter from '@koa/router';
-import { context, trace, Span } from '@opentelemetry/api';
+import { context, trace, Span, SpanKind } from '@opentelemetry/api';
 import { NodeTracerProvider } from '@opentelemetry/sdk-trace-node';
 import { AsyncHooksContextManager } from '@opentelemetry/context-async-hooks';
+import * as testUtils from '@opentelemetry/contrib-test-utils';
 import {
   InMemorySpanExporter,
   SimpleSpanProcessor,
@@ -709,4 +710,36 @@ describe('Koa Instrumentation', () => {
       );
     });
   });
+
+  it('should work with ESM usage', async () => {
+    await testUtils.runTestFixture({
+      cwd: __dirname,
+      argv: ['fixtures/use-koa.mjs'],
+      env: {
+        NODE_OPTIONS:
+          '--experimental-loader=@opentelemetry/instrumentation/hook.mjs',
+        NODE_NO_WARNINGS: '1',
+      },
+      checkResult: (err, stdout, stderr) => {
+        assert.ifError(err);
+      },
+      checkCollector: (collector: testUtils.TestCollector) => {
+        // use-koa.mjs creates a Koa app with a 'GET /post/:id' endpoint and
+        // a `simpleMiddleware`, then makes a single 'GET /post/0' request. We
+        // expect to see spans like this:
+        //    span 'GET /post/:id'
+        //    `- span 'middleware - simpleMiddleware'
+        //       `- span 'router - /post/:id'
+        const spans = collector.sortedSpans;
+        assert.strictEqual(spans[0].name, 'GET /post/:id');
+        assert.strictEqual(spans[0].kind, SpanKind.CLIENT);
+        assert.strictEqual(spans[1].name, 'middleware - simpleMiddleware');
+        assert.strictEqual(spans[1].kind, SpanKind.SERVER);
+        assert.strictEqual(spans[1].parentSpanId, spans[0].spanId);
+        assert.strictEqual(spans[2].name, 'router - /post/:id');
+        assert.strictEqual(spans[2].kind, SpanKind.SERVER);
+        assert.strictEqual(spans[2].parentSpanId, spans[1].spanId);
+      },
+    });
+  });
 });