diff --git a/plugins/node/opentelemetry-plugin-koa/src/koa.ts b/plugins/node/opentelemetry-plugin-koa/src/koa.ts index ebed303cb5..175cdcf1e7 100644 --- a/plugins/node/opentelemetry-plugin-koa/src/koa.ts +++ b/plugins/node/opentelemetry-plugin-koa/src/koa.ts @@ -116,8 +116,7 @@ export class KoaPlugin extends BasePlugin { if (middlewareLayer[kLayerPatched] === true) return middlewareLayer; middlewareLayer[kLayerPatched] = true; this._logger.debug('patching Koa middleware layer'); - - return (context: KoaContext, next: koa.Next) => { + return async (context: KoaContext, next: koa.Next) => { if (this._tracer.getCurrentSpan() === undefined) { return middlewareLayer(context, next); } @@ -132,7 +131,7 @@ export class KoaPlugin extends BasePlugin { }); const startTime = hrTime(); - const result = middlewareLayer(context, next); + const result = await middlewareLayer(context, next); span.end(startTime); return result; }; diff --git a/plugins/node/opentelemetry-plugin-koa/test/koa.test.ts b/plugins/node/opentelemetry-plugin-koa/test/koa.test.ts index e22b775923..ef15767339 100644 --- a/plugins/node/opentelemetry-plugin-koa/test/koa.test.ts +++ b/plugins/node/opentelemetry-plugin-koa/test/koa.test.ts @@ -83,6 +83,13 @@ describe('Koa Plugin - Core Tests', () => { return next(); }; + const asyncMiddleware: koa.Middleware = async (ctx, next) => { + const start = Date.now(); + await next(); + const ms = Date.now() - start; + console.log(`${ctx.method} ${ctx.url} - ${ms}ms`); + }; + describe('Instrumenting core middleware calls', () => { it('should create a child span for middlewares', async () => { const rootSpan = tracer.startSpan('rootSpan'); @@ -151,6 +158,41 @@ describe('Koa Plugin - Core Tests', () => { assert.strictEqual(res, 'test'); server.close(); }); + + it('should handle async middleware functions', async () => { + const rootSpan = tracer.startSpan('rootSpan'); + const app = new koa(); + app.use((ctx, next) => tracer.withSpan(rootSpan, next)); + app.use(asyncMiddleware); + + const server = http.createServer(app.callback()); + await new Promise(resolve => server.listen(0, resolve)); + const port = (server.address() as AddressInfo).port; + assert.strictEqual(memoryExporter.getFinishedSpans().length, 0); + + await tracer.withSpan(rootSpan, async () => { + await httpRequest.get(`http://localhost:${port}`); + rootSpan.end(); + + const requestHandlerSpan = memoryExporter + .getFinishedSpans() + .find(span => span.name.includes('asyncMiddleware')); + assert.notStrictEqual(requestHandlerSpan, undefined); + assert.strictEqual( + requestHandlerSpan?.attributes[AttributeNames.COMPONENT], + KoaComponentName + ); + assert.strictEqual( + requestHandlerSpan?.attributes[AttributeNames.KOA_TYPE], + KoaLayerType.MIDDLEWARE + ); + const exportedRootSpan = memoryExporter + .getFinishedSpans() + .find(span => span.name === 'rootSpan'); + assert.notStrictEqual(exportedRootSpan, undefined); + }); + server.close(); + }); }); describe('Disabling plugin', () => {