Add request-id as header to start/end invocation (#129)

* Adding aws request id to headers on start and end invocation requests

* linted

* unbroke unit tests

* added request id to expected headers in start invocation request unit tests

* adding logs

* fixed assignment to headers in send start invocation

Author: rithikanarayan

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

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
 

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}"

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,15 +72,17 @@
     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
   end
 
   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