datadog-lambda 2.20.0 → 2.22.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 7190d039044c87a9835b29f812e930cc6c3fee9cdc5348226df7fe603b15a7bf
-  data.tar.gz: 35c495079f2b934ba036008084765d92a376f6924f7a4c5468ad34303ee1912c
+  metadata.gz: 1ff16978e295234189602cab59d2c4f2adbff91edc9bcb9a52bbde91fb744369
+  data.tar.gz: 4e89c094a7c0fb2e6161f7179266831df23dd4f665ccaafcc9b1d233fcce0860
 SHA512:
-  metadata.gz: c51f7c325e147d68a713d763c9550f2bf11b3ba9f7492f48c4b5195302f4c5fd4c0da20a9893ea19c27be24c4adf8740186f394250f830cfe8b6d8890c861d7b
-  data.tar.gz: 58f67b709d96a5b9e8ce628c4c1d0c02385508de007772a2560da26dc8cc7e90d23562063e3a8e73a5c90f73e71fd7c0cc0fac305c36cca6d901977d39c3fd6f
+  metadata.gz: 3738d873ac92162f632b6457ad1d188397eec4a3a59d9a523fcdf9b2f53eadd559f92f388f8d5d4a96da6cc38024124efb828b714c0cd44ed050a347eb246150
+  data.tar.gz: 37cfd0a36ae52012ebdc14cab31b1d22b20e968413ef136010fa21f0ee6be2b965257b2066d5d93ed7c3eabdf9989a2a40c84e5341f48aa1bb1ae0e0cffb29d9

data/lib/datadog/lambda/trace/constants.rb

@@ -17,9 +17,11 @@ module Datadog
     DD_TRACE_ID_HEADER = 'x-datadog-trace-id'
     DD_PARENT_ID_HEADER = 'x-datadog-parent-id'
     DD_SAMPLING_PRIORITY_HEADER = 'x-datadog-sampling-priority'
+    DD_ORIGIN = 'x-datadog-origin'
     DD_XRAY_SUBSEGMENT_NAME = 'datadog-metadata'
     DD_XRAY_SUBSEGMENT_KEY = 'trace'
     DD_XRAY_SUBSEGMENT_NAMESPACE = 'datadog'
+    DD_TRACE_MANAGED_SERVICES = 'DD_TRACE_MANAGED_SERVICES'
     SOURCE_XRAY = 'XRAY'
     SOURCE_EVENT = 'EVENT'
     XRAY_ENV_VAR = '_X_AMZN_TRACE_ID'

data/lib/datadog/lambda/trace/listener.rb

@@ -16,6 +16,7 @@ module Datadog
   module Trace
     # TraceListener tracks tracing context information
     class Listener
+      @trace = nil
       def initialize(handler_name:, function_name:, patch_http:,
                      merge_xray_traces:)
         @handler_name = handler_name
@@ -25,18 +26,11 @@ module Datadog
         Datadog::Trace.patch_http if patch_http
       end
 
-      def on_start(event:)
-        trace_context = Datadog::Trace.extract_trace_context(event,
-                                                             @merge_xray_traces)
+      # rubocop:disable Metrics/AbcSize
+      def on_start(event:, request_context:, cold_start:)
+        trace_context = Datadog::Trace.extract_trace_context(event, @merge_xray_traces)
         Datadog::Trace.trace_context = trace_context
         Datadog::Utils.logger.debug "extracted trace context #{trace_context}"
