datadog-lambda 2.19.0 → 2.21.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 14af71dd28cbd28bb7245d9a265c1ef5aae17cc7fca08cdfeec920a0df15863f
-  data.tar.gz: ddaef0b0f205fd9fabdcd1ee2d6663e3f47f1b80e5a7640bb8b0906600807a9d
+  metadata.gz: 6c867a4e0858e0ba773829993e4c60b2d84b26f8afe4e69a766ac75f6ef454c5
+  data.tar.gz: feaa85e2d5d998b32602b74958792a656f692eab4950fdc7da0cccb7d12542bd
 SHA512:
-  metadata.gz: 7e91c96b8bd3c20e420b94621df6eb260e3baaa310fc131600e6343663a6b8f613091a5f7b7beba460cce0c3c4c62b2a932e18c83f47ce2848c4a89b8c135053
-  data.tar.gz: 9dbf1189862c063bdf8eaa361fcebcbca5ee857bb8196b8d420d5a0f9151e7468f971d3628a68baf8182247a57227366ca6387f98523ccb6a86b2ad2b748536d
+  metadata.gz: cd5062d00ae5f88a8c5021f6c76c3b87fa9236d035508e04dbd33113cf51a5b7bec006060ab6dc12ea27da46feaf78ae3eb98d9bd4c907b5a972c7176ca40ee7
+  data.tar.gz: 823cc2efc44cfefd3e0829ca038e966f6b999fd5c2648131a91a2b8b805dbd8ae5aeb5cdba1c14d3c8579aafffb719679b8e1bdd60306f37c80c3bc176f154d2

data/lib/datadog/lambda/metrics.rb

@@ -0,0 +1,70 @@
+# frozen_string_literal: true
+
+#
+# Unless explicitly stated otherwise all files in this repository are licensed
+# under the Apache License Version 2.0.
+#
+# This product includes software developed at Datadog (https://www.datadoghq.com/).
+# Copyright 2023 Datadog, Inc.
+#
+
+require 'datadog/statsd'
+require 'time'
+
+module Datadog
+  # Metrics module contains the singleton class Client to send custom
+  # metrics to Datadog.
+  module Metrics
+    URL = 'localhost'
+    PORT = '8125'
+    # Client is a singleton class that instantiates a Datadog::Statsd
+    # client to send metrics to the Datadog Extension if present.
+    class Client
+      private_class_method :new
+
+      def self.instance
+        @instance ||= new
+      end
+
+      def distribution(name, value, time: nil, **tags)
+        tag_list = get_tags(**tags)
+
+        if Datadog::Utils.extension_running?
+          begin
+            @statsd.distribution(name, value, tags: tag_list)
+          rescue StandardError => e
+            Datadog::Utils.logger.warning "error sending metric to the extension: #{e}"
+          end
+        else
+          time ||= Time.now
+          time_ms = time.to_i
+
+          metric = { e: time_ms, m: name, t: tag_list, v: value }
+          puts metric.to_json
+        end
+      end
+
+      def close
+        @statsd&.close
+      end
+
+      def get_tags(**tags)
+        tag_list = ["dd_lambda_layer:datadog-ruby#{Datadog::Lambda.dd_lambda_layer_tag}"]
+        tags.each do |tag|
+          tag_list << "#{tag[0]}:#{tag[1]}"
+        end
+
+        tag_list
+      end
+
+      private_class_method def initialize
+        if Datadog::Utils.extension_running?
+          Datadog::Utils.logger.debug 'sending metrics through extension'
+          @statsd = Datadog::Statsd.new(URL, PORT, single_thread: true)
+        else
+          Datadog::Utils.logger.debug 'metrics are going to be handled by the forwarder'
+        end
+      end
+    end
+  end
+end

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

@@ -20,6 +20,7 @@ module Datadog
     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,14 @@ 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 = Datadog::Tracing.trace('aws.lambda', **options)
+        Datadog::Utils.send_start_invocation_request(event: event)
+      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

@@ -12,22 +12,60 @@ require 'net/http'
 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'
 
-    def self.extension_running
-      return false unless File.exist?(EXTENSION_PATH)
+    START_INVOCATION_PATH = '/lambda/start-invocation'
+    END_INVOCATION_PATH = '/lambda/end-invocation'
 
