import { Instrumentation , TraceId , Tracer } from 'zipkin' ; import { CoreOptions } from 'request' ; import { formatUrl , IRequestContext , IRequestHeaders , LogFunction , } from '@creditkarma/thrift-server-core' ; import { IRequest , IRequestResponse , IThriftClientFilter , IThriftRequest , NextFunction , } from '@creditkarma/thrift-client' ; import { addL5Dheaders , containsZipkinHeaders , getTracerForService , hasL5DHeader , IZipkinClientOptions , traceIdForHeaders , traceIdFromTraceId , } from '@creditkarma/zipkin-core' ; function applyL5DHeaders ( requestHeaders , headers ) { if ( hasL5DHeader ( requestHeaders ) ) { return addL5Dheaders ( headers ) ; } else { return headers ; } } function readRequestContext ( requestHeaders , tracer , logger ? ) { if ( containsZipkinHeaders ( requestHeaders ) ) { return { traceId : traceIdForHeaders ( requestHeaders ) , headers : requestHeaders , logger , } ; } else { const rootId = tracer . createRootId ( ) ; return { traceId : { traceId : rootId . traceId , spanId : rootId . spanId , parentId : rootId . parentId , sampled : rootId . sampled . getOrElse ( false ) , } , headers : { } , logger , } ; } } function readRequestHeaders ( request ) { if ( request . context && request . context . headers ) { return request . context . headers ; } else { return { } ; } } export function ThriftClientZipkinFilter < Context extends IRequest > ( { localServiceName , remoteServiceName , tracerConfig = { } , } ) { const serviceName = remoteServiceName || localServiceName ; const tracer = getTracerForService ( serviceName , tracerConfig ) ; const instrumentation = new Instrumentation . HttpClient ( { tracer , serviceName : localServiceName , remoteServiceName , } ) ; return { methods : [ ] , handler ( request , next ) { const requestHeaders = readRequestHeaders ( request ) ; const requestContext = readRequestContext ( requestHeaders , tracer , tracerConfig . logger ) ; if ( requestContext . traceId !== undefined ) { tracer . setId ( traceIdFromTraceId ( requestContext . traceId ) ) ; return tracer . scoped ( ( ) => { const updatedHeaders = instrumentation . recordRequest ( { headers : { } } , formatUrl ( request . uri ) , request . methodName || '' ) . headers ; const traceId = tracer . id ; const withLD5Headers = applyL5DHeaders ( requestHeaders , updatedHeaders ) ; return next ( request . data , { headers : withLD5Headers } ) . then ( ( res ) => { instrumentation . recordResponse ( traceId as any , `${ res . statusCode }` ) ; return res ; } , ( err ) => { instrumentation . recordError ( traceId as any , err ) ; return Promise . reject ( err ) ; } ) ; } ) ; } else { return next ( ) ; } } , } ; }