-      rescue StandardError => e
-        Datadog::Utils.logger.error "couldn't read tracing context #{e}"
-      end
-
-      def on_end; end
-
-      def on_wrap(request_context:, cold_start:, &block)
         options = get_option_tags(
           request_context: request_context,
           cold_start: cold_start
@@ -48,9 +42,19 @@ module Datadog
         options[:service] = 'aws.lambda'
         options[:span_type] = 'serverless'
         Datadog::Trace.apply_datadog_trace_context(Datadog::Trace.trace_context)
-        Datadog::Trace.wrap_datadog(options) do
-          block.call
-        end
+
+        trace_digest = Datadog::Utils.send_start_invocation_request(event: event)
+        # 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
+
+        @trace = Datadog::Tracing.trace('aws.lambda', **options)
+      end
+      # rubocop:enable Metrics/AbcSize
+
+      def on_end(response:)
+        Datadog::Utils.send_end_invocation_request(response: response)
+        @trace&.finish
       end
 
       private

data/lib/datadog/lambda/utils/extension.rb

@@ -8,16 +8,24 @@
 # Copyright 2023 Datadog, Inc.
 #
 require 'net/http'
+require 'datadog/tracing/contrib/http/distributed/fetcher'
 
 module Datadog
   # Utils contains utility functions shared between modules
   module Utils
-    AGENT_URL = 'http://127.0.0.1:8124'
-    HELLO_PATH = '/lambda/hello'
-    EXTENSION_CHECK_URI = URI(AGENT_URL + HELLO_PATH).freeze
     EXTENSION_PATH = '/opt/extensions/datadog-agent'
+    EXTENSION_BASE_URL = 'http://127.0.0.1:8124'
 
-    @is_extension_running = nil
+    START_INVOCATION_PATH = '/lambda/start-invocation'
+    END_INVOCATION_PATH = '/lambda/end-invocation'
+
+    START_INVOCATION_URI = URI(EXTENSION_BASE_URL + START_INVOCATION_PATH).freeze
+    END_INVOCATION_URI = URI(EXTENSION_BASE_URL + END_INVOCATION_PATH).freeze
+
+    # Internal communications use Datadog tracing headers
+    PROPAGATOR = Tracing::Distributed::Datadog.new(
+      fetcher: Tracing::Contrib::HTTP::Distributed::Fetcher
+    )
 
     def self.extension_running?
       return @is_extension_running unless @is_extension_running.nil?
@@ -26,16 +34,44 @@ module Datadog
     end
 
     def self.check_extension_running
-      return false unless File.exist?(EXTENSION_PATH)
+      File.exist?(EXTENSION_PATH)
+    end
 
-      begin
-        Net::HTTP.get(EXTENSION_CHECK_URI)
-      rescue StandardError => e
-        Datadog::Utils.logger.debug "extension is not running, returned with error #{e}"
-        return false
+    def self.send_start_invocation_request(event:)
+      return unless extension_running?
+
+      response = Net::HTTP.post(START_INVOCATION_URI, event.to_json, request_headers)
+      # Add origin, since tracer expects it for extraction
+      response[Datadog::Trace::DD_ORIGIN] = 'lambda'
+
+      PROPAGATOR.extract(response)
+    rescue StandardError => e
+      Datadog::Utils.logger.debug "failed on start invocation request to extension: #{e}"
+    end
+
+    def self.send_end_invocation_request(response:)
+      return unless extension_running?
+
+      request = Net::HTTP::Post.new(END_INVOCATION_URI)
+      request.body = response.to_json
+      request[Datadog::Core::Transport::Ext::HTTP::HEADER_DD_INTERNAL_UNTRACED_REQUEST] = 1
+
+      trace_digest = Datadog::Tracing.active_trace&.to_digest
+
+      PROPAGATOR.inject!(trace_digest, request)
+      Net::HTTP.start(END_INVOCATION_URI.host, END_INVOCATION_URI.port) do |http|
+        http.request(request)
       end
+    rescue StandardError => e
+      Datadog::Utils.logger.debug "failed on end invocation request to extension: #{e}"
+    end
 
-      true
+    def self.request_headers
+      {
+        # Header used to avoid tracing requests that are internal to
+        # Datadog products.
+        Datadog::Core::Transport::Ext::HTTP::HEADER_DD_INTERNAL_UNTRACED_REQUEST.to_sym => 'true'
+      }
     end
   end
 end

data/lib/datadog/lambda/version.rb

@@ -12,7 +12,7 @@ module Datadog
   module Lambda
     module VERSION
       MAJOR = 2
-      MINOR = 20
+      MINOR = 22
       PATCH = 0
       PRE = nil
 

data/lib/datadog/lambda.rb

@@ -21,6 +21,7 @@ module Datadog
   # Instruments AWS Lambda functions with Datadog distributed tracing and
   # custom metrics
   module Lambda
+    @response = nil
     @is_cold_start = true
     @patch_http = true
     @metrics_client = Metrics::Client.instance
@@ -30,7 +31,7 @@ module Datadog
     # See https://github.com/DataDog/dd-trace-rb/blob/master/docs/GettingStarted.md#quickstart-for-ruby-applications
     def self.configure_apm
       require 'datadog/tracing'
-      require 'ddtrace/transport/io'
+      require 'datadog/tracing/transport/io'
 
       @patch_http = false
       # Needed to keep trace flushes on a single line
@@ -39,10 +40,13 @@ module Datadog
       Datadog.configure do |c|
         unless Datadog::Utils.extension_running?
           c.tracing.writer = Datadog::Tracing::SyncWriter.new(
-            transport: Datadog::Transport::IO.default
+            transport: Datadog::Tracing::Transport::IO.default
           )
         end
         c.tags = { "_dd.origin": 'lambda' }
+        # Enable AWS SDK instrumentation
+        c.tracing.instrument :aws if trace_managed_services?
+
         yield(c) if block_given?
       end
     end
@@ -54,28 +58,25 @@ module Datadog
     def self.wrap(event, context, &block)
       Datadog::Utils.update_log_level
       @listener ||= initialize_listener
-      @listener.on_start(event: event)
       record_enhanced('invocations', context)
       begin
         cold = @is_cold_start
-        res = @listener.on_wrap(request_context: context, cold_start: cold) do
-          block.call
-        end
+        @listener&.on_start(event: event, request_context: context, cold_start: cold)
+        @response = block.call
       rescue StandardError => e
         record_enhanced('errors', context)
         raise e
       ensure
-        @listener.on_end
+        @listener&.on_end(response: @response)
         @is_cold_start = false
         @metrics_client.close
       end
-      res
+      @response
     end
 
     # Gets the current tracing context
     def self.trace_context
-      context = Hash[Datadog::Trace.trace_context]
-      context
+      Hash[Datadog::Trace.trace_context]
     end
 
     # Send a custom distribution metric
@@ -166,6 +167,15 @@ module Datadog
       dd_enhanced_metrics.downcase == 'true'
     end
 
+    # Read DD_TRACE_MANAGED_SERVICES environment variable
+    # @return [boolean] true if we should trace AWS services
+    def self.trace_managed_services?
+      dd_trace_managed_services = ENV[Trace::DD_TRACE_MANAGED_SERVICES]
+      return true if dd_trace_managed_services.nil?
+
+      dd_trace_managed_services.downcase == 'true'
+    end
+
     def self.initialize_listener
       handler = ENV['_HANDLER'].nil? ? 'handler' : ENV['_HANDLER']
       function = ENV['AWS_LAMBDA_FUNCTION_NAME']
@@ -176,10 +186,18 @@ module Datadog
         Datadog::Utils.logger.debug("Setting merge traces #{merge_xray_traces}")
       end
 
-      Trace::Listener.new(handler_name: handler,
-                          function_name: function,
-                          patch_http: @patch_http,
-                          merge_xray_traces: merge_xray_traces)
+      # Only initialize listener if Tracing enabled.
+      unless Datadog::Tracing.enabled?
+        Datadog::Utils.logger.debug 'dd-trace unavailable'
+        return nil
+      end
+
+      Trace::Listener.new(
+        handler_name: handler,
+        function_name: function,
+        patch_http: @patch_http,
+        merge_xray_traces: merge_xray_traces
+      )
     end
   end
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: datadog-lambda
 version: !ruby/object:Gem::Version
-  version: 2.20.0
+  version: 2.22.0
 platform: ruby
 authors:
 - Datadog, Inc.
 autorequire: 
 bindir: exe
 cert_chain: []
-date: 2023-08-08 00:00:00.000000000 Z
+date: 2023-12-01 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: aws-xray-sdk
@@ -44,14 +44,14 @@ dependencies:
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: 1.12.0
+        version: 1.17.0
   type: :development
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: 1.12.0
+        version: 1.17.0
 - !ruby/object:Gem::Dependency
   name: rake
   requirement: !ruby/object:Gem::Requirement