import { appendThriftObject , formatUrl , getProtocol , getTransport , Int64 , IProtocolConstructor , IRequestContext , ITransportConstructor , LogFunction , MessageType , ProtocolType , readThriftObject , TProtocol , TransportType , TTransport , } from '@creditkarma/thrift-server-core' ; import { getTracerForService , IZipkinTracerConfig , traceIdFromTraceId , ZipkinHeaders , } from '@creditkarma/zipkin-core' ; import { Instrumentation , TraceId , Tracer } from 'zipkin' ; import { IRequestResponse , IThriftClientFilterConfig , IThriftRequest , NextFunction , } from '@creditkarma/thrift-client' ; import * as TTwitter from '../ttwitter/com/creditkarma/finagle/thrift' ; import { defaultLogger } from './logger' ; export { TTwitter } ; export interface IClientId { name : string ; } export interface ITTwitterFileterOptions { localServiceName : string ; remoteServiceName : string ; destHeader ? : string ; isUpgraded ? : boolean ; clientId ? : IClientId ; transportType ? : TransportType ; protocolType ? : ProtocolType ; logger ? : LogFunction ; tracerConfig ? : IZipkinTracerConfig ; } const CAN_TRACE_METHOD_NAME = '__can__finagle__trace__v3__' ; function readRequestContext ( context , tracer , logger ) { if ( context !== undefined && context . traceId !== undefined ) { return { traceId : context . traceId , headers : { } , logger , } ; } else { const traceId = tracer . createRootId ( ) ; return { traceId : { traceId : traceId . traceId , spanId : traceId . spanId , parentId : traceId . parentId , sampled : traceId . sampled . getOrElse ( false ) , } , headers : { } , logger , } ; } } function upgradeRequest ( transportType = 'buffered' , protocolType = 'binary' ) { const Transport = getTransport ( transportType ) ; const Protocol = getProtocol ( protocolType ) ; const options = new TTwitter . ConnectionOptions ( ) ; const writer = new Transport ( ) ; const output = new Protocol ( writer ) ; output . writeMessageBegin ( CAN_TRACE_METHOD_NAME , MessageType . CALL , 0 ) ; options . write ( output ) ; output . writeMessageEnd ( ) ; return output . flush ( ) ; } export function ThriftClientTTwitterFilter < T > ( { localServiceName , remoteServiceName , destHeader = remoteServiceName , isUpgraded = true , clientId , transportType = 'buffered' , protocolType = 'binary' , logger = defaultLogger , tracerConfig , } ) { let hasUpgraded = false ; let upgradeRequested = false ; return { async handler ( request , next ) { if ( isUpgraded ) { function sendUpgradedRequest ( ) { logger ( [ 'info' , 'ThriftClientTTwitterFilter' ] , 'TTwitter upgraded' ) ; const tracer = getTracerForService ( localServiceName , tracerConfig ) ; const instrumentation = new Instrumentation . HttpClient ( { tracer , serviceName : localServiceName , remoteServiceName , } ) ; const requestContext = readRequestContext ( request , tracer , logger ) ; if ( requestContext . traceId !== undefined ) { tracer . setId ( traceIdFromTraceId ( requestContext . traceId ) ) ; return tracer . scoped ( ( ) => { const { headers , } = instrumentation . recordRequest ( { headers : { } } , formatUrl ( request . uri ) , request . methodName || 'post' ) ; const normalHeaders = Object . keys ( headers ) . reduce ( ( acc , name ) => { acc [ name . toLowerCase ( ) ] = headers [ name ] ; return acc ; } , { } ) ; const requestHeader = { trace_id : new Int64 ( `0x${ normalHeaders [ ZipkinHeaders . TraceId ] }` ) , span_id : new Int64 ( `0x${ normalHeaders [ ZipkinHeaders . SpanId ] }` ) , parent_span_id : normalHeaders [ ZipkinHeaders . ParentId ] !== undefined ? new Int64 ( `0x${ normalHeaders [ ZipkinHeaders . ParentId ] }` ) : undefined , sampled : normalHeaders [ ZipkinHeaders . Sampled ] === '1' , client_id : clientId !== undefined ? new TTwitter . ClientId ( clientId ) : undefined , contexts : [ ] , dest : destHeader , delegations : [ ] , } ; return appendThriftObject < TTwitter . IResponseHeaderArgs > ( requestHeader , request . data , TTwitter . RequestHeaderCodec , transportType , protocolType ) . then ( ( extended ) => { return next ( extended , request . context ) . then ( ( res ) : Promise < IRequestResponse > => { return readThriftObject < TTwitter . IResponseHeader > ( res . body , TTwitter . ResponseHeaderCodec , transportType , protocolType ) . then ( ( result ) => { return { statusCode : res . statusCode , headers : { thriftContext : result [ 0 ] , } , body : result [ 1 ] , } ; } , ( err ) => { logger ( [ 'warn' , 'ThriftClientTTwitterFilter' , ] , `Error reading context from Thrift response: ${ err . message }` ) ; return res ; } ) ; } ) ; } ) ; } ) ; } else { return next ( ) ; } } if ( hasUpgraded ) { return sendUpgradedRequest ( ) ; } else if ( upgradeRequested ) { return next ( ) ; } else { logger ( [ 'info' , 'ThriftClientTTwitterFilter' ] , 'Requesting TTwitter upgrade' ) ; upgradeRequested = true ; return next ( upgradeRequest ( ) , request . context ) . then ( ( upgradeResponse ) => { hasUpgraded = true ; return sendUpgradedRequest ( ) ; } , ( err ) => { logger ( [ 'info' ] , `Downgrading TTwitter request: ${ err . message }` ) ; return next ( ) ; } ) ; } } else { return next ( ) ; } } , } ; }