feat: added tracing for koa-router functions and routes

Author: carolinee21

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

@@ -20,6 +20,8 @@ import * as shimmer from 'shimmer';
 import { Parameters, KoaMiddleware, KoaContext, KoaComponentName } from './types';
 import { VERSION } from './version';
 import { getMiddlewareMetadata } from './utils';
+// import { Layer } from '@koa/router';
+import Router = require('@koa/router');
 
 
 /** Koa instrumentation plugin for OpenTelemetry */
@@ -63,38 +65,77 @@ export class KoaPlugin extends BasePlugin<typeof koa> {
         middlewareFunction : KoaMiddleware,
         ...args: Parameters<typeof original>
       ) {
-
-        var oldMiddleware = middlewareFunction;
-        var patchedFunction = plugin._patchLayer(oldMiddleware);
+        console.log("mw name: " + middlewareFunction.name);
+        var isRouterDispatch = middlewareFunction.name == "dispatch";
+        isRouterDispatch = (middlewareFunction as any).router;
+      var patchedFunction;
+        if (isRouterDispatch) {
+          patchedFunction = plugin._patchRoutes(middlewareFunction);
+        } else {
+          patchedFunction = plugin._patchLayer(middlewareFunction, false);         
+        }
         const route = original.apply(this, [patchedFunction]);
-        // var lastLayer = route.middleware[route.middleware.length - 1];
 
         return route;
+
         // tslint:disable-next-line:no-any
       } as any;
 
   }
 
-  private _patchLayer (middlewareLayer: KoaMiddleware) {
+  private _patchLayer (middlewareLayer: KoaMiddleware, isRouter: boolean, layerName?: string) {
     const plugin = this;
+    console.log(".....setting up span patch for: " + layerName ?? middlewareLayer.name);
     const patchedLayer = (context: KoaContext, next: koa.Next) => {
         const currentSpan = this._tracer.getCurrentSpan();
         if (!currentSpan) {
             console.log('--- No current span');
         } 
 
         const metadata = getMiddlewareMetadata(context);
-        const span = plugin._tracer.startSpan(metadata.name, {
+        var spanName = layerName ?? middlewareLayer.name;
+        if (!spanName) {
+          spanName = metadata.name;
+        }
+        const span = plugin._tracer.startSpan(spanName, {
           attributes: metadata.attributes,
         });
-
         var result = middlewareLayer(context, next);
         span.end();
         return result;
-
     }
     return patchedLayer;
   }
+  
+
+  private _patchRoutes (dispatchLayer: KoaMiddleware) {
+    const plugin = this;
+    console.log(".....set up");
+
+    var smth = (dispatchLayer as any);
+    var router = smth.router as Router;
+    
+    var routesStack = router.stack;
+    for (var i = 0; i < routesStack.length; i++) {
+      var pathLayer : Router.Layer = routesStack[i];
+      var path = pathLayer.path;
+      var pathStack = pathLayer.stack;
+      for (var j = 0; j < pathStack.length; j++) {
+        var pop = pathStack[j];
+        var newpop = plugin._patchLayer(pop, true, path);
+        pathStack[j] = newpop;
+      }
+    }
+    
+    const dispatcher = (context: KoaContext, next: koa.Next) => {
+      console.log(".....dispatching");
+      var result = smth(context, next);
+      return result;
+    }
+    return dispatcher;
+
+  }
+  
 }
 
 export const plugin = new KoaPlugin(KoaPlugin.component);

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

@@ -13,13 +13,23 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-import { Middleware, ParameterizedContext, DefaultContext } from 'koa';
+import { Middleware, ParameterizedContext} from 'koa';
+import { RouterParamContext} from '@koa/router';
 
 export type Parameters<T> = T extends (...args: infer T) => any ? T : unknown[];
 
-export type KoaMiddleware = Middleware<ParameterizedContext<any, DefaultContext>>;
+export type KoaMiddleware = Middleware<any, RouterParamContext>;
 
-export type KoaContext = ParameterizedContext<ParameterizedContext<any, DefaultContext>, DefaultContext>;
+// export type KoaContext = ParameterizedContext<ParameterizedContext<any, DefaultContext>, DefaultContext>;
+export type KoaContext = ParameterizedContext<ParameterizedContext<any, any>, any>;
+
+export type KoaLayer = {
+    handle: Function;
+    name: string;
+    params: { [key: string]: string };
+    path: string;
+    regexp: RegExp;
+  };
 
 export enum AttributeNames {
     COMPONENT = 'component',
@@ -35,3 +45,33 @@ export enum AttributeNames {
 
 export const KoaComponentName : string = 'koa';
 
+/*
+
+Argument of type 
+'Middleware<ParameterizedContext<any, RouterParamContext<any, {}>>>' 
+is not assignable to parameter of type 
+'Middleware<ParameterizedContext<ParameterizedContext<any, DefaultContext>, DefaultContext>>'.
+
+
+
+Type 
+'ParameterizedContext<ParameterizedContext<any, DefaultContext>, DefaultContext>' 
+is not assignable to type 
+'ParameterizedContext<any, RouterParamContext<any, {}>>'.
+Type 'ExtendableContext & { state: ParameterizedContext<any, DefaultContext>; } & DefaultContext' is missing the following properties from type 'RouterParamContext<any, {}>': params, router, _matchedRoute, _matchedRouteNamez
+
+
+
+Argument of type 
+'Middleware<ParameterizedContext<any, RouterParamContext<any, {}>>>' 
+is not assignable to parameter of type 
+'Middleware<ParameterizedContext<ParameterizedContext<string, DefaultContext>, DefaultContext>>'.
+  
+
+Type 
+  'ParameterizedContext<ParameterizedContext<string, DefaultContext>, DefaultContext>' 
+  is not assignable to type 
+  'ParameterizedContext<any, RouterParamContext<any, {}>>'.
+
+
+*/