diff --git a/lib/datadog/lambda.rb b/lib/datadog/lambda.rb index b23acd0..4eee6d4 100644 --- a/lib/datadog/lambda.rb +++ b/lib/datadog/lambda.rb @@ -66,7 +66,7 @@ def self.wrap(event, context, &block) record_enhanced('errors', context) raise e ensure - @listener&.on_end(response: @response) + @listener&.on_end(response: @response, request_context: context) @is_cold_start = false @metrics_client.close end diff --git a/lib/datadog/lambda/trace/listener.rb b/lib/datadog/lambda/trace/listener.rb index bff01da..e4ac564 100644 --- a/lib/datadog/lambda/trace/listener.rb +++ b/lib/datadog/lambda/trace/listener.rb @@ -42,7 +42,7 @@ def on_start(event:, request_context:, cold_start:) options[:service] = 'aws.lambda' options[:type] = 'serverless' - trace_digest = Datadog::Utils.send_start_invocation_request(event:) + trace_digest = Datadog::Utils.send_start_invocation_request(event:, request_context:) # Only continue trace from a new one if it exist, or else, # it will create a new trace, which is not ideal here. options[:continue_from] = trace_digest if trace_digest @@ -53,8 +53,8 @@ def on_start(event:, request_context:, cold_start:) end # rubocop:enable Metrics/AbcSize - def on_end(response:) - Datadog::Utils.send_end_invocation_request(response:, span_id: @trace.id) + def on_end(response:, request_context:) + Datadog::Utils.send_end_invocation_request(response:, span_id: @trace.id, request_context:) @trace&.finish end diff --git a/lib/datadog/lambda/utils/extension.rb b/lib/datadog/lambda/utils/extension.rb index 9351ab8..efdf03c 100644 --- a/lib/datadog/lambda/utils/extension.rb +++ b/lib/datadog/lambda/utils/extension.rb @@ -22,6 +22,8 @@ module Utils DD_SPAN_ID_HEADER = 'x-datadog-span-id' DD_PARENT_ID_HEADER = Datadog::Tracing::Distributed::Datadog::PARENT_ID_KEY + LAMBDA_RUNTIME_AWS_REQUEST_HEADER_ID = 'lambda-runtime-aws-request-id' + START_INVOCATION_URI = URI(EXTENSION_BASE_URL + START_INVOCATION_PATH).freeze END_INVOCATION_URI = URI(EXTENSION_BASE_URL + END_INVOCATION_PATH).freeze @@ -40,10 +42,13 @@ def self.check_extension_running File.exist?(EXTENSION_PATH) end - def self.send_start_invocation_request(event:) + def self.send_start_invocation_request(event:, request_context:) return unless extension_running? - response = Net::HTTP.post(START_INVOCATION_URI, event.to_json, request_headers) + headers = request_headers + headers[LAMBDA_RUNTIME_AWS_REQUEST_HEADER_ID] = request_context.aws_request_id + response = Net::HTTP.post(START_INVOCATION_URI, event.to_json, headers) + # Add origin, since tracer expects it for extraction response[Datadog::Trace::DD_ORIGIN] = 'lambda' @@ -53,7 +58,7 @@ def self.send_start_invocation_request(event:) end # rubocop:disable Metrics/AbcSize - def self.send_end_invocation_request(response:, span_id:) + def self.send_end_invocation_request(response:, span_id:, request_context:) return unless extension_running? request = Net::HTTP::Post.new(END_INVOCATION_URI) @@ -66,6 +71,9 @@ def self.send_end_invocation_request(response:, span_id:) # Propagator doesn't inject span_id, so we do it manually # It is needed for the extension to take this span id request[DD_SPAN_ID_HEADER] = span_id.to_s + + request[LAMBDA_RUNTIME_AWS_REQUEST_HEADER_ID] = request_context.aws_request_id + # Remove Parent ID if it is the same as the Span ID request.delete(DD_PARENT_ID_HEADER) if request[DD_PARENT_ID_HEADER] == span_id.to_s Datadog::Utils.logger.debug "End invocation request headers: #{request.to_hash}" diff --git a/test/datadog/lambda/utils/extension.spec.rb b/test/datadog/lambda/utils/extension.spec.rb index 19e7dae..eca4a4c 100644 --- a/test/datadog/lambda/utils/extension.spec.rb +++ b/test/datadog/lambda/utils/extension.spec.rb @@ -4,6 +4,7 @@ require 'datadog/lambda' require 'net/http' +require_relative '../../lambdacontextversion' describe Datadog::Utils do let(:headers) do @@ -25,6 +26,7 @@ describe '#send_start_invocation_request' do context 'when extension is running' do + ctx = LambdaContextVersion.new before(:each) do # Stub the extension_running? method to return true allow(Datadog::Utils).to receive(:extension_running?).and_return(true) @@ -39,11 +41,13 @@ it 'applies trace context from extension' do # Stub POST request to return a trace context + all_headers = Datadog::Utils.request_headers + all_headers['lambda-runtime-aws-request-id'] = ctx.aws_request_id expect(Net::HTTP).to receive(:post) - .with(Datadog::Utils::START_INVOCATION_URI, 'null', Datadog::Utils.request_headers) { headers } + .with(Datadog::Utils::START_INVOCATION_URI, 'null', all_headers) { headers } # Call the start request with an empty event - digest = Datadog::Utils.send_start_invocation_request(event: nil) + digest = Datadog::Utils.send_start_invocation_request(event: nil, request_context: ctx) expect(digest.trace_id.to_s).to eq('4110911582297405557') expect(digest.span_id.to_s).to eq('797643193680388254') @@ -52,11 +56,13 @@ it 'skips applying trace context when headers are not present' do # Stub POST request to return a trace context + all_headers = Datadog::Utils.request_headers + all_headers['lambda-runtime-aws-request-id'] = ctx.aws_request_id expect(Net::HTTP).to receive(:post) - .with(Datadog::Utils::START_INVOCATION_URI, 'null', Datadog::Utils.request_headers) { {} } + .with(Datadog::Utils::START_INVOCATION_URI, 'null', all_headers) { {} } # Call the start request with an empty event - Datadog::Utils.send_start_invocation_request(event: nil) + Datadog::Utils.send_start_invocation_request(event: nil, request_context: ctx) digest = Datadog::Tracing.active_trace.to_digest @@ -66,8 +72,9 @@ end context 'when extension is not running' do + ctx = LambdaContextVersion.new it 'does nothing' do - result = Datadog::Utils.send_start_invocation_request(event: nil) + result = Datadog::Utils.send_start_invocation_request(event: nil, request_context: ctx) expect(result).to eq(nil) end end @@ -75,6 +82,7 @@ describe '#send_end_invocation_request' do context 'when extension is running' do + ctx = LambdaContextVersion.new before(:each) do # Stub the extension_running? method to return true allow(Datadog::Utils).to receive(:extension_running?).and_return(true) @@ -92,13 +100,14 @@ allow(Net::HTTP).to receive(:post) { nil } # Call the start request with an empty event - Datadog::Utils.send_end_invocation_request(response: nil, span_id: nil) + Datadog::Utils.send_end_invocation_request(response: nil, span_id: nil, request_context: ctx) end end context 'when extension is not running' do + ctx = LambdaContextVersion.new it 'does nothing' do - result = Datadog::Utils.send_end_invocation_request(response: nil, span_id: nil) + result = Datadog::Utils.send_end_invocation_request(response: nil, span_id: nil, request_context: ctx) expect(result).to eq(nil) end end