fix(koa): ignore generator-based Koa middleware (open-telemetry#1119)

plugins/node/opentelemetry-instrumentation-koa/README.md

@@ -38,7 +38,9 @@ registerInstrumentations({
 });
 ```
 
-See [`examples`](https://github.com/open-telemetry/opentelemetry-js-contrib/tree/main/plugins/node/opentelemetry-instrumentation-koa/examples) for a short example using both Koa and @koa/router
+See [`examples`](https://github.com/open-telemetry/opentelemetry-js-contrib/tree/main/plugins/node/opentelemetry-instrumentation-koa/examples) for a short example using both Koa and @koa/router.
+
+Note that generator-based middleware are deprecated and won't be instrumented.
 
 ### Koa Instrumentation Options
 

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

@@ -147,7 +147,17 @@ export class KoaInstrumentation extends InstrumentationBase<typeof koa> {
       isLayerIgnored(layerType, this.getConfig())
     )
       return middlewareLayer;
+
+    if (
+      middlewareLayer.constructor.name === 'GeneratorFunction' ||
+      middlewareLayer.constructor.name === 'AsyncGeneratorFunction'
+    ) {
+      api.diag.debug('ignoring generator-based Koa middleware layer');
+      return middlewareLayer;
+    }
+
     middlewareLayer[kLayerPatched] = true;
+
     api.diag.debug('patching Koa middleware layer');
     return async (context: KoaContext, next: koa.Next) => {
       const parent = api.trace.getSpan(api.context.active());

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

@@ -92,34 +92,44 @@ describe('Koa Instrumentation', () => {
     server.close();
   });
 
-  const simpleResponse: koa.Middleware = async (ctx, next) => {
-    ctx.body = 'test';
-    await next();
-  };
-
-  const customMiddleware: koa.Middleware = async (ctx, next) => {
-    for (let i = 0; i < 1000000; i++) {
-      continue;
-    }
-    await next();
-  };
-
-  const spanCreateMiddleware: koa.Middleware = async (ctx, next) => {
-    const span = tracer.startSpan('foo');
-    span.end();
-    await next();
-  };
-
-  const asyncMiddleware: koa.Middleware = async (ctx, next) => {
-    const start = Date.now();
-    await next();
-    const ms = Date.now() - start;
-    ctx.body = `${ctx.method} ${ctx.url} - ${ms}ms`;
-  };
-
-  const failingMiddleware: koa.Middleware = async (_ctx, _next) => {
-    throw new Error('I failed!');
-  };
+  const simpleResponse: () => koa.Middleware = () =>
+    async function simpleResponse(ctx, next) {
+      ctx.body = 'test';
+      await next();
+    };
+
+  const customMiddleware: () => koa.Middleware = () =>
+    async function customMiddleware(ctx, next) {
+      for (let i = 0; i < 1000000; i++) {
+        continue;
+      }
+      await next();
+    };
+
+  const spanCreateMiddleware: () => koa.Middleware = () =>
+    async function spanCreateMiddleware(ctx, next) {
+      const span = tracer.startSpan('foo');
+      span.end();
+      await next();
+    };
+
+  const asyncMiddleware: () => koa.Middleware = () =>
+    async function asyncMiddleware(ctx, next) {
+      const start = Date.now();
+      await next();
+      const ms = Date.now() - start;
+      ctx.body = `${ctx.method} ${ctx.url} - ${ms}ms`;
+    };
+
+  const failingMiddleware: () => koa.Middleware = () =>
+    async function failingMiddleware(_ctx, _next) {
+      throw new Error('I failed!');
+    };
+
+  const generatorMiddleware: () => koa.Middleware = () =>
+    function* generatorMiddleware(next) {
+      yield next;
+    };
 
   describe('Instrumenting @koa/router calls', () => {
     it('should create a child span for middlewares', async () => {
@@ -331,9 +341,9 @@ describe('Koa Instrumentation', () => {
       app.use((ctx, next) =>
         context.with(trace.setSpan(context.active(), rootSpan), next)
       );
-      app.use(customMiddleware);
-      app.use(simpleResponse);
-      app.use(spanCreateMiddleware);
+      app.use(customMiddleware());
+      app.use(simpleResponse());
+      app.use(spanCreateMiddleware());
 
       await context.with(
         trace.setSpan(context.active(), rootSpan),
@@ -384,8 +394,8 @@ describe('Koa Instrumentation', () => {
     });
 
     it('should not create span if there is no parent span', async () => {
-      app.use(customMiddleware);
-      app.use(simpleResponse);
+      app.use(customMiddleware());
+      app.use(simpleResponse());
 
       const res = await httpRequest.get(`http://localhost:${port}`);
       assert.strictEqual(memoryExporter.getFinishedSpans().length, 0);
@@ -397,7 +407,7 @@ describe('Koa Instrumentation', () => {
       app.use((ctx, next) =>
         context.with(trace.setSpan(context.active(), rootSpan), next)
       );
-      app.use(asyncMiddleware);
+      app.use(asyncMiddleware());
 
       await context.with(
         trace.setSpan(context.active(), rootSpan),
@@ -423,12 +433,72 @@ describe('Koa Instrumentation', () => {
       );
     });
 
+    it('should not instrument generator middleware functions', async () => {
+      const rootSpan = tracer.startSpan('rootSpan');
+      app.use((_ctx, next) =>
+        context.with(trace.setSpan(context.active(), rootSpan), next)
+      );
+
+      app.use(generatorMiddleware());
+      app.use(simpleResponse());
+
+      await context.with(
+        trace.setSpan(context.active(), rootSpan),
+        async () => {
+          await httpRequest.get(`http://localhost:${port}`);
+          rootSpan.end();
+          assert.deepStrictEqual(memoryExporter.getFinishedSpans().length, 2);
+
+          const simpleResponseSpan = memoryExporter
+            .getFinishedSpans()
+            .find(span => span.name.includes('simpleResponse'));
+          assert.notStrictEqual(simpleResponseSpan, undefined);
+
+          const exportedRootSpan = memoryExporter
+            .getFinishedSpans()
+            .find(span => span.name === 'rootSpan');
+          assert.notStrictEqual(exportedRootSpan, undefined);
+        }
+      );
+    });
+
+    it('should not instrument the same middleware more than once', async () => {
+      const sameAsyncMiddleware = asyncMiddleware();
+
+      const rootSpan = tracer.startSpan('rootSpan');
+      app.use((_ctx, next) =>
+        context.with(trace.setSpan(context.active(), rootSpan), next)
+      );
+
+      app.use(sameAsyncMiddleware);
+      app.use(sameAsyncMiddleware);
+
+      await context.with(
+        trace.setSpan(context.active(), rootSpan),
+        async () => {
+          await httpRequest.get(`http://localhost:${port}`);
+          rootSpan.end();
+          assert.deepStrictEqual(memoryExporter.getFinishedSpans().length, 2);
+
+          const asyncMiddlewareSpan = memoryExporter
+            .getFinishedSpans()
+            .find(span => span.name.includes('asyncMiddleware'));
+          assert.notStrictEqual(asyncMiddlewareSpan, undefined);
+
+          const exportedRootSpan = memoryExporter
+            .getFinishedSpans()
+            .find(span => span.name === 'rootSpan');
+          assert.notStrictEqual(exportedRootSpan, undefined);
+        }
+      );
+    });
+
     it('should propagate exceptions in the middleware while marking the span with an exception', async () => {
       const rootSpan = tracer.startSpan('rootSpan');
       app.use((_ctx, next) =>
         context.with(trace.setSpan(context.active(), rootSpan), next)
       );
-      app.use(failingMiddleware);
+      app.use(failingMiddleware());
       const res = await httpRequest.get(`http://localhost:${port}`);
       assert.deepStrictEqual(res, 'Internal Server Error');
 
@@ -573,7 +643,7 @@ describe('Koa Instrumentation', () => {
     it('should not create new spans', async () => {
       plugin.disable();
       const rootSpan = tracer.startSpan('rootSpan');
-      app.use(customMiddleware);
+      app.use(customMiddleware());
 
       await context.with(
         trace.setSpan(context.active(), rootSpan),