import * as api from '@opentelemetry/api' ; import { isWrapped , InstrumentationBase , InstrumentationNodeModuleDefinition , } from '@opentelemetry/instrumentation' ; import type * as koa from 'koa' ; import { KoaMiddleware , KoaContext , KoaComponentName , kLayerPatched , KoaLayerType , KoaInstrumentationConfig , } from './types' ; import { AttributeNames } from './enums/AttributeNames' ; import { VERSION } from './version' ; import { getMiddlewareMetadata , isLayerIgnored } from './utils' ; import { getRPCMetadata , RPCType , setRPCMetadata } from '@opentelemetry/core' ; /** Koa instrumentation for OpenTelemetry */ export class KoaInstrumentation extends InstrumentationBase < typeof koa > { static readonly component = KoaComponentName ; constructor ( config ? ) { super ( '@opentelemetry/instrumentation-koa' , VERSION , config ) ; } protected init ( ) { return new InstrumentationNodeModuleDefinition < typeof koa > ( 'koa' , [ '^2.0.0' ] , moduleExports => { if ( moduleExports == null ) { return moduleExports ; } api . diag . debug ( 'Patching Koa' ) ; if ( isWrapped ( moduleExports . prototype . use ) ) { this . _unwrap ( moduleExports . prototype , 'use' ) ; } this . _wrap ( moduleExports . prototype , 'use' , this . _getKoaUsePatch . bind ( this ) ) ; return moduleExports ; } , moduleExports => { api . diag . debug ( 'Unpatching Koa' ) ; if ( isWrapped ( moduleExports . prototype . use ) ) { this . _unwrap ( moduleExports . prototype , 'use' ) ; } } ) ; } param { KoaMiddleware } middleware private _getKoaUsePatch ( original ) { const plugin = this ; return function use ( this , middlewareFunction ) { let patchedFunction ; if ( middlewareFunction . router ) { patchedFunction = plugin . _patchRouterDispatch ( middlewareFunction ) ; } else { patchedFunction = plugin . _patchLayer ( middlewareFunction , false ) ; } return original . apply ( this , [ patchedFunction ] ) ; } ; } param { KoaMiddleware } dispatchLayer private _patchRouterDispatch ( dispatchLayer ) { api . diag . debug ( 'Patching @koa/router dispatch' ) ; const router = dispatchLayer . router ; const routesStack = router ?. stack ?? [ ] ; for ( const pathLayer of routesStack ) { const path = pathLayer . path ; const pathStack = pathLayer . stack ; for ( let j = 0 ; j < pathStack . length ; j ++ ) { const routedMiddleware = pathStack [ j ] ; pathStack [ j ] = this . _patchLayer ( routedMiddleware , true , path ) ; } } return dispatchLayer ; } param { KoaMiddleware } middlewareLayer param { boolean } isRouter param { string ? } layerPath private _patchLayer ( middlewareLayer , isRouter , layerPath ? ) { const layerType = isRouter ? KoaLayerType . ROUTER : KoaLayerType . MIDDLEWARE ; if ( middlewareLayer [ kLayerPatched ] === true || isLayerIgnored ( layerType , this . _config ) ) return middlewareLayer ; middlewareLayer [ kLayerPatched ] = true ; api . diag . debug ( 'patching Koa middleware layer' ) ; return async ( context , next ) => { const parent = api . trace . getSpan ( api . context . active ( ) ) ; if ( parent === undefined ) { return middlewareLayer ( context , next ) ; } const metadata = getMiddlewareMetadata ( context , middlewareLayer , isRouter , layerPath ) ; const span = this . tracer . startSpan ( metadata . name , { attributes : metadata . attributes , } ) ; const rpcMetadata = getRPCMetadata ( api . context . active ( ) ) ; if ( metadata . attributes [ AttributeNames . KOA_TYPE ] === KoaLayerType . ROUTER && rpcMetadata ?. type === RPCType . HTTP ) { rpcMetadata . span . updateName ( `${ context . method } ${ context . _matchedRoute }` ) ; } let newContext = api . trace . setSpan ( api . context . active ( ) , span ) ; if ( rpcMetadata ?. type === RPCType . HTTP ) { newContext = setRPCMetadata ( newContext , Object . assign ( rpcMetadata , { route : context . _matchedRoute } ) ) ; } return api . context . with ( newContext , async ( ) => { try { return await middlewareLayer ( context , next ) ; } catch ( err ) { span . recordException ( err ) ; throw err ; } finally { span . end ( ) ; } } ) ; } ; } }