import { Context } from "aws-lambda" ; import { addLambdaFunctionTagsToXray , TraceContext , extractTraceContext , readStepFunctionContextFromEvent , StepFunctionContext , } from "./context" ; import { patchHttp , unpatchHttp } from "./patch-http" ; import { TraceContextService } from "./trace-context-service" ; import { extractTriggerTags } from "./trigger" ; import { logDebug } from "../utils" ; import { didFunctionColdStart } from "../utils/cold-start" ; import { datadogLambdaVersion } from "../constants" ; import { Source , ddtraceVersion } from "./constants" ; import { patchConsole } from "./patch-console" ; import { SpanContext , TraceOptions , TracerWrapper } from "./tracer-wrapper" ; export type TraceExtractor = ( event : any , context : Context ) => TraceContext ; export interface TraceConfig { default autoPatchHTTP : boolean ; /**
     * Whether to capture the lambda payload and response in Datadog.
     */ captureLambdaPayload : boolean ; /**
     * Whether to automatically patch console.log with Datadog's tracing ids.
     */ injectLogContext : boolean ; default mergeDatadogXrayTraces : boolean ; /**
     * Custom trace extractor function
     */ traceExtractor ? : TraceExtractor ; } export class TraceListener { private contextService : TraceContextService ; private context ? : Context ; private stepFunctionContext ? : StepFunctionContext ; private tracerWrapper : TracerWrapper ; public triggerTags ? : { [ key ] : string ; } ; public get currentTraceHeaders ( ) { return this . contextService . currentTraceHeaders ; } public get currentSpan ( ) { return this . tracerWrapper . currentSpan ; } constructor ( private config , private handlerName ) { this . tracerWrapper = new TracerWrapper ( ) ; this . contextService = new TraceContextService ( this . tracerWrapper ) ; } public onStartInvocation ( event , context ) { const tracerInitialized = this . tracerWrapper . isTracerAvailable ; if ( this . config . injectLogContext ) { patchConsole ( console , this . contextService ) ; logDebug ( "Patched console output with trace context" ) ; } else { logDebug ( "Didn't patch console output with trace context" ) ; } if ( this . config . autoPatchHTTP && ! tracerInitialized ) { logDebug ( "Patching HTTP libraries" ) ; patchHttp ( this . contextService ) ; } else { logDebug ( "Not patching HTTP libraries" , { autoPatchHTTP : this . config . autoPatchHTTP , tracerInitialized } ) ; } this . context = context ; this . triggerTags = extractTriggerTags ( event , context ) ; this . contextService . rootTraceContext = extractTraceContext ( event , context , this . config . traceExtractor ) ; this . stepFunctionContext = readStepFunctionContextFromEvent ( event ) ; } public async onCompleteInvocation ( ) { if ( this . triggerTags ) { addLambdaFunctionTagsToXray ( this . triggerTags ) ; } const tracerInitialized = this . tracerWrapper . isTracerAvailable ; if ( this . config . autoPatchHTTP && ! tracerInitialized ) { logDebug ( "Unpatching HTTP libraries" ) ; unpatchHttp ( ) ; } } public onWrap < T = ( ... args : any [ ] ) => any > ( func ) { let parentSpanContext = null ; if ( this . contextService . traceSource === Source . Event || this . config . mergeDatadogXrayTraces ) { const rootTraceHeaders = this . contextService . rootTraceHeaders ; parentSpanContext = this . tracerWrapper . extract ( rootTraceHeaders ) ; logDebug ( "Attempting to find parent for the aws.lambda span" ) ; } else { logDebug ( "Didn't attempt to find parent for aws.lambda span" , { mergeDatadogXrayTraces : this . config . mergeDatadogXrayTraces , traceSource : this . contextService . traceSource , } ) ; } const options = { } ; if ( this . context ) { logDebug ( "Creating the aws.lambda span" ) ; const functionArn = ( this . context . invokedFunctionArn ?? "" ) . toLowerCase ( ) ; const tk = functionArn . split ( ":" ) ; options . tags = { cold_start : didFunctionColdStart ( ) , function_arn : tk . length > 7 ? tk . slice ( 0 , 7 ) . join ( ":" ) : functionArn , function_version : tk . length > 7 ? tk [ 7 ] : "$LATEST" , request_id : this . context . awsRequestId , resource_names : this . context . functionName , functionname : this . context ?. functionName ?. toLowerCase ( ) , datadog_lambda : datadogLambdaVersion , dd_trace : ddtraceVersion , } ; if ( ( this . contextService . traceSource === Source . Xray && this . config . mergeDatadogXrayTraces ) || this . contextService . traceSource === Source . Event ) { options . tags [ "_dd.parent_source" ] = this . contextService . traceSource ; } if ( this . triggerTags ) { options . tags = { ... options . tags , ... this . triggerTags } ; } } if ( this . stepFunctionContext ) { logDebug ( "Applying step function context to the aws.lambda span" ) ; options . tags = { ... options . tags , ... this . stepFunctionContext , } ; } if ( parentSpanContext !== null ) { options . childOf = parentSpanContext ; } options . type = "serverless" ; options . service = "aws.lambda" ; if ( this . context ) { options . resource = this . context . functionName ; } return this . tracerWrapper . wrap ( "aws.lambda" , options , func ) ; } }