feat(koa): add requestHook support (open-telemetry#1099)

Co-authored-by: Amir Blum <[email protected]>
Co-authored-by: Valentin Marchaud <[email protected]>

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

@@ -45,12 +45,30 @@ See [`examples/koa`](https://github.com/open-telemetry/opentelemetry-js-contrib/
 | Options | Type | Example | Description |
 | ------- | ---- | ------- | ----------- |
 | `ignoreLayersType`| `KoaLayerType[]` | `['middleware']` | Ignore layers of specified type. |
+| `requestHook` | `KoaRequestCustomAttributeFunction` | `(span, info) => {}` | Function for adding custom attributes to Koa middleware layers. Receives params: `Span, KoaRequestInfo`. |
 
 `ignoreLayersType` accepts an array of `KoaLayerType` which can take the following string values:
 
 - `router`,
 - `middleware`.
 
+#### Using `requestHook`
+
+Instrumentation configuration accepts a custom "hook" function which will be called for every instrumented Koa middleware layer involved in a request. Custom attributes can be set on the span or run any custom logic per layer.
+
+```javascript
+import { KoaInstrumentation } from "@opentelemetry/instrumentation-koa"
+
+const koaInstrumentation = new KoaInstrumentation({
+  requestHook: function (span: Span, info: KoaRequestInfo) {
+    span.setAttribute(
+      'http.method',
+      info.context.request.method
+    )
+  }
+});
+```
+
 ## Koa Packages
 
 This package provides automatic tracing for middleware added using either the core [`koa`](https://github.com/koajs/koa) package or the [`@koa/router`](https://github.com/koajs/router) package.

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

@@ -55,11 +55,13 @@
     "@opentelemetry/sdk-trace-node": "^1.3.1",
     "@types/mocha": "7.0.2",
     "@types/node": "16.11.21",
+    "@types/sinon": "10.0.9",
     "gts": "3.1.0",
     "koa": "2.13.1",
     "mocha": "7.2.0",
     "nyc": "15.1.0",
     "rimraf": "3.0.2",
+    "sinon": "14.0.0",
     "test-all-versions": "5.0.1",
     "ts-mocha": "10.0.0",
     "typescript": "4.3.5"

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

@@ -19,6 +19,7 @@ import {
   isWrapped,
   InstrumentationBase,
   InstrumentationNodeModuleDefinition,
+  safeExecuteInTheMiddle,
 } from '@opentelemetry/instrumentation';
 
 import type * as koa from 'koa';
@@ -38,9 +39,22 @@ import { getRPCMetadata, RPCType, setRPCMetadata } from '@opentelemetry/core';
 /** Koa instrumentation for OpenTelemetry */
 export class KoaInstrumentation extends InstrumentationBase<typeof koa> {
   static readonly component = KoaComponentName;
-  constructor(config?: KoaInstrumentationConfig) {
-    super('@opentelemetry/instrumentation-koa', VERSION, config);
+  constructor(config: KoaInstrumentationConfig = {}) {
+    super(
+      '@opentelemetry/instrumentation-koa',
+      VERSION,
+      Object.assign({}, config)
+    );
+  }
+
+  override setConfig(config: KoaInstrumentationConfig = {}) {
+    this._config = Object.assign({}, config);
+  }
+
+  override getConfig(): KoaInstrumentationConfig {
+    return this._config as KoaInstrumentationConfig;
   }
+
   protected init() {
     return new InstrumentationNodeModuleDefinition<typeof koa>(
       'koa',
@@ -130,7 +144,7 @@ export class KoaInstrumentation extends InstrumentationBase<typeof koa> {
     // Skip patching layer if its ignored in the config
     if (
       middlewareLayer[kLayerPatched] === true ||
-      isLayerIgnored(layerType, this._config)
+      isLayerIgnored(layerType, this.getConfig())
     )
       return middlewareLayer;
     middlewareLayer[kLayerPatched] = true;
@@ -168,6 +182,20 @@ export class KoaInstrumentation extends InstrumentationBase<typeof koa> {
           Object.assign(rpcMetadata, { route: context._matchedRoute })
         );
       }
+
+      if (this.getConfig().requestHook) {
+        safeExecuteInTheMiddle(
+          () =>
+            this.getConfig().requestHook!(span, { context, middlewareLayer }),
+          e => {
+            if (e) {
+              api.diag.error('koa instrumentation: request hook failed', e);
+            }
+          },
+          true
+        );
+      }
+
       return api.context.with(newContext, async () => {
         try {
           return await middlewareLayer(context, next);

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

@@ -16,6 +16,7 @@
 import type { Middleware, ParameterizedContext, DefaultState } from 'koa';
 import type { RouterParamContext } from '@koa/router';
 import type * as Router from '@koa/router';
+import { Span } from '@opentelemetry/api';
 import { InstrumentationConfig } from '@opentelemetry/instrumentation';
 
 /**
@@ -31,12 +32,28 @@ export type KoaMiddleware = Middleware<DefaultState, KoaContext> & {
 
 export type KoaContext = ParameterizedContext<DefaultState, RouterParamContext>;
 
+export type KoaRequestInfo = {
+  context: KoaContext;
+  middlewareLayer: KoaMiddleware;
+};
+
+/**
+ * Function that can be used to add custom attributes to the current span
+ * @param span - The Express middleware layer span.
+ * @param context - The current KoaContext.
+ */
+export interface KoaRequestCustomAttributeFunction {
+  (span: Span, info: KoaRequestInfo): void;
+}
+
 /**
  * Options available for the Koa Instrumentation (see [documentation](https://github.com/open-telemetry/opentelemetry-js/tree/main/packages/opentelemetry-Instrumentation-koa#koa-Instrumentation-options))
  */
 export interface KoaInstrumentationConfig extends InstrumentationConfig {
   /** Ignore specific layers based on their type */
   ignoreLayersType?: KoaLayerType[];
+  /** Function for adding custom attributes to each middleware layer span */
+  requestHook?: KoaRequestCustomAttributeFunction;
 }
 
 export enum KoaLayerType {

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

@@ -15,7 +15,7 @@
  */
 
 import * as KoaRouter from '@koa/router';
-import { context, trace } from '@opentelemetry/api';
+import { context, trace, Span } from '@opentelemetry/api';
 import { NodeTracerProvider } from '@opentelemetry/sdk-trace-node';
 import { AsyncHooksContextManager } from '@opentelemetry/context-async-hooks';
 import {
@@ -30,8 +30,9 @@ const plugin = new KoaInstrumentation();
 import * as assert from 'assert';
 import * as koa from 'koa';
 import * as http from 'http';
+import * as sinon from 'sinon';
 import { AddressInfo } from 'net';
-import { KoaLayerType } from '../src/types';
+import { KoaLayerType, KoaRequestInfo } from '../src/types';
 import { AttributeNames } from '../src/enums/AttributeNames';
 import { RPCType, setRPCMetadata } from '@opentelemetry/core';
 
@@ -459,6 +460,115 @@ describe('Koa Instrumentation', () => {
     });
   });
 
+  describe('Using requestHook', () => {
+    it('should ignore requestHook which throws exception', async () => {
+      const rootSpan = tracer.startSpan('rootSpan');
+      const rpcMetadata = { type: RPCType.HTTP, span: rootSpan };
+      app.use((ctx, next) =>
+        context.with(
+          setRPCMetadata(
+            trace.setSpan(context.active(), rootSpan),
+            rpcMetadata
+          ),
+          next
+        )
+      );
+
+      const requestHook = sinon.spy((span: Span, info: KoaRequestInfo) => {
+        span.setAttribute(
+          SemanticAttributes.HTTP_METHOD,
+          info.context.request.method
+        );
+
+        throw Error('error thrown in requestHook');
+      });
+
+      plugin.setConfig({
+        requestHook,
+      });
+
+      const router = new KoaRouter();
+      router.get('/post/:id', ctx => {
+        ctx.body = `Post id: ${ctx.params.id}`;
+      });
+
+      app.use(router.routes());
+
+      await context.with(
+        trace.setSpan(context.active(), rootSpan),
+        async () => {
+          await httpRequest.get(`http://localhost:${port}/post/0`);
+          rootSpan.end();
+
+          assert.deepStrictEqual(memoryExporter.getFinishedSpans().length, 2);
+          const requestHandlerSpan = memoryExporter
+            .getFinishedSpans()
+            .find(span => span.name.includes('router - /post/:id'));
+          assert.notStrictEqual(requestHandlerSpan, undefined);
+          assert.strictEqual(
+            requestHandlerSpan?.attributes['http.method'],
+            'GET'
+          );
+
+          sinon.assert.threw(requestHook);
+        }
+      );
+    });
+
+    it('should call requestHook when set in config', async () => {
+      const rootSpan = tracer.startSpan('rootSpan');
+      const rpcMetadata = { type: RPCType.HTTP, span: rootSpan };
+      app.use((ctx, next) =>
+        context.with(
+          setRPCMetadata(
+            trace.setSpan(context.active(), rootSpan),
+            rpcMetadata
+          ),
+          next
+        )
+      );
+
+      const requestHook = sinon.spy((span: Span, info: KoaRequestInfo) => {
+        span.setAttribute('http.method', info.context.request.method);
+        span.setAttribute('app.env', info.context.app.env);
+      });
+
+      plugin.setConfig({
+        requestHook,
+      });
+
+      const router = new KoaRouter();
+      router.get('/post/:id', ctx => {
+        ctx.body = `Post id: ${ctx.params.id}`;
+      });
+
+      app.use(router.routes());
+
+      await context.with(
+        trace.setSpan(context.active(), rootSpan),
+        async () => {
+          await httpRequest.get(`http://localhost:${port}/post/0`);
+          rootSpan.end();
+
+          assert.deepStrictEqual(memoryExporter.getFinishedSpans().length, 2);
+          const requestHandlerSpan = memoryExporter
+            .getFinishedSpans()
+            .find(span => span.name.includes('router - /post/:id'));
+          assert.notStrictEqual(requestHandlerSpan, undefined);
+          sinon.assert.calledOnce(requestHook);
+          assert.strictEqual(
+            requestHandlerSpan?.attributes['http.method'],
+            'GET'
+          );
+          assert.strictEqual(
+            requestHandlerSpan?.attributes['app.env'],
+            'development'
+          );
+        }
+      );
+    });
+  });
+
   describe('Disabling koa instrumentation', () => {
     it('should not create new spans', async () => {
       plugin.disable();