-      begin
-        Net::HTTP.get(EXTENSION_CHECK_URI)
-      rescue StandardError => e
-        Datadog::Utils.logger.debug "extension is not running, returned with error #{e}"
-        return false
+    START_INVOCATION_URI = URI(EXTENSION_BASE_URL + START_INVOCATION_PATH).freeze
+    END_INVOCATION_URI = URI(EXTENSION_BASE_URL + END_INVOCATION_PATH).freeze
+
+    def self.extension_running?
+      return @is_extension_running unless @is_extension_running.nil?
+
+      @is_extension_running = check_extension_running
+    end
+
+    def self.check_extension_running
+      File.exist?(EXTENSION_PATH)
+    end
+
+    def self.send_start_invocation_request(event:)
+      return unless extension_running?
+
+      response = Net::HTTP.post(START_INVOCATION_URI, event.to_json, request_headers)
+
+      trace_digest = Tracing::Propagation::HTTP.extract(response)
+      # Only continue trace from a new one if it exist, or else,
+      # it will create a new trace, which is not ideal here.
+      Tracing.continue_trace!(trace_digest) if trace_digest
+    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 = Datadog::Tracing.active_trace
+      Tracing::Propagation::HTTP.inject!(trace, 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 = 19
+      MINOR = 21
       PATCH = 0
       PRE = nil
 

data/lib/datadog/lambda.rb

@@ -9,39 +9,44 @@
 #
 # rubocop:disable Metrics/ModuleLength
 
+require 'datadog/lambda/metrics'
 require 'datadog/lambda/trace/listener'
 require 'datadog/lambda/utils/logger'
 require 'datadog/lambda/utils/extension'
 require 'datadog/lambda/trace/patch_http'
 require 'json'
-require 'time'
 require 'datadog/lambda/version'
 
 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
 
     # Configures Datadog's APM tracer with lambda specific defaults.
     # Same options can be given as Datadog.configure in tracer
     # 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
       $stdout.sync = true
 
       Datadog.configure do |c|
-        unless Datadog::Utils.extension_running
+        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
@@ -53,27 +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
@@ -85,15 +88,7 @@ module Datadog
       raise 'name must be a string' unless name.is_a?(String)
       raise 'value must be a number' unless value.is_a?(Numeric)
 
-      time ||= Time.now
-      time_ms = time.to_f.to_i
-
-      tag_list = ["dd_lambda_layer:datadog-ruby#{dd_lambda_layer_tag}"]
-      tags.each do |tag|
-        tag_list.push("#{tag[0]}:#{tag[1]}")
-      end
-      metric = { e: time_ms, m: name, t: tag_list, v: value }
-      puts metric.to_json
+      @metrics_client.distribution(name, value, time: time, **tags)
     end
 
     def self.dd_lambda_layer_tag
@@ -172,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']
@@ -182,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.19.0
+  version: 2.21.0
 platform: ruby
 authors:
 - Datadog, Inc.
 autorequire: 
 bindir: exe
 cert_chain: []
-date: 2023-05-10 00:00:00.000000000 Z
+date: 2023-10-16 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: aws-xray-sdk
@@ -24,20 +24,34 @@ dependencies:
     - - "~>"
       - !ruby/object:Gem::Version
         version: 0.11.3
+- !ruby/object:Gem::Dependency
+  name: dogstatsd-ruby
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '5.0'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '5.0'
 - !ruby/object:Gem::Dependency
   name: ddtrace
   requirement: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: 1.11.0
+        version: 1.15.0
   type: :development
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: 1.11.0
+        version: 1.15.0
 - !ruby/object:Gem::Dependency
   name: rake
   requirement: !ruby/object:Gem::Requirement
@@ -119,6 +133,7 @@ extensions: []
 extra_rdoc_files: []
 files:
 - lib/datadog/lambda.rb
+- lib/datadog/lambda/metrics.rb
 - lib/datadog/lambda/trace/constants.rb
 - lib/datadog/lambda/trace/context.rb
 - lib/datadog/lambda/trace/ddtrace.rb